explicitly close socat socket (#12632)

* explicitly close socat socket * source can be closed multiple times because AirbyteSource is an AutoCloseable close should be idempotent though so calling it multiple times is fine
airbytehq · May 23, 2022 · a25233c · a25233c
1 parent 50e223b
commit a25233c
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 3 deletions.
diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/DefaultReplicationWorker.java b/airbyte-workers/src/main/java/io/airbyte/workers/DefaultReplicationWorker.java
@@ -300,6 +300,13 @@ private static Runnable getReplicationRunnable(final AirbyteSource source,
             if (recordsRead % 1000 == 0) {
               LOGGER.info("Records read: {} ({})", recordsRead, FileUtils.byteCountToDisplaySize(messageTracker.getTotalBytesEmitted()));
             }
+          } else {
+            LOGGER.info("Source has no more messages, closing connection.");
+            try {
+              source.close();
+            } catch (final Exception e) {
+              throw new SourceException("Source cannot be stopped!", e);
+            }
           }
         }
         LOGGER.info("Total records read: {} ({})", recordsRead, FileUtils.byteCountToDisplaySize(messageTracker.getTotalBytesEmitted()));

diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/process/KubePodProcess.java b/airbyte-workers/src/main/java/io/airbyte/workers/process/KubePodProcess.java
@@ -454,7 +454,7 @@ public KubePodProcess(final boolean isOrchestrator,
     final Container relayStdout = new ContainerBuilder()
         .withName("relay-stdout")
         .withImage(socatImage)
-        .withCommand("sh", "-c", String.format("cat %s | socat -d -d - TCP:%s:%s", STDOUT_PIPE_FILE, processRunnerHost, stdoutLocalPort))
+        .withCommand("sh", "-c", String.format("cat %s | socat -d -d -t 60 - TCP:%s:%s", STDOUT_PIPE_FILE, processRunnerHost, stdoutLocalPort))
         .withVolumeMounts(pipeVolumeMount, terminationVolumeMount)
         .withResources(sidecarResources)
         .withImagePullPolicy(sidecarImagePullPolicy)
@@ -463,7 +463,7 @@ public KubePodProcess(final boolean isOrchestrator,
     final Container relayStderr = new ContainerBuilder()
         .withName("relay-stderr")
         .withImage(socatImage)
-        .withCommand("sh", "-c", String.format("cat %s | socat -d -d - TCP:%s:%s", STDERR_PIPE_FILE, processRunnerHost, stderrLocalPort))
+        .withCommand("sh", "-c", String.format("cat %s | socat -d -d -t 60 - TCP:%s:%s", STDERR_PIPE_FILE, processRunnerHost, stderrLocalPort))
         .withVolumeMounts(pipeVolumeMount, terminationVolumeMount)
         .withResources(sidecarResources)
         .withImagePullPolicy(sidecarImagePullPolicy)

diff --git a/airbyte-workers/src/test/java/io/airbyte/workers/DefaultReplicationWorkerTest.java b/airbyte-workers/src/test/java/io/airbyte/workers/DefaultReplicationWorkerTest.java
@@ -11,6 +11,7 @@
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
@@ -133,7 +134,7 @@ void test() throws Exception {
     verify(destination).start(destinationConfig, jobRoot);
     verify(destination).accept(RECORD_MESSAGE1);
     verify(destination).accept(RECORD_MESSAGE2);
-    verify(source).close();
+    verify(source, atLeastOnce()).close();
     verify(destination).close();
   }