instrument netty rejected execution handlers

Author: richardstartin

dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/RejectedExecutionHandlerInstrumentation.java

@@ -1,8 +1,10 @@
 package datadog.trace.instrumentation.java.concurrent;
 
 import static datadog.trace.agent.tooling.bytebuddy.matcher.DDElementMatchers.hasInterface;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf;
 import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.cancelTask;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.nameEndsWith;
 import static net.bytebuddy.matcher.ElementMatchers.named;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 
@@ -34,7 +36,8 @@ public ElementMatcher<? super TypeDescription> typeMatcher() {
             "java.util.concurrent.ThreadPoolExecutor$DiscardPolicy",
             "java.util.concurrent.ThreadPoolExecutor$DiscardOldestPolicy",
             "java.util.concurrent.ThreadPoolExecutor$CallerRunsPolicy")
-        .or(hasInterface(named("java.util.concurrent.RejectedExecutionHandler")));
+        .or(hasInterface(named("java.util.concurrent.RejectedExecutionHandler")))
+        .or(hasInterface(nameEndsWith("netty.util.concurrent.RejectedExecutionHandler")));
   }
 
   @Override
@@ -50,9 +53,10 @@ public Map<String, String> contextStore() {
   public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
     return Collections.singletonMap(
         isMethod()
-            .and(named("rejectedExecution"))
-            .and(takesArgument(0, named("java.lang.Runnable")))
-            .and(takesArgument(1, named("java.util.concurrent.ThreadPoolExecutor"))),
+            // JDK or netty
+            .and(namedOneOf("rejectedExecution", "rejected"))
+            // must not constrain or use second parameter
+            .and(takesArgument(0, named("java.lang.Runnable"))),
         getClass().getName() + "$Reject");
   }
 

dd-java-agent/instrumentation/java-concurrent/src/test/groovy/RejectedExecutionTest.groovy

@@ -1,6 +1,8 @@
 import datadog.trace.agent.test.AgentTestRunner
 import datadog.trace.api.DDId
 import datadog.trace.core.DDSpan
+import io.netty.util.concurrent.DefaultEventExecutor
+import io.netty.util.concurrent.DefaultThreadFactory
 
 import java.util.concurrent.ArrayBlockingQueue
 import java.util.concurrent.CountDownLatch
@@ -134,20 +136,7 @@ class RejectedExecutionTest extends AgentTestRunner {
     testClosure()
 
     then:
-    TEST_WRITER.waitForTraces(1)
-    TEST_WRITER.size() == 1
-    List<DDSpan> trace = TEST_WRITER.get(0)
-    trace.size() == 2
-    DDSpan parent = trace.find {
-      it.getOperationName() == "parent"
-    }
-    DDSpan child = trace.find {
-      it.getOperationName() == "asyncChild"
-    }
-    parent != null
-    child != null
-    parent.getParentId() == DDId.ZERO
-    parent.getSpanId() == child.getParentId()
+    expectParentChildTrace()
 
 
     where:
@@ -164,6 +153,54 @@ class RejectedExecutionTest extends AgentTestRunner {
     ]
   }
 
+  def "trace reported with swallowing netty rejected execution handler" () {
+    setup:
+    DefaultEventExecutor executor = new DefaultEventExecutor(null, new DefaultThreadFactory(DefaultEventExecutor),
+    1, handler)
+    CountDownLatch latch = new CountDownLatch(1)
+    // this task will block the executor thread ensuring the next task gets enqueued
+    executor.submit({
+      latch.await()
+    })
+    // this will sit in the queue
+    executor.submit({})
+
+    when:
+    runUnderTrace("parent") {
+      activeScope().setAsyncPropagation(true)
+      // must be rejected because the queue will be full until some
+      // time after the first task is released
+      executor.submit((Runnable) new JavaAsyncChild(true, false))
+      latch.countDown()
+    }
+
+    then:
+    expectParentChildTrace()
+
+    where:
+    handler << [
+      new SwallowingRejectedExecutionHandler()
+    ]
+  }
+
+  def expectParentChildTrace() {
+    TEST_WRITER.waitForTraces(1)
+    TEST_WRITER.size() == 1
+    List<DDSpan> trace = TEST_WRITER.get(0)
+    assert trace.size() == 2
+    DDSpan parent = trace.find {
+      it.getOperationName() == "parent"
+    }
+    DDSpan child = trace.find {
+      it.getOperationName() == "asyncChild"
+    }
+    assert parent != null
+    assert child != null
+    assert parent.getParentId() == DDId.ZERO
+    assert parent.getSpanId() == child.getParentId()
+    return true
+  }
+
   def setupShutdownExecutor(RejectedExecutionHandler rejectedExecutionHandler) {
     ExecutorService pool = new ThreadPoolExecutor(1,
       1,

dd-java-agent/instrumentation/java-concurrent/src/test/java/SwallowingRejectedExecutionHandler.java

@@ -1,7 +1,12 @@
+import io.netty.util.concurrent.SingleThreadEventExecutor;
 import java.util.concurrent.RejectedExecutionHandler;
 import java.util.concurrent.ThreadPoolExecutor;
 
-public class SwallowingRejectedExecutionHandler implements RejectedExecutionHandler {
+public class SwallowingRejectedExecutionHandler
+    implements RejectedExecutionHandler, io.netty.util.concurrent.RejectedExecutionHandler {
   @Override
   public void rejectedExecution(Runnable r, ThreadPoolExecutor executor) {}
+
+  @Override
+  public void rejected(Runnable task, SingleThreadEventExecutor executor) {}
 }