diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/UserAgent.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/UserAgent.java index 6aca696ec..35ccbc4d3 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/UserAgent.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/UserAgent.java @@ -6,6 +6,9 @@ import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.Objects; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; import java.util.regex.Pattern; import java.util.stream.Collectors; @@ -14,8 +17,8 @@ public class UserAgent { private static String productVersion = "0.0.0"; private static class Info { - private String key; - private String value; + private final String key; + private final String value; public Info(String key, String value) { this.key = key; @@ -29,9 +32,29 @@ public String getKey() { public String getValue() { return value; } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof Info)) { + return false; + } + Info info = (Info) o; + return key.equals(info.key) && value.equals(info.value); + } + + @Override + public int hashCode() { + return Objects.hash(key, value); + } } - private static final ArrayList otherInfo = new ArrayList<>(); + // Callers such as JDBC drivers re-register the same entries on every connection, and asString() + // runs on every request, so entries are deduplicated and read without locking. + // Package-private for testing. + static final Set otherInfo = ConcurrentHashMap.newKeySet(); // TODO: check if reading from // /META-INF/maven/com.databricks/databrics-sdk-java/pom.properties @@ -96,9 +119,7 @@ public static void withPartner(String partner) { public static void withOtherInfo(String key, String value) { matchAlphanum(key); matchAlphanumOrSemVer(value); - synchronized (otherInfo) { - otherInfo.add(new Info(key, value)); - } + otherInfo.add(new Info(key, value)); } private static String osName() { @@ -137,13 +158,10 @@ public static String asString() { if (!metaHarness.isEmpty()) { segments.add(String.format("meta-harness/%s", metaHarness)); } - // Concurrent iteration over ArrayList must be guarded with synchronized. - synchronized (otherInfo) { - segments.addAll( - otherInfo.stream() - .map(e -> String.format("%s/%s", e.getKey(), e.getValue())) - .collect(Collectors.toSet())); - } + segments.addAll( + otherInfo.stream() + .map(e -> String.format("%s/%s", e.getKey(), e.getValue())) + .collect(Collectors.toSet())); return segments.stream().collect(Collectors.joining(" ")); } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/UserAgentLoadTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/UserAgentLoadTest.java index c07e2b64c..6af1e7ed8 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/UserAgentLoadTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/UserAgentLoadTest.java @@ -21,12 +21,13 @@ public void testAsStringConcurrent() throws InterruptedException, ExecutionExcep List> futures = new ArrayList<>(); int successCount = 0; int failureCount = 0; + int otherInfoSizeBefore = UserAgent.otherInfo.size(); // Add some user agent info Callable task = () -> { try { - UserAgent.withOtherInfo("key1", "value1"); + UserAgent.withOtherInfo("load-test", "1.0.0"); UserAgent.asString(); return true; } catch (Exception e) { @@ -60,5 +61,7 @@ public void testAsStringConcurrent() throws InterruptedException, ExecutionExcep // Optionally, you can assert that there were no failures assertEquals(0, failureCount); + // Concurrent registrations of the same entry add it once. + assertEquals(otherInfoSizeBefore + 1, UserAgent.otherInfo.size()); } } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/UserAgentTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/UserAgentTest.java index dd7794ea8..6ec448e24 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/UserAgentTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/UserAgentTest.java @@ -57,6 +57,23 @@ public void testUserAgentWithOtherInfo() { Assertions.assertTrue(userAgent.contains("key2/value2")); } + @Test + public void testUserAgentWithOtherInfoDeduplicated() { + int sizeBefore = UserAgent.otherInfo.size(); + for (int i = 0; i < 1000; i++) { + UserAgent.withOtherInfo("dedup", "1.0.0"); + } + Assertions.assertEquals(sizeBefore + 1, UserAgent.otherInfo.size()); + + // Same key with a different value is a separate entry. + UserAgent.withOtherInfo("dedup", "2.0.0"); + Assertions.assertEquals(sizeBefore + 2, UserAgent.otherInfo.size()); + + String userAgent = UserAgent.asString(); + Assertions.assertTrue(userAgent.contains("dedup/1.0.0")); + Assertions.assertTrue(userAgent.contains("dedup/2.0.0")); + } + @Test public void testUserAgentWithInvalidKey() { Assertions.assertThrows(