Merge pull request #262 from alcaeus/fix-generator-memory-leak

Fix memory leak due to partially iterated generators

lib/Alcaeus/MongoDbAdapter/AbstractCursor.php

@@ -52,7 +52,7 @@ abstract class AbstractCursor
     protected $db;
 
     /**
-     * @var \Iterator
+     * @var CursorIterator
      */
     protected $iterator;
 
@@ -292,9 +292,8 @@ abstract class AbstractCursor
     protected function ensureIterator()
     {
         if ($this->iterator === null) {
-            // MongoDB\Driver\Cursor needs to be wrapped into a \Generator so that a valid \Iterator with working implementations of
-            // next, current, valid, key and rewind is returned. These methods don't work if we wrap the Cursor inside an \IteratorIterator
             $this->iterator = $this->wrapTraversable($this->ensureCursor());
+            $this->iterator->rewind();
         }
 
         return $this->iterator;
@@ -302,13 +301,11 @@ abstract class AbstractCursor
 
     /**
      * @param \Traversable $traversable
-     * @return \Generator
+     * @return CursorIterator
      */
     protected function wrapTraversable(\Traversable $traversable)
     {
-        foreach ($traversable as $key => $value) {
-            yield $key => $value;
-        }
+        return new CursorIterator($traversable);
     }
 
     /**

lib/Alcaeus/MongoDbAdapter/CursorIterator.php

@@ -0,0 +1,38 @@
+<?php
+
+namespace Alcaeus\MongoDbAdapter;
+
+use IteratorIterator;
+use MongoDB\BSON\ObjectID;
+use Traversable;
+
+/**
+ * @internal
+ */
+final class CursorIterator extends IteratorIterator
+{
+    /** @var bool */
+    private $useIdAsKey;
+
+    public function __construct(Traversable $iterator, $useIdAsKey = false)
+    {
+        parent::__construct($iterator);
+
+        $this->useIdAsKey = $useIdAsKey;
+    }
+
+    public function key()
+    {
+        if (!$this->useIdAsKey) {
+            return parent::key();
+        }
+
+        $current = $this->current();
+
+        if (!isset($current->_id) || (is_object($current->_id) && !$current->_id instanceof ObjectID)) {
+            return parent::key();
+        }
+
+        return (string) $current->_id;
+    }
+}

lib/Mongo/MongoCursor.php

@@ -18,6 +18,7 @@ if (class_exists('MongoCursor', false)) {
 }
 
 use Alcaeus\MongoDbAdapter\AbstractCursor;
+use Alcaeus\MongoDbAdapter\CursorIterator;
 use Alcaeus\MongoDbAdapter\TypeConverter;
 use Alcaeus\MongoDbAdapter\ExceptionConverter;
 use MongoDB\Driver\Cursor;
@@ -472,16 +473,11 @@ class MongoCursor extends AbstractCursor implements Iterator, Countable, MongoCu
 
     /**
      * @param \Traversable $traversable
-     * @return \Generator
+     * @return CursorIterator
      */
     protected function wrapTraversable(\Traversable $traversable)
     {
-        foreach ($traversable as $key => $value) {
-            if (isset($value->_id) && ($value->_id instanceof \MongoDB\BSON\ObjectID || !is_object($value->_id))) {
-                $key = (string) $value->_id;
-            }
-            yield $key => $value;
-        }
+        return new CursorIterator($traversable, true);
     }
 
     /**