Cleanup. Rev version.

tryangul · tryangul · commit 1367b4f48d3a · 2025-09-24T11:45:52.000-07:00
diff --git a/airbyte-cdk/bulk/changelog.md b/airbyte-cdk/bulk/changelog.md
@@ -1,3 +1,7 @@
+## Version 0.1.41
+
+Datflow Load CDK: Fixes hang when one of many parallel pipelines fails. Organizes thread pools. 
+
 ## Version 0.1.40
 
 Add gradle task to bump CDK version + add changelog entry
@@ -6,6 +10,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
diff --git a/airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/dataflow/config/DispatcherBeanFactory.kt b/airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/dataflow/config/DispatcherBeanFactory.kt
@@ -13,9 +13,7 @@ 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.
- */
+/** The dispatchers (think views of thread pools) and static scopes we use for dataflow. */
 @Factory
 class DispatcherBeanFactory {
     @Named("pipelineRunnerDispatcher")
diff --git a/airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/dataflow/pipeline/PipelineRunner.kt b/airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/dataflow/pipeline/PipelineRunner.kt
@@ -56,7 +56,8 @@ class PipelineRunner(
         reconciler.flushCompleteStates()
 
         if (store.hasStates()) {
-            val stateException = 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
         }
diff --git a/airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/file/ClientSocket.kt b/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 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)
     }
 }
diff --git a/airbyte-cdk/bulk/core/load/src/test/kotlin/io/airbyte/cdk/load/dataflow/pipeline/DataFlowPipelineTest.kt b/airbyte-cdk/bulk/core/load/src/test/kotlin/io/airbyte/cdk/load/dataflow/pipeline/DataFlowPipelineTest.kt
@@ -34,7 +34,7 @@ class DataFlowPipelineTest {
         val testScope = TestScope(this.testScheduler)
         val aggregationDispatcher = StandardTestDispatcher(testScope.testScheduler)
         val flushDispatcher = StandardTestDispatcher(testScope.testScheduler)
-        
+
         // Given
         val initialIO = mockk<DataFlowStageIO>()
         val input = flowOf(initialIO)
@@ -68,7 +68,7 @@ class DataFlowPipelineTest {
 
         // When
         pipeline.run()
-        
+
         // Advance the test scheduler to process all coroutines
         testScope.testScheduler.advanceUntilIdle()
 
@@ -81,4 +81,4 @@ class DataFlowPipelineTest {
             completionHandler.apply(null)
         }
     }
-}
+}
diff --git a/airbyte-cdk/bulk/core/load/src/test/kotlin/io/airbyte/cdk/load/dataflow/pipeline/PipelineRunnerTest.kt b/airbyte-cdk/bulk/core/load/src/test/kotlin/io/airbyte/cdk/load/dataflow/pipeline/PipelineRunnerTest.kt
@@ -53,8 +53,15 @@ class PipelineRunnerTest {
     fun `run should execute all pipelines concurrently`() = runTest {
         // Given
         val pipelineScope = Fixtures.testScope(this.coroutineContext)
-        val runner = PipelineRunner(reconciler, store, listOf(pipeline1, pipeline2, pipeline3), inputStreams, pipelineScope)
-        
+        val runner =
+            PipelineRunner(
+                reconciler,
+                store,
+                listOf(pipeline1, pipeline2, pipeline3),
+                inputStreams,
+                pipelineScope
+            )
+
         coEvery { pipeline1.run() } coAnswers { delay(100) }
         coEvery { pipeline2.run() } coAnswers { delay(50) }
         coEvery { pipeline3.run() } coAnswers { delay(75) }
@@ -79,13 +86,15 @@ class PipelineRunnerTest {
         val exception = RuntimeException("Pipeline failed")
         coEvery { pipeline1.run() } throws exception
         val failingScope = Fixtures.testScope(this.coroutineContext)
-        val failingRunner = PipelineRunner(reconciler, store, listOf(pipeline1), inputStreams, failingScope)
+        val failingRunner =
+            PipelineRunner(reconciler, store, listOf(pipeline1), inputStreams, failingScope)
 
         // When/Then
-        val thrownException = assertThrows<RuntimeException> {
-            failingRunner.run()
-            testScheduler.advanceUntilIdle()
-        }
+        val thrownException =
+            assertThrows<RuntimeException> {
+                failingRunner.run()
+                testScheduler.advanceUntilIdle()
+            }
 
         assertEquals("Pipeline failed", thrownException.message)
 
@@ -99,7 +108,8 @@ class PipelineRunnerTest {
         // Given
         coEvery { pipeline1.run() } just Runs
         val singlePipelineScope = Fixtures.testScope(this.coroutineContext)
-        val singlePipelineRunner = PipelineRunner(reconciler, store, listOf(pipeline1), inputStreams, singlePipelineScope)
+        val singlePipelineRunner =
+            PipelineRunner(reconciler, store, listOf(pipeline1), inputStreams, singlePipelineScope)
 
         // When
         singlePipelineRunner.run()
@@ -132,7 +142,8 @@ class PipelineRunnerTest {
         // Given
         coEvery { pipeline1.run() } just Runs
         val singlePipelineScope = Fixtures.testScope(this.coroutineContext)
-        val singlePipelineRunner = PipelineRunner(reconciler, store, listOf(pipeline1), inputStreams, singlePipelineScope)
+        val singlePipelineRunner =
+            PipelineRunner(reconciler, store, listOf(pipeline1), inputStreams, singlePipelineScope)
 
         // When
         singlePipelineRunner.run()
@@ -151,13 +162,15 @@ class PipelineRunnerTest {
         coEvery { pipeline1.run() } just Runs
         coEvery { reconciler.disable() } throws RuntimeException("Failed to disable")
         val singlePipelineScope = Fixtures.testScope(this.coroutineContext)
-        val singlePipelineRunner = PipelineRunner(reconciler, store, listOf(pipeline1), inputStreams, singlePipelineScope)
+        val singlePipelineRunner =
+            PipelineRunner(reconciler, store, listOf(pipeline1), inputStreams, singlePipelineScope)
 
         // When/Then
-        val exception = assertThrows<RuntimeException> { 
-            singlePipelineRunner.run()
-            testScheduler.advanceUntilIdle()
-        }
+        val exception =
+            assertThrows<RuntimeException> {
+                singlePipelineRunner.run()
+                testScheduler.advanceUntilIdle()
+            }
 
         assertEquals("Failed to disable", exception.message)
 
@@ -172,13 +185,15 @@ class PipelineRunnerTest {
         coEvery { pipeline1.run() } just Runs
         every { reconciler.flushCompleteStates() } throws RuntimeException("Failed to flush")
         val singlePipelineScope = Fixtures.testScope(this.coroutineContext)
-        val singlePipelineRunner = PipelineRunner(reconciler, store, listOf(pipeline1), inputStreams, singlePipelineScope)
+        val singlePipelineRunner =
+            PipelineRunner(reconciler, store, listOf(pipeline1), inputStreams, singlePipelineScope)
 
         // When/Then
-        val exception = assertThrows<RuntimeException> { 
-            singlePipelineRunner.run()
-            testScheduler.advanceUntilIdle()
-        }
+        val exception =
+            assertThrows<RuntimeException> {
+                singlePipelineRunner.run()
+                testScheduler.advanceUntilIdle()
+            }
 
         assertEquals("Failed to flush", exception.message)
 
@@ -193,28 +208,30 @@ class PipelineRunnerTest {
         // Given
         val exception = RuntimeException("Pipeline failed")
         coEvery { pipeline1.run() } throws exception
-        
+
         val failingScope = Fixtures.testScope(this.coroutineContext)
-        val failingRunner = PipelineRunner(reconciler, store, listOf(pipeline1), inputStreams, failingScope)
+        val failingRunner =
+            PipelineRunner(reconciler, store, listOf(pipeline1), inputStreams, failingScope)
 
         // When
-        assertThrows<RuntimeException> { 
+        assertThrows<RuntimeException> {
             failingRunner.run()
             testScheduler.advanceUntilIdle()
         }
-        
+
         // Then - verify that closeAll was called due to exception handler
         verify(atLeast = 1) { inputStreams.closeAll() }
     }
 
     @Test
     fun `run should handle large number of pipelines`() = runTest {
         // Given
-        val pipelines = (1..32).map {
-            val pipeline = mockk<DataFlowPipeline>()
-            coEvery { pipeline.run() } just Runs
-            pipeline
-        }
+        val pipelines =
+            (1..32).map {
+                val pipeline = mockk<DataFlowPipeline>()
+                coEvery { pipeline.run() } just Runs
+                pipeline
+            }
 
         val largeScope = Fixtures.testScope(this.coroutineContext)
         val largeRunner = PipelineRunner(reconciler, store, pipelines, inputStreams, largeScope)
@@ -237,28 +254,32 @@ class PipelineRunnerTest {
         var pipeline2Complete = false
         var reconcilerDisabled = false
 
-        coEvery { pipeline1.run() } coAnswers {
-            delay(100)
-            pipeline1Complete = true
-        }
-        coEvery { pipeline2.run() } coAnswers {
-            delay(50)
-            pipeline2Complete = true
-        }
-        coEvery { reconciler.disable() } answers {
-            reconcilerDisabled = true
-            assertTrue(
-                pipeline1Complete,
-                "Pipeline 1 should be complete before disabling reconciler"
-            )
-            assertTrue(
-                pipeline2Complete,
-                "Pipeline 2 should be complete before disabling reconciler"
-            )
-        }
+        coEvery { pipeline1.run() } coAnswers
+            {
+                delay(100)
+                pipeline1Complete = true
+            }
+        coEvery { pipeline2.run() } coAnswers
+            {
+                delay(50)
+                pipeline2Complete = true
+            }
+        coEvery { reconciler.disable() } answers
+            {
+                reconcilerDisabled = true
+                assertTrue(
+                    pipeline1Complete,
+                    "Pipeline 1 should be complete before disabling reconciler"
+                )
+                assertTrue(
+                    pipeline2Complete,
+                    "Pipeline 2 should be complete before disabling reconciler"
+                )
+            }
 
         val twoScope = Fixtures.testScope(this.coroutineContext)
-        val twoRunner = PipelineRunner(reconciler, store, listOf(pipeline1, pipeline2), inputStreams, twoScope)
+        val twoRunner =
+            PipelineRunner(reconciler, store, listOf(pipeline1, pipeline2), inputStreams, twoScope)
 
         // When
         twoRunner.run()
@@ -269,39 +290,49 @@ class PipelineRunnerTest {
     }
 
     @Test
-    fun `run should throw IllegalStateException when unflushed states exist at sync end`() = runTest {
-        // Given
-        coEvery { pipeline1.run() } just Runs
-        every { store.hasStates() } returns true // Simulate unflushed states exist
-
-        val singlePipelineScope = Fixtures.testScope(this.coroutineContext)
-        val singlePipelineRunner = PipelineRunner(reconciler, store, listOf(pipeline1), inputStreams, singlePipelineScope)
-
-        // When/Then
-        val exception = assertThrows<IllegalStateException> { 
-            singlePipelineRunner.run()
-            testScheduler.advanceUntilIdle()
+    fun `run should throw IllegalStateException when unflushed states exist at sync end`() =
+        runTest {
+            // Given
+            coEvery { pipeline1.run() } just Runs
+            every { store.hasStates() } returns true // Simulate unflushed states exist
+
+            val singlePipelineScope = Fixtures.testScope(this.coroutineContext)
+            val singlePipelineRunner =
+                PipelineRunner(
+                    reconciler,
+                    store,
+                    listOf(pipeline1),
+                    inputStreams,
+                    singlePipelineScope
+                )
+
+            // When/Then
+            val exception =
+                assertThrows<IllegalStateException> {
+                    singlePipelineRunner.run()
+                    testScheduler.advanceUntilIdle()
+                }
+
+            // Then
+            assertEquals("Sync completed, but unflushed states were detected.", exception.message)
+
+            // Verify all operations completed before the check
+            coVerify(exactly = 1) { reconciler.run() }
+            coVerify(exactly = 1) { pipeline1.run() }
+            coVerify(exactly = 1) { reconciler.disable() }
+            coVerify(exactly = 1) { reconciler.flushCompleteStates() }
+            coVerify(exactly = 1) { store.hasStates() }
         }
 
-        // Then
-        assertEquals("Sync completed, but unflushed states were detected.", exception.message)
-
-        // Verify all operations completed before the check
-        coVerify(exactly = 1) { reconciler.run() }
-        coVerify(exactly = 1) { pipeline1.run() }
-        coVerify(exactly = 1) { reconciler.disable() }
-        coVerify(exactly = 1) { reconciler.flushCompleteStates() }
-        coVerify(exactly = 1) { store.hasStates() }
-    }
-
     @Test
     fun `run should not throw exception when all states are flushed`() = runTest {
         // Given
         coEvery { pipeline1.run() } just Runs
         every { store.hasStates() } returns false // All states are flushed
 
         val singlePipelineScope = Fixtures.testScope(this.coroutineContext)
-        val singlePipelineRunner = PipelineRunner(reconciler, store, listOf(pipeline1), inputStreams, singlePipelineScope)
+        val singlePipelineRunner =
+            PipelineRunner(reconciler, store, listOf(pipeline1), inputStreams, singlePipelineScope)
 
         // When - should complete without exception
         singlePipelineRunner.run()
diff --git a/airbyte-cdk/bulk/core/load/src/test/kotlin/io/airbyte/cdk/load/dataflow/state/StateReconcilerTest.kt b/airbyte-cdk/bulk/core/load/src/test/kotlin/io/airbyte/cdk/load/dataflow/state/StateReconcilerTest.kt
@@ -50,7 +50,7 @@ class StateReconcilerTest {
     fun setUp() {
         testScope = TestScope(StandardTestDispatcher())
         reconcilerScope = CoroutineScope(testScope.coroutineContext)
-        
+
         stateReconciler =
             StateReconciler(
                 stateStore,
@@ -146,7 +146,7 @@ class StateReconcilerTest {
         every { emittedStatsStore.getStats() } returns statsList
 
         every { consumer.accept(any<AirbyteMessage>()) } just Runs
-        
+
         // Create a new reconciler with the test scope for this test
         val localReconciler =
             StateReconciler(
@@ -156,7 +156,7 @@ class StateReconcilerTest {
                 this.backgroundScope,
                 interval.toJavaDuration(),
             )
-        
+
         // When
         localReconciler.run()
 
diff --git a/airbyte-cdk/bulk/version.properties b/airbyte-cdk/bulk/version.properties
@@ -1 +1 @@
-version=0.1.40
+version=0.1.41

Original file line number	Diff line number	Diff line change
`@@ -56,7 +56,8 @@ class PipelineRunner(`
`56`	`56`	`reconciler.flushCompleteStates()`
`57`	`57`
`58`	`58`	`if (store.hasStates()) {`
`59`		`- val stateException = IllegalStateException("Sync completed, but unflushed states were detected.")`
	`59`	`+ val stateException =`
	`60`	`+ IllegalStateException("Sync completed, but unflushed states were detected.")`
`60`	`61`	`log.info { "Destination Pipeline Completed — Exceptionally: $stateException" }`
`61`	`62`	`throw stateException`
`62`	`63`	`}`