Provide Serializable Interface for Creating WorkerRuns

airbyte-config/models/build.gradle

@@ -21,4 +21,5 @@ jsonSchema2Pojo {
     generateBuilders = true
     includeConstructors = false
     includeSetters = true
+    serializable = true
 }

airbyte-scheduler/src/main/java/io/airbyte/scheduler/JobSubmitter.java

@@ -38,6 +38,8 @@
 import io.airbyte.config.persistence.ConfigNotFoundException;
 import io.airbyte.config.persistence.ConfigRepository;
 import io.airbyte.scheduler.persistence.JobPersistence;
+import io.airbyte.scheduler.worker_run.SchedulerWorkerRunWithEnvironmentFactory;
+import io.airbyte.scheduler.worker_run.WorkerRun;
 import io.airbyte.validation.json.JsonValidationException;
 import io.airbyte.workers.WorkerConstants;
 import java.io.IOException;
@@ -58,16 +60,16 @@ public class JobSubmitter implements Runnable {
   private final ExecutorService threadPool;
   private final JobPersistence persistence;
   private final ConfigRepository configRepository;
-  private final WorkerRunFactory workerRunFactory;
+  private final SchedulerWorkerRunWithEnvironmentFactory schedulerWorkerRunWithEnvironmentFactory;
 
   public JobSubmitter(final ExecutorService threadPool,
                       final JobPersistence persistence,
                       final ConfigRepository configRepository,
-                      final WorkerRunFactory workerRunFactory) {
+                      final SchedulerWorkerRunWithEnvironmentFactory schedulerWorkerRunWithEnvironmentFactory) {
     this.threadPool = threadPool;
     this.persistence = persistence;
     this.configRepository = configRepository;
-    this.workerRunFactory = workerRunFactory;
+    this.schedulerWorkerRunWithEnvironmentFactory = schedulerWorkerRunWithEnvironmentFactory;
   }
 
   @Override
@@ -91,7 +93,7 @@ public void run() {
 
   @VisibleForTesting
   void submitJob(Job job) {
-    final WorkerRun workerRun = workerRunFactory.create(job);
+    final WorkerRun workerRun = schedulerWorkerRunWithEnvironmentFactory.create(job);
     // we need to know the attempt number before we begin the job lifecycle. thus we state what the
     // attempt number should be. if it is not, that the lifecycle will fail. this should not happen as
     // long as job submission for a single job is single threaded. this is a compromise to allow the job

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

@@ -38,6 +38,7 @@
 import io.airbyte.db.Databases;
 import io.airbyte.scheduler.persistence.DefaultJobPersistence;
 import io.airbyte.scheduler.persistence.JobPersistence;
+import io.airbyte.scheduler.worker_run.SchedulerWorkerRunWithEnvironmentFactory;
 import io.airbyte.workers.process.DockerProcessBuilderFactory;
 import io.airbyte.workers.process.KubeProcessBuilderFactory;
 import io.airbyte.workers.process.ProcessBuilderFactory;
@@ -93,11 +94,11 @@ public SchedulerApp(Path workspaceRoot,
   public void start() throws IOException {
     final ExecutorService workerThreadPool = Executors.newFixedThreadPool(MAX_WORKERS, THREAD_FACTORY);
     final ScheduledExecutorService scheduledPool = Executors.newSingleThreadScheduledExecutor();
-    final WorkerRunFactory workerRunFactory = new WorkerRunFactory(workspaceRoot, pbf);
+    final SchedulerWorkerRunWithEnvironmentFactory workerRunWithEnvironmentFactory = new SchedulerWorkerRunWithEnvironmentFactory(workspaceRoot, pbf);
 
     final JobRetrier jobRetrier = new JobRetrier(jobPersistence, Instant::now);
     final JobScheduler jobScheduler = new JobScheduler(jobPersistence, configRepository);
-    final JobSubmitter jobSubmitter = new JobSubmitter(workerThreadPool, jobPersistence, configRepository, workerRunFactory);
+    final JobSubmitter jobSubmitter = new JobSubmitter(workerThreadPool, jobPersistence, configRepository, workerRunWithEnvironmentFactory);
 
     Map<String, String> mdc = MDC.getCopyOfContextMap();
 

airbyte-scheduler/src/main/java/io/airbyte/scheduler/worker_run/BaseWorkerRunFactory.java

@@ -0,0 +1,72 @@
+/*
+ * 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.scheduler.worker_run;
+
+import io.airbyte.config.JobOutput;
+import io.airbyte.workers.Worker;
+import io.airbyte.workers.process.AirbyteIntegrationLauncher;
+import io.airbyte.workers.process.IntegrationLauncher;
+import io.airbyte.workers.process.ProcessBuilderFactory;
+import java.io.Serializable;
+import java.nio.file.Path;
+
+/**
+ * This classes exists to make testing WorkerRunFactories easier. For testing the extending class
+ * can use the protected constructor to pass in mocks. In production, the public constructor should
+ * be preferred.
+ *
+ * @param <T> Input config type
+ */
+public abstract class BaseWorkerRunFactory<T extends Serializable> implements WorkerRunFactory<T> {
+
+  final IntegrationLauncherFactory integrationLauncherFactory;
+  final WorkerRunCreator workerRunCreator;
+
+  public BaseWorkerRunFactory() {
+    this(AirbyteIntegrationLauncher::new, WorkerRun::new);
+  }
+
+  BaseWorkerRunFactory(IntegrationLauncherFactory integrationLauncherFactory, WorkerRunCreator workerRunCreator) {
+    this.integrationLauncherFactory = integrationLauncherFactory;
+    this.workerRunCreator = workerRunCreator;
+  }
+
+  // exists to make testing easier.
+  @FunctionalInterface
+  interface IntegrationLauncherFactory {
+
+    IntegrationLauncher create(long jobId, int attempt, final String image, ProcessBuilderFactory pbf);
+
+  }
+
+  // exists to make testing easier.
+  @FunctionalInterface
+  interface WorkerRunCreator {
+
+    <T> WorkerRun create(Path jobRoot, T input, Worker<T, JobOutput> worker);
+
+  }
+
+}