Refactor SubmitBatchRequests to use Coroutines

src/main/kotlin/org/wfanet/measurement/reporting/service/api/SubmitBatchRequests.kt

@@ -16,11 +16,13 @@
 
 package org.wfanet.measurement.reporting.service.api
 
-import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.Deferred
+import kotlinx.coroutines.async
+import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.asFlow
-import kotlinx.coroutines.flow.flatMapConcat
 import kotlinx.coroutines.flow.flow
+import kotlinx.coroutines.sync.Semaphore
+import kotlinx.coroutines.sync.withPermit
 
 class BatchRequestException(message: String? = null, cause: Throwable? = null) :
   Exception(message, cause)
@@ -47,20 +49,35 @@ fun <T> Flow<T>.chunked(chunkSize: Int): Flow<List<T>> {
   }
 }
 
-/** Submits multiple RPCs by dividing the input items to batches. */
-@OptIn(ExperimentalCoroutinesApi::class) // For `flatMapConcat`.
+/**
+ * Submits multiple RPCs by dividing the input items to batches.
+ *
+ * @return [Flow] that emits [List]s containing the results of the multiple RPCs.
+ */
 suspend fun <ITEM, RESP, RESULT> submitBatchRequests(
-  items: Flow<ITEM>,
+  items: Collection<ITEM>,
   limit: Int,
   callRpc: suspend (List<ITEM>) -> RESP,
   parseResponse: (RESP) -> List<RESULT>,
-): Flow<RESULT> {
+): Flow<List<RESULT>> {
   if (limit <= 0) {
     throw BatchRequestException(
       "Invalid limit",
       IllegalArgumentException("The size limit of a batch must be greater than 0."),
     )
   }
 
-  return items.chunked(limit).flatMapConcat { batch -> parseResponse(callRpc(batch)).asFlow() }
+  // For network requests, the number of concurrent coroutines needs to be capped. To be on the safe
+  // side, a low number is chosen.
+  val batchSemaphore = Semaphore(3)
+  return flow {
+    coroutineScope {
+      val deferred: List<Deferred<List<RESULT>>> =
+        items.chunked(limit).map { batch: List<ITEM> ->
+          async { batchSemaphore.withPermit { parseResponse(callRpc(batch)) } }
+        }
+
+      deferred.forEach { emit(it.await()) }
+    }
+  }
 }

src/main/kotlin/org/wfanet/measurement/reporting/service/api/v2alpha/MetricsService.kt

@@ -47,9 +47,6 @@ import kotlinx.coroutines.asExecutor
 import kotlinx.coroutines.async
 import kotlinx.coroutines.awaitAll
 import kotlinx.coroutines.coroutineScope
-import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.asFlow
-import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.toList
 import kotlinx.coroutines.withContext
 import org.jetbrains.annotations.BlockingExecutor
@@ -308,7 +305,7 @@ class MetricsService(
       val measurementConsumer: MeasurementConsumer = getMeasurementConsumer(principal)
 
       // Gets all external IDs of primitive reporting sets from the metric list.
-      val externalPrimitiveReportingSetIds: Flow<String> =
+      val externalPrimitiveReportingSetIds: List<String> =
         internalMetricsList
           .flatMap { internalMetric ->
             internalMetric.weightedMeasurementsList.flatMap { weightedMeasurement ->
@@ -318,24 +315,27 @@ class MetricsService(
             }
           }
           .distinct()
-          .asFlow()
 
       val callBatchGetInternalReportingSetsRpc:
         suspend (List<String>) -> BatchGetReportingSetsResponse =
         { items ->
           batchGetInternalReportingSets(principal.resourceKey.measurementConsumerId, items)
         }
 
-      val internalPrimitiveReportingSetMap: Map<String, InternalReportingSet> =
+      val internalPrimitiveReportingSetMap: Map<String, InternalReportingSet> = buildMap {
         submitBatchRequests(
             externalPrimitiveReportingSetIds,
             BATCH_GET_REPORTING_SETS_LIMIT,
             callBatchGetInternalReportingSetsRpc,
           ) { response: BatchGetReportingSetsResponse ->
             response.reportingSetsList
           }
-          .toList()
-          .associateBy { it.externalReportingSetId }
+          .collect { reportingSets: List<InternalReportingSet> ->
+            for (reportingSet in reportingSets) {
+              computeIfAbsent(reportingSet.externalReportingSetId) { reportingSet }
+            }
+          }
+      }
 
       val dataProviderNames = mutableSetOf<String>()
       for (internalPrimitiveReportingSet in internalPrimitiveReportingSetMap.values) {
@@ -372,14 +372,16 @@ class MetricsService(
           batchCreateCmmsMeasurements(principal, items)
         }
 
-      val cmmsMeasurements: Flow<Measurement> =
+      val cmmsMeasurements: List<Measurement> =
         submitBatchRequests(
-          cmmsCreateMeasurementRequests.asFlow(),
-          BATCH_KINGDOM_MEASUREMENTS_LIMIT,
-          callBatchCreateMeasurementsRpc,
-        ) { response: BatchCreateMeasurementsResponse ->
-          response.measurementsList
-        }
+            cmmsCreateMeasurementRequests,
+            BATCH_KINGDOM_MEASUREMENTS_LIMIT,
+            callBatchCreateMeasurementsRpc,
+          ) { response: BatchCreateMeasurementsResponse ->
+            response.measurementsList
+          }
+          .toList()
+          .flatten()
 
       // Set CMMS measurement IDs.
       val callBatchSetCmmsMeasurementIdsRpc:
@@ -400,7 +402,7 @@ class MetricsService(
         ) { response: BatchSetCmmsMeasurementIdsResponse ->
           response.measurementsList
         }
-        .toList()
+        .collect {}
     }
 
     /** Sets a batch of CMMS [MeasurementIds] to the [InternalMeasurement] table. */
@@ -800,13 +802,13 @@ class MetricsService(
                 batchSetInternalMeasurementResults(items, apiAuthenticationKey, principal)
               }
             submitBatchRequests(
-                measurementsList.asFlow(),
+                measurementsList,
                 BATCH_SET_MEASUREMENT_RESULTS_LIMIT,
                 callBatchSetInternalMeasurementResultsRpc,
               ) { response: BatchSetCmmsMeasurementResultsResponse ->
                 response.measurementsList
               }
-              .toList()
+              .collect {}
 
             anyUpdate = true
           }
@@ -823,13 +825,13 @@ class MetricsService(
                 )
               }
             submitBatchRequests(
-                measurementsList.asFlow(),
+                measurementsList,
                 BATCH_SET_MEASUREMENT_FAILURES_LIMIT,
                 callBatchSetInternalMeasurementFailuresRpc,
               ) { response: BatchSetCmmsMeasurementFailuresResponse ->
                 response.measurementsList
               }
-              .toList()
+              .collect {}
 
             anyUpdate = true
           }
@@ -926,13 +928,14 @@ class MetricsService(
         }
 
       return submitBatchRequests(
-          measurementNames.asFlow(),
+          measurementNames,
           BATCH_KINGDOM_MEASUREMENTS_LIMIT,
           callBatchGetMeasurementsRpc,
         ) { response: BatchGetMeasurementsResponse ->
           response.measurementsList
         }
         .toList()
+        .flatten()
     }
 
     /** Batch get CMMS measurements. */

src/main/kotlin/org/wfanet/measurement/reporting/service/api/v2alpha/ReportSchedulesService.kt

@@ -33,7 +33,6 @@ import java.time.ZonedDateTime
 import java.time.temporal.TemporalAdjusters
 import java.time.zone.ZoneRulesException
 import kotlin.math.min
-import kotlinx.coroutines.flow.asFlow
 import org.projectnessie.cel.Env
 import org.wfanet.measurement.api.v2alpha.DataProvider
 import org.wfanet.measurement.api.v2alpha.DataProviderKey
@@ -673,31 +672,32 @@ class ReportSchedulesService(
       while (externalReportingSetIdSet.isNotEmpty()) {
         retrievedExternalReportingSetIdSet.addAll(externalReportingSetIdSet)
 
-        submitBatchRequests(
-            externalReportingSetIdSet.asFlow(),
-            BATCH_GET_REPORTING_SETS_LIMIT,
-            callRpc,
-          ) { response ->
+        submitBatchRequests(externalReportingSetIdSet, BATCH_GET_REPORTING_SETS_LIMIT, callRpc) {
+            response ->
             externalReportingSetIdSet.clear()
             response.reportingSetsList
           }
-          .collect {
-            if (it.hasComposite()) {
-              val lhsExternalReportingSetId = it.composite.lhs.externalReportingSetId
-              if (lhsExternalReportingSetId.isNotEmpty()) {
-                if (!retrievedExternalReportingSetIdSet.contains(lhsExternalReportingSetId)) {
-                  externalReportingSetIdSet.add(lhsExternalReportingSetId)
+          .collect { internalReportingSets: List<InternalReportingSet> ->
+            for (internalReportingSet in internalReportingSets) {
+              if (internalReportingSet.hasComposite()) {
+                val lhsExternalReportingSetId =
+                  internalReportingSet.composite.lhs.externalReportingSetId
+                if (lhsExternalReportingSetId.isNotEmpty()) {
+                  if (!retrievedExternalReportingSetIdSet.contains(lhsExternalReportingSetId)) {
+                    externalReportingSetIdSet.add(lhsExternalReportingSetId)
+                  }
                 }
-              }
 
-              val rhsExternalReportingSetId = it.composite.rhs.externalReportingSetId
-              if (rhsExternalReportingSetId.isNotEmpty()) {
-                if (!retrievedExternalReportingSetIdSet.contains(rhsExternalReportingSetId)) {
-                  externalReportingSetIdSet.add(rhsExternalReportingSetId)
+                val rhsExternalReportingSetId =
+                  internalReportingSet.composite.rhs.externalReportingSetId
+                if (rhsExternalReportingSetId.isNotEmpty()) {
+                  if (!retrievedExternalReportingSetIdSet.contains(rhsExternalReportingSetId)) {
+                    externalReportingSetIdSet.add(rhsExternalReportingSetId)
+                  }
                 }
               }
+              reportingSets.add(internalReportingSet)
             }
-            reportingSets.add(it)
           }
       }