refactor: Extract listResources utility function for handling pagination

src/main/kotlin/org/wfanet/measurement/api/v2alpha/BUILD.bazel

@@ -52,7 +52,7 @@ kt_jvm_library(
     deps = [
         "//src/main/kotlin/org/wfanet/measurement/api/v2alpha:context_keys",
         "//src/main/kotlin/org/wfanet/measurement/api/v2alpha:measurement_principal",
-        "//src/main/kotlin/org/wfanet/measurement/common/api/grpc",
+        "//src/main/kotlin/org/wfanet/measurement/common/api/grpc:akid_principal_server_interceptor",
         "//src/main/kotlin/org/wfanet/measurement/common/grpc:context",
         "//src/main/kotlin/org/wfanet/measurement/common/identity",
         "//src/main/proto/wfa/measurement/api/v2alpha:duchy_kt_jvm_proto",

src/main/kotlin/org/wfanet/measurement/common/api/grpc/BUILD.bazel

@@ -3,7 +3,7 @@ load("@wfa_rules_kotlin_jvm//kotlin:defs.bzl", "kt_jvm_library")
 package(default_visibility = ["//visibility:public"])
 
 kt_jvm_library(
-    name = "grpc",
+    name = "akid_principal_server_interceptor",
     srcs = ["AkidPrincipalServerInterceptor.kt"],
     deps = [
         "//src/main/kotlin/org/wfanet/measurement/common/api:principal",
@@ -14,3 +14,13 @@ kt_jvm_library(
         "@wfa_common_jvm//src/main/kotlin/org/wfanet/measurement/common/grpc",
     ],
 )
+
+kt_jvm_library(
+    name = "list_resources",
+    srcs = ["ListResources.kt"],
+    deps = [
+        "@wfa_common_jvm//imports/java/com/google/protobuf",
+        "@wfa_common_jvm//imports/kotlin/kotlinx/coroutines:core",
+        "@wfa_rules_kotlin_jvm//imports/io/gprc/kotlin:stub",
+    ],
+)

src/main/kotlin/org/wfanet/measurement/common/api/grpc/ListResources.kt

@@ -0,0 +1,92 @@
+/*
+ * Copyright 2024 The Cross-Media Measurement Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.wfanet.measurement.common.api.grpc
+
+import com.google.protobuf.Message
+import io.grpc.kotlin.AbstractCoroutineStub
+import kotlin.coroutines.coroutineContext
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.ensureActive
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.asFlow
+import kotlinx.coroutines.flow.flattenConcat
+import kotlinx.coroutines.flow.flow
+import kotlinx.coroutines.flow.map
+
+/** A [List] of resources from a paginated List method. */
+data class ResourceList<T : Message>(
+  val resources: List<T>,
+  /**
+   * A token that can be sent on subsequent requests to retrieve the next page. If empty, there are
+   * no subsequent pages.
+   */
+  val nextPageToken: String,
+) : List<T> by resources
+
+/**
+ * Lists resources from a paginated List method on this stub.
+ *
+ * @param pageToken page token for initial request
+ * @param list function which calls the appropriate List method on the stub
+ */
+fun <T : Message, S : AbstractCoroutineStub<S>> S.listResources(
+  pageToken: String = "",
+  list: suspend S.(pageToken: String) -> ResourceList<T>,
+): Flow<ResourceList<T>> =
+  listResources(Int.MAX_VALUE, pageToken) { nextPageToken, _ -> list(nextPageToken) }
+
+/**
+ * Lists resources from a paginated List method on this stub.
+ *
+ * @param limit maximum number of resources to emit
+ * @param pageToken page token for initial request
+ * @param list function which calls the appropriate List method on the stub, returning no more than
+ *   the specified remaining number of resources
+ */
+fun <T : Message, S : AbstractCoroutineStub<S>> S.listResources(
+  limit: Int,
+  pageToken: String = "",
+  list: suspend S.(pageToken: String, remaining: Int) -> ResourceList<T>,
+): Flow<ResourceList<T>> {
+  require(limit > 0) { "limit must be positive" }
+  return flow {
+    var remaining: Int = limit
+    var nextPageToken = pageToken
+
+    while (true) {
+      coroutineContext.ensureActive()
+
+      val resourceList: ResourceList<T> = list(nextPageToken, remaining)
+      require(resourceList.size <= remaining) {
+        "List call must ensure that limit is not exceeded. " +
+          "Returned ${resourceList.size} items when only $remaining were remaining"
+      }
+      emit(resourceList)
+
+      remaining -= resourceList.size
+      nextPageToken = resourceList.nextPageToken
+      if (nextPageToken.isEmpty() || remaining == 0) {
+        break
+      }
+    }
+  }
+}
+
+/** @see [flattenConcat] */
+@ExperimentalCoroutinesApi // Overloads experimental `flattenConcat` function.
+fun <T : Message> Flow<ResourceList<T>>.flattenConcat(): Flow<T> =
+  map { it.asFlow() }.flattenConcat()

src/main/kotlin/org/wfanet/measurement/integration/common/BUILD.bazel

@@ -319,6 +319,7 @@ kt_jvm_library(
     ],
     deps = [
         ":in_process_cmms_components",
+        "//src/main/kotlin/org/wfanet/measurement/common/api/grpc:list_resources",
         "//src/main/kotlin/org/wfanet/measurement/kingdom/batch:measurement_system_prober",
         "//src/main/kotlin/org/wfanet/measurement/kingdom/deploy/common/service:data_services",
         "@wfa_common_jvm//imports/java/com/google/common/truth",

src/main/kotlin/org/wfanet/measurement/integration/common/InProcessMeasurementSystemProberIntegrationTest.kt

@@ -22,6 +22,9 @@ import java.io.File
 import java.nio.file.Paths
 import java.time.Clock
 import java.time.Duration
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.toList
 import kotlinx.coroutines.runBlocking
 import org.junit.After
 import org.junit.Before
@@ -30,13 +33,15 @@ import org.junit.Rule
 import org.junit.Test
 import org.wfanet.measurement.api.v2alpha.DataProvidersGrpcKt.DataProvidersCoroutineStub
 import org.wfanet.measurement.api.v2alpha.EventGroupsGrpcKt.EventGroupsCoroutineStub
-import org.wfanet.measurement.api.v2alpha.ListMeasurementsResponse
 import org.wfanet.measurement.api.v2alpha.Measurement
 import org.wfanet.measurement.api.v2alpha.MeasurementConsumersGrpcKt.MeasurementConsumersCoroutineStub
 import org.wfanet.measurement.api.v2alpha.MeasurementsGrpcKt.MeasurementsCoroutineStub
 import org.wfanet.measurement.api.v2alpha.RequisitionsGrpcKt.RequisitionsCoroutineStub
 import org.wfanet.measurement.api.v2alpha.listMeasurementsRequest
 import org.wfanet.measurement.api.withAuthenticationKey
+import org.wfanet.measurement.common.api.grpc.ResourceList
+import org.wfanet.measurement.common.api.grpc.flattenConcat
+import org.wfanet.measurement.common.api.grpc.listResources
 import org.wfanet.measurement.common.getRuntimePath
 import org.wfanet.measurement.common.identity.withPrincipalName
 import org.wfanet.measurement.common.testing.ProviderRule
@@ -129,33 +134,32 @@ abstract class InProcessMeasurementSystemProberIntegrationTest(
     assertThat(measurements.size).isEqualTo(1)
   }
 
+  @OptIn(ExperimentalCoroutinesApi::class) // For `flattenConcat`.
   private suspend fun listMeasurements(): List<Measurement> {
-    var nextPageToken = ""
     val measurementConsumerData = inProcessCmmsComponents.getMeasurementConsumerData()
 
-    do {
-      val response: ListMeasurementsResponse =
-        try {
-          publicMeasurementsClient
-            .withAuthenticationKey(measurementConsumerData.apiAuthenticationKey)
-            .listMeasurements(
-              listMeasurementsRequest {
-                parent = measurementConsumerData.name
-                pageToken = nextPageToken
-              }
-            )
-        } catch (e: StatusException) {
-          throw Exception(
-            "Unable to list measurements for measurement consumer ${measurementConsumerData.name}",
-            e,
-          )
+    val measurementLists: Flow<ResourceList<Measurement>> =
+      publicMeasurementsClient
+        .withAuthenticationKey(measurementConsumerData.apiAuthenticationKey)
+        .listResources { pageToken ->
+          val response =
+            try {
+              listMeasurements(
+                listMeasurementsRequest {
+                  parent = measurementConsumerData.name
+                  this.pageToken = pageToken
+                }
+              )
+            } catch (e: StatusException) {
+              throw Exception(
+                "Unable to list measurements for measurement consumer ${measurementConsumerData.name}",
+                e,
+              )
+            }
+          ResourceList(response.measurementsList, response.nextPageToken)
         }
-      if (response.measurementsList.isNotEmpty()) {
-        return response.measurementsList
-      }
-      nextPageToken = response.nextPageToken
-    } while (nextPageToken.isNotEmpty())
-    return emptyList()
+
+    return measurementLists.flattenConcat().toList()
   }
 
   companion object {

src/main/kotlin/org/wfanet/measurement/kingdom/batch/BUILD.bazel

@@ -47,6 +47,7 @@ kt_jvm_library(
         "//src/main/kotlin/org/wfanet/measurement/api:api_key_constants",
         "//src/main/kotlin/org/wfanet/measurement/api/v2alpha:packed_messages",
         "//src/main/kotlin/org/wfanet/measurement/api/v2alpha:resource_key",
+        "//src/main/kotlin/org/wfanet/measurement/common/api/grpc:list_resources",
         "//src/main/proto/wfa/measurement/api/v2alpha:data_provider_kt_jvm_proto",
         "//src/main/proto/wfa/measurement/api/v2alpha:data_providers_service_kt_jvm_grpc_proto",
         "//src/main/proto/wfa/measurement/api/v2alpha:event_group_kt_jvm_proto",