add cleanup iterator

1 files changed, 25 insertions, 0 deletions

diff --git a/core/src/main/scala/spark/util/CleanupIterator.scala b/core/src/main/scala/spark/util/CleanupIterator.scala
new file mode 100644
index 0000000000..d2093c0230
--- /dev/null
+++ b/core/src/main/scala/spark/util/CleanupIterator.scala
@@ -0,0 +1,25 @@
+package spark.util
+
+/**
+ * Wrapper around an iterator which calls a cleanup method when its finished iterating through its elements
+ */
+abstract class CleanupIterator[+A, +I <: Iterator[A]](sub: I) extends Iterator[A]{
+  def next = sub.next
+  def hasNext = {
+    val r = sub.hasNext
+    if (!r) {
+      cleanup
+    }
+    r
+  }
+
+  def cleanup
+}
+
+object CleanupIterator {
+  def apply[A, I <: Iterator[A]](sub: I, cleanupFunction: => Unit) : CleanupIterator[A,I] = {
+    new CleanupIterator[A,I](sub) {
+      def cleanup = cleanupFunction
+    }
+  }
+}
\ No newline at end of file