diff options
7 files changed, 343 insertions, 0 deletions
diff --git a/contrib/agent/README.md b/contrib/agent/README.md index ed0824ac..35dcc9db 100644 --- a/contrib/agent/README.md +++ b/contrib/agent/README.md @@ -20,6 +20,12 @@ The context of the caller of [Executor#execute](https://docs.oracle.com/javase/8 is automatically propagated to the submitted Runnable. +### Automatic context propagation for Threads + +The context of the caller of [Thread#start](https://docs.oracle.com/javase/8/docs/api/java/lang/Thread.html#start--) +is automatically propagated to the new thread. + + ## Design Ideas We see tracing as a cross-cutting concern which the *OpenCensus Agent for Java* weaves into diff --git a/contrib/agent/src/integration-test/java/io/opencensus/contrib/agent/instrumentation/ThreadInstrumentationTest.java b/contrib/agent/src/integration-test/java/io/opencensus/contrib/agent/instrumentation/ThreadInstrumentationTest.java new file mode 100644 index 00000000..f0d105ac --- /dev/null +++ b/contrib/agent/src/integration-test/java/io/opencensus/contrib/agent/instrumentation/ThreadInstrumentationTest.java @@ -0,0 +1,132 @@ +/* + * Copyright 2017, Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * http://www.apache.org/licenses/LICENSE-2.0 + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.opencensus.contrib.agent.instrumentation; + +import static com.google.common.truth.Truth.assertThat; + +import io.grpc.Context; +import java.util.concurrent.Executor; +import java.util.concurrent.atomic.AtomicBoolean; +import org.junit.After; +import org.junit.Test; + +/** + * Integration tests for {@link ThreadInstrumentation}. + * + * <p>The integration tests are executed in a separate JVM that has the OpenCensus agent enabled + * via the {@code -javaagent} command line option. + */ +public class ThreadInstrumentationTest { + + private static final Context.Key<String> KEY = Context.key("mykey"); + + private Context previousContext; + + @After + public void afterMethod() { + Context.current().detach(previousContext); + } + + @Test(timeout = 5000) + public void start_Runnable() throws Exception { + final Context context = Context.current().withValue(KEY, "myvalue"); + previousContext = context.attach(); + + final AtomicBoolean tested = new AtomicBoolean(false); + + Thread thread = new Thread(new Runnable() { + @Override + public void run() { + assertThat(Context.current()).isSameAs(context); + assertThat(KEY.get()).isEqualTo("myvalue"); + tested.set(true); + } + }); + + thread.start(); + thread.join(); + + assertThat(tested.get()).isTrue(); + } + + @Test(timeout = 5000) + public void start_Subclass() throws Exception { + final Context context = Context.current().withValue(KEY, "myvalue"); + previousContext = context.attach(); + + final AtomicBoolean tested = new AtomicBoolean(false); + + class MyThread extends Thread { + + @Override + public void run() { + assertThat(Context.current()).isSameAs(context); + assertThat(KEY.get()).isEqualTo("myvalue"); + tested.set(true); + } + } + + Thread thread = new MyThread(); + + thread.start(); + thread.join(); + + assertThat(tested.get()).isTrue(); + } + + /** + * Tests that the automatic context propagation added by {@link ThreadInstrumentation} does not + * interfere with the automatically propagated context from Executor#execute. + */ + @Test(timeout = 5000) + public void start_automaticallyWrappedRunnable() throws Exception { + final Context context = Context.current().withValue(KEY, "myvalue"); + previousContext = context.attach(); + + final AtomicBoolean tested = new AtomicBoolean(false); + + Executor newThreadExecutor = new Executor() { + @Override + public void execute(Runnable command) { + // Attach a new context before starting a new thread. This new context will be propagated to + // the new thread as in #start_Runnable. However, since the Runnable has been wrapped in a + // different context (by automatic instrumentation of Executor#execute), that context will + // be attached when executing the Runnable. + Context context2 = Context.current().withValue(KEY, "wrong context"); + context2.attach(); + Thread thread = new Thread(command); + thread.start(); + try { + thread.join(); + } catch (InterruptedException ex) { + Thread.currentThread().interrupt(); + } + context2.detach(context); + } + }; + + newThreadExecutor.execute(new Runnable() { + @Override + public void run() { + // Assert that the automatic context propagation added by ThreadInstrumentation did not + // interfere with the automatically propagated context from Executor#execute. + assertThat(Context.current()).isSameAs(context); + assertThat(KEY.get()).isEqualTo("myvalue"); + tested.set(true); + } + }); + + assertThat(tested.get()).isTrue(); + } +} diff --git a/contrib/agent/src/main/java/io/opencensus/contrib/agent/bootstrap/ContextManager.java b/contrib/agent/src/main/java/io/opencensus/contrib/agent/bootstrap/ContextManager.java index 89a2f025..dccaf35e 100644 --- a/contrib/agent/src/main/java/io/opencensus/contrib/agent/bootstrap/ContextManager.java +++ b/contrib/agent/src/main/java/io/opencensus/contrib/agent/bootstrap/ContextManager.java @@ -70,4 +70,25 @@ public final class ContextManager { public static Runnable wrapInCurrentContext(Runnable runnable) { return contextStrategy.wrapInCurrentContext(runnable); } + + /** + * Saves the context that is associated with the current scope. + * + * <p>The context will be attached when entering the specified thread's {@link Thread#run()} + * method. + * + * @param thread a {@link Thread} object + */ + public static void saveContextForThread(Thread thread) { + contextStrategy.saveContextForThread(thread); + } + + /** + * Attaches the context that was previously saved for the specified thread. + * + * @param thread a {@link Thread} object + */ + public static void attachContextForThread(Thread thread) { + contextStrategy.attachContextForThread(thread); + } } diff --git a/contrib/agent/src/main/java/io/opencensus/contrib/agent/bootstrap/ContextStrategy.java b/contrib/agent/src/main/java/io/opencensus/contrib/agent/bootstrap/ContextStrategy.java index 1ab5154c..1d1cf738 100644 --- a/contrib/agent/src/main/java/io/opencensus/contrib/agent/bootstrap/ContextStrategy.java +++ b/contrib/agent/src/main/java/io/opencensus/contrib/agent/bootstrap/ContextStrategy.java @@ -26,4 +26,21 @@ public interface ContextStrategy { * @return the wrapped {@link Runnable} object */ Runnable wrapInCurrentContext(Runnable runnable); + + /** + * Saves the context that is associated with the current scope. + * + * <p>The context will be attached when entering the specified thread's {@link Thread#run()} + * method. + * + * @param thread a {@link Thread} object + */ + void saveContextForThread(Thread thread); + + /** + * Attaches the context that was previously saved for the specified thread. + * + * @param thread a {@link Thread} object + */ + void attachContextForThread(Thread thread); } diff --git a/contrib/agent/src/main/java/io/opencensus/contrib/agent/instrumentation/ContextStrategyImpl.java b/contrib/agent/src/main/java/io/opencensus/contrib/agent/instrumentation/ContextStrategyImpl.java index 616d3d43..534d0352 100644 --- a/contrib/agent/src/main/java/io/opencensus/contrib/agent/instrumentation/ContextStrategyImpl.java +++ b/contrib/agent/src/main/java/io/opencensus/contrib/agent/instrumentation/ContextStrategyImpl.java @@ -13,8 +13,11 @@ package io.opencensus.contrib.agent.instrumentation; +import com.google.common.cache.Cache; +import com.google.common.cache.CacheBuilder; import io.grpc.Context; import io.opencensus.contrib.agent.bootstrap.ContextStrategy; +import java.lang.ref.WeakReference; /** * Implementation of {@link ContextStrategy} for accessing and manipulating the @@ -22,8 +25,48 @@ import io.opencensus.contrib.agent.bootstrap.ContextStrategy; */ final class ContextStrategyImpl implements ContextStrategy { + /** + * Thread-safe mapping of {@link Thread}s to {@link Context}s, used for tunneling the caller's + * {@link Context} of {@link Thread#start()} to {@link Thread#run()}. + * + * <p>A thread is inserted into this map when {@link Thread#start()} is called, and removed when + * {@link Thread#run()} is called. + * + * <p>NB: {@link Thread#run()} is not guaranteed to be called after {@link Thread#start()}, for + * example when attempting to start a thread a second time. Therefore, threads are wrapped in + * {@link WeakReference}s so that this map does not prevent the garbage collection of otherwise + * unreferenced threads. Unreferenced threads will be automatically removed from the map by the + * routine cleanup of the underlying {@link Cache} implementation. + * + * <p>NB: A side-effect of {@link CacheBuilder#weakKeys()} is the use of identity ({@code ==}) + * comparison to determine equality of threads. Identity comparison is required here because + * subclasses of {@link Thread} might override {@link Object#hashCode()} and {@link + * Object#equals(java.lang.Object)} with potentially broken implementations. + * + * <p>NB: Using thread IDs as keys was considered: It's unclear how to safely detect and cleanup + * otherwise unreferenced threads IDs from the map. + */ + private final Cache<Thread, Context> savedContexts + = CacheBuilder.newBuilder().weakKeys().build(); + @Override public Runnable wrapInCurrentContext(Runnable runnable) { return Context.current().wrap(runnable); } + + @Override + public void saveContextForThread(Thread thread) { + savedContexts.put(thread, Context.current()); + } + + @Override + public void attachContextForThread(Thread thread) { + if (Thread.currentThread() == thread) { + Context context = savedContexts.getIfPresent(thread); + if (context != null) { + savedContexts.invalidate(thread); + context.attach(); + } + } + } } diff --git a/contrib/agent/src/main/java/io/opencensus/contrib/agent/instrumentation/ThreadInstrumentation.java b/contrib/agent/src/main/java/io/opencensus/contrib/agent/instrumentation/ThreadInstrumentation.java new file mode 100644 index 00000000..0b6b70ab --- /dev/null +++ b/contrib/agent/src/main/java/io/opencensus/contrib/agent/instrumentation/ThreadInstrumentation.java @@ -0,0 +1,107 @@ +/* + * Copyright 2017, Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * http://www.apache.org/licenses/LICENSE-2.0 + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.opencensus.contrib.agent.instrumentation; + +import static com.google.common.base.Preconditions.checkNotNull; +import static net.bytebuddy.matcher.ElementMatchers.isSubTypeOf; +import static net.bytebuddy.matcher.ElementMatchers.named; + +import com.google.auto.service.AutoService; +import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; +import io.opencensus.contrib.agent.bootstrap.ContextManager; +import net.bytebuddy.agent.builder.AgentBuilder; +import net.bytebuddy.asm.Advice; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.dynamic.DynamicType; +import net.bytebuddy.matcher.ElementMatcher; +import net.bytebuddy.utility.JavaModule; + +/** + * Propagates the context of the caller of {@link Thread#start} to the new thread, just like the + * Microsoft .Net Framework propagates the <a + * href="https://msdn.microsoft.com/en-us/library/system.threading.executioncontext(v=vs.110).aspx">System.Threading.ExecutionContext</a>. + * + * <p>NB: A similar effect could be achieved with {@link InheritableThreadLocal}, but the semantics + * are different: {@link InheritableThreadLocal} inherits values when the thread object is + * initialized as opposed to when {@link Thread#start()} is called. + */ +@AutoService(Instrumenter.class) +public final class ThreadInstrumentation implements Instrumenter { + + @Override + public AgentBuilder instrument(AgentBuilder agentBuilder) { + checkNotNull(agentBuilder, "agentBuilder"); + + return agentBuilder + .type(createMatcher()) + .transform(createTransformer()); + } + + private static class Transformer implements AgentBuilder.Transformer { + + @Override + public DynamicType.Builder<?> transform(DynamicType.Builder<?> builder, + TypeDescription typeDescription, ClassLoader classLoader, JavaModule module) { + return builder + .visit(Advice.to(Start.class).on(named("start"))) + .visit(Advice.to(Run.class).on(named("run"))); + } + } + + private static ElementMatcher.Junction<TypeDescription> createMatcher() { + // TODO(stschmidt): Exclude known call sites that already propagate the context. + + return isSubTypeOf(Thread.class); + } + + private static AgentBuilder.Transformer createTransformer() { + return new Transformer(); + } + + private static class Start { + + /** + * Saves the context that is associated with the current scope. + * + * <p>The context will be attached when entering the thread's {@link Thread#run()} method. + * + * <p>NB: This method is never called as is. Instead, Byte Buddy copies the method's bytecode + * into Thread#start. + * + * @see Advice + */ + @Advice.OnMethodEnter + @SuppressFBWarnings("UPM_UNCALLED_PRIVATE_METHOD") + private static void enter(@Advice.This Thread thread) { + ContextManager.saveContextForThread(thread); + } + } + + private static class Run { + + /** + * Attaches the context that was previously saved for this thread. + * + * <p>NB: This method is never called as is. Instead, Byte Buddy copies the method's bytecode + * into Thread#run. + * + * @see Advice + */ + @Advice.OnMethodEnter + @SuppressFBWarnings("UPM_UNCALLED_PRIVATE_METHOD") + private static void enter(@Advice.This Thread thread) { + ContextManager.attachContextForThread(thread); + } + } +} diff --git a/contrib/agent/src/test/java/io/opencensus/contrib/agent/bootstrap/ContextManagerTest.java b/contrib/agent/src/test/java/io/opencensus/contrib/agent/bootstrap/ContextManagerTest.java index c072392f..cdf96b52 100644 --- a/contrib/agent/src/test/java/io/opencensus/contrib/agent/bootstrap/ContextManagerTest.java +++ b/contrib/agent/src/test/java/io/opencensus/contrib/agent/bootstrap/ContextManagerTest.java @@ -42,6 +42,9 @@ public class ContextManagerTest { @Mock private Runnable runnable; + @Mock + private Thread thread; + @Test public void setContextStrategy_already_initialized() { exception.expect(IllegalStateException.class); @@ -55,4 +58,18 @@ public class ContextManagerTest { Mockito.verify(mockContextStrategy).wrapInCurrentContext(runnable); } + + @Test + public void saveContextForThread() { + ContextManager.saveContextForThread(thread); + + Mockito.verify(mockContextStrategy).saveContextForThread(thread); + } + + @Test + public void attachContextForThread() { + ContextManager.attachContextForThread(thread); + + Mockito.verify(mockContextStrategy).attachContextForThread(thread); + } } |