fix: Execute exchange tasks as new coroutines.

src/main/kotlin/org/wfanet/panelmatch/client/deploy/ExchangeWorkflowDaemon.kt

@@ -15,17 +15,15 @@
 package org.wfanet.panelmatch.client.deploy
 
 import java.time.Clock
-import kotlinx.coroutines.Job
-import kotlinx.coroutines.currentCoroutineContext
+import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.isActive
 import kotlinx.coroutines.job
+import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
-import org.wfanet.measurement.common.logAndSuppressExceptionSuspend
 import org.wfanet.measurement.common.throttler.Throttler
 import org.wfanet.panelmatch.client.common.Identity
 import org.wfanet.panelmatch.client.exchangetasks.ExchangeTaskMapper
 import org.wfanet.panelmatch.client.launcher.ApiClient
-import org.wfanet.panelmatch.client.launcher.ExchangeStepLauncher
 import org.wfanet.panelmatch.client.launcher.ExchangeStepValidatorImpl
 import org.wfanet.panelmatch.client.launcher.ExchangeTaskExecutor
 import org.wfanet.panelmatch.client.storage.PrivateStorageSelector
@@ -36,6 +34,7 @@ import org.wfanet.panelmatch.client.storage.StorageDetailsProvider
 import org.wfanet.panelmatch.common.ExchangeDateKey
 import org.wfanet.panelmatch.common.Timeout
 import org.wfanet.panelmatch.common.certificates.CertificateManager
+import org.wfanet.panelmatch.common.loggerFor
 import org.wfanet.panelmatch.common.secrets.SecretMap
 import org.wfanet.panelmatch.common.storage.StorageFactory
 
@@ -112,35 +111,69 @@ abstract class ExchangeWorkflowDaemon : Runnable {
   override fun run() = runBlocking { runSuspending() }
 
   suspend fun runSuspending() {
-    val exchangeStepLauncher =
-      ExchangeStepLauncher(apiClient = apiClient, taskLauncher = stepExecutor)
     when (runMode) {
-      RunMode.DAEMON -> runDaemon(exchangeStepLauncher)
-      RunMode.CRON_JOB -> runCronJob(exchangeStepLauncher)
+      RunMode.DAEMON -> runDaemon()
+      RunMode.CRON_JOB -> runCronJob()
     }
   }
 
-  /** Runs [exchangeStepLauncher] in an infinite loop. */
-  protected open suspend fun runDaemon(exchangeStepLauncher: ExchangeStepLauncher) {
-    throttler.loopOnReady {
-      // All errors thrown inside the loop should be suppressed such that the daemon doesn't crash.
-      logAndSuppressExceptionSuspend { exchangeStepLauncher.findAndRunExchangeStep() }
+  /**
+   * Claims exchange steps and executes them in an infinite loop. Claimed steps are launched as
+   * child coroutines to allow multiple steps to execute concurrently.
+   */
+  protected open suspend fun runDaemon() = coroutineScope {
+    while (coroutineContext.isActive) {
+      val step =
+        throttler.onReady {
+          try {
+            apiClient.claimExchangeStep()
+          } catch (e: Exception) {
+            logger.severe("Failed to claim exchange step: $e")
+            null
+          }
+        }
+
+      if (step != null) {
+        launch {
+          try {
+            stepExecutor.execute(step)
+          } catch (e: Exception) {
+            logger.severe("Failed to execute exchange step: $e")
+          }
+        }
+      }
     }
   }
 
   /**
-   * Runs [exchangeStepLauncher] in a loop until there are no remaining tasks and all launched tasks
-   * have completed.
+   * Claims exchange steps and executes them until all available steps are exhausted, then returns.
+   * Claimed steps are launched as child coroutines to allow multiple steps to execute concurrently.
    */
-  protected open suspend fun runCronJob(exchangeStepLauncher: ExchangeStepLauncher) {
-    val activeJobs = mutableListOf<Job>()
-    do {
-      activeJobs.removeIf { !it.isActive }
-      val job = logAndSuppressExceptionSuspend { exchangeStepLauncher.findAndRunExchangeStep() }
-      if (job != null) {
-        activeJobs += job
+  protected open suspend fun runCronJob() = coroutineScope {
+    while (coroutineContext.isActive) {
+      val step =
+        throttler.onReady {
+          try {
+            apiClient.claimExchangeStep()
+          } catch (e: Exception) {
+            logger.severe("Failed to claim exchange step: $e")
+            null
+          }
+        }
+
+      if (step != null) {
+        launch {
+          try {
+            stepExecutor.execute(step)
+          } catch (e: Exception) {
+            logger.severe("Failed to execute exchange step: $e")
+          }
+        }
+      } else if (coroutineContext.job.children.none { it.isActive }) {
+        logger.info("All available steps executed; shutting down.")
+        break
       }
-    } while (currentCoroutineContext().isActive && activeJobs.isNotEmpty())
+    }
   }
 
   enum class RunMode {
@@ -153,4 +186,8 @@ abstract class ExchangeWorkflowDaemon : Runnable {
      */
     CRON_JOB,
   }
+
+  companion object {
+    private val logger by loggerFor()
+  }
 }

src/main/kotlin/org/wfanet/panelmatch/client/launcher/ExchangeStepExecutor.kt

@@ -14,10 +14,8 @@
 
 package org.wfanet.panelmatch.client.launcher
 
-import kotlinx.coroutines.Job
-
 /** Executes [ApiClient.ClaimedExchangeStep]s. */
 interface ExchangeStepExecutor {
-  /** Executes [exchangeStep] in a new coroutine and returns the running [Job]. */
-  suspend fun execute(exchangeStep: ApiClient.ClaimedExchangeStep): Job
+  /** Executes [exchangeStep]. */
+  suspend fun execute(exchangeStep: ApiClient.ClaimedExchangeStep)
 }

src/main/kotlin/org/wfanet/panelmatch/client/launcher/ExchangeTaskExecutor.kt

@@ -17,14 +17,9 @@ package org.wfanet.panelmatch.client.launcher
 import com.google.protobuf.ByteString
 import java.util.logging.Level
 import java.util.logging.Logger
-import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineName
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.Job
-import kotlinx.coroutines.cancel
-import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
 import org.wfanet.measurement.storage.StorageClient
 import org.wfanet.measurement.storage.StorageClient.Blob
 import org.wfanet.panelmatch.client.common.ExchangeContext
@@ -54,12 +49,11 @@ class ExchangeTaskExecutor(
   private val privateStorageSelector: PrivateStorageSelector,
   private val exchangeTaskMapper: ExchangeTaskMapper,
   private val validator: ExchangeStepValidator,
-  private val dispatcher: CoroutineDispatcher = Dispatchers.Default,
 ) : ExchangeStepExecutor {
 
-  override suspend fun execute(exchangeStep: ApiClient.ClaimedExchangeStep): Job = coroutineScope {
+  override suspend fun execute(exchangeStep: ApiClient.ClaimedExchangeStep) {
     val attemptKey = exchangeStep.attemptKey
-    launch(dispatcher + CoroutineName(attemptKey.toString()) + TaskLog(attemptKey.toString())) {
+    withContext(CoroutineName(attemptKey.toString()) + TaskLog(attemptKey.toString())) {
       try {
         val validatedStep = validator.validate(exchangeStep)
         val context =
@@ -79,7 +73,6 @@ class ExchangeTaskExecutor(
             else -> ExchangeStepAttempt.State.FAILED
           }
         markAsFinished(attemptKey, attemptState)
-        cancel("Task failed and reported back to Kingdom. Cancelling task scope.", e)
       }
     }
   }

src/test/kotlin/org/wfanet/panelmatch/client/launcher/BUILD.bazel

@@ -18,28 +18,6 @@ kt_jvm_test(
     ],
 )
 
-kt_jvm_test(
-    name = "ExchangeStepLauncherTest",
-    timeout = "short",
-    srcs = ["ExchangeStepLauncherTest.kt"],
-    test_class = "org.wfanet.panelmatch.client.launcher.ExchangeStepLauncherTest",
-    deps = [
-        "//src/main/kotlin/org/wfanet/panelmatch/client/common",
-        "//src/main/kotlin/org/wfanet/panelmatch/client/launcher",
-        "//src/main/kotlin/org/wfanet/panelmatch/common",
-        "//src/main/kotlin/org/wfanet/panelmatch/common/testing",
-        "//src/main/proto/wfa/panelmatch/client/internal:exchange_workflow_kt_jvm_proto",
-        "@wfa_common_jvm//imports/java/com/google/common/truth",
-        "@wfa_common_jvm//imports/java/com/google/common/truth/extensions/proto",
-        "@wfa_common_jvm//imports/java/org/junit",
-        "@wfa_common_jvm//imports/java/org/mockito",
-        "@wfa_common_jvm//imports/kotlin/kotlinx/coroutines:core",
-        "@wfa_common_jvm//imports/kotlin/org/mockito/kotlin",
-        "@wfa_common_jvm//src/main/kotlin/org/wfanet/measurement/common",
-        "@wfa_measurement_system//src/main/kotlin/org/wfanet/measurement/common/api:resource_key",
-    ],
-)
-
 kt_jvm_test(
     name = "ExchangeStepValidatorImplTest",
     timeout = "short",