Actually apply Socat CPU override for experimentation (#7534)

airbytehq · Jun 28, 2023 · 217bbd9 · 217bbd9
1 parent 8d04286
commit 217bbd9
Show file tree

Hide file tree

Showing 9 changed files with 52 additions and 23 deletions.
diff --git a/...s-worker/src/main/java/io/airbyte/workers/orchestrator/KubeOrchestratorHandleFactory.java b/...s-worker/src/main/java/io/airbyte/workers/orchestrator/KubeOrchestratorHandleFactory.java
@@ -4,15 +4,11 @@
 
 package io.airbyte.workers.orchestrator;
 
-import static io.airbyte.config.EnvConfigs.SOCAT_KUBE_CPU_LIMIT;
-import static io.airbyte.config.EnvConfigs.SOCAT_KUBE_CPU_REQUEST;
-
 import com.google.common.annotations.VisibleForTesting;
 import io.airbyte.commons.functional.CheckedSupplier;
 import io.airbyte.commons.temporal.TemporalUtils;
 import io.airbyte.config.ReplicationOutput;
 import io.airbyte.config.StandardSyncInput;
-import io.airbyte.featureflag.ConcurrentSocatResources;
 import io.airbyte.featureflag.Connection;
 import io.airbyte.featureflag.ContainerOrchestratorDevImage;
 import io.airbyte.featureflag.ContainerOrchestratorJavaOpts;
@@ -27,7 +23,6 @@
 import io.airbyte.workers.sync.ReplicationLauncherWorker;
 import io.micronaut.context.annotation.Requires;
 import io.micronaut.context.annotation.Value;
-import io.micronaut.core.util.StringUtils;
 import io.temporal.activity.ActivityExecutionContext;
 import jakarta.inject.Named;
 import jakarta.inject.Singleton;
@@ -85,7 +80,8 @@ public CheckedSupplier<Worker<StandardSyncInput, ReplicationOutput>, Exception>
         activityContext,
         serverPort,
         temporalUtils,
-        workerConfigs);
+        workerConfigs,
+        featureFlagClient);
   }
 
   /**
@@ -114,14 +110,6 @@ static ContainerOrchestratorConfig injectContainerOrchestratorConfig(final Featu
       envMap.put("JAVA_OPTS", injectedJavaOpts);
     }
 
-    // Allow for the override of the socat pod CPU resources as part of the concurrent source read
-    // experimentation
-    final String socatResources = client.stringVariation(ConcurrentSocatResources.INSTANCE, new Connection(connectionId));
-    if (StringUtils.isNotEmpty(socatResources)) {
-      envMap.put(SOCAT_KUBE_CPU_LIMIT, socatResources);
-      envMap.put(SOCAT_KUBE_CPU_REQUEST, socatResources);
-    }
-
     // This is messy because the ContainerOrchestratorConfig is immutable, so we alwasy have to create
     // an
     // entirely new object.

diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/DbtLauncherWorker.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/DbtLauncherWorker.java
@@ -10,6 +10,7 @@
 import io.airbyte.commons.json.Jsons;
 import io.airbyte.commons.temporal.TemporalUtils;
 import io.airbyte.config.OperatorDbtInput;
+import io.airbyte.featureflag.FeatureFlagClient;
 import io.airbyte.persistence.job.models.IntegrationLauncherConfig;
 import io.airbyte.persistence.job.models.JobRunConfig;
 import io.airbyte.workers.ContainerOrchestratorConfig;
@@ -35,7 +36,8 @@ public DbtLauncherWorker(final UUID connectionId,
                            final ContainerOrchestratorConfig containerOrchestratorConfig,
                            final Supplier<ActivityExecutionContext> activityContext,
                            final Integer serverPort,
-                           final TemporalUtils temporalUtils) {
+                           final TemporalUtils temporalUtils,
+                           final FeatureFlagClient featureFlagClient) {
     super(
         connectionId,
         DBT,
@@ -50,6 +52,7 @@ public DbtLauncherWorker(final UUID connectionId,
         serverPort,
         temporalUtils,
         workerConfigs,
+        featureFlagClient,
         // Custom connector does not use Dbt at this moment, thus this flag for runnning job under
         // isolated pool can be set to false.
         false);

diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/LauncherWorker.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/LauncherWorker.java
@@ -4,6 +4,8 @@
 
 package io.airbyte.workers.sync;
 
+import static io.airbyte.config.EnvConfigs.SOCAT_KUBE_CPU_LIMIT;
+import static io.airbyte.config.EnvConfigs.SOCAT_KUBE_CPU_REQUEST;
 import static io.airbyte.metrics.lib.ApmTraceConstants.Tags.PROCESS_EXIT_VALUE_KEY;
 import static io.airbyte.metrics.lib.ApmTraceConstants.WORKER_OPERATION_NAME;
 import static io.airbyte.workers.process.Metadata.CONNECTION_ID_LABEL_KEY;
@@ -15,6 +17,9 @@
 import io.airbyte.commons.temporal.TemporalUtils;
 import io.airbyte.commons.temporal.sync.OrchestratorConstants;
 import io.airbyte.config.ResourceRequirements;
+import io.airbyte.featureflag.ConcurrentSocatResources;
+import io.airbyte.featureflag.Connection;
+import io.airbyte.featureflag.FeatureFlagClient;
 import io.airbyte.metrics.lib.ApmTraceUtils;
 import io.airbyte.persistence.job.models.JobRunConfig;
 import io.airbyte.workers.ContainerOrchestratorConfig;
@@ -31,6 +36,7 @@
 import io.fabric8.kubernetes.api.model.DeletionPropagation;
 import io.fabric8.kubernetes.api.model.Pod;
 import io.fabric8.kubernetes.client.KubernetesClientException;
+import io.micronaut.core.util.StringUtils;
 import io.temporal.activity.ActivityExecutionContext;
 import java.nio.file.Path;
 import java.time.Duration;
@@ -44,6 +50,8 @@
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import lombok.extern.slf4j.Slf4j;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  * Coordinates configuring and managing the state of an async process. This is tied to the (job_id,
@@ -55,6 +63,8 @@
 @Slf4j
 public abstract class LauncherWorker<INPUT, OUTPUT> implements Worker<INPUT, OUTPUT> {
 
+  private static final Logger LOGGER = LoggerFactory.getLogger(LauncherWorker.class);
+
   private static final Duration MAX_DELETION_TIMEOUT = Duration.ofSeconds(45);
 
   /**
@@ -79,6 +89,7 @@ public abstract class LauncherWorker<INPUT, OUTPUT> implements Worker<INPUT, OUT
   private final boolean isCustomConnector;
   private final AtomicBoolean cancelled = new AtomicBoolean(false);
   private AsyncOrchestratorPodProcess process;
+  private final FeatureFlagClient featureFlagClient;
 
   public LauncherWorker(final UUID connectionId,
                         final String application,
@@ -92,6 +103,7 @@ public LauncherWorker(final UUID connectionId,
                         final Integer serverPort,
                         final TemporalUtils temporalUtils,
                         final WorkerConfigs workerConfigs,
+                        final FeatureFlagClient featureFlagClient,
                         final boolean isCustomConnector) {
 
     this.connectionId = connectionId;
@@ -106,6 +118,7 @@ public LauncherWorker(final UUID connectionId,
     this.serverPort = serverPort;
     this.temporalUtils = temporalUtils;
     this.workerConfigs = workerConfigs;
+    this.featureFlagClient = featureFlagClient;
     this.isCustomConnector = isCustomConnector;
 
     // Generate a random UUID to unique identify the pod process
@@ -130,6 +143,15 @@ public OUTPUT run(final INPUT input, final Path jobRoot) throws WorkerException
         // Merge in the env from the ContainerOrchestratorConfig
         containerOrchestratorConfig.environmentVariables().entrySet().stream().forEach(e -> envMap.putIfAbsent(e.getKey(), e.getValue()));
 
+        // Allow for the override of the socat pod CPU resources as part of the concurrent source read
+        // experimentation
+        final String socatResources = featureFlagClient.stringVariation(ConcurrentSocatResources.INSTANCE, new Connection(connectionId));
+        if (StringUtils.isNotEmpty(socatResources)) {
+          LOGGER.info("Overriding Socat CPU limit and request to {}.", socatResources);
+          envMap.put(SOCAT_KUBE_CPU_LIMIT, socatResources);
+          envMap.put(SOCAT_KUBE_CPU_REQUEST, socatResources);
+        }
+
         final Map<String, String> fileMap = new HashMap<>(additionalFileMap);
         fileMap.putAll(Map.of(
             OrchestratorConstants.INIT_FILE_APPLICATION, application,

diff --git a/...yte-commons-worker/src/main/java/io/airbyte/workers/sync/NormalizationLauncherWorker.java b/...yte-commons-worker/src/main/java/io/airbyte/workers/sync/NormalizationLauncherWorker.java
@@ -11,6 +11,7 @@
 import io.airbyte.commons.temporal.TemporalUtils;
 import io.airbyte.config.NormalizationInput;
 import io.airbyte.config.NormalizationSummary;
+import io.airbyte.featureflag.FeatureFlagClient;
 import io.airbyte.persistence.job.models.IntegrationLauncherConfig;
 import io.airbyte.persistence.job.models.JobRunConfig;
 import io.airbyte.workers.ContainerOrchestratorConfig;
@@ -36,7 +37,8 @@ public NormalizationLauncherWorker(final UUID connectionId,
                                      final ContainerOrchestratorConfig containerOrchestratorConfig,
                                      final Supplier<ActivityExecutionContext> activityContext,
                                      final Integer serverPort,
-                                     final TemporalUtils temporalUtils) {
+                                     final TemporalUtils temporalUtils,
+                                     final FeatureFlagClient featureFlagClient) {
     super(
         connectionId,
         NORMALIZATION,
@@ -51,6 +53,7 @@ public NormalizationLauncherWorker(final UUID connectionId,
         serverPort,
         temporalUtils,
         workerConfigs,
+        featureFlagClient,
         // Normalization process will happen only on a fixed set of connectors,
         // thus they are not going to be run under custom connectors. Setting this to false.
         false);

diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/ReplicationLauncherWorker.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/ReplicationLauncherWorker.java
@@ -12,6 +12,7 @@
 import io.airbyte.config.ReplicationOutput;
 import io.airbyte.config.ResourceRequirements;
 import io.airbyte.config.StandardSyncInput;
+import io.airbyte.featureflag.FeatureFlagClient;
 import io.airbyte.persistence.job.models.IntegrationLauncherConfig;
 import io.airbyte.persistence.job.models.JobRunConfig;
 import io.airbyte.workers.ContainerOrchestratorConfig;
@@ -42,7 +43,8 @@ public ReplicationLauncherWorker(final UUID connectionId,
                                    final Supplier<ActivityExecutionContext> activityContext,
                                    final Integer serverPort,
                                    final TemporalUtils temporalUtils,
-                                   final WorkerConfigs workerConfigs) {
+                                   final WorkerConfigs workerConfigs,
+                                   final FeatureFlagClient featureFlagClient) {
     super(
         connectionId,
         REPLICATION,
@@ -58,6 +60,7 @@ public ReplicationLauncherWorker(final UUID connectionId,
         serverPort,
         temporalUtils,
         workerConfigs,
+        featureFlagClient,
         sourceLauncherConfig.getIsCustomConnector() || destinationLauncherConfig.getIsCustomConnector());
   }
 

diff --git a/airbyte-featureflag/src/main/kotlin/FlagDefinitions.kt b/airbyte-featureflag/src/main/kotlin/FlagDefinitions.kt
@@ -64,7 +64,7 @@ object RefreshSchemaPeriod : Temporary<Int>(key= "refreshSchema.period.hours", d
 
 object ConcurrentSourceStreamRead : Temporary<Boolean>(key = "concurrent.source.stream.read", default = false)
 
-object ConcurrentSocatResources : Temporary<String>(key = "concurrent.socat.resources", default = "0.5")
+object ConcurrentSocatResources : Temporary<String>(key = "concurrent.socat.resources", default = "")
 
 object ReplicationWorkerImpl : Permanent<String>(key = "platform.replication-worker-impl", default = "default")
 

diff --git a/...workers/src/main/java/io/airbyte/workers/temporal/sync/DbtTransformationActivityImpl.java b/...workers/src/main/java/io/airbyte/workers/temporal/sync/DbtTransformationActivityImpl.java
@@ -22,6 +22,7 @@
 import io.airbyte.config.ResourceRequirements;
 import io.airbyte.config.helpers.LogConfigs;
 import io.airbyte.config.persistence.split_secrets.SecretsHydrator;
+import io.airbyte.featureflag.FeatureFlagClient;
 import io.airbyte.metrics.lib.ApmTraceUtils;
 import io.airbyte.metrics.lib.MetricClientFactory;
 import io.airbyte.metrics.lib.OssMetricsRegistry;
@@ -67,6 +68,7 @@ public class DbtTransformationActivityImpl implements DbtTransformationActivity
   private final AirbyteConfigValidator airbyteConfigValidator;
   private final TemporalUtils temporalUtils;
   private final AirbyteApiClient airbyteApiClient;
+  private final FeatureFlagClient featureFlagClient;
 
   public DbtTransformationActivityImpl(@Named("containerOrchestratorConfig") final Optional<ContainerOrchestratorConfig> containerOrchestratorConfig,
                                        final WorkerConfigsProvider workerConfigsProvider,
@@ -79,7 +81,8 @@ public DbtTransformationActivityImpl(@Named("containerOrchestratorConfig") final
                                        @Value("${micronaut.server.port}") final Integer serverPort,
                                        final AirbyteConfigValidator airbyteConfigValidator,
                                        final TemporalUtils temporalUtils,
-                                       final AirbyteApiClient airbyteApiClient) {
+                                       final AirbyteApiClient airbyteApiClient,
+                                       final FeatureFlagClient featureFlagClient) {
     this.containerOrchestratorConfig = containerOrchestratorConfig;
     this.workerConfigsProvider = workerConfigsProvider;
     this.processFactory = processFactory;
@@ -92,6 +95,7 @@ public DbtTransformationActivityImpl(@Named("containerOrchestratorConfig") final
     this.airbyteConfigValidator = airbyteConfigValidator;
     this.temporalUtils = temporalUtils;
     this.airbyteApiClient = airbyteApiClient;
+    this.featureFlagClient = featureFlagClient;
   }
 
   @Trace(operationName = ACTIVITY_TRACE_OPERATION_NAME)
@@ -173,7 +177,8 @@ private CheckedSupplier<Worker<OperatorDbtInput, Void>, Exception> getContainerL
         containerOrchestratorConfig.get(),
         activityContext,
         serverPort,
-        temporalUtils);
+        temporalUtils,
+        featureFlagClient);
   }
 
 }
diff --git a/...yte-workers/src/main/java/io/airbyte/workers/temporal/sync/NormalizationActivityImpl.java b/...yte-workers/src/main/java/io/airbyte/workers/temporal/sync/NormalizationActivityImpl.java
@@ -29,6 +29,7 @@
 import io.airbyte.config.StandardSyncOutput;
 import io.airbyte.config.helpers.LogConfigs;
 import io.airbyte.config.persistence.split_secrets.SecretsHydrator;
+import io.airbyte.featureflag.FeatureFlagClient;
 import io.airbyte.metrics.lib.ApmTraceUtils;
 import io.airbyte.metrics.lib.MetricClientFactory;
 import io.airbyte.metrics.lib.OssMetricsRegistry;
@@ -76,6 +77,7 @@ public class NormalizationActivityImpl implements NormalizationActivity {
   private final AirbyteConfigValidator airbyteConfigValidator;
   private final TemporalUtils temporalUtils;
   private final AirbyteApiClient airbyteApiClient;
+  private final FeatureFlagClient featureFlagClient;
 
   private static final String V1_NORMALIZATION_MINOR_VERSION = "3";
 
@@ -90,7 +92,8 @@ public NormalizationActivityImpl(@Named("containerOrchestratorConfig") final Opt
                                    @Value("${micronaut.server.port}") final Integer serverPort,
                                    final AirbyteConfigValidator airbyteConfigValidator,
                                    final TemporalUtils temporalUtils,
-                                   final AirbyteApiClient airbyteApiClient) {
+                                   final AirbyteApiClient airbyteApiClient,
+                                   final FeatureFlagClient featureFlagClient) {
     this.containerOrchestratorConfig = containerOrchestratorConfig;
     this.workerConfigsProvider = workerConfigsProvider;
     this.processFactory = processFactory;
@@ -103,6 +106,7 @@ public NormalizationActivityImpl(@Named("containerOrchestratorConfig") final Opt
     this.airbyteConfigValidator = airbyteConfigValidator;
     this.temporalUtils = temporalUtils;
     this.airbyteApiClient = airbyteApiClient;
+    this.featureFlagClient = featureFlagClient;
   }
 
   @Trace(operationName = ACTIVITY_TRACE_OPERATION_NAME)
@@ -265,7 +269,8 @@ private CheckedSupplier<Worker<NormalizationInput, NormalizationSummary>, Except
         containerOrchestratorConfig.get(),
         activityContext,
         serverPort,
-        temporalUtils);
+        temporalUtils,
+        featureFlagClient);
   }
 
 }
diff --git a/flags.yml b/flags.yml
@@ -22,6 +22,6 @@ flags:
   - name: concurrent.source.stream.read
     serve: false
   - name: concurrent.socat.resources
-    serve: "0.5"
+    serve: ""
   - name: platform.add-scheduling-jitter
     serve: false