Use resolved address for peer.hostname when available without hitting the cache

Author: amarziali

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/BaseDecorator.java

@@ -2,13 +2,12 @@
 
 import static datadog.trace.api.cache.RadixTreeCache.PORTS;
 import static datadog.trace.api.cache.RadixTreeCache.UNSET_PORT;
+import static datadog.trace.bootstrap.instrumentation.java.net.HostNameResolver.hostName;
 
 import datadog.context.ContextScope;
 import datadog.trace.api.Config;
 import datadog.trace.api.DDTags;
 import datadog.trace.api.Functions;
-import datadog.trace.api.cache.DDCache;
-import datadog.trace.api.cache.DDCaches;
 import datadog.trace.api.cache.QualifiedClassNameCache;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -40,8 +39,6 @@ public String apply(Class<?> clazz) {
           },
           Functions.PrefixJoin.of("."));
 
-  private static final DDCache<String, String> HOSTNAME_CACHE = DDCaches.newFixedSizeCache(64);
-
   protected final boolean traceAnalyticsEnabled;
   protected final Double traceAnalyticsSampleRate;
 
@@ -200,11 +197,4 @@ public CharSequence className(final Class<?> clazz) {
     String simpleName = clazz.getSimpleName();
     return simpleName.isEmpty() ? CLASS_NAMES.getClassName(clazz) : simpleName;
   }
-
-  private static String hostName(InetAddress remoteAddress, String ip) {
-    if (null != ip) {
-      return HOSTNAME_CACHE.computeIfAbsent(ip, _ip -> remoteAddress.getHostName());
-    }
-    return remoteAddress.getHostName();
-  }
 }

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/net/HostNameResolver.java

@@ -0,0 +1,72 @@
+package datadog.trace.bootstrap.instrumentation.java.net;
+
+import datadog.trace.api.cache.DDCache;
+import datadog.trace.api.cache.DDCaches;
+import datadog.trace.util.MethodHandles;
+import java.lang.invoke.MethodHandle;
+import java.net.InetAddress;
+
+public final class HostNameResolver {
+  private static final MethodHandle HOLDER_GET;
+  private static final MethodHandle HOSTNAME_GET;
+
+  private static final DDCache<String, String> HOSTNAME_CACHE = DDCaches.newFixedSizeCache(64);
+
+  static {
+    MethodHandle holderTmp = null, hostnameTmp = null;
+    try {
+      final ClassLoader cl = HostNameResolver.class.getClassLoader();
+      final MethodHandles methodHandles = new MethodHandles(cl);
+
+      final Class<?> holderClass =
+          Class.forName("java.net.InetAddress$InetAddressHolder", false, cl);
+      holderTmp = methodHandles.method(InetAddress.class, "holder");
+      if (holderTmp != null) {
+        hostnameTmp = methodHandles.method(holderClass, "getHostName");
+      }
+    } catch (Throwable ignored) {
+      holderTmp = null;
+    } finally {
+      if (holderTmp != null && hostnameTmp != null) {
+        HOLDER_GET = holderTmp;
+        HOSTNAME_GET = hostnameTmp;
+      } else {
+        HOLDER_GET = null;
+        HOSTNAME_GET = null;
+      }
+    }
+  }
+
+  private HostNameResolver() {}
+
+  static String getAlreadyResolvedHostName(InetAddress address) {
+    if (HOLDER_GET == null) {
+      return null;
+    }
+    try {
+      final Object holder = HOLDER_GET.invoke(address);
+      return (String) HOSTNAME_GET.invoke(holder);
+    } catch (final Throwable ignored) {
+    }
+    return null;
+  }
+
+  private static String fromCache(InetAddress remoteAddress, String ip) {
+    final String alreadyResolved = HostNameResolver.getAlreadyResolvedHostName(remoteAddress);
+    if (alreadyResolved != null) {
+      return alreadyResolved;
+    }
+    if (null != ip) {
+      return HOSTNAME_CACHE.computeIfAbsent(ip, _ip -> remoteAddress.getHostName());
+    }
+    return remoteAddress.getHostName();
+  }
+
+  public static String hostName(InetAddress address, String ip) {
+    final String alreadyResolved = getAlreadyResolvedHostName(address);
+    if (alreadyResolved != null) {
+      return alreadyResolved;
+    }
+    return fromCache(address, ip);
+  }
+}

dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/instrumentation/java/net/HostNameResolverTest.groovy

@@ -0,0 +1,45 @@
+package datadog.trace.bootstrap.instrumentation.java.net
+
+import datadog.trace.test.util.DDSpecification
+
+class HostNameResolverTest extends DDSpecification {
+  def "should directly get the hostname for already resolved address #address"() {
+    given:
+    def host = HostNameResolver.getAlreadyResolvedHostName(address)
+    expect:
+    host == expected
+    where:
+    address                                                           | expected
+    new Inet4Address("test", InetAddress.getLocalHost().getAddress()) | "test"
+    new Inet6Address("test", InetAddress.getLocalHost().getAddress()) | "test"
+  }
+
+  def "should return null when directly get the address for unresolved #address"() {
+    given:
+    def host = HostNameResolver.getAlreadyResolvedHostName(address)
+    expect:
+    host == null
+    where:
+    address                                                           | _
+    InetAddress.getByAddress(InetAddress.getLocalHost().getAddress()) | _
+    new Inet6Address(null, InetAddress.getLocalHost().getAddress())   | _
+  }
+
+  def "should use the cache for unresolved addresses"() {
+    setup:
+    def inet1 = Mock(InetAddress)
+    def inet2 = Mock(InetAddress)
+    when:
+    def address = new InetSocketAddress(inet1, 666)
+    def host = HostNameResolver.hostName(address.getAddress(), "127.0.0.1")
+    then:
+    host == "somehost"
+    1 * inet1.getHostName() >> "somehost"
+    when:
+    address = new InetSocketAddress(inet2, 666)
+    host = HostNameResolver.hostName(address.getAddress(), "127.0.0.1")
+    then:
+    0 * inet2.getHostName()
+    host == "somehost"
+  }
+}