add network config to titus request (#1185)

* add network config to titus request

* address comments

* remove unnecessary files

* address param location

* address nit comments

genie-web/src/main/java/com/netflix/genie/web/agent/launchers/dtos/TitusBatchJobRequest.java

@@ -71,6 +71,8 @@ public class TitusBatchJobRequest {
     @NotNull
     @NonNull
     private JobGroupInfo jobGroupInfo;
+    @Nullable
+    private NetworkConfiguration networkConfiguration;
 
     /**
      * Titus job owner POJO.
@@ -190,6 +192,15 @@ public static class Batch {
         private long runtimeLimitSec;
     }
 
+    /**
+     * Titus job network configuration.
+     */
+    @Data
+    @Builder
+    public static class NetworkConfiguration {
+        private String networkMode;
+    }
+
     /**
      * Titus job disruption budget.
      */

genie-web/src/main/java/com/netflix/genie/web/agent/launchers/impl/TitusAgentLauncherImpl.java

@@ -22,6 +22,7 @@
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.node.JsonNodeFactory;
 import com.github.benmanes.caffeine.cache.Cache;
+import com.google.common.base.Strings;
 import com.netflix.genie.common.internal.dtos.ComputeResources;
 import com.netflix.genie.common.internal.dtos.Image;
 import com.netflix.genie.common.internal.tracing.brave.BraveTracePropagator;
@@ -102,6 +103,13 @@ public class TitusAgentLauncherImpl implements AgentLauncher {
             .map(s -> placeholders.getOrDefault(s, s))
             .collect(Collectors.toList());
     private static final Logger LOG = LoggerFactory.getLogger(TitusAgentLauncherImpl.class);
+    // For Titus network mode details, see
+    // https://github.com/Netflix/titus-api-definitions/blob/master/doc/titus-v3-spec.md#networkconfigurationnetworkmode
+    private static final String TITUS_NETWORK_MODE_IPV4 = "Ipv4Only";
+    private static final String TITUS_NETWORK_MODE_DUAL_STACK = "Ipv6AndIpv4";
+    private static final String TITUS_NETWORK_MODE_DUAL_STACK_FALLBACK = "Ipv6AndIpv4Fallback";
+    private static final String TITUS_NETWORK_MODE_IPV6 = "Ipv6Only";
+    private static final String TITUS_NETWORK_MODE_HIGH_SCALE = "HighScale";
 
     private final RestTemplate restTemplate;
     private final RetryTemplate retryTemplate;
@@ -261,7 +269,7 @@ private TitusBatchJobRequest createJobRequest(final ResolvedJob resolvedJob) thr
 
         final Map<String, String> jobAttributes = this.createJobAttributes(jobId, resolvedJob);
 
-        final TitusBatchJobRequest request = TitusBatchJobRequest.builder()
+        final TitusBatchJobRequest.TitusBatchJobRequestBuilder requestBuilder = TitusBatchJobRequest.builder()
             .owner(
                 TitusBatchJobRequest.Owner
                     .builder()
@@ -339,14 +347,39 @@ private TitusBatchJobRequest createJobRequest(final ResolvedJob resolvedJob) thr
                     .detail(this.titusAgentLauncherProperties.getDetail())
                     .sequence(this.titusAgentLauncherProperties.getSequence())
                     .build()
-            )
-            .build();
+            );
+
+        final Optional<String> networkConfiguration = validateNetworkConfiguration(this.environment.getProperty(
+            TitusAgentLauncherProperties.CONTAINER_NETWORK_MODE,
+            String.class));
+        networkConfiguration.ifPresent(config -> requestBuilder.networkConfiguration(
+                TitusBatchJobRequest.NetworkConfiguration.builder().networkMode(config).build()));
+
+        final TitusBatchJobRequest request = requestBuilder.build();
 
         // Run the request through the security adapter to add any necessary context
         this.jobRequestAdapter.modifyJobRequest(request, resolvedJob);
         return request;
     }
 
+    private Optional<String> validateNetworkConfiguration(@Nullable final String networkConfig) {
+        if (Strings.isNullOrEmpty(networkConfig)) {
+            return Optional.empty();
+        }
+
+        switch (networkConfig) {
+            case TITUS_NETWORK_MODE_IPV4:
+            case TITUS_NETWORK_MODE_DUAL_STACK:
+            case TITUS_NETWORK_MODE_DUAL_STACK_FALLBACK:
+            case TITUS_NETWORK_MODE_IPV6:
+            case TITUS_NETWORK_MODE_HIGH_SCALE:
+                return Optional.of(networkConfig);
+            default:
+                return Optional.empty();
+        }
+
+    }
+
     /**
      * Helper method to avoid runtime errors if for some reason the DataSize converters aren't loaded.
      *

genie-web/src/main/java/com/netflix/genie/web/properties/TitusAgentLauncherProperties.java

@@ -169,6 +169,11 @@ public class TitusAgentLauncherProperties {
      */
     public static final String AGENT_IMAGE_KEY_PROPERTY = PREFIX + ".agentImageKey";
 
+    /**
+     * The property for titus container network mode.
+     */
+    public static final String CONTAINER_NETWORK_MODE = PREFIX + ".networkMode";
+
     /**
      * Whether the Titus Agent Launcher is enabled.
      */

genie-web/src/test/groovy/com/netflix/genie/web/agent/launchers/impl/TitusAgentLauncherImplSpec.groovy

@@ -605,6 +605,39 @@ class TitusAgentLauncherImplSpec extends Specification {
         requestCapture.getContainer().getAttributes().get(prop2Key) == prop2Value
     }
 
+    def "Check titus container network mode"() {
+        TitusBatchJobResponse response = toTitusResponse("{ \"id\" : \"" + TITUS_JOB_ID + "\" }")
+        TitusBatchJobRequest requestCapture
+
+        when:
+        this.environment.withProperty(TitusAgentLauncherProperties.CONTAINER_NETWORK_MODE, inputNetworkMode)
+        this.launcher.launchAgent(this.resolvedJob, null)
+
+        then:
+        1 * this.restTemplate.postForObject(TITUS_ENDPOINT, _ as TitusBatchJobRequest, TitusBatchJobResponse.class) >> {
+            args ->
+                requestCapture = args[1] as TitusBatchJobRequest
+                return response
+        }
+        1 * this.adapter.modifyJobRequest(_ as TitusBatchJobRequest, this.resolvedJob)
+        requestCapture != null
+        if (requestNetworkMode != null) {
+            requestCapture.getNetworkConfiguration() != null
+            requestCapture.getNetworkConfiguration().getNetworkMode() == requestNetworkMode
+        } else {
+            requestCapture.getNetworkConfiguration() == null
+        }
+
+        where:
+        inputNetworkMode      | requestNetworkMode
+        "Ipv4Only"            | "Ipv4Only"
+        "Ipv6AndIpv4"         | "Ipv6AndIpv4"
+        "Ipv6AndIpv4Fallback" | "Ipv6AndIpv4Fallback"
+        "Ipv6Only"            | "Ipv6Only"
+        "HighScale"           | "HighScale"
+        "NonSense"            | null
+    }
+
     def "Retry policy works as expected"() {
         def retryCodes = EnumSet.of(HttpStatus.REQUEST_TIMEOUT, HttpStatus.SERVICE_UNAVAILABLE)
         def maxRetries = 2