fix: inefficient form metadata query leading to excessive memory usage (

#5763)
opengovsg · Feb 14, 2023 · 38abbc9 · 38abbc9
1 parent a3c3374
commit 38abbc9
Showing 1 changed file with 53 additions and 51 deletions.
diff --git a/src/app/models/submission.server.model.ts b/src/app/models/submission.server.model.ts
@@ -1,6 +1,5 @@
 import moment from 'moment-timezone'
 import mongoose, { Mongoose, QueryCursor, Schema } from 'mongoose'
-import type { FixedLengthArray } from 'type-fest'
 
 import {
   FormAuthType,
@@ -274,10 +273,7 @@ EncryptSubmissionSchema.statics.findSingleMetadata = function (
  * Unexported as the type is only used in {@see findAllMetadataByFormId} for
  * now.
  */
-type MetadataAggregateResult = {
-  pageResults: Pick<ISubmissionSchema, '_id' | 'created'>[]
-  allResults: FixedLengthArray<{ count: number }, 1> | []
-}
+type MetadataAggregateResult = Pick<ISubmissionSchema, '_id' | 'created'>
 
 EncryptSubmissionSchema.statics.findAllMetadataByFormId = function (
   formId: string,
@@ -294,54 +290,60 @@ EncryptSubmissionSchema.statics.findAllMetadataByFormId = function (
 }> {
   const numToSkip = (page - 1) * pageSize
 
-  return (
-    this.aggregate()
-      .match({
-        // Casting to ObjectId as Mongoose does not cast pipeline stages.
-        // See https://mongoosejs.com/docs/api.html#aggregate_Aggregate.
+  // return documents within the page
+  const pageResults = this.aggregate<MetadataAggregateResult>([
+    {
+      $match: {
         form: mongoose.Types.ObjectId(formId),
         submissionType: SubmissionType.Encrypt,
-      })
-      .sort({ created: -1 })
-      .facet({
-        pageResults: [
-          { $skip: numToSkip },
-          { $limit: pageSize },
-          { $project: { _id: 1, created: 1 } },
-        ],
-        allResults: [
-          { $group: { _id: null, count: { $sum: 1 } } },
-          { $project: { _id: 0 } },
-        ],
-      })
-      // prevents out-of-memory for large search results (max 100MB).
-      .allowDiskUse(true)
-      .then((result: MetadataAggregateResult[]) => {
-        const [{ pageResults, allResults }] = result
-        const [numResults] = allResults
-        const count = numResults?.count ?? 0
-
-        let currentNumber = count - numToSkip
-
-        const metadata = pageResults.map((data) => {
-          const metadataEntry: StorageModeSubmissionMetadata = {
-            number: currentNumber,
-            refNo: data._id,
-            submissionTime: moment(data.created)
-              .tz('Asia/Singapore')
-              .format('Do MMM YYYY, h:mm:ss a'),
-          }
-
-          currentNumber--
-          return metadataEntry
-        })
-
-        return {
-          metadata,
-          count,
-        }
-      })
-  )
+      },
+    },
+    {
+      $sort: {
+        created: -1,
+      },
+    },
+    {
+      $skip: numToSkip,
+    },
+    {
+      $limit: pageSize,
+    },
+    {
+      $project: {
+        _id: 1,
+        created: 1,
+      },
+    },
+  ])
+
+  const count =
+    this.countDocuments({
+      form: mongoose.Types.ObjectId(formId),
+      submissionType: SubmissionType.Encrypt,
+    }).exec() ?? 0
+
+  return Promise.all([pageResults, count]).then(([result, count]) => {
+    let currentNumber = count - numToSkip
+
+    const metadata = result.map((data) => {
+      const metadataEntry: StorageModeSubmissionMetadata = {
+        number: currentNumber,
+        refNo: data._id,
+        submissionTime: moment(data.created)
+          .tz('Asia/Singapore')
+          .format('Do MMM YYYY, h:mm:ss a'),
+      }
+
+      currentNumber--
+      return metadataEntry
+    })
+
+    return {
+      metadata,
+      count,
+    }
+  })
 }
 
 EncryptSubmissionSchema.statics.getSubmissionCursorByFormId = function (