revert: "fix: detect source container OOM by monitoring pipe closure" (#18388)

Author: benmoriceau

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

@@ -35,25 +35,11 @@ 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()
 
@@ -118,55 +104,14 @@ 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() {
-      if (outputStreamExhausted) {
-        return true
-      }
-      val hasNext = messageIterator.hasNext()
-      if (!hasNext) {
-        outputStreamExhausted = true
-        return true
-      }
-      return false
-    }
+    get() = !messageIterator.hasNext() && containerIOHandle.exitCodeExists()
 
   override val exitValue: Int
-    /**
-     * 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
-    }
+    get() = containerIOHandle.getExitCode()
 
   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,23 +29,9 @@ 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()
@@ -86,55 +72,14 @@ 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.
-     *
-     * We consider the source finished in one case:
-     * 1. The iterator has no more messages
+     * Short circuit early to avoid checking the underlying process. note: hasNext is blocking.
      */
-    get() {
-      if (outputStreamExhausted) {
-        return true
-      }
-      val hasNext = messageIterator.hasNext()
-      if (!hasNext) {
-        outputStreamExhausted = true
-        return true
-      }
-      return false
-    }
+    get() = !messageIterator.hasNext() && containerIOHandle.exitCodeExists()
 
   override val exitValue: Int
-    /**
-     * 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
-    }
+    get() = containerIOHandle.getExitCode()
 
   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,8 +213,7 @@ internal class LocalContainerAirbyteDestinationTest {
   }
 
   @Test
-  internal fun testDestinationIsFinishedWithExitCode() {
-    // Test normal finish with exit code file present
+  internal fun testDestinationIsFinished() {
     val iterator =
       mockk<Iterator<AirbyteMessage>> {
         every { hasNext() } returns false
@@ -236,31 +235,15 @@ internal class LocalContainerAirbyteDestinationTest {
       jobRoot = jobRoot,
     )
     assertEquals(true, 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 true
+    assertEquals(false, destination.isFinished)
 
-    val destination =
-      LocalContainerAirbyteDestination(
-        streamFactory = streamFactory,
-        messageMetricsTracker = messageMetricsTracker,
-        messageWriterFactory = messageWriterFactory,
-        containerIOHandle = containerIOHandle,
-        containerLogMdcBuilder = containerLogMdcBuilder,
-        destinationTimeoutMonitor = destinationTimeoutMonitor,
-      )
+    every { iterator.hasNext() } returns false
+    exitValueFile.delete()
+    assertEquals(false, destination.isFinished)
 
-    destination.start(
-      destinationConfig = workerDestinationConfig,
-      jobRoot = jobRoot,
-    )
+    every { iterator.hasNext() } returns true
     assertEquals(false, destination.isFinished)
   }
 
@@ -321,7 +304,6 @@ internal class LocalContainerAirbyteDestinationTest {
     val mockedContainerIOHandle =
       mockk<ContainerIOHandle> {
         every { terminate() } returns true
-        every { exitCodeExists() } returns true
         every { getExitCode() } returns exitValue
       }
 
@@ -351,7 +333,6 @@ 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>()
@@ -390,7 +371,6 @@ internal class LocalContainerAirbyteDestinationTest {
     val mockedContainerIOHandle =
       mockk<ContainerIOHandle> {
         every { terminate() } returns true
-        every { exitCodeExists() } returns true
         every { getExitCode() } returns exitValue
       }
     every { messageMetricsTracker.flushDestReadCountMetric() } returns Unit
@@ -419,7 +399,6 @@ internal class LocalContainerAirbyteDestinationTest {
     val mockedContainerIOHandle =
       mockk<ContainerIOHandle> {
         every { terminate() } returns false
-        every { exitCodeExists() } returns true
         every { getExitCode() } returns exitValue
       }
     every { messageMetricsTracker.flushDestReadCountMetric() } returns Unit
@@ -448,7 +427,6 @@ internal class LocalContainerAirbyteDestinationTest {
     val mockedContainerIOHandle =
       mockk<ContainerIOHandle> {
         every { terminate() } returns false
-        every { exitCodeExists() } returns true
         every { getExitCode() } returns exitValue
       }
     every { messageMetricsTracker.flushDestReadCountMetric() } returns Unit