Fix Ec2HostLocationSupplier using a different format to Cassandra

atlasdb-cassandra/src/main/java/com/palantir/atlasdb/keyvalue/cassandra/pool/Ec2HostLocationSupplier.java

@@ -19,6 +19,7 @@
 import java.io.IOException;
 import java.util.function.Supplier;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.palantir.logsafe.Preconditions;
 
 import okhttp3.OkHttpClient;
@@ -46,13 +47,32 @@ public HostLocation get() {
             Preconditions.checkState(response.isSuccessful(),
                     "Getting AWS host metadata was not successful");
 
-            String responseString = response.body().string();
-            String datacenter = responseString.substring(0, responseString.length() - 2);
-            String rack = responseString.substring(responseString.length() - 1);
-
-            return HostLocation.of(datacenter, rack);
+            return parseHostLocation(response.body().string());
         } catch (IOException e) {
             throw new RuntimeException(e);
         }
     }
+
+    @VisibleForTesting
+    static HostLocation parseHostLocation(String az) {
+        // Code is copied from Cassandra's Ec2Snitch. The result of this parsing must match Cassandra's as closely as
+        // possible, as the strings are later matched exactly.
+
+        String ec2region;
+        String ec2zone;
+
+        // Split "us-east-1a" or "asia-1a" into "us-east"/"1a" and "asia"/"1a".
+        String[] splits = az.split("-");
+        ec2zone = splits[splits.length - 1];
+
+        // hack for CASSANDRA-4026
+        ec2region = az.substring(0, az.length() - 1);
+        if (ec2region.endsWith("1")) {
+            ec2region = az.substring(0, az.length() - 3);
+        }
+
+        return HostLocation.of(ec2region, ec2zone);
+
+
+    }
 }

atlasdb-cassandra/src/test/java/com/palantir/atlasdb/keyvalue/cassandra/pool/HostLocationSupplierTest.java

@@ -81,4 +81,10 @@ public void shouldReturnEmptyLocationFromEc2Exception() {
         assertThat(hostLocationSupplier.get()).isNotPresent();
     }
 
+    @Test
+    public void shouldReturnHostLocationInCassandraStyle() {
+        HostLocation awsLocation = HostLocation.of("us-east", "1a");
+        assertThat(Ec2HostLocationSupplier.parseHostLocation("us-east-1a")).isEqualTo(awsLocation);
+    }
+
 }

changelog/@unreleased/pr-4283.v2.yml

@@ -0,0 +1,18 @@
+type: fix
+fix:
+  description: |-
+    Fix Ec2HostLocationSupplier using a different format to Cassandra
+
+    **Goals (and why)**:
+    Modify Ec2HostLocationSupplier to emulate Cassandra's Ec2Snitch logic (as it is later comparing against parsed info from that class).
+
+    **Implementation Description (bullets)**:
+    The response string from the AWS endpoint is now parsed using code copied from Cassandra's implementation (or as close as possible)
+
+    **Testing (What was existing testing like?  What have you done to improve it?)**:
+    Added a test to make sure that the parsing gives the expected result
+
+    **Priority (whenever / two weeks / yesterday)**:
+    Today
+  links:
+  - https://github.com/palantir/atlasdb/pull/4283