Add support for maxTimeMS mongoOption (#3018)

Author: TylerBrock

spec/MongoStorageAdapter.spec.js

@@ -60,6 +60,43 @@ describe_only_db('mongo')('MongoStorageAdapter', () => {
     });
   });
 
+  it('find succeeds when query is within maxTimeMS', (done) => {
+    const maxTimeMS = 250;
+    let adapter = new MongoStorageAdapter({
+      uri: databaseURI,
+      mongoOptions: { maxTimeMS },
+    });
+    adapter.createObject('Foo', { fields: {} }, { objectId: 'abcde' })
+    .then(() => adapter._rawFind('Foo', { '$where': `sleep(${maxTimeMS / 2})` }))
+    .then(
+      () => done(),
+      (err) => {
+        done.fail(`maxTimeMS should not affect fast queries ${err}`);
+      }
+    );
+  })
+
+  it('find fails when query exceeds maxTimeMS', (done) => {
+    const maxTimeMS = 250;
+    let adapter = new MongoStorageAdapter({
+      uri: databaseURI,
+      mongoOptions: { maxTimeMS },
+    });
+    adapter.createObject('Foo', { fields: {} }, { objectId: 'abcde' })
+    .then(() => adapter._rawFind('Foo', { '$where': `sleep(${maxTimeMS * 2})` }))
+    .then(
+      () => {
+        done.fail('Find succeeded despite taking too long!');
+      },
+      (err) => {
+        expect(err.name).toEqual('MongoError');
+        expect(err.code).toEqual(50);
+        expect(err.message).toEqual('operation exceeded time limit');
+        done();
+      }
+    );
+  });
+
   it('stores pointers with a _p_ prefix', (done) => {
     let obj = {
       objectId: 'bar',

src/Adapters/Storage/Mongo/MongoCollection.js

@@ -13,8 +13,8 @@ export default class MongoCollection {
   // none, then build the geoindex.
   // This could be improved a lot but it's not clear if that's a good
   // idea. Or even if this behavior is a good idea.
-  find(query, { skip, limit, sort, keys } = {}) {
-    return this._rawFind(query, { skip, limit, sort, keys })
+  find(query, { skip, limit, sort, keys, maxTimeMS } = {}) {
+    return this._rawFind(query, { skip, limit, sort, keys, maxTimeMS })
       .catch(error => {
         // Check for "no geoindex" error
         if (error.code != 17007 && !error.message.match(/unable to find index for .geoNear/)) {
@@ -30,22 +30,29 @@ export default class MongoCollection {
         index[key] = '2d';
         return this._mongoCollection.createIndex(index)
           // Retry, but just once.
-          .then(() => this._rawFind(query, { skip, limit, sort, keys }));
+          .then(() => this._rawFind(query, { skip, limit, sort, keys, maxTimeMS }));
       });
   }
 
-  _rawFind(query, { skip, limit, sort, keys } = {}) {
+  _rawFind(query, { skip, limit, sort, keys, maxTimeMS } = {}) {
     let findOperation = this._mongoCollection
       .find(query, { skip, limit, sort })
 
     if (keys) {
       findOperation = findOperation.project(keys);
     }
+
+    if (maxTimeMS) {
+      findOperation = findOperation.maxTimeMS(maxTimeMS);
+    }
+
     return findOperation.toArray();
   }
 
-  count(query, { skip, limit, sort } = {}) {
-    return this._mongoCollection.count(query, { skip, limit, sort });
+  count(query, { skip, limit, sort, maxTimeMS } = {}) {
+    let countOperation = this._mongoCollection.count(query, { skip, limit, sort, maxTimeMS });
+
+    return countOperation;
   }
 
   insertOne(object) {

src/Adapters/Storage/Mongo/MongoStorageAdapter.js

@@ -92,6 +92,9 @@ export class MongoStorageAdapter {
     this._uri = uri;
     this._collectionPrefix = collectionPrefix;
     this._mongoOptions = mongoOptions;
+
+    // MaxTimeMS is not a global MongoDB client option, it is applied per operation.
+    this._maxTimeMS = mongoOptions.maxTimeMS;
   }
 
   connect() {
@@ -329,7 +332,13 @@ export class MongoStorageAdapter {
       return memo;
     }, {});
     return this._adaptiveCollection(className)
-    .then(collection => collection.find(mongoWhere, { skip, limit, sort: mongoSort, keys: mongoKeys }))
+    .then(collection => collection.find(mongoWhere, {
+      skip,
+      limit,
+      sort: mongoSort,
+      keys: mongoKeys,
+      maxTimeMS: this._maxTimeMS,
+    }))
     .then(objects => objects.map(object => mongoObjectToParseObject(className, object, schema)))
   }
 
@@ -358,14 +367,18 @@ export class MongoStorageAdapter {
 
   // Used in tests
   _rawFind(className, query) {
-    return this._adaptiveCollection(className).then(collection => collection.find(query));
+    return this._adaptiveCollection(className).then(collection => collection.find(query, {
+      maxTimeMS: this._maxTimeMS,
+    }));
   }
 
-  // Executs a count.
+  // Executes a count.
   count(className, schema, query) {
     schema = convertParseSchemaToMongoSchema(schema);
     return this._adaptiveCollection(className)
-    .then(collection => collection.count(transformWhere(className, query, schema)));
+    .then(collection => collection.count(transformWhere(className, query, schema), {
+      maxTimeMS: this._maxTimeMS,
+    }));
   }
 
   performInitialization() {