fix: detect source container OOM by monitoring pipe closure (#18381)

Co-authored-by: Claude <[email protected]>

Author: benmoriceau

airbyte-container-orchestrator/src/main/kotlin/io/airbyte/container/orchestrator/worker/io/LocalContainerAirbyteDestination.kt

@@ -35,11 +35,25 @@ class LocalContainerAirbyteDestination(
   private val containerIOHandle: ContainerIOHandle,
   private val containerLogMdcBuilder: MdcScope.Builder,
   private val flushImmediately: Boolean = false,
+  private val exitCodeWaitSeconds: Long = EXIT_CODE_WAIT_SECONDS,
 ) : AirbyteDestination {
+  companion object {
+    // Fallback exit code when the destination container is killed (e.g., OOM) and no exit code file is written
+    const val FALLBACK_EXIT_CODE = 1
+
+    // Time to wait for exit code file after pipe closes (to handle race condition between pipe close and exit code write)
+    const val EXIT_CODE_WAIT_SECONDS = 10L
+  }
+
   private val inputHasEnded = AtomicBoolean(false)
   private lateinit var messageIterator: Iterator<AirbyteMessage>
   private lateinit var writer: AirbyteMessageBufferedWriter<AirbyteMessage>
 
+  // Tracks when the output stream iterator is exhausted (pipes closed).
+  // This is important for detecting OOM kills where the container dies without writing an exit code file.
+  @Volatile
+  private var outputStreamExhausted = false
+
   override fun close() {
     emitDestinationMessageCountMetrics()
 
@@ -104,14 +118,55 @@ class LocalContainerAirbyteDestination(
   }
 
   override val isFinished: Boolean
-    /*
+    /**
      * As this check is done on every message read, it is important for this operation to be efficient.
      * Short circuit early to avoid checking the underlying process. Note: hasNext is blocking.
+     *
+     * We consider the destination finished in one case:
+     * 1. The iterator has no more messages
      */
-    get() = !messageIterator.hasNext() && containerIOHandle.exitCodeExists()
+    get() {
+      if (outputStreamExhausted) {
+        return true
+      }
+      val hasNext = messageIterator.hasNext()
+      if (!hasNext) {
+        outputStreamExhausted = true
+        return true
+      }
+      return false
+    }
 
   override val exitValue: Int
-    get() = containerIOHandle.getExitCode()
+    /**
+     * Returns the exit code of the destination container.
+     * If no exit code file exists but the output stream was exhausted (pipes closed),
+     * this waits briefly for the exit code file to appear (to handle the race condition
+     * between pipe close and exit code file write), then returns a fallback exit code
+     * if the file still doesn't exist.
+     */
+    get() {
+      if (containerIOHandle.exitCodeExists()) {
+        return containerIOHandle.getExitCode()
+      }
+      if (outputStreamExhausted) {
+        // Wait for exit code file - there's a race condition where the pipe closes
+        // before the shell script writes the exit code file
+        try {
+          Thread.sleep(exitCodeWaitSeconds * 1000)
+        } catch (e: InterruptedException) {
+          Thread.currentThread().interrupt()
+        }
+        if (containerIOHandle.exitCodeExists()) {
+          return containerIOHandle.getExitCode()
+        }
+        logger.error {
+          "No exit code file found after waiting ${exitCodeWaitSeconds}s. Container may have been killed (OOM or other signal). Returning fallback exit code $FALLBACK_EXIT_CODE."
+        }
+        return FALLBACK_EXIT_CODE
+      }
+      return containerIOHandle.getExitCode() // This will throw IllegalStateException as expected
+    }
 
   override fun attemptRead(): Optional<AirbyteMessage> {
     val m = if (messageIterator.hasNext()) messageIterator.next() else null

airbyte-container-orchestrator/src/main/kotlin/io/airbyte/container/orchestrator/worker/io/LocalContainerAirbyteSource.kt

@@ -29,9 +29,23 @@ class LocalContainerAirbyteSource(
   private val messageMetricsTracker: MessageMetricsTracker,
   private val containerIOHandle: ContainerIOHandle,
   private val containerLogMdcBuilder: MdcScope.Builder,
+  private val exitCodeWaitSeconds: Long = EXIT_CODE_WAIT_SECONDS,
 ) : AirbyteSource {
+  companion object {
+    // Fallback exit code when the source container is killed (e.g., OOM) and no exit code file is written
+    const val FALLBACK_EXIT_CODE = 1
+
+    // Time to wait for exit code file after pipe closes (to handle race condition between pipe close and exit code write)
+    const val EXIT_CODE_WAIT_SECONDS = 10L
+  }
+
   private lateinit var messageIterator: Iterator<AirbyteMessage>
 
+  // Tracks when the output stream iterator is exhausted (pipes closed).
+  // This is important for detecting OOM kills where the container dies without writing an exit code file.
+  @Volatile
+  private var outputStreamExhausted = false
+
   override fun close() {
     messageMetricsTracker.flushSourceReadCountMetric()
     val terminationResult = containerIOHandle.terminate()
@@ -72,14 +86,55 @@ class LocalContainerAirbyteSource(
   }
 
   override val isFinished: Boolean
-    /*
+    /**
      * As this check is done on every message read, it is important for this operation to be efficient.
-     * Short circuit early to avoid checking the underlying process. note: hasNext is blocking.
+     * Short circuit early to avoid checking the underlying process. Note: hasNext is blocking.
+     *
+     * We consider the source finished in one case:
+     * 1. The iterator has no more messages
      */
-    get() = !messageIterator.hasNext() && containerIOHandle.exitCodeExists()
+    get() {
+      if (outputStreamExhausted) {
+        return true
+      }
+      val hasNext = messageIterator.hasNext()
+      if (!hasNext) {
+        outputStreamExhausted = true
+        return true
+      }
+      return false
+    }
 
   override val exitValue: Int
-    get() = containerIOHandle.getExitCode()
+    /**
+     * Returns the exit code of the source container.
+     * If no exit code file exists but the output stream was exhausted (pipes closed),
+     * this waits briefly for the exit code file to appear (to handle the race condition
+     * between pipe close and exit code file write), then returns a fallback exit code
+     * if the file still doesn't exist.
+     */
+    get() {
+      if (containerIOHandle.exitCodeExists()) {
+        return containerIOHandle.getExitCode()
+      }
+      if (outputStreamExhausted) {
+        // Wait for exit code file - there's a race condition where the pipe closes
+        // before the shell script writes the exit code file
+        try {
+          Thread.sleep(exitCodeWaitSeconds * 1000)
+        } catch (e: InterruptedException) {
+          Thread.currentThread().interrupt()
+        }
+        if (containerIOHandle.exitCodeExists()) {
+          return containerIOHandle.getExitCode()
+        }
+        logger.error {
+          "No exit code file found after waiting ${exitCodeWaitSeconds}s. Container may have been killed (OOM or other signal). Returning fallback exit code $FALLBACK_EXIT_CODE."
+        }
+        return FALLBACK_EXIT_CODE
+      }
+      return containerIOHandle.getExitCode() // This will throw IllegalStateException as expected
+    }
 
   override fun attemptRead(): Optional<AirbyteMessage> {
     val m = if (messageIterator.hasNext()) messageIterator.next() else null

airbyte-container-orchestrator/src/test/kotlin/io/airbyte/container/orchestrator/worker/io/LocalContainerAirbyteDestinationTest.kt

@@ -213,7 +213,8 @@ internal class LocalContainerAirbyteDestinationTest {
   }
 
   @Test
-  internal fun testDestinationIsFinished() {
+  internal fun testDestinationIsFinishedWithExitCode() {
+    // Test normal finish with exit code file present
     val iterator =
       mockk<Iterator<AirbyteMessage>> {
         every { hasNext() } returns false
@@ -235,15 +236,31 @@ internal class LocalContainerAirbyteDestinationTest {
       jobRoot = jobRoot,
     )
     assertEquals(true, destination.isFinished)
+  }
 
-    every { iterator.hasNext() } returns true
-    assertEquals(false, destination.isFinished)
+  @Test
+  internal fun testDestinationIsNotFinishedWhileHasMessages() {
+    // Test that destination is not finished while there are still messages
+    val iterator =
+      mockk<Iterator<AirbyteMessage>> {
+        every { hasNext() } returns true
+      }
+    every { stream.iterator() } returns iterator
 
-    every { iterator.hasNext() } returns false
-    exitValueFile.delete()
-    assertEquals(false, destination.isFinished)
+    val destination =
+      LocalContainerAirbyteDestination(
+        streamFactory = streamFactory,
+        messageMetricsTracker = messageMetricsTracker,
+        messageWriterFactory = messageWriterFactory,
+        containerIOHandle = containerIOHandle,
+        containerLogMdcBuilder = containerLogMdcBuilder,
+        destinationTimeoutMonitor = destinationTimeoutMonitor,
+      )
 
-    every { iterator.hasNext() } returns true
+    destination.start(
+      destinationConfig = workerDestinationConfig,
+      jobRoot = jobRoot,
+    )
     assertEquals(false, destination.isFinished)
   }
 
@@ -304,6 +321,7 @@ internal class LocalContainerAirbyteDestinationTest {
     val mockedContainerIOHandle =
       mockk<ContainerIOHandle> {
         every { terminate() } returns true
+        every { exitCodeExists() } returns true
         every { getExitCode() } returns exitValue
       }
 
@@ -333,6 +351,7 @@ internal class LocalContainerAirbyteDestinationTest {
       mockk<ContainerIOHandle> {
         every { getErrInputStream() } returns mockk<InputStream>()
         every { terminate() } returns true
+        every { exitCodeExists() } returns true
         every { getExitCode() } returns exitValue
         every { getInputStream() } returns mockk<InputStream>()
         every { getOutputStream() } returns mockk<OutputStream>()
@@ -371,6 +390,7 @@ internal class LocalContainerAirbyteDestinationTest {
     val mockedContainerIOHandle =
       mockk<ContainerIOHandle> {
         every { terminate() } returns true
+        every { exitCodeExists() } returns true
         every { getExitCode() } returns exitValue
       }
     every { messageMetricsTracker.flushDestReadCountMetric() } returns Unit
@@ -399,6 +419,7 @@ internal class LocalContainerAirbyteDestinationTest {
     val mockedContainerIOHandle =
       mockk<ContainerIOHandle> {
         every { terminate() } returns false
+        every { exitCodeExists() } returns true
         every { getExitCode() } returns exitValue
       }
     every { messageMetricsTracker.flushDestReadCountMetric() } returns Unit
@@ -427,6 +448,7 @@ internal class LocalContainerAirbyteDestinationTest {
     val mockedContainerIOHandle =
       mockk<ContainerIOHandle> {
         every { terminate() } returns false
+        every { exitCodeExists() } returns true
         every { getExitCode() } returns exitValue
       }
     every { messageMetricsTracker.flushDestReadCountMetric() } returns Unit