change flag to enable_dualstack, refactor ipv6 setup to separate meth…

…ods, address other code improvement suggestions

framework/infrastructure/gcp/compute.py

@@ -39,7 +39,6 @@ class ComputeV1(
     _WAIT_FOR_BACKEND_SLEEP_SEC = 4
     _WAIT_FOR_OPERATION_SEC = 60 * 10
     gfe_debug_header: Optional[str]
-    add_dualstack_support: bool = False
 
     @dataclasses.dataclass(frozen=True)
     class GcpResource:
@@ -62,12 +61,10 @@ def __init__(
         api_manager: gcp.api.GcpApiManager,
         project: str,
         gfe_debug_header: Optional[str] = None,
-        add_dualstack_support: bool = False,
         version: str = "v1",
     ):
         super().__init__(api_manager.compute(version), project)
         self.gfe_debug_header = gfe_debug_header
-        self.add_dualstack_support = add_dualstack_support
 
     class HealthCheckProtocol(enum.Enum):
         TCP = enum.auto()
@@ -154,6 +151,7 @@ def create_backend_service_traffic_director(
         subset_size: Optional[int] = None,
         locality_lb_policies: Optional[List[dict]] = None,
         outlier_detection: Optional[dict] = None,
+        enable_dualstack: bool = False,
     ) -> "GcpResource":
         if not isinstance(protocol, self.BackendServiceProtocol):
             raise TypeError(f"Unexpected Backend Service protocol: {protocol}")
@@ -165,7 +163,7 @@ def create_backend_service_traffic_director(
         }
         # If add dualstack support is specified True, config the backend service
         # to support IPv6
-        if self.add_dualstack_support:
+        if enable_dualstack:
             body["ipAddressSelectionPolicy"] = "PREFER_IPV6"
         # If affinity header is specified, config the backend service to support
         # affinity, and set affinity header to the one given.

framework/infrastructure/traffic_director.py

@@ -116,7 +116,6 @@ def __init__(
             project,
             version=compute_api_version,
             gfe_debug_header=xds_flags.GFE_DEBUG_HEADER.value,
-            add_dualstack_support=xds_flags.ADD_DUALSTACK_SUPPORT.value,
         )
 
         # Settings
@@ -132,14 +131,14 @@ def __init__(
         self.firewall_rule: Optional[GcpResource] = None
         self.firewall_rule_ipv6: Optional[GcpResource] = None
         self.target_proxy: Optional[GcpResource] = None
-        self.target_proxy_v6: Optional[GcpResource] = None
+        self.target_proxy_ipv6: Optional[GcpResource] = None
         # TODO(sergiitk): remove this flag once target proxy resource loaded
         self.target_proxy_is_http: bool = False
         self.alternative_target_proxy: Optional[GcpResource] = None
         self.forwarding_rule: Optional[GcpResource] = None
-        self.forwarding_rule_v6: Optional[GcpResource] = None
+        self.forwarding_rule_ipv6: Optional[GcpResource] = None
         self.alternative_forwarding_rule: Optional[GcpResource] = None
-        self.add_dualstack_support = xds_flags.ADD_DUALSTACK_SUPPORT.value
+        self.enable_dualstack = xds_flags.ENABLE_DUALSTACK.value
 
         # Backends.
         self.backends = set()
@@ -177,11 +176,9 @@ def setup_routing_rule_map_for_grpc(self, service_host, service_port):
         self.create_target_proxy()
         self.create_forwarding_rule(service_port)
 
-        if self.add_dualstack_support:
-            self.create_target_proxy(ipv6=True)
-            self.create_forwarding_rule(
-                service_port, ipv6=True, ip_address="::"
-            )
+        if self.enable_dualstack:
+            self.create_target_proxy_ipv6()
+            self.create_forwarding_rule_ipv6(service_port)
 
     def cleanup(self, *, force=False):
         # Cleanup in the reverse order of creation
@@ -190,9 +187,8 @@ def cleanup(self, *, force=False):
         self.delete_alternative_forwarding_rule(force=force)
         self.delete_target_http_proxy(force=force)
         self.delete_target_grpc_proxy(force=force)
-        if self.add_dualstack_support:
-            self.delete_target_http_proxy(force=force, ipv6=True)
-            self.delete_forwarding_rule(force=force, ipv6=True)
+        if self.enable_dualstack:
+            self.delete_target_proxy_ipv6(force=force)
         self.delete_alternative_target_grpc_proxy(force=force)
         self.delete_url_map(force=force)
         self.delete_alternative_url_map(force=force)
@@ -261,6 +257,7 @@ def create_backend_service(
             affinity_header=affinity_header,
             locality_lb_policies=locality_lb_policies,
             outlier_detection=outlier_detection,
+            enable_dualstack=self.enable_dualstack,
         )
         self.backend_service = resource
         self.backend_service_protocol = protocol
@@ -352,7 +349,10 @@ def create_alternative_backend_service(
             'Creating %s Alternative Backend Service "%s"', protocol.name, name
         )
         resource = self.compute.create_backend_service_traffic_director(
-            name, health_check=self.health_check, protocol=protocol
+            name,
+            health_check=self.health_check,
+            protocol=protocol,
+            enable_dualstack=self.enable_dualstack,
         )
         self.alternative_backend_service = resource
         self.alternative_backend_service_protocol = protocol
@@ -433,6 +433,7 @@ def create_affinity_backend_service(
             health_check=self.health_check,
             protocol=protocol,
             affinity_header=TEST_AFFINITY_METADATA_KEY,
+            enable_dualstack=self.enable_dualstack,
         )
         self.affinity_backend_service = resource
         self.affinity_backend_service_protocol = protocol
@@ -608,12 +609,8 @@ def delete_alternative_url_map(self, force=False):
         self.compute.delete_url_map(name)
         self.url_map = None
 
-    def create_target_proxy(self, ipv6=False):
-        if ipv6:
-            name = self.make_resource_name(self.TARGET_PROXY_NAME_IPV6)
-        else:
-            name = self.make_resource_name(self.TARGET_PROXY_NAME)
-
+    def create_target_proxy(self):
+        name = self.make_resource_name(self.TARGET_PROXY_NAME)
         if self.backend_service_protocol is BackendServiceProtocol.GRPC:
             target_proxy_type = "GRPC"
             create_proxy_fn = self.compute.create_target_grpc_proxy
@@ -631,58 +628,58 @@ def create_target_proxy(self, ipv6=False):
             target_proxy_type,
             self.url_map.name,
         )
+        self.target_proxy = create_proxy_fn(name, self.url_map)
 
-        if ipv6:
-            self.target_proxy_v6 = create_proxy_fn(name, self.url_map)
-        else:
-            self.target_proxy = create_proxy_fn(name, self.url_map)
+    def create_target_proxy_ipv6(self):
+        name = self.make_resource_name(self.TARGET_PROXY_NAME_IPV6)
+        target_proxy_type = "HTTP"  # TODO: Support GRPC target proxy as well
+        create_proxy_fn = self.compute.create_target_http_proxy
 
-    def delete_target_grpc_proxy(self, force=False, ipv6=False):
+        logger.info(
+            'Creating target %s proxy "%s" to URL map %s',
+            name,
+            target_proxy_type,
+            self.url_map.name,
+        )
+        self.target_proxy_ipv6 = create_proxy_fn(name, self.url_map)
+
+    def delete_target_grpc_proxy(self, force=False):
         if force:
-            if ipv6:
-                name = self.make_resource_name(self.TARGET_PROXY_NAME_IPV6)
-            else:
-                name = self.make_resource_name(self.TARGET_PROXY_NAME)
-        elif ipv6:
-            if self.target_proxy_v6:
-                name = self.target_proxy_v6.name
+            name = self.make_resource_name(self.TARGET_PROXY_NAME)
         elif self.target_proxy:
             name = self.target_proxy.name
         else:
             return
-
         logger.info('Deleting Target GRPC proxy "%s"', name)
         self.compute.delete_target_grpc_proxy(name)
-        if ipv6:
-            self.target_proxy_v6 = None
-        else:
-            self.target_proxy = None
-
+        self.target_proxy = None
         self.target_proxy_is_http = False
 
-    def delete_target_http_proxy(self, force=False, ipv6=False):
+    def delete_target_http_proxy(self, force=False):
         if force:
-            if ipv6:
-                name = self.make_resource_name(self.TARGET_PROXY_NAME_IPV6)
-            else:
-                name = self.make_resource_name(self.TARGET_PROXY_NAME)
-        elif ipv6:
-            if self.target_proxy_v6 and self.target_proxy_is_http:
-                name = self.target_proxy_v6.name
+            name = self.make_resource_name(self.TARGET_PROXY_NAME)
         elif self.target_proxy and self.target_proxy_is_http:
             name = self.target_proxy.name
         else:
             return
-
         logger.info('Deleting HTTP Target proxy "%s"', name)
         self.compute.delete_target_http_proxy(name)
-        if ipv6:
-            self.target_proxy_v6 = None
-        else:
-            self.target_proxy = None
-
+        self.target_proxy = None
         self.target_proxy_is_http = False
 
+    def delete_target_proxy_ipv6(self, force=False):
+        if force:
+            name = self.make_resource_name(self.TARGET_PROXY_NAME_IPV6)
+        elif self.target_proxy_ipv6:
+            name = self.target_proxy_ipv6.name
+        else:
+            return
+        logger.info('Deleting Target HTTP proxy "%s"', name)
+        self.compute.delete_target_http_proxy(
+            name
+        )  # TODO: Support GRPC target proxy as well
+        self.target_proxy_ipv6 = None
+
     def create_alternative_target_proxy(self):
         name = self.make_resource_name(self.ALTERNATIVE_TARGET_PROXY_NAME)
         if self.backend_service_protocol is BackendServiceProtocol.GRPC:
@@ -724,60 +721,59 @@ def find_unused_forwarding_rule_port(
         # TODO(sergiitk): custom exception
         raise RuntimeError("Couldn't find unused forwarding rule port")
 
-    def create_forwarding_rule(
-        self, src_port: int, ipv6=False, ip_address="0.0.0.0"
-    ):
-        if ipv6:
-            name = self.make_resource_name(self.FORWARDING_RULE_NAME_IPV6)
-            target_proxy = self.target_proxy_v6
-        else:
-            name = self.make_resource_name(self.FORWARDING_RULE_NAME)
-            target_proxy = self.target_proxy
+    def create_forwarding_rule(self, src_port: int):
+        name = self.make_resource_name(self.FORWARDING_RULE_NAME)
+        src_port = int(src_port)
+        logging.info(
+            'Creating forwarding rule "%s" in network "%s": 0.0.0.0:%s -> %s',
+            name,
+            self.network,
+            src_port,
+            self.target_proxy.url,
+        )
+        resource = self.compute.create_forwarding_rule(
+            name, src_port, self.target_proxy, self.network_url
+        )
+        self.forwarding_rule = resource
 
+        return resource
+
+    def create_forwarding_rule_ipv6(self, src_port: int):
+        name = self.make_resource_name(self.FORWARDING_RULE_NAME_IPV6)
         src_port = int(src_port)
         logging.info(
-            'Creating forwarding rule "%s" in network "%s": %s:%s -> %s',
+            'Creating forwarding rule "%s" in network "%s": [::]:%s -> %s',
             name,
             self.network,
             src_port,
-            ip_address,
-            target_proxy.url,
+            self.target_proxy_ipv6.url,
         )
         resource = self.compute.create_forwarding_rule(
             name,
             src_port,
-            target_proxy,
+            self.target_proxy_ipv6,
             self.network_url,
-            ip_address=ip_address,
+            ip_address="::",
         )
-
-        if ipv6:
-            self.forwarding_rule_v6 = resource
-        else:
-            self.forwarding_rule = resource
+        self.forwarding_rule_v6 = resource
 
         return resource
 
-    def delete_forwarding_rule(self, force=False, ipv6=False):
+    def delete_forwarding_rule(self, force=False):
         if force:
-            if ipv6:
-                name = self.make_resource_name(self.FORWARDING_RULE_NAME_IPV6)
-            else:
-                name = self.make_resource_name(self.FORWARDING_RULE_NAME)
-        elif ipv6:
-            if self.forwarding_rule_v6:
-                name = self.forwarding_rule_v6.name
+            name = self.make_resource_name(self.FORWARDING_RULE_NAME)
+            name_ipv6 = self.make_resource_name(self.FORWARDING_RULE_NAME_IPV6)
         elif self.forwarding_rule:
             name = self.forwarding_rule.name
+            name_ipv6 = self.forwarding_rule_v6.name
         else:
             return
-
         logger.info('Deleting Forwarding rule "%s"', name)
         self.compute.delete_forwarding_rule(name)
-        if ipv6:
-            self.forwarding_rule_v6 = None
-        else:
-            self.forwarding_rule = None
+        self.forwarding_rule = None
+        logger.info('Deleting Forwarding rule "%s"', name_ipv6)
+        self.compute.delete_forwarding_rule(name_ipv6)
+        self.forwarding_rule_ipv6 = None
 
     def create_alternative_forwarding_rule(
         self, src_port: int, ip_address="0.0.0.0"

framework/test_app/runners/k8s/k8s_base_runner.py

@@ -1000,10 +1000,14 @@ def _wait_pod_started(self, name, **kwargs) -> k8s.V1Pod:
         pod = self.k8s_namespace.get_pod(name)
 
         if hasattr(pod.status, "pod_ip_s"):  # if running with dualstack support
-            pod_ip_s = pod.status.pod_ip_s
+            logger.info(
+                "Pod %s ready, IPs: %s", pod.metadata.name, pod.status.pod_ip_s
+            )
         else:
-            pod_ip_s = pod.status.pod_ip
-        logger.info("Pod %s ready, IP: %s", pod.metadata.name, pod_ip_s)
+            logger.info(
+                "Pod %s ready, IP: %s", pod.metadata.name, pod.status.pod_ip
+            )
+
         return pod
 
     def _pod_started_logic(self, pod: k8s.V1Pod) -> bool:

framework/xds_flags.py

@@ -190,10 +190,10 @@
     help="Whether to enable GFE debug headers and what value to use.",
 )
 
-ADD_DUALSTACK_SUPPORT = flags.DEFINE_bool(
-    "add_dualstack_support",
+ENABLE_DUALSTACK = flags.DEFINE_bool(
+    "ENABLE_DUALSTACK",
     default=False,
-    help=(f"Add support for Dual Stack resources to the framework."),
+    help=("Enable support for Dual Stack resources to the framework."),
 )
 
 

kubernetes-manifests/client.deployment.yaml

@@ -66,7 +66,7 @@ spec:
               value: "true"
             - name: GRPC_EXPERIMENTAL_XDS_ENABLE_OVERRIDE_HOST
               value: "true"
-            % if add_dualstack_support:
+            % if enable_dualstack:
             - name: GRPC_EXPERIMENTAL_ENABLE_NEW_PICK_FIRST
               value: "true"
             - name: GRPC_EXPERIMENTAL_XDS_DUALSTACK_ENDPOINTS

kubernetes-manifests/server.service.yaml

@@ -10,7 +10,7 @@ metadata:
     cloud.google.com/neg: '{"exposed_ports": {"${test_port}":{"name":"${neg_name}"}}}'
 spec:
   type: ClusterIP
-  % if add_dualstack_support:
+  % if enable_dualstack:
   ipFamilyPolicy: PreferDualStack
   ipFamilies:
     - IPv6