Fix janus_gc_deleted_batches_total (#3160)

* Add batch_aggregations rows to make test fail * GC: Fix deleted batch counter updates
divviup · May 23, 2024 · d39af88 · d39af88
1 parent fc0b52e
commit d39af88
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 16 deletions.
diff --git a/aggregator_core/src/datastore.rs b/aggregator_core/src/datastore.rs
@@ -4854,25 +4854,29 @@ deleted_aggregate_share_jobs AS (
       AND (LOWER(batch_interval) < $2
         OR (aggregate_share_jobs.batch_identifier = batches_to_delete.batch_identifier
           AND aggregate_share_jobs.aggregation_param = batches_to_delete.aggregation_param))
+),
+deleted_batch_aggregations AS (
+    DELETE FROM batch_aggregations
+    USING batches_to_delete
+    WHERE task_id = $1
+      AND batch_aggregations.batch_identifier = batches_to_delete.batch_identifier
+      AND batch_aggregations.aggregation_param = batches_to_delete.aggregation_param
 )
-DELETE FROM batch_aggregations
-USING batches_to_delete
-WHERE task_id = $1
-  AND batch_aggregations.batch_identifier = batches_to_delete.batch_identifier
-  AND batch_aggregations.aggregation_param = batches_to_delete.aggregation_param",
+SELECT COUNT(1) AS batch_count FROM batches_to_delete",
             )
             .await?;
-        self.execute(
-            &stmt,
-            &[
-                /* task_id */ &task_info.pkey,
-                /* threshold */
-                &task_info.report_expiry_threshold(&self.clock.now().as_naive_date_time()?)?,
-                /* limit */ &i64::try_from(limit)?,
-            ],
-        )
-        .await
-        .map_err(Into::into)
+        let row = self
+            .query_one(
+                &stmt,
+                &[
+                    /* task_id */ &task_info.pkey,
+                    /* threshold */
+                    &task_info.report_expiry_threshold(&self.clock.now().as_naive_date_time()?)?,
+                    /* limit */ &i64::try_from(limit)?,
+                ],
+            )
+            .await?;
+        row.get_bigint_and_convert("batch_count")
     }
 
     /// Retrieve all global HPKE keypairs.

diff --git a/aggregator_core/src/datastore/tests.rs b/aggregator_core/src/datastore/tests.rs
@@ -6125,6 +6125,23 @@ async fn delete_expired_collection_artifacts(ephemeral_datastore: EphemeralDatas
             },
         );
         tx.put_batch_aggregation(&batch_aggregation).await.unwrap();
+        for ord in 1..8 {
+            let batch_aggregation = BatchAggregation::<0, Q, dummy::Vdaf>::new(
+                *task.id(),
+                batch_identifier.clone(),
+                dummy::AggregationParam(0),
+                ord,
+                client_timestamp_interval,
+                BatchAggregationState::Aggregating {
+                    aggregate_share: None,
+                    report_count: 0,
+                    checksum: ReportIdChecksum::default(),
+                    aggregation_jobs_created: 0,
+                    aggregation_jobs_terminated: 0,
+                },
+            );
+            tx.put_batch_aggregation(&batch_aggregation).await.unwrap();
+        }
 
         if task.role() == &Role::Leader {
             let collection_job = CollectionJob::<0, Q, dummy::Vdaf>::new(