Support filtering shared class-info by arbitrary predicates (#16)

For example, filtering by class-loader sub-graphs

Author: mcculls

utils/src/main/java/datadog/instrument/utils/ClassInfoCache.java

@@ -8,6 +8,7 @@
 
 import java.util.Arrays;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.IntPredicate;
 
 /**
  * Shares class information from multiple classloaders in a single cache.
@@ -132,6 +133,47 @@ public T find(CharSequence className, int classLoaderKeyId) {
     }
   }
 
+  /**
+   * Finds information for the given class-name, filtered by class-loader key.
+   *
+   * @param className the class-name
+   * @param classLoaderKeyFilter the filter for the class-loader key
+   * @return information shared under the class-name and filtered class-loader
+   * @see ClassLoaderIndex#getClassLoaderKeyId(ClassLoader)
+   */
+  @SuppressWarnings("unchecked")
+  public T find(CharSequence className, IntPredicate classLoaderKeyFilter) {
+    final int hash = className.hashCode();
+    final SharedInfo[] shared = this.shared;
+    final int slotMask = this.slotMask;
+
+    // try to find matching slot, rehashing after each attempt
+    for (int i = 1, h = hash; true; i++, h = rehash(h)) {
+      int slot = slotMask & h;
+      SharedInfo existing = shared[slot];
+      if (existing != null) {
+        if (existing.className.contentEquals(className)) {
+          // apply filter to class-loader key, -1 always matches
+          if (existing.classLoaderKeyId < 0
+              || classLoaderKeyFilter.test(existing.classLoaderKeyId)) {
+            // use global TICKS as a substitute for access time
+            // TICKS is only incremented in 'share' for performance reasons
+            existing.accessed = TICKS.get();
+            return (T) existing.classInfo;
+          }
+          // fall-through and quit; name matched but class-loader didn't
+        } else if (i < MAX_HASH_ATTEMPTS) {
+          continue; // rehash and try again
+        }
+      }
+      // quit search when:
+      // * we find an empty slot (we know there won't be further info)
+      // * we find a slot with the same name but different class-loader
+      // * we've exhausted all hash attempts
+      return null;
+    }
+  }
+
   /**
    * Shares information for the given class-name, under the given class-loader key.
    *

utils/src/test/java/datadog/instrument/utils/ClassInfoCacheTest.java

@@ -8,6 +8,7 @@
 
 import java.util.HashSet;
 import java.util.Set;
+import java.util.function.IntPredicate;
 import org.junit.jupiter.api.Test;
 
 class ClassInfoCacheTest {
@@ -19,6 +20,11 @@ void basicOperation() {
     ClassLoader myCL = newCL();
     ClassLoader notMyCL = newCL();
 
+    int myCLKey = ClassLoaderIndex.getClassLoaderKeyId(myCL);
+
+    IntPredicate myCLFilter = sameCLKey(myCLKey);
+    IntPredicate notMyCLFilter = myCLFilter.negate();
+
     assertNull(cache.find("example.test.MyGlobalClass"));
     assertNull(cache.find("example.test.MyLocalClass"));
     assertNull(cache.find("example.test.NotMyClass"));
@@ -31,6 +37,14 @@ void basicOperation() {
     assertNull(cache.find("example.test.MyLocalClass", notMyCL));
     assertNull(cache.find("example.test.NotMyClass", notMyCL));
 
+    assertNull(cache.find("example.test.MyGlobalClass", myCLFilter));
+    assertNull(cache.find("example.test.MyLocalClass", myCLFilter));
+    assertNull(cache.find("example.test.NotMyClass", myCLFilter));
+
+    assertNull(cache.find("example.test.MyGlobalClass", notMyCLFilter));
+    assertNull(cache.find("example.test.MyLocalClass", notMyCLFilter));
+    assertNull(cache.find("example.test.NotMyClass", notMyCLFilter));
+
     cache.share("example.test.MyGlobalClass", "my global data");
     cache.share("example.test.MyLocalClass", "my local data", myCL);
 
@@ -46,6 +60,14 @@ void basicOperation() {
     assertNull(cache.find("example.test.MyLocalClass", notMyCL));
     assertNull(cache.find("example.test.NotMyClass", notMyCL));
 
+    assertEquals("my global data", cache.find("example.test.MyGlobalClass", myCLFilter));
+    assertEquals("my local data", cache.find("example.test.MyLocalClass", myCLFilter));
+    assertNull(cache.find("example.test.NotMyClass", myCLFilter));
+
+    assertEquals("my global data", cache.find("example.test.MyGlobalClass", notMyCLFilter));
+    assertNull(cache.find("example.test.MyLocalClass", notMyCLFilter));
+    assertNull(cache.find("example.test.NotMyClass", notMyCLFilter));
+
     cache.clear();
 
     assertNull(cache.find("example.test.MyGlobalClass"));
@@ -59,6 +81,14 @@ void basicOperation() {
     assertNull(cache.find("example.test.MyGlobalClass", notMyCL));
     assertNull(cache.find("example.test.MyLocalClass", notMyCL));
     assertNull(cache.find("example.test.NotMyClass", notMyCL));
+
+    assertNull(cache.find("example.test.MyGlobalClass", myCLFilter));
+    assertNull(cache.find("example.test.MyLocalClass", myCLFilter));
+    assertNull(cache.find("example.test.NotMyClass", myCLFilter));
+
+    assertNull(cache.find("example.test.MyGlobalClass", notMyCLFilter));
+    assertNull(cache.find("example.test.MyLocalClass", notMyCLFilter));
+    assertNull(cache.find("example.test.NotMyClass", notMyCLFilter));
   }
 
   @Test
@@ -111,12 +141,20 @@ void overflow() {
     for (int i = 0; i < 300; i++) {
       if (overwritten.contains(i)) {
         assertNull(cache.find("example.MyClass" + i), "rem " + i);
+        assertNull(cache.find("example.MyClass" + i, i), "rem " + i);
+        assertNull(cache.find("example.MyClass" + i, sameCLKey(i)), "rem " + i);
       } else {
         assertEquals(i, cache.find("example.MyClass" + i), "add " + i);
+        assertEquals(i, cache.find("example.MyClass" + i, i), "add " + i);
+        assertEquals(i, cache.find("example.MyClass" + i, sameCLKey(i)), "add " + i);
       }
     }
   }
 
+  private static IntPredicate sameCLKey(int clKey) {
+    return k -> k == clKey;
+  }
+
   private static ClassLoader newCL() {
     return new ClassLoader() {};
   }