Test bucket report row sampling and handle empty samples

Author: bean1352

modules/module-mongodb-storage/src/storage/implementation/MongoSyncBucketStorage.ts

@@ -576,40 +576,52 @@ export abstract class MongoSyncBucketStorage
   /**
    * Estimate a bucket's live rows from a sample of its operations.
    *
-   * `pipelinePrefix` must select the bucket's operations (and, when `sampled`, randomly down-sample them) and
-   * yield documents with top-level `op`, `table` and `row_id` fields. Returns the distinct row count (exact
-   * when the whole bucket was read, otherwise estimated via {@link storage.estimateDistinctRows}); fragmentation is
-   * then `operations / rows`.
+   * `buildPrefix(applySample)` returns a pipeline prefix that selects the bucket's operations (down-sampled
+   * when `applySample` is true) and yields documents with top-level `op`, `table` and `row_id` fields.
+   * Returns the distinct row count (exact when the whole bucket was read, otherwise estimated via
+   * {@link storage.estimateDistinctRows}); fragmentation is then `operations / rows`.
    */
   protected async estimateRowsFromOperationSample<T extends mongo.Document>(
     collection: mongo.Collection<T>,
-    pipelinePrefix: mongo.Document[],
+    buildPrefix: (applySample: boolean) => mongo.Document[],
     operations: number,
     sampled: boolean
   ): Promise<BucketRowEstimate> {
-    const pipeline: mongo.Document[] = [
-      ...pipelinePrefix,
-      {
-        $facet: {
-          sampledOps: [{ $count: 'count' }],
-          distinctRows: [
-            { $match: { op: { $in: ['PUT', 'REMOVE'] } } },
-            { $group: { _id: { table: '$table', row_id: '$row_id' } } },
-            { $count: 'count' }
-          ]
+    const runCounts = async (applySample: boolean) => {
+      const pipeline: mongo.Document[] = [
+        ...buildPrefix(applySample),
+        {
+          $facet: {
+            sampledOps: [{ $count: 'count' }],
+            distinctRows: [
+              { $match: { op: { $in: ['PUT', 'REMOVE'] } } },
+              { $group: { _id: { table: '$table', row_id: '$row_id' } } },
+              { $count: 'count' }
+            ]
+          }
         }
-      }
-    ];
-
-    type FacetResult = { sampledOps: { count: number }[]; distinctRows: { count: number }[] };
-    const [result] = await collection
-      .aggregate<FacetResult>(pipeline, { allowDiskUse: false, maxTimeMS: storage.BUCKET_REPORT_TIMEOUT_MS })
-      .toArray();
+      ];
+      type FacetResult = { sampledOps: { count: number }[]; distinctRows: { count: number }[] };
+      const [result] = await collection
+        .aggregate<FacetResult>(pipeline, { allowDiskUse: false, maxTimeMS: storage.BUCKET_REPORT_TIMEOUT_MS })
+        .toArray();
+      return {
+        sampledOps: result?.sampledOps[0]?.count ?? 0,
+        distinctRows: result?.distinctRows[0]?.count ?? 0
+      };
+    };
 
-    const sampledOps = result?.sampledOps[0]?.count ?? 0;
-    const distinctRows = result?.distinctRows[0]?.count ?? 0;
-    if (sampledOps == 0 || distinctRows == 0) {
-      // Nothing row-bearing was sampled (e.g. a bucket of only MOVE/CLEAR ops): treat as fully fragmented.
+    let { sampledOps, distinctRows } = await runCounts(sampled);
+    if (sampled && sampledOps == 0) {
+      // A document-level `$sampleRate` can select nothing when a bucket spans very few storage documents
+      // (v3 batches operations into a document). Fall back to an exact read so the bucket is not reported as
+      // zero rows. This reads the whole bucket only in the rare empty-sample case, which cannot happen for a
+      // bucket large enough to span many documents.
+      distinctRows = (await runCounts(false)).distinctRows;
+      return { rows: distinctRows, estimated: false };
+    }
+    if (distinctRows == 0) {
+      // Nothing row-bearing was found (e.g. a bucket of only MOVE/CLEAR ops): treat as fully fragmented.
       return { rows: 0, estimated: sampled };
     }
     if (!sampled) {

modules/module-mongodb-storage/src/storage/implementation/v1/MongoSyncBucketStorageV1.ts

@@ -204,22 +204,25 @@ export class MongoSyncBucketStorageV1 extends MongoSyncBucketStorage {
   protected estimateBucketRows(candidate: TopBucketCandidate): Promise<BucketRowEstimate> {
     // v1/v2 store one document per operation, so a bucket's ops are an id-prefix range that can be sampled directly.
     const sampled = this.shouldSampleBucketRows(candidate.operations);
-    // Range-match on the whole `_id` (g, b, o) so the {_id} index is used; a dotted `{'_id.g','_id.b'}` match
-    // cannot use the compound-object index and would scan the whole collection per bucket.
-    const prefix: mongo.Document[] = [
-      {
-        $match: {
-          _id: idPrefixFilter<{ g: number; b: string; o: unknown }>(
-            { g: this.replicationStreamId, b: candidate.bucket },
-            ['o']
-          )
+    const buildPrefix = (applySample: boolean): mongo.Document[] => {
+      // Range-match on the whole `_id` (g, b, o) so the {_id} index is used; a dotted `{'_id.g','_id.b'}` match
+      // cannot use the compound-object index and would scan the whole collection per bucket.
+      const prefix: mongo.Document[] = [
+        {
+          $match: {
+            _id: idPrefixFilter<{ g: number; b: string; o: unknown }>(
+              { g: this.replicationStreamId, b: candidate.bucket },
+              ['o']
+            )
+          }
         }
+      ];
+      if (applySample) {
+        prefix.push({ $match: { $sampleRate: this.bucketRowSampleRate(candidate.operations) } });
       }
-    ];
-    if (sampled) {
-      prefix.push({ $match: { $sampleRate: this.bucketRowSampleRate(candidate.operations) } });
-    }
-    return this.estimateRowsFromOperationSample(this.db.bucketDataV1, prefix, candidate.operations, sampled);
+      return prefix;
+    };
+    return this.estimateRowsFromOperationSample(this.db.bucketDataV1, buildPrefix, candidate.operations, sampled);
   }
 
   protected createMongoParameterCompactor(

modules/module-mongodb-storage/src/storage/implementation/v3/MongoSyncBucketStorageV3.ts

@@ -214,16 +214,19 @@ export class MongoSyncBucketStorageV3 extends MongoSyncBucketStorage {
     // Sample whole batch documents, then unwind to operation level so the shared estimator sees one doc per op.
     const sampled = this.shouldSampleBucketRows(candidate.operations);
     const collection = this.db.bucketData(this.replicationStreamId, candidate.defId!);
-    // Range-match on the whole `_id` (b, o) so the {_id} index is used; a dotted `{'_id.b': ...}` match
-    // cannot use the compound-object index and would scan the whole collection per bucket.
-    const prefix: mongo.Document[] = [
-      { $match: { _id: idPrefixFilter<{ b: string; o: unknown }>({ b: candidate.bucket }, ['o']) } }
-    ];
-    if (sampled) {
-      prefix.push({ $match: { $sampleRate: this.bucketRowSampleRate(candidate.operations) } });
-    }
-    prefix.push({ $unwind: '$ops' }, { $replaceRoot: { newRoot: '$ops' } });
-    return this.estimateRowsFromOperationSample(collection, prefix, candidate.operations, sampled);
+    const buildPrefix = (applySample: boolean): mongo.Document[] => {
+      // Range-match on the whole `_id` (b, o) so the {_id} index is used; a dotted `{'_id.b': ...}` match
+      // cannot use the compound-object index and would scan the whole collection per bucket.
+      const prefix: mongo.Document[] = [
+        { $match: { _id: idPrefixFilter<{ b: string; o: unknown }>({ b: candidate.bucket }, ['o']) } }
+      ];
+      if (applySample) {
+        prefix.push({ $match: { $sampleRate: this.bucketRowSampleRate(candidate.operations) } });
+      }
+      prefix.push({ $unwind: '$ops' }, { $replaceRoot: { newRoot: '$ops' } });
+      return prefix;
+    };
+    return this.estimateRowsFromOperationSample(collection, buildPrefix, candidate.operations, sampled);
   }
 
   protected createMongoParameterCompactor(

packages/service-core-tests/src/tests/register-bucket-report-tests.ts

@@ -210,4 +210,60 @@ bucket_definitions:
     expect(report.totals.bucketCount).toEqual(2);
     expect(report.totals).toMatchObject({ operations: 3, estimated: false });
   });
+
+  test('samples the row count for a bucket above the sampling threshold', async () => {
+    await using factory = await generateStorageFactory();
+    const { stream, content } = await test_utils.deploySyncRules(
+      factory,
+      updateSyncRulesFromYaml(GLOBAL_SYNC_RULES, { storageVersion })
+    );
+    const bucketStorage = factory.getInstance(stream);
+
+    await using writer = await bucketStorage.createWriter(test_utils.BATCH_OPTIONS);
+    const testTable = await test_utils.resolveTestTable(writer, 'test', ['id'], config);
+    await writer.markAllSnapshotDone('1/1');
+
+    // 50 rows, each updated 25 times, is 1,300 operations against 50 live rows. That is past the 1,000
+    // operation threshold, so the report samples the operation history rather than reading it in full and
+    // the row count comes back as an estimate. The value per update varies so no two writes are identical.
+    // Each round is flushed separately so the operations span many storage documents, as they would in real
+    // replication (some backends batch operations per document, and a sample must see more than one).
+    const rowCount = 50;
+    const updatesPerRow = 25;
+    for (let row = 0; row < rowCount; row++) {
+      await writer.save({
+        sourceTable: testTable,
+        tag: storage.SaveOperationTag.INSERT,
+        after: { id: `r${row}` },
+        afterReplicaId: test_utils.rid(`r${row}`)
+      });
+    }
+    await writer.commit('1/1');
+    await writer.flush();
+    for (let update = 0; update < updatesPerRow; update++) {
+      for (let row = 0; row < rowCount; row++) {
+        await writer.save({
+          sourceTable: testTable,
+          tag: storage.SaveOperationTag.UPDATE,
+          after: { id: `r${row}`, value: `v${update}` },
+          afterReplicaId: test_utils.rid(`r${row}`)
+        });
+      }
+      await writer.commit('1/1');
+      await writer.flush();
+    }
+
+    const bucket = test_utils.bucketRequest(content, 'global[]').bucket;
+    const report = await getReport(bucketStorage);
+    const stats = report.buckets.find((b) => b.bucket === bucket)!;
+
+    // The operation count is exact (read from bucket_state); the row count is a sampled estimate.
+    expect(stats.operations).toEqual(rowCount + rowCount * updatesPerRow);
+    expect(stats.rowsEstimated).toEqual(true);
+    // The sample covers enough of a bucket this fragmented to recover the 50 live rows within a small margin.
+    expect(stats.rows).toBeGreaterThanOrEqual(45);
+    expect(stats.rows).toBeLessThanOrEqual(55);
+    // Fragmentation is operations / rows, so a heavily updated bucket reads well above 1.
+    expect(stats.fragmentation).toBeGreaterThan(10);
+  });
 }