[WIP] Bulk load cdk: add checkpointing test

airbyte-cdk/bulk/core/base/src/main/kotlin/io/airbyte/cdk/AirbyteConnectorRunner.kt

@@ -58,7 +58,7 @@ sealed class AirbyteConnectorRunner(
     val testBeanDefinitions: Array<out RuntimeBeanDefinition<*>>,
     val testProperties: Map<String, String> = emptyMap(),
 ) {
-    val envs: Array<String> = arrayOf(Environment.CLI, connectorType)
+        val envs: Array<String> = arrayOf(Environment.CLI, connectorType)
 
     inline fun <reified R : Runnable> run() {
         val picocliCommandLineFactory = PicocliCommandLineFactory(this)

airbyte-cdk/bulk/core/load/build.gradle

@@ -23,6 +23,7 @@ dependencies {
     testImplementation("org.jetbrains.kotlinx:kotlinx-coroutines-test:1.8.1")
     implementation "org.jetbrains.kotlin:kotlin-reflect:2.0.20"
     testFixturesImplementation "uk.org.webcompere:system-stubs-jupiter:2.1.7"
+    testFixturesImplementation "org.jetbrains.kotlinx:kotlinx-coroutines-test:1.8.1"
 }
 
 task integrationTest(type: Test) {

airbyte-cdk/bulk/core/load/src/testFixtures/kotlin/io/airbyte/cdk/load/test/util/IntegrationTest.kt

@@ -23,7 +23,7 @@ import java.time.format.DateTimeFormatter
 import java.util.concurrent.atomic.AtomicBoolean
 import javax.inject.Inject
 import kotlin.test.fail
-import kotlinx.coroutines.async
+import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
 import org.apache.commons.lang3.RandomStringUtils
 import org.junit.jupiter.api.AfterEach
@@ -150,7 +150,7 @@ abstract class IntegrationTest(
                 catalog.asProtocolObject(),
             )
         return runBlocking {
-            val destinationCompletion = async { destination.run() }
+            val destinationCompletion = launch { destination.run() }
             messages.forEach { destination.sendMessage(it.asProtocolMessage()) }
             if (streamStatus != null) {
                 catalog.streams.forEach {
@@ -175,7 +175,7 @@ abstract class IntegrationTest(
                 }
             }
             destination.shutdown()
-            destinationCompletion.await()
+            destinationCompletion.join()
             destination.readMessages()
         }
     }

airbyte-cdk/bulk/core/load/src/testFixtures/kotlin/io/airbyte/cdk/load/test/util/destination_process/DestinationProcess.kt

@@ -29,6 +29,9 @@ interface DestinationProcess {
     suspend fun run()
 
     fun sendMessage(message: AirbyteMessage)
+    fun sendMessages(vararg messages: AirbyteMessage) {
+        messages.forEach { sendMessage(it) }
+    }
 
     /** Return all messages the destination emitted since the last call to [readMessages]. */
     fun readMessages(): List<AirbyteMessage>

airbyte-cdk/bulk/core/load/src/testFixtures/kotlin/io/airbyte/cdk/load/test/util/destination_process/DockerizedDestination.kt

@@ -24,6 +24,7 @@ import java.time.Clock
 import java.util.Locale
 import java.util.Scanner
 import javax.inject.Singleton
+import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.launch
@@ -183,6 +184,11 @@ class DockerizedDestination(
                 getMdcScope().use { process.errorReader().forEachLine { logger.error { it } } }
                 stderrDrained.complete(Unit)
             }
+        }.invokeOnCompletion { cause: Throwable? ->
+            // TODO does this work correctly?
+            if (cause is CancellationException) {
+                process.destroy()
+            }
         }
     }
 
@@ -196,15 +202,20 @@ class DockerizedDestination(
     }
 
     override suspend fun shutdown() {
+        println("EDGAO DEBUG entered shutdown")
         withContext(Dispatchers.IO) {
+            println("EDGAO DEBUG closing stdin")
             destinationStdin.close()
-            // The old cdk had a 1-minute timeout here. That seems... weird?
-            // We can just rely on the junit timeout, presumably?
-            process.waitFor()
             // Wait for ourselves to drain stdout/stderr. Otherwise we won't capture
             // all the destination output (logs/trace messages).
+            println("EDGAO DEBUG draining streams")
             stdoutDrained.join()
             stderrDrained.join()
+            // The old cdk had a 1-minute timeout here. That seems... weird?
+            // We can just rely on the junit timeout, presumably?
+            println("EDGAO DEBUG waiting for process exit")
+            process.waitFor()
+            println("EDGAO DEBUG grabbing exit code")
             val exitCode = process.exitValue()
             if (exitCode != 0) {
                 // Hey look, it's possible to extract the error from a failed destination process!
@@ -214,8 +225,10 @@ class DockerizedDestination(
                         .traces()
                         .filter { it.type == AirbyteTraceMessage.Type.ERROR }
                         .map { it.error }
+                println("EDGAO DEBUG exiting uncleanly: $exitCode, $filteredTraces")
                 throw DestinationUncleanExitException(exitCode, filteredTraces)
             }
+            println("EDGAO DEBUG exited with exit code: $exitCode")
         }
     }
 }

airbyte-cdk/bulk/core/load/src/testFixtures/kotlin/io/airbyte/cdk/load/test/util/destination_process/NonDockerizedDestination.kt

@@ -15,6 +15,8 @@ import java.io.PipedInputStream
 import java.io.PipedOutputStream
 import java.io.PrintWriter
 import javax.inject.Singleton
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.withContext
 
 class NonDockerizedDestination(
     command: String,
@@ -53,7 +55,9 @@ class NonDockerizedDestination(
     }
 
     override suspend fun run() {
-        destination.run()
+        withContext(Dispatchers.IO) {
+            destination.run()
+        }
     }
 
     override fun sendMessage(message: AirbyteMessage) {