Add AzureSnitch

As we were implementing the snitch itself, we noticed that the constructors
for cloud-based snitches are unnecessarilly complicated and
we took the opportunity to make them simpler.

patch by Stefan Miklosovic; reviewed by German Eichberger and Jacek Lewandowski for CASSANDRA-18646

CHANGES.txt

@@ -1,4 +1,5 @@
 5.0
+ * Add AzureSnitch (CASSANDRA-18646)
  * Implementation of the Unified Compaction Strategy as described in CEP-26 (CASSANDRA-18397)
  * Upgrade commons cli to 1.5.0 (CASSANDRA-18659)
  * Disable the deprecated keyspace/table thresholds and convert them to guardrails (CASSANDRA-18617)

NEWS.txt

@@ -172,6 +172,7 @@ New features
     - `cassandra-stress` has a new option called '-jmx' which enables a user to pass username and password to JMX (CASSANDRA-18544)
     - It is possible to read all credentials for `cassandra-stress` from a file via option `-credentials-file` (CASSANDRA-18544)
     - nodetool info displays bootstrap state a node is in as well as if it was decommissioned or if it failed to decommission (CASSANDRA-18555)
+    - Added snitch for Microsoft Azure of name AzureSnitch (CASSANDRA-18646)
 
 Upgrading
 ---------

conf/cassandra-rackdc.properties

@@ -43,6 +43,11 @@ rack=rack1
 # it is refreshed. Defaults to 21600. Can not be smaller than 30 and bigger than 21600. Has to be an integer.
 # ec2_metadata_token_ttl_seconds=21600
 
+# AzureSnitch
+# Options are:
+# Version of API to talk to. When not set, defaults to '2021-12-13'.
+# azure_api_version=2021-12-13
+
 # For all cloud-based snitches, there are following options available:
 #
 # Property to change metadata service url for a cloud-based snitch. The endpoint of a particular

conf/cassandra.yaml

@@ -1316,6 +1316,11 @@ slow_query_log_timeout: 500ms
 #    Snitch for getting dc and rack of a node from metadata service of Alibaba cloud.
 #    This snitch that assumes an ECS region is a DC and an ECS availability_zone is a rack.
 #
+# AzureSnitch:
+#    Gets datacenter from 'location' and rack from 'zone' fields of 'compute' object
+#    from instance metadata service. If the availability zone is not enabled, it will use the fault
+#    domain and get its respective value.
+#
 # CloudstackSnitch:
 #    A snitch that assumes a Cloudstack Zone follows the typical convention 
 #    country-location-az and uses a country/location tuple as a datacenter

src/java/org/apache/cassandra/locator/AbstractCloudMetadataServiceConnector.java

@@ -45,8 +45,11 @@ abstract class AbstractCloudMetadataServiceConnector
     protected final String metadataServiceUrl;
     protected final int requestTimeoutMs;
 
-    public AbstractCloudMetadataServiceConnector(SnitchProperties properties)
+    private final SnitchProperties properties;
+
+    public AbstractCloudMetadataServiceConnector(SnitchProperties snitchProperties)
     {
+        this.properties = snitchProperties;
         String parsedMetadataServiceUrl = properties.get(METADATA_URL_PROPERTY, null);
 
         try
@@ -65,7 +68,6 @@ abstract class AbstractCloudMetadataServiceConnector
                                              ex);
         }
 
-
         String metadataRequestTimeout = properties.get(METADATA_REQUEST_TIMEOUT_PROPERTY, DEFAULT_METADATA_REQUEST_TIMEOUT);
 
         try
@@ -80,6 +82,11 @@ abstract class AbstractCloudMetadataServiceConnector
         }
     }
 
+    public SnitchProperties getProperties()
+    {
+        return properties;
+    }
+
     public final String apiCall(String query) throws IOException
     {
         return apiCall(metadataServiceUrl, query, "GET", ImmutableMap.of(), 200);

src/java/org/apache/cassandra/locator/AbstractCloudMetadataServiceSnitch.java

@@ -40,24 +40,20 @@ abstract class AbstractCloudMetadataServiceSnitch extends AbstractNetworkTopolog
     static final String DEFAULT_RACK = "UNKNOWN-RACK";
 
     protected final AbstractCloudMetadataServiceConnector connector;
-    protected final SnitchProperties snitchProperties;
 
     private final String localRack;
     private final String localDc;
 
     private Map<InetAddressAndPort, Map<String, String>> savedEndpoints;
 
-    public AbstractCloudMetadataServiceSnitch(AbstractCloudMetadataServiceConnector connector,
-                                              SnitchProperties snitchProperties,
-                                              Pair<String, String> dcAndRack)
+    public AbstractCloudMetadataServiceSnitch(AbstractCloudMetadataServiceConnector connector, Pair<String, String> dcAndRack)
     {
         this.connector = connector;
-        this.snitchProperties = snitchProperties;
         this.localDc = dcAndRack.left;
         this.localRack = dcAndRack.right;
 
         logger.info(format("%s using datacenter: %s, rack: %s, connector: %s, properties: %s",
-                           getClass().getName(), getLocalDatacenter(), getLocalRack(), connector, snitchProperties));
+                           getClass().getName(), getLocalDatacenter(), getLocalRack(), connector, connector.getProperties()));
     }
 
     @Override

src/java/org/apache/cassandra/locator/AlibabaCloudSnitch.java

@@ -19,8 +19,6 @@ package org.apache.cassandra.locator;
 
 import java.io.IOException;
 
-import com.google.common.collect.ImmutableMap;
-
 import org.apache.cassandra.locator.AbstractCloudMetadataServiceConnector.DefaultCloudMetadataServiceConnector;
 
 import static org.apache.cassandra.locator.AbstractCloudMetadataServiceConnector.METADATA_URL_PROPERTY;
@@ -44,14 +42,13 @@ public class AlibabaCloudSnitch extends AbstractCloudMetadataServiceSnitch
 
     public AlibabaCloudSnitch(SnitchProperties properties) throws IOException
     {
-        this(properties, new DefaultCloudMetadataServiceConnector(properties.putIfAbsent(METADATA_URL_PROPERTY,
-                                                                                         DEFAULT_METADATA_SERVICE_URL)));
+        this(new DefaultCloudMetadataServiceConnector(properties.putIfAbsent(METADATA_URL_PROPERTY,
+                                                                             DEFAULT_METADATA_SERVICE_URL)));
     }
 
-    public AlibabaCloudSnitch(SnitchProperties properties, AbstractCloudMetadataServiceConnector connector) throws IOException
+    public AlibabaCloudSnitch(AbstractCloudMetadataServiceConnector connector) throws IOException
     {
-        super(connector, properties, SnitchUtils.parseDcAndRack(connector.apiCall(ZONE_NAME_QUERY_URL,
-                                                                                  ImmutableMap.of()),
-                                                                properties.getDcSuffix()));
+        super(connector, SnitchUtils.parseDcAndRack(connector.apiCall(ZONE_NAME_QUERY_URL),
+                                                    connector.getProperties().getDcSuffix()));
     }
 }

src/java/org/apache/cassandra/locator/AzureSnitch.java

@@ -0,0 +1,102 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.cassandra.locator;
+
+import java.io.IOException;
+
+import com.google.common.collect.ImmutableMap;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import org.apache.cassandra.locator.AbstractCloudMetadataServiceConnector.DefaultCloudMetadataServiceConnector;
+import org.apache.cassandra.utils.JsonUtils;
+import org.apache.cassandra.utils.Pair;
+
+import static java.lang.String.format;
+import static org.apache.cassandra.locator.AbstractCloudMetadataServiceConnector.METADATA_URL_PROPERTY;
+
+/**
+ * AzureSnitch will resolve datacenter and rack by calling {@code /metadata/instance/compute} endpoint returning
+ * the response in JSON format for API version {@code 2021-12-13}. The version of API is configurable via property
+ * {@code azure_api_version} in cassandra-rackdc.properties.
+ * <p>
+ * A datacenter is resolved from {@code location} field and a rack is resolved by looking into {@code zone} field first.
+ * When zone is not set, or it is empty string, it will look into {@code platformFaultDomain} field. Such resolved
+ * value is prepended by {@code rack-} string.
+ */
+public class AzureSnitch extends AbstractCloudMetadataServiceSnitch
+{
+    static final String DEFAULT_METADATA_SERVICE_URL = "http://169.254.169.254";
+    static final String METADATA_QUERY_TEMPLATE = "/metadata/instance/compute?api-version=%s&format=json";
+    static final String METADATA_HEADER = "Metadata";
+    static final String API_VERSION_PROPERTY_KEY = "azure_api_version";
+    static final String DEFAULT_API_VERSION = "2021-12-13";
+
+    public AzureSnitch() throws IOException
+    {
+        this(new SnitchProperties());
+    }
+
+    public AzureSnitch(SnitchProperties properties) throws IOException
+    {
+        this(new DefaultCloudMetadataServiceConnector(properties.putIfAbsent(METADATA_URL_PROPERTY,
+                                                                             DEFAULT_METADATA_SERVICE_URL)));
+    }
+
+    public AzureSnitch(AbstractCloudMetadataServiceConnector connector) throws IOException
+    {
+        super(connector, resolveDcAndRack(connector));
+    }
+
+    private static Pair<String, String> resolveDcAndRack(AbstractCloudMetadataServiceConnector connector) throws IOException
+    {
+        String apiVersion = connector.getProperties().get(API_VERSION_PROPERTY_KEY, DEFAULT_API_VERSION);
+        String response = connector.apiCall(format(METADATA_QUERY_TEMPLATE, apiVersion), ImmutableMap.of(METADATA_HEADER, "true"));
+        JsonNode jsonNode = JsonUtils.JSON_OBJECT_MAPPER.readTree(response);
+
+        JsonNode location = jsonNode.get("location");
+        JsonNode zone = jsonNode.get("zone");
+        JsonNode platformFaultDomain = jsonNode.get("platformFaultDomain");
+
+        String datacenter;
+        String rack;
+
+        if (location == null || location.isNull() || location.asText().isEmpty())
+            datacenter = DEFAULT_DC;
+        else
+            datacenter = location.asText();
+
+        if (zone == null || zone.isNull() || zone.asText().isEmpty())
+        {
+            if (platformFaultDomain == null || platformFaultDomain.isNull() || platformFaultDomain.asText().isEmpty())
+            {
+                rack = DEFAULT_RACK;
+            }
+            else
+            {
+                rack = platformFaultDomain.asText();
+            }
+        }
+        else
+        {
+            rack = zone.asText();
+        }
+
+        return Pair.create(datacenter + connector.getProperties().getDcSuffix(), "rack-" + rack);
+    }
+}

src/java/org/apache/cassandra/locator/CloudstackSnitch.java

@@ -27,11 +27,11 @@ import java.util.regex.Pattern;
 import org.apache.cassandra.exceptions.ConfigurationException;
 import org.apache.cassandra.io.util.File;
 import org.apache.cassandra.io.util.FileReader;
+import org.apache.cassandra.locator.AbstractCloudMetadataServiceConnector.DefaultCloudMetadataServiceConnector;
 import org.apache.cassandra.utils.JVMStabilityInspector;
 import org.apache.cassandra.utils.Pair;
 
 import static org.apache.cassandra.locator.AbstractCloudMetadataServiceConnector.METADATA_URL_PROPERTY;
-import static org.apache.cassandra.locator.CloudstackSnitch.CloudstackConnector.ZONE_NAME_QUERY_URI;
 
 /**
  * A snitch that assumes a Cloudstack Zone follows the typical convention
@@ -43,6 +43,8 @@ import static org.apache.cassandra.locator.CloudstackSnitch.CloudstackConnector.
 @Deprecated
 public class CloudstackSnitch extends AbstractCloudMetadataServiceSnitch
 {
+    static final String ZONE_NAME_QUERY_URI = "/latest/meta-data/availability-zone";
+
     private static final String[] LEASE_FILES =
     {
     "file:///var/lib/dhcp/dhclient.eth0.leases",
@@ -56,26 +58,16 @@ public class CloudstackSnitch extends AbstractCloudMetadataServiceSnitch
 
     public CloudstackSnitch(SnitchProperties snitchProperties) throws IOException
     {
-        this(snitchProperties, new CloudstackConnector(snitchProperties.putIfAbsent(METADATA_URL_PROPERTY, csMetadataEndpoint())));
+        this(new DefaultCloudMetadataServiceConnector(snitchProperties.putIfAbsent(METADATA_URL_PROPERTY, csMetadataEndpoint())));
     }
 
-    public CloudstackSnitch(SnitchProperties properties, AbstractCloudMetadataServiceConnector connector) throws IOException
+    public CloudstackSnitch(AbstractCloudMetadataServiceConnector connector) throws IOException
     {
-        super(connector, properties, resolveDcAndRack(connector));
+        super(connector, resolveDcAndRack(connector));
         logger.warn("{} is deprecated and not actively maintained. It will be removed in the next " +
                     "major version of Cassandra.", CloudstackSnitch.class.getName());
     }
 
-    static class CloudstackConnector extends AbstractCloudMetadataServiceConnector
-    {
-        static final String ZONE_NAME_QUERY_URI = "/latest/meta-data/availability-zone";
-
-        protected CloudstackConnector(SnitchProperties properties)
-        {
-            super(properties);
-        }
-    }
-
     private static Pair<String, String> resolveDcAndRack(AbstractCloudMetadataServiceConnector connector) throws IOException
     {
         String zone = connector.apiCall(ZONE_NAME_QUERY_URI);

src/java/org/apache/cassandra/locator/Ec2MultiRegionSnitch.java

@@ -55,12 +55,12 @@ public class Ec2MultiRegionSnitch extends Ec2Snitch
 
     public Ec2MultiRegionSnitch(SnitchProperties props) throws IOException, ConfigurationException
     {
-        this(props, Ec2MetadataServiceConnector.create(props));
+        this(Ec2MetadataServiceConnector.create(props));
     }
 
-    Ec2MultiRegionSnitch(SnitchProperties props, AbstractCloudMetadataServiceConnector connector) throws IOException
+    Ec2MultiRegionSnitch(AbstractCloudMetadataServiceConnector connector) throws IOException
     {
-        super(props, connector);
+        super(connector);
         InetAddress localPublicAddress = InetAddress.getByName(connector.apiCall(PUBLIC_IP_QUERY));
         logger.info("EC2Snitch using publicIP as identifier: {}", localPublicAddress);
         localPrivateAddress = connector.apiCall(PRIVATE_IP_QUERY);

src/java/org/apache/cassandra/locator/Ec2Snitch.java

@@ -63,22 +63,22 @@ public class Ec2Snitch extends AbstractCloudMetadataServiceSnitch
 
     public Ec2Snitch(SnitchProperties props) throws IOException, ConfigurationException
     {
-        this(props, Ec2MetadataServiceConnector.create(props));
+        this(Ec2MetadataServiceConnector.create(props));
     }
 
-    Ec2Snitch(SnitchProperties props, AbstractCloudMetadataServiceConnector connector) throws IOException
+    Ec2Snitch(AbstractCloudMetadataServiceConnector connector) throws IOException
     {
-        super(connector, props, getDcAndRack(props, connector));
-        usingLegacyNaming = isUsingLegacyNaming(props);
+        super(connector, getDcAndRack(connector));
+        usingLegacyNaming = isUsingLegacyNaming(connector.getProperties());
     }
 
-    private static Pair<String, String> getDcAndRack(SnitchProperties props, AbstractCloudMetadataServiceConnector connector) throws IOException
+    private static Pair<String, String> getDcAndRack(AbstractCloudMetadataServiceConnector connector) throws IOException
     {
         String az = connector.apiCall(ZONE_NAME_QUERY);
 
         // if using the full naming scheme, region name is created by removing letters from the
         // end of the availability zone and zone is the full zone name
-        boolean usingLegacyNaming = isUsingLegacyNaming(props);
+        boolean usingLegacyNaming = isUsingLegacyNaming(connector.getProperties());
         String region;
         String localDc;
         String localRack;
@@ -101,7 +101,7 @@ public class Ec2Snitch extends AbstractCloudMetadataServiceSnitch
             localRack = az;
         }
 
-        localDc = region.concat(props.getDcSuffix());
+        localDc = region.concat(connector.getProperties().getDcSuffix());
 
         return Pair.create(localDc, localRack);
     }

src/java/org/apache/cassandra/locator/GoogleCloudSnitch.java

@@ -41,14 +41,14 @@ public class GoogleCloudSnitch extends AbstractCloudMetadataServiceSnitch
 
     public GoogleCloudSnitch(SnitchProperties properties) throws IOException
     {
-        this(properties, new DefaultCloudMetadataServiceConnector(properties.putIfAbsent(METADATA_URL_PROPERTY,
-                                                                                         DEFAULT_METADATA_SERVICE_URL)));
+        this(new DefaultCloudMetadataServiceConnector(properties.putIfAbsent(METADATA_URL_PROPERTY,
+                                                                             DEFAULT_METADATA_SERVICE_URL)));
     }
 
-    public GoogleCloudSnitch(SnitchProperties properties, AbstractCloudMetadataServiceConnector connector) throws IOException
+    public GoogleCloudSnitch(AbstractCloudMetadataServiceConnector connector) throws IOException
     {
-        super(connector, properties, SnitchUtils.parseDcAndRack(connector.apiCall(ZONE_NAME_QUERY_URL,
-                                                                                  ImmutableMap.of("Metadata-Flavor", "Google")),
-                                                                properties.getDcSuffix()));
+        super(connector, SnitchUtils.parseDcAndRack(connector.apiCall(ZONE_NAME_QUERY_URL,
+                                                                      ImmutableMap.of("Metadata-Flavor", "Google")),
+                                                    connector.getProperties().getDcSuffix()));
     }
 }

test/unit/org/apache/cassandra/locator/AlibabaCloudSnitchTest.java

@@ -20,7 +20,6 @@ package org.apache.cassandra.locator;
 import java.io.IOException;
 import java.util.EnumMap;
 import java.util.Map;
-import java.util.Properties;
 
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
@@ -33,16 +32,19 @@ import org.apache.cassandra.exceptions.ConfigurationException;
 import org.apache.cassandra.gms.ApplicationState;
 import org.apache.cassandra.gms.Gossiper;
 import org.apache.cassandra.gms.VersionedValue;
+import org.apache.cassandra.locator.AbstractCloudMetadataServiceConnector.DefaultCloudMetadataServiceConnector;
 import org.apache.cassandra.service.StorageService;
+import org.apache.cassandra.utils.Pair;
 
 import static org.apache.cassandra.ServerTestUtils.cleanup;
 import static org.apache.cassandra.ServerTestUtils.mkdirs;
 import static org.apache.cassandra.config.CassandraRelevantProperties.GOSSIP_DISABLE_THREAD_VALIDATION;
+import static org.apache.cassandra.locator.AbstractCloudMetadataServiceConnector.METADATA_URL_PROPERTY;
+import static org.apache.cassandra.locator.AlibabaCloudSnitch.DEFAULT_METADATA_SERVICE_URL;
 import static org.junit.Assert.assertEquals;
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anyMap;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.spy;
 
 public class AlibabaCloudSnitchTest
 {
@@ -66,10 +68,12 @@ public class AlibabaCloudSnitchTest
     {
         az = "cn-hangzhou-f";
 
-        AbstractCloudMetadataServiceConnector mock = mock(AbstractCloudMetadataServiceConnector.class);
-        when(mock.apiCall(any(), anyMap())).thenReturn(az);
+        DefaultCloudMetadataServiceConnector spiedConnector = spy(new DefaultCloudMetadataServiceConnector(
+        new SnitchProperties(Pair.create(METADATA_URL_PROPERTY, DEFAULT_METADATA_SERVICE_URL))));
 
-        AlibabaCloudSnitch snitch = new AlibabaCloudSnitch(new SnitchProperties(new Properties()), mock);
+        doReturn(az).when(spiedConnector).apiCall(any());
+
+        AlibabaCloudSnitch snitch = new AlibabaCloudSnitch(spiedConnector);
         InetAddressAndPort local = InetAddressAndPort.getByName("127.0.0.1");
         InetAddressAndPort nonlocal = InetAddressAndPort.getByName("127.0.0.7");
 
@@ -90,11 +94,12 @@ public class AlibabaCloudSnitchTest
     public void testNewRegions() throws IOException, ConfigurationException
     {
         az = "us-east-1a";
+        DefaultCloudMetadataServiceConnector spiedConnector = spy(new DefaultCloudMetadataServiceConnector(
+        new SnitchProperties(Pair.create(METADATA_URL_PROPERTY, DEFAULT_METADATA_SERVICE_URL))));
 
-        AbstractCloudMetadataServiceConnector mock = mock(AbstractCloudMetadataServiceConnector.class);
-        when(mock.apiCall(any(), anyMap())).thenReturn(az);
+        doReturn(az).when(spiedConnector).apiCall(any());
 
-        AlibabaCloudSnitch snitch = new AlibabaCloudSnitch(new SnitchProperties(new Properties()), mock);
+        AlibabaCloudSnitch snitch = new AlibabaCloudSnitch(spiedConnector);
         InetAddressAndPort local = InetAddressAndPort.getByName("127.0.0.1");
         assertEquals("us-east", snitch.getDatacenter(local));
         assertEquals("1a", snitch.getRack(local));

test/unit/org/apache/cassandra/locator/AzureConnectorMockingTest.java

@@ -0,0 +1,91 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.cassandra.locator;
+
+import java.util.Properties;
+
+import org.junit.Rule;
+import org.junit.Test;
+
+import com.github.tomakehurst.wiremock.junit.WireMockRule;
+import org.apache.cassandra.locator.AbstractCloudMetadataServiceConnector.DefaultCloudMetadataServiceConnector;
+
+import static com.github.tomakehurst.wiremock.client.WireMock.aResponse;
+import static com.github.tomakehurst.wiremock.client.WireMock.equalTo;
+import static com.github.tomakehurst.wiremock.client.WireMock.get;
+import static com.github.tomakehurst.wiremock.client.WireMock.urlEqualTo;
+import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.wireMockConfig;
+import static java.lang.String.format;
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.cassandra.locator.AbstractCloudMetadataServiceConnector.METADATA_URL_PROPERTY;
+import static org.apache.cassandra.locator.AzureSnitch.API_VERSION_PROPERTY_KEY;
+import static org.apache.cassandra.locator.AzureSnitch.METADATA_HEADER;
+import static org.apache.cassandra.locator.AzureSnitch.METADATA_QUERY_TEMPLATE;
+import static org.junit.Assert.assertEquals;
+
+public class AzureConnectorMockingTest
+{
+    @Rule
+    public final WireMockRule service = new WireMockRule(wireMockConfig().bindAddress("127.0.0.1").port(8080));
+
+    @Test
+    public void testConnector() throws Throwable
+    {
+        String RESPONSE = "{\"location\": \"PolandCentral\", \"zone\": \"1\", \"platformFaultDomain\": \"5\"}";
+
+        service.stubFor(get(urlEqualTo(format(METADATA_QUERY_TEMPLATE, "2021-12-13")))
+                        .withHeader(METADATA_HEADER, equalTo("true"))
+                        .willReturn(aResponse().withBody(RESPONSE)
+                                               .withStatus(200)
+                                               .withHeader("Content-Type", "application/json; charset=utf-8")
+                                               .withHeader("Content-Length", String.valueOf(RESPONSE.getBytes(UTF_8).length))));
+
+        Properties p = new Properties();
+        p.setProperty(METADATA_URL_PROPERTY, "http://127.0.0.1:8080");
+
+        SnitchProperties snitchProperties = new SnitchProperties(p);
+        AzureSnitch azureSnitch = new AzureSnitch(new DefaultCloudMetadataServiceConnector(snitchProperties));
+
+        assertEquals("rack-1", azureSnitch.getLocalRack());
+        assertEquals("PolandCentral", azureSnitch.getLocalDatacenter());
+    }
+
+    @Test
+    public void testMissingZone() throws Throwable
+    {
+        String RESPONSE = "{\"location\": \"PolandCentral\", \"zone\": \"\", \"platformFaultDomain\": \"5\"}";
+
+        service.stubFor(get(urlEqualTo(format(METADATA_QUERY_TEMPLATE, "2021-12-14")))
+                        .withHeader(METADATA_HEADER, equalTo("true"))
+                        .willReturn(aResponse().withBody(RESPONSE)
+                                               .withStatus(200)
+                                               .withHeader("Content-Type", "application/json; charset=utf-8")
+                                               .withHeader("Content-Length", String.valueOf(RESPONSE.getBytes(UTF_8).length))));
+
+        Properties p = new Properties();
+        p.setProperty(METADATA_URL_PROPERTY, "http://127.0.0.1:8080");
+        p.setProperty(API_VERSION_PROPERTY_KEY, "2021-12-14");
+
+        SnitchProperties snitchProperties = new SnitchProperties(p);
+        AzureSnitch azureSnitch = new AzureSnitch(new DefaultCloudMetadataServiceConnector(snitchProperties));
+
+        assertEquals("rack-5", azureSnitch.getLocalRack());
+        assertEquals("PolandCentral", azureSnitch.getLocalDatacenter());
+    }
+}

test/unit/org/apache/cassandra/locator/CloudstackSnitchTest.java

@@ -21,7 +21,6 @@ package org.apache.cassandra.locator;
 import java.io.IOException;
 import java.util.EnumMap;
 import java.util.Map;
-import java.util.Properties;
 
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
@@ -34,15 +33,18 @@ import org.apache.cassandra.exceptions.ConfigurationException;
 import org.apache.cassandra.gms.ApplicationState;
 import org.apache.cassandra.gms.Gossiper;
 import org.apache.cassandra.gms.VersionedValue;
+import org.apache.cassandra.locator.AbstractCloudMetadataServiceConnector.DefaultCloudMetadataServiceConnector;
 import org.apache.cassandra.service.StorageService;
+import org.apache.cassandra.utils.Pair;
 
 import static org.apache.cassandra.ServerTestUtils.cleanup;
 import static org.apache.cassandra.ServerTestUtils.mkdirs;
 import static org.apache.cassandra.config.CassandraRelevantProperties.GOSSIP_DISABLE_THREAD_VALIDATION;
+import static org.apache.cassandra.locator.AbstractCloudMetadataServiceConnector.METADATA_URL_PROPERTY;
 import static org.junit.Assert.assertEquals;
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.spy;
 
 public class CloudstackSnitchTest
 {
@@ -65,10 +67,13 @@ public class CloudstackSnitchTest
     public void testRacks() throws IOException, ConfigurationException
     {
         az = "ch-gva-1";
-        CloudstackSnitch.CloudstackConnector mock = mock(CloudstackSnitch.CloudstackConnector.class);
-        when(mock.apiCall(any())).thenReturn(az);
 
-        CloudstackSnitch snitch = new CloudstackSnitch(new SnitchProperties(new Properties()), mock);
+        DefaultCloudMetadataServiceConnector spiedConnector = spy(new DefaultCloudMetadataServiceConnector(
+        new SnitchProperties(Pair.create(METADATA_URL_PROPERTY, "http://127.0.0.1"))));
+
+        doReturn(az).when(spiedConnector).apiCall(any());
+
+        CloudstackSnitch snitch = new CloudstackSnitch(spiedConnector);
         InetAddressAndPort local = InetAddressAndPort.getByName("127.0.0.1");
         InetAddressAndPort nonlocal = InetAddressAndPort.getByName("127.0.0.7");
 
@@ -88,15 +93,19 @@ public class CloudstackSnitchTest
     @Test
     public void testNewRegions() throws IOException, ConfigurationException
     {
-        az = "ch-gva-1";
-        CloudstackSnitch.CloudstackConnector mock = mock(CloudstackSnitch.CloudstackConnector.class);
-        when(mock.apiCall(any())).thenReturn(az);
-        CloudstackSnitch snitch = new CloudstackSnitch(new SnitchProperties(new Properties()), mock);
+        az = "us-east-1a";
+
+        DefaultCloudMetadataServiceConnector spiedConnector = spy(new DefaultCloudMetadataServiceConnector(
+        new SnitchProperties(Pair.create(METADATA_URL_PROPERTY, "http://127.0.0.1"))));
+
+        doReturn(az).when(spiedConnector).apiCall(any());
+
+        CloudstackSnitch snitch = new CloudstackSnitch(spiedConnector);
 
         InetAddressAndPort local = InetAddressAndPort.getByName("127.0.0.1");
 
-        assertEquals("ch-gva", snitch.getDatacenter(local));
-        assertEquals("1", snitch.getRack(local));
+        assertEquals("us-east", snitch.getDatacenter(local));
+        assertEquals("1a", snitch.getRack(local));
     }
 
     @AfterClass

test/unit/org/apache/cassandra/locator/Ec2SnitchTest.java

@@ -22,7 +22,6 @@ import java.util.Collections;
 import java.util.EnumMap;
 import java.util.HashSet;
 import java.util.Map;
-import java.util.Properties;
 import java.util.Set;
 
 import org.junit.AfterClass;
@@ -48,7 +47,10 @@ import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.when;
 
 public class Ec2SnitchTest
@@ -85,51 +87,55 @@ public class Ec2SnitchTest
     public void testLegacyRac() throws Exception
     {
         Ec2MetadataServiceConnector connectorMock = mock(Ec2MetadataServiceConnector.class);
-        when(connectorMock.apiCall(anyString())).thenReturn("us-east-1d");
-        Ec2Snitch snitch = new Ec2Snitch(legacySnitchProps, connectorMock);
+        doReturn("us-east-1d").when(connectorMock).apiCall(anyString());
+        doReturn(legacySnitchProps).when(connectorMock).getProperties();
+        Ec2Snitch snitch = new Ec2Snitch(connectorMock);
         testLegacyRacInternal(snitch);
     }
 
     @Test
     public void testMultiregionLegacyRac() throws Exception
     {
-        Ec2MetadataServiceConnector multiRegionConnectorMock = mock(Ec2MetadataServiceConnector.class);
-        when(multiRegionConnectorMock.apiCall(anyString())).then((Answer<String>) invocation -> {
+        Ec2MetadataServiceConnector spy = spy(Ec2MetadataServiceConnector.create(legacySnitchProps));
+        doReturn(legacySnitchProps).when(spy).getProperties();
+        doAnswer((Answer<String>) invocation -> {
             String query = invocation.getArgument(0);
             return (PUBLIC_IP_QUERY.equals(query) || PRIVATE_IP_QUERY.equals(query)) ? "127.0.0.1" : "us-east-1d";
-        });
+        }).when(spy).apiCall(anyString());
 
-        Ec2Snitch snitch = new Ec2MultiRegionSnitch(legacySnitchProps, multiRegionConnectorMock);
+        Ec2Snitch snitch = new Ec2MultiRegionSnitch(spy);
         testLegacyRacInternal(snitch);
     }
 
     @Test
     public void testLegacyNewRegions() throws Exception
     {
-        Ec2MetadataServiceConnector connectorMock = mock(Ec2MetadataServiceConnector.class);
-        when(connectorMock.apiCall(anyString())).thenReturn("us-east-2d");
-        testLegacyNewRegionsInternal(new Ec2Snitch(legacySnitchProps, connectorMock));
+        Ec2MetadataServiceConnector spy = spy(Ec2MetadataServiceConnector.create(legacySnitchProps));
+        doReturn(legacySnitchProps).when(spy).getProperties();
+        doReturn("us-east-2d").when(spy).apiCall(anyString());
+        testLegacyNewRegionsInternal(new Ec2Snitch(spy));
     }
 
     @Test
     public void testLegacyMultiRegionNewRegions() throws Exception
     {
-        Ec2MetadataServiceConnector multiRegionConnectorMock = mock(Ec2MetadataServiceConnector.class);
-        when(multiRegionConnectorMock.apiCall(anyString())).then((Answer<String>) invocation -> {
+        Ec2MetadataServiceConnector spy = spy(Ec2MetadataServiceConnector.create(legacySnitchProps));
+        doReturn(legacySnitchProps).when(spy).getProperties();
+        doAnswer((Answer<String>) invocation -> {
             String query = invocation.getArgument(0);
             return (PUBLIC_IP_QUERY.equals(query) || PRIVATE_IP_QUERY.equals(query)) ? "127.0.0.1" : "us-east-2d";
-        });
+        }).when(spy).apiCall(anyString());
 
-        testLegacyNewRegionsInternal(new Ec2MultiRegionSnitch(legacySnitchProps, multiRegionConnectorMock));
+        testLegacyNewRegionsInternal(new Ec2MultiRegionSnitch(spy));
     }
 
-
     @Test
     public void testFullNamingScheme() throws Exception
     {
         Ec2MetadataServiceConnector connectorMock = mock(Ec2MetadataServiceConnector.class);
         when(connectorMock.apiCall(anyString())).thenReturn("us-east-2d");
-        Ec2Snitch snitch = new Ec2Snitch(new SnitchProperties(new Properties()), connectorMock);
+        when(connectorMock.getProperties()).thenReturn(new SnitchProperties());
+        Ec2Snitch snitch = new Ec2Snitch(connectorMock);
 
         InetAddressAndPort local = InetAddressAndPort.getByName("127.0.0.1");
 
@@ -137,6 +143,7 @@ public class Ec2SnitchTest
         assertEquals("us-east-2d", snitch.getRack(local));
 
         Ec2MetadataServiceConnector multiRegionConnectorMock = mock(Ec2MetadataServiceConnector.class);
+        when(connectorMock.getProperties()).thenReturn(new SnitchProperties());
         when(multiRegionConnectorMock.apiCall(anyString())).then((Answer<String>) invocation -> {
             String query = invocation.getArgument(0);
             return (PUBLIC_IP_QUERY.equals(query) || PRIVATE_IP_QUERY.equals(query)) ? "127.0.0.1" : "us-east-2d";

test/unit/org/apache/cassandra/locator/GoogleCloudSnitchTest.java

@@ -22,7 +22,6 @@ package org.apache.cassandra.locator;
 import java.io.IOException;
 import java.util.EnumMap;
 import java.util.Map;
-import java.util.Properties;
 
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
@@ -35,16 +34,20 @@ import org.apache.cassandra.exceptions.ConfigurationException;
 import org.apache.cassandra.gms.ApplicationState;
 import org.apache.cassandra.gms.Gossiper;
 import org.apache.cassandra.gms.VersionedValue;
+import org.apache.cassandra.locator.AbstractCloudMetadataServiceConnector.DefaultCloudMetadataServiceConnector;
 import org.apache.cassandra.service.StorageService;
+import org.apache.cassandra.utils.Pair;
 
 import static org.apache.cassandra.ServerTestUtils.cleanup;
 import static org.apache.cassandra.ServerTestUtils.mkdirs;
 import static org.apache.cassandra.config.CassandraRelevantProperties.GOSSIP_DISABLE_THREAD_VALIDATION;
+import static org.apache.cassandra.locator.AbstractCloudMetadataServiceConnector.METADATA_URL_PROPERTY;
+import static org.apache.cassandra.locator.AlibabaCloudSnitch.DEFAULT_METADATA_SERVICE_URL;
 import static org.junit.Assert.assertEquals;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyMap;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.spy;
 
 public class GoogleCloudSnitchTest
 {
@@ -66,10 +69,12 @@ public class GoogleCloudSnitchTest
     {
         String az = "us-central1-a";
 
-        AbstractCloudMetadataServiceConnector mock = mock(AbstractCloudMetadataServiceConnector.class);
-        when(mock.apiCall(any(), anyMap())).thenReturn(az);
+        DefaultCloudMetadataServiceConnector spiedConnector = spy(new DefaultCloudMetadataServiceConnector(
+        new SnitchProperties(Pair.create(METADATA_URL_PROPERTY, DEFAULT_METADATA_SERVICE_URL))));
 
-        GoogleCloudSnitch snitch = new GoogleCloudSnitch(new SnitchProperties(new Properties()), mock);
+        doReturn(az).when(spiedConnector).apiCall(any(), anyMap());
+
+        GoogleCloudSnitch snitch = new GoogleCloudSnitch(spiedConnector);
         InetAddressAndPort local = InetAddressAndPort.getByName("127.0.0.1");
         InetAddressAndPort nonlocal = InetAddressAndPort.getByName("127.0.0.7");
 
@@ -90,10 +95,13 @@ public class GoogleCloudSnitchTest
     public void testNewRegions() throws IOException, ConfigurationException
     {
         String az = "asia-east1-a";
-        AbstractCloudMetadataServiceConnector mock = mock(AbstractCloudMetadataServiceConnector.class);
-        when(mock.apiCall(any(), anyMap())).thenReturn(az);
 
-        GoogleCloudSnitch snitch = new GoogleCloudSnitch(new SnitchProperties(new Properties()), mock);
+        DefaultCloudMetadataServiceConnector spiedConnector = spy(new DefaultCloudMetadataServiceConnector(
+        new SnitchProperties(Pair.create(METADATA_URL_PROPERTY, DEFAULT_METADATA_SERVICE_URL))));
+
+        doReturn(az).when(spiedConnector).apiCall(any(), anyMap());
+
+        GoogleCloudSnitch snitch = new GoogleCloudSnitch(spiedConnector);
         InetAddressAndPort local = InetAddressAndPort.getByName("127.0.0.1");
         assertEquals("asia-east1", snitch.getDatacenter(local));
         assertEquals("a", snitch.getRack(local));