Switching to use SpanConcurrentMap in the concurrency plugin (#2173)

Author: eyalkoren

CHANGELOG.asciidoc

@@ -53,6 +53,7 @@ exclude known static files in all applications might be changed to `/app/static/
 ** `DynamicTransformer.Accessor.get().ensureInstrumented` is now `DynamicTransformer.ensureInstrumented` - {pull}#2164[#2164]
 * Switching last instrumentations (`trace_methods`, sparkjava, JDK `HttpServer` and Struts 2) to
 `TracerAwareInstrumentation` - {pull}#2170[#2170]
+* Replace concurrency plugin maps to `SpanConcurrentHashMap` ones - {pull}#2173[#2173]
 
 [[release-notes-1.x]]
 === Java Agent version 1.x

apm-agent-plugins/apm-java-concurrent-plugin/src/main/java/co/elastic/apm/agent/concurrent/JavaConcurrent.java

@@ -18,12 +18,12 @@
  */
 package co.elastic.apm.agent.concurrent;
 
+import co.elastic.apm.agent.collections.WeakConcurrentProviderImpl;
 import co.elastic.apm.agent.impl.Tracer;
 import co.elastic.apm.agent.impl.transaction.AbstractSpan;
 import co.elastic.apm.agent.sdk.DynamicTransformer;
 import co.elastic.apm.agent.sdk.ElasticApmInstrumentation;
 import co.elastic.apm.agent.sdk.state.GlobalState;
-import co.elastic.apm.agent.sdk.weakconcurrent.WeakConcurrent;
 import co.elastic.apm.agent.sdk.weakconcurrent.WeakMap;
 
 import javax.annotation.Nullable;
@@ -41,7 +41,8 @@
 @GlobalState
 public class JavaConcurrent {
 
- private static final WeakMap<Object, AbstractSpan<?>> contextMap = WeakConcurrent.buildMap();
+ private static final WeakMap<Object, AbstractSpan<?>> contextMap = WeakConcurrentProviderImpl.createWeakSpanMap();
+
  private static final List<Class<? extends ElasticApmInstrumentation>> RUNNABLE_CALLABLE_FJTASK_INSTRUMENTATION = Collections.
  <Class<? extends ElasticApmInstrumentation>>singletonList(RunnableCallableForkJoinTaskInstrumentation.class);
  static final ThreadLocal<Boolean> needsContext = new ThreadLocal<>();
@@ -57,10 +58,7 @@ public class JavaConcurrent {
  }
 
  private static void removeContext(Object o) {
- AbstractSpan<?> context = contextMap.remove(o);
- if (context != null) {
- context.decrementReferences();
- }
+ contextMap.remove(o);
  }
 
  private static boolean shouldAvoidContextPropagation(@Nullable Object executable) {
@@ -69,21 +67,34 @@ private static boolean shouldAvoidContextPropagation(@Nullable Object executable
  needsContext.get() == Boolean.FALSE;
  }
 
+ /**
+ * Retrieves the context mapped to the provided task and activates it on the current thread.
+ * It is the responsibility of the caller to deactivate the returned context at the right time.
+ * If the mapped context is already the active span of this thread, this method returns {@code null}.
+ * @param o a task for which running there may be a context to activate
+ * @param tracer the tracer
+ * @return the context mapped to the provided task or {@code null} if such does not exist or if the mapped context
+ * is already the active one on the current thread.
+ */
  @Nullable
  public static AbstractSpan<?> restoreContext(Object o, Tracer tracer) {
  // When an Executor executes directly on the current thread we need to enable this thread for context propagation again
  needsContext.set(Boolean.TRUE);
- AbstractSpan<?> context = contextMap.remove(o);
+
+ // we cannot remove yet, as this decrements the reference count, which may cause already ended spans to be recycled ahead of time
+ AbstractSpan<?> context = contextMap.get(o);
  if (context == null) {
  return null;
  }
- if (tracer.getActive() != context) {
- context.activate();
- context.decrementReferences();
- return context;
- } else {
- context.decrementReferences();
- return null;
+
+ try {
+ if (tracer.getActive() != context) {
+ return context.activate();
+ } else {
+ return null;
+ }
+ } finally {
+ contextMap.remove(o);
  }
  }
 
@@ -110,7 +121,6 @@ public static Runnable withContext(@Nullable Runnable runnable, Tracer tracer) {
  private static void captureContext(Object task, AbstractSpan<?> active) {
  DynamicTransformer.ensureInstrumented(task.getClass(), RUNNABLE_CALLABLE_FJTASK_INSTRUMENTATION);
  contextMap.put(task, active);
- active.incrementReferences();
  // Do no discard branches leading to async operations so not to break span references
  active.setNonDiscardable();
  }

apm-agent-plugins/apm-java-concurrent-plugin/src/test/java/co/elastic/apm/agent/concurrent/ScheduledExecutorServiceTest.java

@@ -56,20 +56,53 @@ public ScheduledFuture<?> schedule(Runnable command, long delay, TimeUnit unit)
 
  @AfterEach
  void tearDown() {
- assertThat(tracer.getActive()).isEqualTo(transaction);
- transaction.deactivate().end();
+ if (tracer.getActive() != null) {
+ transaction.deactivate().end();
+ }
  }
 
  @Test
  void testScheduleCallable() throws Exception {
  final ScheduledFuture<? extends AbstractSpan<?>> future = scheduler.schedule(() -> tracer.getActive(), 0, TimeUnit.SECONDS);
  assertThat(future.get()).isEqualTo(transaction);
+ assertThat(tracer.getActive()).isEqualTo(transaction);
  }
 
  @Test
  void testScheduleRunnable() throws Exception {
  AtomicReference<AbstractSpan<?>> ref = new AtomicReference<>();
  scheduler.schedule(() -> ref.set(tracer.getActive()), 0, TimeUnit.SECONDS).get();
  assertThat(ref.get()).isEqualTo(transaction);
+ assertThat(tracer.getActive()).isEqualTo(transaction);
+ }
+
+ @Test
+ void testScheduleCallable_delayAndEndTransaction() throws Exception {
+ final ScheduledFuture<? extends AbstractSpan<?>> scheduledFuture = scheduler.schedule(() -> tracer.getActive(), 50, TimeUnit.MILLISECONDS);
+ verifyEndedTransactionIsStillReferenced(scheduledFuture);
+ assertThat(scheduledFuture.get()).isEqualTo(transaction);
+ }
+
+ @Test
+ void testScheduleRunnable_delayAndEndTransaction() throws Exception {
+ AtomicReference<AbstractSpan<?>> ref = new AtomicReference<>();
+ ScheduledFuture<?> scheduledTaskFuture = scheduler.schedule(() -> ref.set(tracer.getActive()), 50, TimeUnit.MILLISECONDS);
+ verifyEndedTransactionIsStillReferenced(scheduledTaskFuture);
+ assertThat(ref.get()).isEqualTo(transaction);
+ }
+
+ private void verifyEndedTransactionIsStillReferenced(ScheduledFuture<?> scheduledTaskFuture) throws InterruptedException, java.util.concurrent.ExecutionException {
+ transaction.deactivate().end();
+ // decrementing references to mock what the real reporter would do
+ transaction.decrementReferences();
+ // make sure the transaction is still referenced and not recycled yet
+ assertThat(transaction.isReferenced()).isTrue();
+ assertThat(transaction.getTraceContext().getId().isEmpty()).isFalse();
+
+ scheduledTaskFuture.get();
+
+ // make sure the transaction is recycled after the task has terminated
+ assertThat(transaction.isReferenced()).isFalse();
+ assertThat(transaction.getTraceContext().getId().isEmpty()).isTrue();
  }
 }