Rbroughan/fix proto input stream hang (#66576)

Co-authored-by: Subodh Kant Chaturvedi <[email protected]>

Author: tryangul

airbyte-cdk/bulk/changelog.md

@@ -1,3 +1,7 @@
+## Version 0.1.42
+
+Datflow Load CDK: Fixes hang when one of many parallel pipelines fails. Organizes thread pools. 
+
 ## Version 0.1.41
 
 **Extract CDK**
@@ -12,6 +16,8 @@ Add gradle task to bump CDK version + add changelog entry
 
 Minor fixes with stream completion logic + proto conversion in Load CDK.
 
+* **Changed:** Minor fixes with stream completion logic + proto conversion
+
 ## Version 0.1.38
 
 Adds stats support for "speed" mode to the Load CDK

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/dataflow/config/ConnectorInputStreams.kt

@@ -11,4 +11,6 @@ import java.io.InputStream
  * subject to Micronaut merging like beans into a single list leading to injecting unexpected extra
  * input streams.
  */
-class ConnectorInputStreams(private val values: List<InputStream>) : List<InputStream> by values
+class ConnectorInputStreams(private val values: List<InputStream>) : List<InputStream> by values {
+    fun closeAll() = values.forEach { it.close() }
+}

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/dataflow/config/DispatcherBeanFactory.kt

@@ -0,0 +1,44 @@
+/*
+ * Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.cdk.load.dataflow.config
+
+import io.micronaut.context.annotation.Factory
+import jakarta.inject.Named
+import jakarta.inject.Singleton
+import java.util.concurrent.Executors
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.asCoroutineDispatcher
+
+/** The dispatchers (think views of thread pools) and static scopes we use for dataflow. */
+@Factory
+class DispatcherBeanFactory {
+    @Named("pipelineRunnerDispatcher")
+    @Singleton
+    fun pipelineRunnerDispatcher() = Dispatchers.Default
+
+    @Named("stateReconcilerDispatcher") @Singleton fun stateReconcilerDispatcher() = Dispatchers.IO
+
+    @Named("aggregationDispatcher")
+    @Singleton
+    fun aggregationDispatcher(
+        @Named("inputStreams") inputStreams: ConnectorInputStreams,
+    ) = Executors.newFixedThreadPool(inputStreams.size).asCoroutineDispatcher()
+
+    @Named("flushDispatcher") @Singleton fun flushDispatcher() = Dispatchers.IO
+
+    @Named("pipelineRunnerScope")
+    @Singleton
+    fun pipelineRunnerDispatcher(
+        @Named("pipelineRunnerDispatcher") dispatcher: CoroutineDispatcher,
+    ) = CoroutineScope(dispatcher)
+
+    @Named("stateReconcilerScope")
+    @Singleton
+    fun stateDispatcher(
+        @Named("stateReconcilerDispatcher") dispatcher: CoroutineDispatcher,
+    ) = CoroutineScope(dispatcher)
+}

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/dataflow/config/InputBeanFactory.kt

@@ -30,6 +30,7 @@ import io.micronaut.context.annotation.Requires
 import io.micronaut.context.annotation.Value
 import jakarta.inject.Named
 import jakarta.inject.Singleton
+import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.flow.Flow
 
 /**
@@ -134,6 +135,8 @@ class InputBeanFactory {
         stateHistogramStore: StateHistogramStore,
         statsStore: CommittedStatsStore,
         memoryAndParallelismConfig: MemoryAndParallelismConfig,
+        @Named("aggregationDispatcher") aggregationDispatcher: CoroutineDispatcher,
+        @Named("flushDispatcher") flushDispatcher: CoroutineDispatcher,
     ): List<DataFlowPipeline> =
         inputFlows.map {
             val aggStore = aggregateStoreFactory.make()
@@ -153,6 +156,8 @@ class InputBeanFactory {
                 state = state,
                 completionHandler = completionHandler,
                 memoryAndParallelismConfig = memoryAndParallelismConfig,
+                aggregationDispatcher = aggregationDispatcher,
+                flushDispatcher = flushDispatcher,
             )
         }
 }

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/dataflow/input/ProtobufDestinationMessageInputFlow.kt

@@ -13,6 +13,10 @@ import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.FlowCollector
 import kotlinx.coroutines.withContext
 
+/**
+ * Performs non-cooperative blocking IO. Does not respond directly to coroutine
+ * CancellationExceptions.
+ */
 class ProtobufDestinationMessageInputFlow(
     private val inputStream: InputStream,
     private val reader: ProtobufDataChannelReader,

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/dataflow/pipeline/DataFlowPipeline.kt

@@ -6,9 +6,7 @@ package io.airbyte.cdk.load.dataflow.pipeline
 
 import io.airbyte.cdk.load.dataflow.config.MemoryAndParallelismConfig
 import io.airbyte.cdk.load.dataflow.stages.AggregateStage
-import java.util.concurrent.Executors
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.asCoroutineDispatcher
+import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.buffer
 import kotlinx.coroutines.flow.flowOn
@@ -24,19 +22,19 @@ class DataFlowPipeline(
     private val state: DataFlowStage,
     private val completionHandler: PipelineCompletionHandler,
     private val memoryAndParallelismConfig: MemoryAndParallelismConfig,
+    private val aggregationDispatcher: CoroutineDispatcher,
+    private val flushDispatcher: CoroutineDispatcher,
 ) {
-    private val customDispatcher = Executors.newSingleThreadExecutor().asCoroutineDispatcher()
-
     suspend fun run() {
         input
             .map(parse::apply)
             .transform { aggregate.apply(it, this) }
             .buffer(capacity = memoryAndParallelismConfig.maxBufferedAggregates)
-            .flowOn(customDispatcher)
+            .flowOn(aggregationDispatcher)
             .map(flush::apply)
             .map(state::apply)
             .onCompletion { completionHandler.apply(it) }
-            .flowOn(Dispatchers.IO)
+            .flowOn(flushDispatcher)
             .collect {}
     }
 }

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/dataflow/pipeline/PipelineCompletionHandler.kt

@@ -23,7 +23,6 @@ class PipelineCompletionHandler(
         cause: Throwable?,
     ) = coroutineScope {
         if (cause != null) {
-            log.error { "Destination Pipeline Completed — Exceptionally" }
             throw cause
         }
 

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/dataflow/pipeline/PipelineRunner.kt

@@ -4,48 +4,84 @@
 
 package io.airbyte.cdk.load.dataflow.pipeline
 
+import io.airbyte.cdk.load.dataflow.config.ConnectorInputStreams
 import io.airbyte.cdk.load.dataflow.state.StateReconciler
 import io.airbyte.cdk.load.dataflow.state.StateStore
 import io.github.oshai.kotlinlogging.KotlinLogging
+import jakarta.inject.Named
 import jakarta.inject.Singleton
 import java.lang.IllegalStateException
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.CoroutineExceptionHandler
 import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.coroutineScope
+import kotlinx.coroutines.ExecutorCoroutineDispatcher
+import kotlinx.coroutines.cancel
 import kotlinx.coroutines.joinAll
 import kotlinx.coroutines.launch
 
+/**
+ * Orchestrates the running of pipelines in parallel, handles and propagates errors and manages the
+ * state reconciler lifecycle.
+ */
 @Singleton
 class PipelineRunner(
     private val reconciler: StateReconciler,
     private val store: StateStore,
-    val pipelines: List<DataFlowPipeline>,
+    private val pipelines: List<DataFlowPipeline>,
+    private val inputStreams: ConnectorInputStreams,
+    @Named("pipelineRunnerScope") private val pipelineScope: CoroutineScope,
+    @Named("aggregationDispatcher") private val aggregationDispatcher: CoroutineDispatcher,
 ) {
     private val log = KotlinLogging.logger {}
 
-    suspend fun run() = coroutineScope {
+    private var terminalException: Throwable? = null
+
+    suspend fun run() {
         log.info { "Destination Pipeline Starting..." }
-        log.info { "Running with ${pipelines.size} input streams..." }
 
-        reconciler.run(CoroutineScope(Dispatchers.IO))
+        log.info { "Starting state reconciler..." }
+        reconciler.run()
 
+        log.info { "Starting ${pipelines.size} pipelines..." }
+        pipelines.map { p -> pipelineScope.launch(exceptionHandler) { p.run() } }.joinAll()
+        log.info { "Individual pipelines complete..." }
+
+        // shutdown the reconciler regardless of success or failure, so we don't hang
         try {
-            pipelines.map { p -> launch { p.run() } }.joinAll()
-            log.info { "Individual pipelines complete..." }
-        } finally {
-            // shutdown the reconciler regardless of success or failure, so we don't hang
-            log.info { "Disabling reconciler..." }
+            log.info { "Disabling state reconciler..." }
             reconciler.disable()
+        } finally {
+            if (aggregationDispatcher is ExecutorCoroutineDispatcher) {
+                aggregationDispatcher.close()
+            }
+        }
+
+        terminalException?.let {
+            log.error(terminalException) { "Destination Pipeline Completed — Exceptionally" }
+            throw it
         }
 
         log.info { "Flushing final states..." }
         reconciler.flushCompleteStates()
 
-        log.info { "Destination Pipeline Completed — Successfully" }
-
         if (store.hasStates()) {
-            log.info { "Unflushed states detected. Failing sync." }
-            throw IllegalStateException("Sync completed, but unflushed states were detected.")
+            val stateException =
+                IllegalStateException("Sync completed, but unflushed states were detected.")
+            log.info { "Destination Pipeline Completed — Exceptionally: $stateException" }
+            throw stateException
         }
+
+        log.info { "Destination Pipeline Completed — Successfully" }
+    }
+
+    // ensure all pipelines close when a single pipeline throws an exception
+    private val exceptionHandler = CoroutineExceptionHandler { context, exception ->
+        log.error { "Caught Pipeline Exception: $exception\n Cancelling Destination Pipeline..." }
+        // close each input stream to cancel any blocking reads on them that would prevent shutdown
+        inputStreams.closeAll()
+        context.cancel()
+        // capture the child exception to re-throw (core-CDK relies on catching thrown exceptions)
+        // exceptions directly thrown from within this block will be suppressed
+        terminalException = exception
     }
 }

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/dataflow/state/StateReconciler.kt

@@ -24,18 +24,18 @@ class StateReconciler(
     private val stateStore: StateStore,
     private val emittedStatsStore: EmittedStatsStore,
     private val consumer: OutputConsumer,
-    @Named("stateReconciliationInterval")
-    reconciliationInterval: Duration?, // only java durations can be injected
+    @Named("stateReconcilerScope") private val scope: CoroutineScope,
+    @Named("stateReconcilerInterval") interval: Duration?, // only java durations can be injected
 ) {
     // allow overriding this for test purposes
-    private val reconciliationInterval = reconciliationInterval?.toKotlinDuration() ?: 30.seconds
+    private val interval = interval?.toKotlinDuration() ?: 30.seconds
     private lateinit var job: Job
 
-    fun run(scope: CoroutineScope) {
+    fun run() {
         job =
             scope.launch {
                 while (true) {
-                    delay(reconciliationInterval)
+                    delay(interval)
                     flushCompleteStates()
                     flushEmittedStats()
                 }

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/file/ClientSocket.kt

@@ -90,20 +90,6 @@ class ClientSocket(
         // will break tests. TODO: Anything else.
         log.info { "Socket file $socketPath connected for reading" }
 
-        val inputStream = Channels.newInputStream(openedSocket).buffered(bufferSizeBytes)
-
-        return SocketInputStream(openedSocket, inputStream)
-    }
-}
-
-class SocketInputStream(
-    private val socketChannel: SocketChannel,
-    private val inputStream: InputStream,
-) : InputStream() {
-    override fun read(): Int = inputStream.read()
-
-    override fun close() {
-        inputStream.close()
-        socketChannel.close()
+        return Channels.newInputStream(openedSocket).buffered(bufferSizeBytes)
     }
 }