Kube: Better Port Abstraction.

airbyte-config/models/src/main/java/io/airbyte/config/EnvConfigs.java

@@ -28,6 +28,7 @@
 import io.airbyte.config.helpers.LogClientSingleton;
 import java.nio.file.Path;
 import java.util.Arrays;
+import java.util.HashSet;
 import java.util.Optional;
 import java.util.Set;
 import java.util.function.Function;
@@ -202,9 +203,11 @@ public String getTemporalHost() {
 
   @Override
   public Set<Integer> getTemporalWorkerPorts() {
-    return Arrays.stream(getEnvOrDefault(TEMPORAL_WORKER_PORTS, "").split(","))
-        .map(Integer::valueOf)
-        .collect(Collectors.toSet());
+    var ports = getEnvOrDefault(TEMPORAL_WORKER_PORTS, "");
+    if (ports.isEmpty()) {
+      return new HashSet<>();
+    }
+    return Arrays.stream(ports.split(",")).map(Integer::valueOf).collect(Collectors.toSet());
   }
 
   @Override

airbyte-scheduler/app/src/main/java/io/airbyte/scheduler/app/SchedulerApp.java

@@ -44,6 +44,7 @@
 import io.airbyte.scheduler.persistence.JobPersistence;
 import io.airbyte.scheduler.persistence.job_tracker.JobTracker;
 import io.airbyte.workers.process.DockerProcessFactory;
+import io.airbyte.workers.process.KubePortManagerSingleton;
 import io.airbyte.workers.process.KubeProcessFactory;
 import io.airbyte.workers.process.ProcessFactory;
 import io.airbyte.workers.process.WorkerHeartbeatServer;
@@ -62,10 +63,8 @@
 import java.time.Instant;
 import java.util.Map;
 import java.util.Optional;
-import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
-import java.util.concurrent.LinkedBlockingDeque;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.TimeUnit;
@@ -177,11 +176,10 @@ private static ProcessFactory getProcessBuilderFactory(Configs configs) throws I
     if (configs.getWorkerEnvironment() == Configs.WorkerEnvironment.KUBERNETES) {
       final ApiClient officialClient = Config.defaultClient();
       final KubernetesClient fabricClient = new DefaultKubernetesClient();
-      final BlockingQueue<Integer> workerPorts = new LinkedBlockingDeque<>(configs.getTemporalWorkerPorts());
       final String localIp = InetAddress.getLocalHost().getHostAddress();
       final String kubeHeartbeatUrl = localIp + ":" + KUBE_HEARTBEAT_PORT;
       LOGGER.info("Using Kubernetes namespace: {}", configs.getKubeNamespace());
-      return new KubeProcessFactory(configs.getKubeNamespace(), officialClient, fabricClient, kubeHeartbeatUrl, workerPorts);
+      return new KubeProcessFactory(configs.getKubeNamespace(), officialClient, fabricClient, kubeHeartbeatUrl);
     } else {
       return new DockerProcessFactory(
           configs.getWorkspaceRoot(),
@@ -222,6 +220,7 @@ public static void main(String[] args) throws IOException, InterruptedException
     final JobNotifier jobNotifier = new JobNotifier(configs.getWebappUrl(), configRepository);
 
     if (configs.getWorkerEnvironment() == Configs.WorkerEnvironment.KUBERNETES) {
+      KubePortManagerSingleton.warnIfInsufficientPorts(Integer.parseInt(configs.getSubmitterNumThreads()));
       Map<String, String> mdc = MDC.getCopyOfContextMap();
       Executors.newSingleThreadExecutor().submit(
           () -> {

airbyte-workers/src/main/java/io/airbyte/workers/WorkerUtils.java

@@ -108,12 +108,12 @@ static void gentleCloseWithHeartbeat(final Process process,
                                        final Duration checkHeartbeatDuration,
                                        final Duration forcedShutdownDuration,
                                        final BiConsumer<Process, Duration> forceShutdown) {
-
     while (process.isAlive() && heartbeatMonitor.isBeating()) {
       try {
         if (new EnvConfigs().getWorkerEnvironment().equals(WorkerEnvironment.KUBERNETES)) {
           LOGGER.debug("Gently closing process {} with heartbeat..", process.info().commandLine().get());
         }
+
         process.waitFor(checkHeartbeatDuration.toMillis(), TimeUnit.MILLISECONDS);
       } catch (InterruptedException e) {
         LOGGER.error("Exception while waiting for process to finish", e);
@@ -125,6 +125,7 @@ static void gentleCloseWithHeartbeat(final Process process,
         if (new EnvConfigs().getWorkerEnvironment().equals(WorkerEnvironment.KUBERNETES)) {
           LOGGER.debug("Gently closing process {} without heartbeat..", process.info().commandLine().get());
         }
+
         process.waitFor(gracefulShutdownDuration.toMillis(), TimeUnit.MILLISECONDS);
       } catch (InterruptedException e) {
         LOGGER.error("Exception during grace period for process to finish. This can happen when cancelling jobs.");
@@ -136,6 +137,7 @@ static void gentleCloseWithHeartbeat(final Process process,
       if (new EnvConfigs().getWorkerEnvironment().equals(WorkerEnvironment.KUBERNETES)) {
         LOGGER.debug("Force shutdown process {}..", process.info().commandLine().get());
       }
+
       forceShutdown.accept(process, forcedShutdownDuration);
     }
   }

airbyte-workers/src/main/java/io/airbyte/workers/process/KubePodProcess.java

@@ -63,7 +63,6 @@
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.function.Consumer;
 import org.apache.commons.io.output.NullOutputStream;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -136,7 +135,6 @@ public class KubePodProcess extends Process {
   private Integer returnCode = null;
   private Long lastStatusCheck = null;
 
-  private final Consumer<Integer> portReleaser;
   private final ServerSocket stdoutServerSocket;
   private final int stdoutLocalPort;
   private final ServerSocket stderrServerSocket;
@@ -286,7 +284,6 @@ private static void waitForInitPodToRun(KubernetesClient client, Pod podDefiniti
 
   public KubePodProcess(ApiClient officialClient,
                         KubernetesClient fabricClient,
-                        Consumer<Integer> portReleaser,
                         String podName,
                         String namespace,
                         String image,
@@ -300,7 +297,6 @@ public KubePodProcess(ApiClient officialClient,
                         final String... args)
       throws IOException, InterruptedException {
     this.fabricClient = fabricClient;
-    this.portReleaser = portReleaser;
     this.stdoutLocalPort = stdoutLocalPort;
     this.stderrLocalPort = stderrLocalPort;
 
@@ -536,12 +532,19 @@ private void close() {
     Exceptions.swallow(this.stdoutServerSocket::close);
     Exceptions.swallow(this.stderrServerSocket::close);
     Exceptions.swallow(this.executorService::shutdownNow);
-    try {
-      portReleaser.accept(stdoutLocalPort);
-      portReleaser.accept(stderrLocalPort);
-    } catch (Exception e) {
-      LOGGER.error("Error releasing ports ", e);
+
+    var stdoutPortReleased = KubePortManagerSingleton.offer(stdoutLocalPort);
+    if (!stdoutPortReleased) {
+      LOGGER.warn("Error while releasing port {} from pod {}. This can cause the scheduler to run out of ports.", stdoutLocalPort,
+          podDefinition.getMetadata().getName());
     }
+
+    var stderrPortReleased = KubePortManagerSingleton.offer(stderrLocalPort);
+    if (!stderrPortReleased) {
+      LOGGER.warn("Error while releasing port {} from pod {}. This can cause the scheduler to run out of ports.", stderrLocalPort,
+          podDefinition.getMetadata().getName());
+    }
+
     LOGGER.debug("Closed {}", podDefinition.getMetadata().getName());
   }
 

airbyte-workers/src/main/java/io/airbyte/workers/process/KubePortManagerSingleton.java

@@ -0,0 +1,81 @@
+/*
+ * MIT License
+ *
+ * Copyright (c) 2020 Airbyte
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+package io.airbyte.workers.process;
+
+import com.google.common.annotations.VisibleForTesting;
+import io.airbyte.config.EnvConfigs;
+import java.util.Set;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.LinkedBlockingDeque;
+import java.util.concurrent.TimeUnit;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Convenience wrapper around a thread-safe BlockingQueue. Keeps track of available ports for Kube
+ * Pod Processes.
+ *
+ * Although this data structure can do without the wrapper class, this class allows easier testing
+ * via the {@link #getNumAvailablePorts()} function.
+ *
+ * The singleton pattern clarifies that only one copy of this class is intended to exists per
+ * scheduler deployment.
+ */
+public class KubePortManagerSingleton {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(KubePortManagerSingleton.class);
+  private static final int MAX_PORTS_PER_WORKERS = 4; // A sync has two workers. Each worker requires 2 ports.
+  private static BlockingQueue<Integer> workerPorts = new LinkedBlockingDeque<>(new EnvConfigs().getTemporalWorkerPorts());
+
+  public static Integer take() throws InterruptedException {
+    return workerPorts.poll(10, TimeUnit.MINUTES);
+  }
+
+  public static boolean offer(Integer port) {
+    if (!workerPorts.contains(port)) {
+      workerPorts.add(port);
+      return true;
+    }
+    return false;
+  }
+
+  public static int getNumAvailablePorts() {
+    return workerPorts.size();
+  }
+
+  public static void warnIfInsufficientPorts(int potentialConcurrentWorkers) {
+    var maxUsedPorts = potentialConcurrentWorkers * MAX_PORTS_PER_WORKERS;
+    if (maxUsedPorts > workerPorts.size()) {
+      LOGGER.warn("{} workers can use up to {} ports at once. Only {} ports allocated. This might result in stuck jobs.",
+          potentialConcurrentWorkers, maxUsedPorts, workerPorts.size());
+    }
+  }
+
+  @VisibleForTesting
+  protected static void setWorkerPorts(Set<Integer> ports) {
+    workerPorts = new LinkedBlockingDeque<>(ports);
+  }
+
+}

airbyte-workers/src/main/java/io/airbyte/workers/process/KubeProcessFactory.java

@@ -31,8 +31,6 @@
 import io.kubernetes.client.openapi.ApiClient;
 import java.nio.file.Path;
 import java.util.Map;
-import java.util.concurrent.BlockingQueue;
-import java.util.function.Consumer;
 import org.apache.commons.lang3.RandomStringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -45,7 +43,6 @@ public class KubeProcessFactory implements ProcessFactory {
   private final ApiClient officialClient;
   private final KubernetesClient fabricClient;
   private final String kubeHeartbeatUrl;
-  private final BlockingQueue<Integer> workerPorts;
 
   /**
    * @param namespace kubernetes namespace where spawned pods will live
@@ -58,13 +55,11 @@ public class KubeProcessFactory implements ProcessFactory {
   public KubeProcessFactory(String namespace,
                             ApiClient officialClient,
                             KubernetesClient fabricClient,
-                            String kubeHeartbeatUrl,
-                            BlockingQueue<Integer> workerPorts) {
+                            String kubeHeartbeatUrl) {
     this.namespace = namespace;
     this.officialClient = officialClient;
     this.fabricClient = fabricClient;
     this.kubeHeartbeatUrl = kubeHeartbeatUrl;
-    this.workerPorts = workerPorts;
   }
 
   @Override
@@ -80,27 +75,18 @@ public Process create(String jobId,
       throws WorkerException {
     try {
       // used to differentiate source and destination processes with the same id and attempt
+
       final String podName = createPodName(imageName, jobId, attempt);
 
-      final int stdoutLocalPort = workerPorts.take();
+      final int stdoutLocalPort = KubePortManagerSingleton.take();
       LOGGER.info("{} stdoutLocalPort = {}", podName, stdoutLocalPort);
 
-      final int stderrLocalPort = workerPorts.take();
+      final int stderrLocalPort = KubePortManagerSingleton.take();
       LOGGER.info("{} stderrLocalPort = {}", podName, stderrLocalPort);
 
-      final Consumer<Integer> portReleaser = port -> {
-        if (!workerPorts.contains(port)) {
-          workerPorts.add(port);
-          LOGGER.info("{} releasing: {}", podName, port);
-        } else {
-          LOGGER.info("{} skipping releasing: {}", podName, port);
-        }
-      };
-
       return new KubePodProcess(
           officialClient,
           fabricClient,
-          portReleaser,
           podName,
           namespace,
           imageName,

airbyte-workers/src/test-integration/java/io/airbyte/workers/process/KubePodProcessIntegrationTest.java

@@ -46,8 +46,6 @@
 import java.util.List;
 import java.util.Optional;
 import java.util.Set;
-import java.util.concurrent.BlockingQueue;
-import java.util.concurrent.LinkedBlockingDeque;
 import org.apache.commons.lang.RandomStringUtils;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
@@ -58,27 +56,25 @@ public class KubePodProcessIntegrationTest {
 
   private static final boolean IS_MINIKUBE = Boolean.parseBoolean(Optional.ofNullable(System.getenv("IS_MINIKUBE")).orElse("false"));
   private List<Integer> openPorts;
-  private List<Integer> openWorkerPorts;
   private int heartbeatPort;
   private String heartbeatUrl;
   private ApiClient officialClient;
   private KubernetesClient fabricClient;
-  private BlockingQueue<Integer> workerPorts;
   private KubeProcessFactory processFactory;
 
   private static WorkerHeartbeatServer server;
 
   @BeforeEach
   public void setup() throws Exception {
     openPorts = new ArrayList<>(getOpenPorts(5));
-    openWorkerPorts = openPorts.subList(1, openPorts.size() - 1);
+    KubePortManagerSingleton.setWorkerPorts(new HashSet<>(openPorts.subList(1, openPorts.size() - 1)));
+
     heartbeatPort = openPorts.get(0);
     heartbeatUrl = getHost() + ":" + heartbeatPort;
 
     officialClient = Config.defaultClient();
     fabricClient = new DefaultKubernetesClient();
-    workerPorts = new LinkedBlockingDeque<>(openWorkerPorts);
-    processFactory = new KubeProcessFactory("default", officialClient, fabricClient, heartbeatUrl, workerPorts);
+    processFactory = new KubeProcessFactory("default", officialClient, fabricClient, heartbeatUrl);
 
     server = new WorkerHeartbeatServer(heartbeatPort);
     server.startBackground();