PHPLIB-501: Use IteratorIterator to prevent memory leak with generators

Author: alcaeus

src/Model/CachingIterator.php

@@ -18,8 +18,8 @@
 namespace MongoDB\Model;
 
 use Countable;
-use Generator;
 use Iterator;
+use IteratorIterator;
 use Traversable;
 use function count;
 use function current;
@@ -41,7 +41,7 @@ class CachingIterator implements Countable, Iterator
     /** @var array */
     private $items = [];
 
-    /** @var Generator */
+    /** @var IteratorIterator */
     private $iterator;
 
     /** @var boolean */
@@ -50,6 +50,9 @@ class CachingIterator implements Countable, Iterator
     /** @var boolean */
     private $iteratorExhausted = false;
 
+    /** @var boolean */
+    private $iteratorPrepared = false;
+
     /**
      * Initialize the iterator and stores the first item in the cache. This
      * effectively rewinds the Traversable and the wrapping Generator, which
@@ -61,8 +64,7 @@ class CachingIterator implements Countable, Iterator
      */
     public function __construct(Traversable $traversable)
     {
-        $this->iterator = $this->wrapTraversable($traversable);
-        $this->storeCurrentItem();
+        $this->iterator = new IteratorIterator($traversable);
     }
 
     /**
@@ -82,6 +84,8 @@ public function count()
      */
     public function current()
     {
+        $this->prepareIterator();
+
         return current($this->items);
     }
 
@@ -91,6 +95,8 @@ public function current()
      */
     public function key()
     {
+        $this->prepareIterator();
+
         return key($this->items);
     }
 
@@ -100,9 +106,15 @@ public function key()
      */
     public function next()
     {
+        $this->prepareIterator();
+
         if (! $this->iteratorExhausted) {
+            $this->iteratorAdvanced = true;
             $this->iterator->next();
+
             $this->storeCurrentItem();
+
+            $this->iteratorExhausted = ! $this->iterator->valid();
         }
 
         next($this->items);
@@ -114,6 +126,8 @@ public function next()
      */
     public function rewind()
     {
+        $this->prepareIterator();
+
         /* If the iterator has advanced, exhaust it now so that future iteration
          * can rely on the cache.
          */
@@ -138,11 +152,24 @@ public function valid()
      */
     private function exhaustIterator()
     {
+        $this->prepareIterator();
+
         while (! $this->iteratorExhausted) {
             $this->next();
         }
     }
 
+    private function prepareIterator()
+    {
+        if ($this->iteratorPrepared) {
+            return;
+        }
+
+        $this->iterator->rewind();
+        $this->iteratorPrepared = true;
+        $this->storeCurrentItem();
+    }
+
     /**
      * Stores the current item in the cache.
      */
@@ -156,20 +183,4 @@ private function storeCurrentItem()
 
         $this->items[$key] = $this->iterator->current();
     }
-
-    /**
-     * Wraps the Traversable with a Generator.
-     *
-     * @param Traversable $traversable
-     * @return Generator
-     */
-    private function wrapTraversable(Traversable $traversable)
-    {
-        foreach ($traversable as $key => $value) {
-            yield $key => $value;
-            $this->iteratorAdvanced = true;
-        }
-
-        $this->iteratorExhausted = true;
-    }
 }