WIP: separate per-store context storage from boxed store-ids that forward to ObjectStoreDispatch

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/BoxedContextStore.java

@@ -0,0 +1,39 @@
+package datadog.trace.bootstrap;
+
+import datadog.instrument.fieldinject.ObjectStoreDispatch;
+import java.util.function.Function;
+
+/** Boxes a store-id and dispatches context operations through {@link ObjectStoreDispatch}. */
+public final class BoxedContextStore implements ContextStore<Object, Object> {
+  final int storeId;
+
+  BoxedContextStore(final int storeId) {
+    this.storeId = storeId;
+  }
+
+  @Override
+  public Object get(final Object key) {
+    return ObjectStoreDispatch.get(key, storeId);
+  }
+
+  @Override
+  public void put(final Object key, final Object context) {
+    ObjectStoreDispatch.put(key, storeId, context);
+  }
+
+  @Override
+  public Object getOrPut(final Object key, final Object context) {
+    return ObjectStoreDispatch.getOrPut(key, storeId, context);
+  }
+
+  @Override
+  public Object getOrCompute(
+      final Object key, final Function<? super Object, Object> contextFactory) {
+    return ObjectStoreDispatch.getOrCompute(key, storeId, contextFactory);
+  }
+
+  @Override
+  public Object remove(final Object key) {
+    return ObjectStoreDispatch.remove(key, storeId);
+  }
+}

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/BoxedContextStores.java

@@ -0,0 +1,133 @@
+package datadog.trace.bootstrap;
+
+import java.util.Arrays;
+import java.util.concurrent.ConcurrentHashMap;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Intentional copy of FieldBackedContextStores but for {@link BoxedContextStore}.
+ *
+ * <p>Only one is loaded at runtime; the explicit field type helps the JIT optimize.
+ */
+public final class BoxedContextStores {
+
+  private static final Logger log = LoggerFactory.getLogger(BoxedContextStores.class);
+
+  // these fields will be accessed directly from field-injected instrumentation
+  public static final BoxedContextStore contextStore0 = new BoxedContextStore(0);
+  public static final BoxedContextStore contextStore1 = new BoxedContextStore(1);
+  public static final BoxedContextStore contextStore2 = new BoxedContextStore(2);
+  public static final BoxedContextStore contextStore3 = new BoxedContextStore(3);
+  public static final BoxedContextStore contextStore4 = new BoxedContextStore(4);
+  public static final BoxedContextStore contextStore5 = new BoxedContextStore(5);
+  public static final BoxedContextStore contextStore6 = new BoxedContextStore(6);
+  public static final BoxedContextStore contextStore7 = new BoxedContextStore(7);
+  public static final BoxedContextStore contextStore8 = new BoxedContextStore(8);
+  public static final BoxedContextStore contextStore9 = new BoxedContextStore(9);
+  public static final BoxedContextStore contextStore10 = new BoxedContextStore(10);
+  public static final BoxedContextStore contextStore11 = new BoxedContextStore(11);
+  public static final BoxedContextStore contextStore12 = new BoxedContextStore(12);
+  public static final BoxedContextStore contextStore13 = new BoxedContextStore(13);
+  public static final BoxedContextStore contextStore14 = new BoxedContextStore(14);
+  public static final BoxedContextStore contextStore15 = new BoxedContextStore(15);
+  public static final BoxedContextStore contextStore16 = new BoxedContextStore(16);
+  public static final BoxedContextStore contextStore17 = new BoxedContextStore(17);
+  public static final BoxedContextStore contextStore18 = new BoxedContextStore(18);
+  public static final BoxedContextStore contextStore19 = new BoxedContextStore(19);
+  public static final BoxedContextStore contextStore20 = new BoxedContextStore(20);
+  public static final BoxedContextStore contextStore21 = new BoxedContextStore(21);
+  public static final BoxedContextStore contextStore22 = new BoxedContextStore(22);
+  public static final BoxedContextStore contextStore23 = new BoxedContextStore(23);
+  public static final BoxedContextStore contextStore24 = new BoxedContextStore(24);
+  public static final BoxedContextStore contextStore25 = new BoxedContextStore(25);
+  public static final BoxedContextStore contextStore26 = new BoxedContextStore(26);
+  public static final BoxedContextStore contextStore27 = new BoxedContextStore(27);
+  public static final BoxedContextStore contextStore28 = new BoxedContextStore(28);
+  public static final BoxedContextStore contextStore29 = new BoxedContextStore(29);
+  public static final BoxedContextStore contextStore30 = new BoxedContextStore(30);
+  public static final BoxedContextStore contextStore31 = new BoxedContextStore(31);
+
+  // keep track of all allocated stores so far
+  private static volatile BoxedContextStore[] stores = {
+    contextStore0,
+    contextStore1,
+    contextStore2,
+    contextStore3,
+    contextStore4,
+    contextStore5,
+    contextStore6,
+    contextStore7,
+    contextStore8,
+    contextStore9,
+    contextStore10,
+    contextStore11,
+    contextStore12,
+    contextStore13,
+    contextStore14,
+    contextStore15,
+    contextStore16,
+    contextStore17,
+    contextStore18,
+    contextStore19,
+    contextStore20,
+    contextStore21,
+    contextStore22,
+    contextStore23,
+    contextStore24,
+    contextStore25,
+    contextStore26,
+    contextStore27,
+    contextStore28,
+    contextStore29,
+    contextStore30,
+    contextStore31
+  };
+
+  private static final ConcurrentHashMap<String, BoxedContextStore> storesByName =
+      new ConcurrentHashMap<>();
+  private static final Object allocationLock = new Object();
+
+  private static int nextStoreId;
+
+  public static int getContextStoreId(String keyClassName, String contextClassName) {
+    String storeName = keyClassName + ';' + contextClassName;
+    BoxedContextStore store = storesByName.get(storeName);
+    if (store == null) {
+      synchronized (allocationLock) {
+        // speculatively create the next store in the sequence and attempt to map this name to it;
+        // if another thread has mapped this name then the store will be kept for the next mapping
+        BoxedContextStore existing =
+            storesByName.putIfAbsent(storeName, store = allocateStore(nextStoreId));
+        if (existing != null) {
+          return existing.storeId;
+        }
+        nextStoreId++;
+      }
+      log.debug(
+          "Allocated ContextStore #{} - instrumentation.target.context={}->{}",
+          store.storeId,
+          keyClassName,
+          contextClassName);
+    }
+    return store.storeId;
+  }
+
+  public static BoxedContextStore getContextStore(final int storeId) {
+    return stores[storeId]; // guaranteed to be populated by getContextStoreId
+  }
+
+  // this method should only be called while holding the allocation lock
+  private static BoxedContextStore allocateStore(int storeId) {
+    BoxedContextStore[] snapshot = stores;
+    if (storeId >= snapshot.length) {
+      stores = snapshot = Arrays.copyOf(snapshot, storeId + 16);
+    }
+    // check for pre-allocated / speculatively allocated stores
+    BoxedContextStore store = snapshot[storeId];
+    if (store == null) {
+      snapshot[storeId] = store = new BoxedContextStore(storeId);
+    }
+    return store;
+  }
+}

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/ContextStores.java

@@ -0,0 +1,39 @@
+package datadog.trace.bootstrap;
+
+import datadog.trace.api.InstrumenterConfig;
+
+/** Allocates {@link ContextStore}s. */
+public final class ContextStores {
+
+  // fast lookup available for fixed number of stores
+  public static final int FAST_STORE_ID_LIMIT = 32;
+
+  private static final boolean MAP_PER_STORE =
+      InstrumenterConfig.get().isRuntimeContextMapPerStore();
+
+  public static final String STORES_DESCRIPTOR =
+      MAP_PER_STORE
+          ? "Ldatadog/trace/bootstrap/FieldBackedContextStores;"
+          : "Ldatadog/trace/bootstrap/BoxedContextStores;";
+
+  public static final String STORE_DESCRIPTOR =
+      MAP_PER_STORE
+          ? "Ldatadog/trace/bootstrap/FieldBackedContextStore;"
+          : "Ldatadog/trace/bootstrap/BoxedContextStore;";
+
+  public static int getContextStoreId(String keyClassName, String contextClassName) {
+    if (MAP_PER_STORE) {
+      return FieldBackedContextStores.getContextStoreId(keyClassName, contextClassName);
+    } else {
+      return BoxedContextStores.getContextStoreId(keyClassName, contextClassName);
+    }
+  }
+
+  public static ContextStore<?, ?> getContextStore(final int storeId) {
+    if (MAP_PER_STORE) {
+      return FieldBackedContextStores.getContextStore(storeId);
+    } else {
+      return BoxedContextStores.getContextStore(storeId);
+    }
+  }
+}

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/FieldBackedContextStore.java

@@ -1,18 +1,13 @@
 package datadog.trace.bootstrap;
 
 import datadog.instrument.fieldinject.KeyWithValue;
-import datadog.instrument.fieldinject.ObjectStoreDispatch;
-import datadog.trace.api.InstrumenterConfig;
 import java.util.function.Function;
 
 /**
  * {@link ContextStore} that attempts to store context in its keys by using bytecode-injected
  * fields. Delegates to a lazy {@link WeakMap} for keys that don't have a field for this store.
  */
 public final class FieldBackedContextStore implements ContextStore<Object, Object> {
-  private static final boolean MAP_PER_STORE =
-      InstrumenterConfig.get().isRuntimeContextMapPerStore();
-
   final int storeId;
 
   FieldBackedContextStore(final int storeId) {
@@ -23,21 +18,17 @@ public final class FieldBackedContextStore implements ContextStore<Object, Objec
   public Object get(final Object key) {
     if (key instanceof KeyWithValue) {
       return ((KeyWithValue) key).$get$__dd_instrument$(storeId);
-    } else if (MAP_PER_STORE) {
-      return weakStore().get(key);
     } else {
-      return ObjectStoreDispatch.get(key, storeId);
+      return weakStore().get(key);
     }
   }
 
   @Override
   public void put(final Object key, final Object context) {
     if (key instanceof KeyWithValue) {
       ((KeyWithValue) key).$put$__dd_instrument$(storeId, context);
-    } else if (MAP_PER_STORE) {
-      weakStore().put(key, context);
     } else {
-      ObjectStoreDispatch.put(key, storeId, context);
+      weakStore().put(key, context);
     }
   }
 
@@ -56,15 +47,14 @@ public Object getOrPut(final Object key, final Object context) {
         }
       }
       return existingContext;
-    } else if (MAP_PER_STORE) {
-      return weakStore().getOrPut(key, context);
     } else {
-      return ObjectStoreDispatch.getOrPut(key, storeId, context);
+      return weakStore().getOrPut(key, context);
     }
   }
 
   @Override
-  public Object getOrCompute(Object key, Function<? super Object, Object> contextFactory) {
+  public Object getOrCompute(
+      final Object key, final Function<? super Object, Object> contextFactory) {
     if (key instanceof KeyWithValue) {
       final KeyWithValue accessor = (KeyWithValue) key;
       Object existingContext = accessor.$get$__dd_instrument$(storeId);
@@ -78,15 +68,13 @@ public Object getOrCompute(Object key, Function<? super Object, Object> contextF
         }
       }
       return existingContext;
-    } else if (MAP_PER_STORE) {
-      return weakStore().getOrCompute(key, contextFactory);
     } else {
-      return ObjectStoreDispatch.getOrCompute(key, storeId, contextFactory);
+      return weakStore().getOrCompute(key, contextFactory);
     }
   }
 
   @Override
-  public Object remove(Object key) {
+  public Object remove(final Object key) {
     if (key instanceof KeyWithValue) {
       final KeyWithValue accessor = (KeyWithValue) key;
       Object existingContext = accessor.$get$__dd_instrument$(storeId);
@@ -99,10 +87,8 @@ public Object remove(Object key) {
         }
       }
       return existingContext;
-    } else if (MAP_PER_STORE) {
-      return weakStore().remove(key);
     } else {
-      return ObjectStoreDispatch.remove(key, storeId);
+      return weakStore().remove(key);
     }
   }
 

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/FieldBackedContextStores.java

@@ -10,9 +10,6 @@ public final class FieldBackedContextStores {
 
   private static final Logger log = LoggerFactory.getLogger(FieldBackedContextStores.class);
 
-  // provide fast lookup for a fixed number of stores
-  public static final int FAST_STORE_ID_LIMIT = 32;
-
   // these fields will be accessed directly from field-injected instrumentation
   public static final FieldBackedContextStore contextStore0 = new FieldBackedContextStore(0);
   public static final FieldBackedContextStore contextStore1 = new FieldBackedContextStore(1);