Specify time limit operation on a mongodb cursor

lib/Doctrine/MongoDB/Cursor.php

@@ -78,6 +78,7 @@ class Cursor implements CursorInterface
     protected $sort;
     protected $tailable;
     protected $timeout;
+    protected $maxTimeMS;
 
     /**
      * Constructor.
@@ -508,6 +509,9 @@ public function recreate()
         if ($this->timeout !== null) {
             $this->mongoCursor->timeout($this->timeout);
         }
+        if ($this->maxTimeMS !== null) {
+            $this->mongoCursor->addOption('$maxTimeMS', $this->maxTimeMS);
+        }
     }
 
     /**
@@ -664,6 +668,20 @@ public function timeout($ms)
         return $this;
     }
 
+    /**
+     * Wrapper method for MongoCursor::maxTimeMS().
+     *
+     * @see http://php.net/manual/en/mongocursor.maxtimems.php
+     * @param integer $ms
+     * @return self
+     */
+    public function maxTimeMS($ms)
+    {
+        $this->maxTimeMS = (integer)$ms;
+        $this->mongoCursor->addOption('$maxTimeMS', (integer)$ms);
+        return $this;
+    }
+
     /**
      * Return the cursor's results as an array.
      *

lib/Doctrine/MongoDB/EagerCursor.php

@@ -440,4 +440,15 @@ public function timeout($ms)
 
         return $this;
     }
+
+
+    /**
+     * {@inheritdoc}
+     */
+    public function maxTimeMS($ms)
+    {
+        $this->cursor->maxTimeMS($ms);
+
+        return $this;
+    }
 }

lib/Doctrine/MongoDB/LoggableCursor.php

@@ -139,4 +139,17 @@ public function sort($fields)
 
         return parent::sort($fields);
     }
+
+    /**
+     * @see Cursor::maxTimeMS()
+     */
+    public function maxTimeMS($ms)
+    {
+        $this->log(array(
+            'maxTimeMS' => true,
+            'maxTimeMSNum' => $ms,
+        ));
+
+        return parent::maxTimeMS($ms);
+    }
 }

lib/Doctrine/MongoDB/Query/Builder.php

@@ -738,6 +738,18 @@ public function hint($index)
         return $this;
     }
 
+    /**
+     * Specifies a cumulative time limit in milliseconds for processing operations on a cursor.
+     *
+     * @param int $ms
+     * @return $this
+     */
+    public function maxTimeMS($ms)
+    {
+        $this->query['maxTimeMS'] = $ms;
+        return $this;
+    }
+
     /**
      * Set the immortal cursor flag.
      *

lib/Doctrine/MongoDB/Query/Query.php

@@ -395,7 +395,7 @@ protected function prepareCursor(Cursor $cursor)
             $cursor->setReadPreference($this->query['readPreference'], $this->query['readPreferenceTags']);
         }
 
-        foreach ($this->getQueryOptions('hint', 'immortal', 'limit', 'skip', 'slaveOkay', 'sort') as $key => $value) {
+        foreach ($this->getQueryOptions('hint', 'immortal', 'limit', 'skip', 'slaveOkay', 'sort', 'maxTimeMS') as $key => $value) {
             $cursor->$key($value);
         }
 

tests/Doctrine/MongoDB/Tests/CursorTest.php

@@ -381,6 +381,43 @@ public function testRecreate()
         $cursor->recreate();
     }
 
+    public function testSetMaxTimeMSWhenRecreateCursor()
+    {
+        $self = $this;
+
+        $setCursorExpectations = function($mongoCursor) use ($self) {
+            $mongoCursor->expects($self->once())
+                ->method('addOption')
+                ->with(
+                    $self->equalTo('$maxTimeMS'),
+                    $self->equalTo(30000)
+                );
+        };
+
+        $mongoCursor = $this->getMockMongoCursor();
+        $recreatedMongoCursor = $this->getMockMongoCursor();
+
+        $setCursorExpectations($mongoCursor);
+        $setCursorExpectations($recreatedMongoCursor);
+
+        $mongoCollection = $this->getMockCollection();
+        $mongoCollection->expects($this->once())
+            ->method('find')
+            ->with(array('x' => 9500), array())
+            ->will($this->returnValue($recreatedMongoCursor));
+
+        $collection = $this->getMockCollection();
+        $collection->expects($this->once())
+            ->method('getMongoCollection')
+            ->will($this->returnValue($mongoCollection));
+
+        $cursor = $this->getTestCursor($collection, $mongoCursor, array('x' => 9500));
+
+        $cursor->maxTimeMS(30000);
+
+        $cursor->recreate();
+    }
+
     private function getMockCollection()
     {
         return $this->getMockBuilder('Doctrine\MongoDB\Collection')

tests/Doctrine/MongoDB/Tests/LoggableCursorTest.php

@@ -23,6 +23,7 @@ public function provideLoggedMethods()
             array('limit'),
             array('hint', array()),
             array('snapshot'),
+            array('maxTimeMS', array())
         );
     }
 }

tests/Doctrine/MongoDB/Tests/Query/QueryTest.php

@@ -286,6 +286,32 @@ public function testUseIdentifierKeys()
         $this->assertSame($cursor, $query->execute());
     }
 
+    public function testSpecifyMaxTimeMSOnCursor()
+    {
+        $cursor = $this->getMockCursor();
+        $collection = $this->getMockCollection();
+
+        $collection->expects($this->once())
+            ->method('find')
+            ->with($this->equalTo(array('foo' => 'bar')))
+            ->will($this->returnValue($cursor));
+
+        $cursor->expects($this->once())
+            ->method('maxTimeMS')
+            ->with($this->equalTo(30000))
+            ->will($this->returnValue($cursor));
+
+        $queryArray = array(
+            'type' => Query::TYPE_FIND,
+            'query' => array('foo' => 'bar'),
+            'maxTimeMS' => 30000
+        );
+
+        $query = new Query($collection, $queryArray, array());
+
+        $this->assertSame($cursor, $query->execute());
+    }
+
     /**
      * @return \Doctrine\MongoDB\Collection
      */