Add a trampoline for tracing-related methods. (#766)

* Add a trampoline for tracing-related methods.

* ./gradlew googleJavaFormat

* Also initialize the trampoline.

* Improve Javadoc. Minor cleanups.

6 files changed, 272 insertions, 1 deletions

diff --git a/contrib/agent/build.gradle b/contrib/agent/build.gradle
index 308aacfd..2b15c22d 100644
--- a/contrib/agent/build.gradle
+++ b/contrib/agent/build.gradle
@@ -34,6 +34,7 @@ def agentRepackaged = "${agentPackage}.deps"
 dependencies {
   compileOnly libraries.auto_service
   compileOnly libraries.grpc_context
+  compileOnly project(':opencensus-api')
   compile libraries.findbugs_annotations
   compile libraries.guava
   compile libraries.byte_buddy
@@ -155,7 +156,9 @@ configurations {
 }
 
 dependencies {
-  integrationTestCompile libraries.grpc_context
+  integrationTestCompile project(':opencensus-api')
+  integrationTestRuntime libraries.grpc_context
+  integrationTestRuntime project(':opencensus-impl')
 }
 
 // Disable checkstyle for integration tests if not java8.
diff --git a/contrib/agent/src/main/java/io/opencensus/contrib/agent/bootstrap/TraceStrategy.java b/contrib/agent/src/main/java/io/opencensus/contrib/agent/bootstrap/TraceStrategy.java
new file mode 100644
index 00000000..e9029554
--- /dev/null
+++ b/contrib/agent/src/main/java/io/opencensus/contrib/agent/bootstrap/TraceStrategy.java
@@ -0,0 +1,39 @@
+/*
+ * Copyright 2017, OpenCensus Authors
+ *
+ * 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.bootstrap;
+
+import com.google.errorprone.annotations.MustBeClosed;
+import io.opencensus.trace.Span;
+import io.opencensus.trace.SpanBuilder;
+import io.opencensus.trace.Tracer;
+import java.io.Closeable;
+
+/** Strategy interface for creating and manipulating trace spans. */
+public interface TraceStrategy {
+
+  /**
+   * Starts a new span and sets it as the current span.
+   *
+   * @param spanName the name of the returned {@link Span}
+   * @return an object that defines a scope where the newly created {@code Span} will be set to the
+   *     current Context
+   * @see Tracer#spanBuilder(java.lang.String)
+   * @see SpanBuilder#startScopedSpan()
+   */
+  @MustBeClosed
+  Closeable startScopedSpan(String spanName);
+}
diff --git a/contrib/agent/src/main/java/io/opencensus/contrib/agent/bootstrap/TraceTrampoline.java b/contrib/agent/src/main/java/io/opencensus/contrib/agent/bootstrap/TraceTrampoline.java
new file mode 100644
index 00000000..7927ce5f
--- /dev/null
+++ b/contrib/agent/src/main/java/io/opencensus/contrib/agent/bootstrap/TraceTrampoline.java
@@ -0,0 +1,92 @@
+/*
+ * Copyright 2017, OpenCensus Authors
+ *
+ * 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.bootstrap;
+
+import com.google.errorprone.annotations.MustBeClosed;
+import java.io.Closeable;
+
+/**
+ * {@code TraceTrampoline} provides methods for creating and manipulating trace spans from
+ * instrumented bytecode.
+ *
+ * <p>{@code TraceTrampoline} avoids tight coupling with the concrete trace API through the {@link
+ * TraceStrategy} interface.
+ *
+ * <p>Both {@link TraceTrampoline} and {@link TraceStrategy} are loaded by the bootstrap classloader
+ * so that they can be used from classes loaded by the bootstrap classloader. A concrete
+ * implementation of {@link TraceStrategy} will be loaded by the system classloader. This allows for
+ * using the same trace API as the instrumented application.
+ *
+ * <p>{@code TraceTrampoline} is implemented as a static class to allow for easy and fast use from
+ * instrumented bytecode. We cannot use dependency injection for the instrumented bytecode.
+ */
+public final class TraceTrampoline {
+
+  // Not synchronized to avoid any synchronization costs after initialization.
+  // The agent is responsible for initializing this once (through #setTraceStrategy) before any
+  // other method of this class is called.
+  private static TraceStrategy traceStrategy;
+
+  private TraceTrampoline() {}
+
+  /**
+   * Sets the concrete strategy for creating and manipulating trace spans.
+   *
+   * <p>NB: The agent is responsible for setting the trace strategy once before any other method of
+   * this class is called.
+   *
+   * @param traceStrategy the concrete strategy for creating and manipulating trace spans
+   */
+  public static void setTraceStrategy(TraceStrategy traceStrategy) {
+    if (TraceTrampoline.traceStrategy != null) {
+      throw new IllegalStateException("traceStrategy was already set");
+    }
+
+    if (traceStrategy == null) {
+      throw new NullPointerException("traceStrategy");
+    }
+
+    TraceTrampoline.traceStrategy = traceStrategy;
+  }
+
+  /**
+   * Starts a new span and sets it as the current span.
+   *
+   * <p>Enters the scope of code where the newly created {@code Span} is in the current Context, and
+   * returns an object that represents that scope. When the returned object is closed, the scope is
+   * exited, the previous Context is restored, and the newly created {@code Span} is ended using
+   * {@link io.opencensus.trace.Span#end}.
+   *
+   * <p>Callers must eventually close the returned object to avoid leaking the Context.
+   *
+   * <p>Supports the try-with-resource idiom.
+   *
+   * <p>NB: The return type of this method is intentionally {@link Closeable} and not the more
+   * specific {@link io.opencensus.common.Scope} because the latter would not be visible from
+   * classes loaded by the bootstrap classloader.
+   *
+   * @param spanName the name of the returned {@link io.opencensus.trace.Span}
+   * @return an object that defines a scope where the newly created {@code Span} will be set to the
+   *     current Context
+   * @see io.opencensus.trace.Tracer#spanBuilder(String)
+   * @see io.opencensus.trace.SpanBuilder#startScopedSpan()
+   */
+  @MustBeClosed
+  public static Closeable startScopedSpan(String spanName) {
+    return traceStrategy.startScopedSpan(spanName);
+  }
+}
diff --git a/contrib/agent/src/main/java/io/opencensus/contrib/agent/instrumentation/TraceStrategyImpl.java b/contrib/agent/src/main/java/io/opencensus/contrib/agent/instrumentation/TraceStrategyImpl.java
new file mode 100644
index 00000000..4c10b118
--- /dev/null
+++ b/contrib/agent/src/main/java/io/opencensus/contrib/agent/instrumentation/TraceStrategyImpl.java
@@ -0,0 +1,41 @@
+/*
+ * Copyright 2017, OpenCensus Authors
+ *
+ * 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 com.google.errorprone.annotations.MustBeClosed;
+import io.opencensus.contrib.agent.bootstrap.TraceStrategy;
+import io.opencensus.trace.Tracing;
+import io.opencensus.trace.samplers.Samplers;
+import java.io.Closeable;
+
+/** Implementation of {@link TraceStrategy} for creating and manipulating trace spans. */
+final class TraceStrategyImpl implements TraceStrategy {
+
+  @MustBeClosed
+  @Override
+  public Closeable startScopedSpan(String spanName) {
+    checkNotNull(spanName, "spanName");
+
+    return Tracing.getTracer()
+        .spanBuilder(spanName)
+        .setSampler(Samplers.alwaysSample())
+        .setRecordEvents(true)
+        .startScopedSpan();
+  }
+}
diff --git a/contrib/agent/src/main/java/io/opencensus/contrib/agent/instrumentation/TraceTrampolineInitializer.java b/contrib/agent/src/main/java/io/opencensus/contrib/agent/instrumentation/TraceTrampolineInitializer.java
new file mode 100644
index 00000000..4d879458
--- /dev/null
+++ b/contrib/agent/src/main/java/io/opencensus/contrib/agent/instrumentation/TraceTrampolineInitializer.java
@@ -0,0 +1,36 @@
+/*
+ * Copyright 2017, OpenCensus Authors
+ *
+ * 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 com.google.auto.service.AutoService;
+import io.opencensus.contrib.agent.bootstrap.TraceStrategy;
+import io.opencensus.contrib.agent.bootstrap.TraceTrampoline;
+import net.bytebuddy.agent.builder.AgentBuilder;
+
+/** Initializes the {@link TraceTrampoline} with a concrete {@link TraceStrategy}. */
+@AutoService(Instrumenter.class)
+public final class TraceTrampolineInitializer implements Instrumenter {
+
+  @Override
+  public AgentBuilder instrument(AgentBuilder agentBuilder) {
+    // TODO(stschmidt): Gracefully handle the case of missing trace API at runtime,
+    // maybe load the missing classes from a JAR that comes with the agent JAR.
+    TraceTrampoline.setTraceStrategy(new TraceStrategyImpl());
+
+    return agentBuilder;
+  }
+}
diff --git a/contrib/agent/src/test/java/io/opencensus/contrib/agent/bootstrap/TraceTrampolineTest.java b/contrib/agent/src/test/java/io/opencensus/contrib/agent/bootstrap/TraceTrampolineTest.java
new file mode 100644
index 00000000..7519fc1d
--- /dev/null
+++ b/contrib/agent/src/test/java/io/opencensus/contrib/agent/bootstrap/TraceTrampolineTest.java
@@ -0,0 +1,60 @@
+/*
+ * Copyright 2017, OpenCensus Authors
+ *
+ * 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.bootstrap;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Mockito.mock;
+
+import java.io.Closeable;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+import org.junit.runner.RunWith;
+import org.mockito.Mockito;
+import org.mockito.junit.MockitoJUnitRunner;
+
+/** Unit tests for {@link TraceTrampoline}. */
+@RunWith(MockitoJUnitRunner.class)
+public class TraceTrampolineTest {
+
+  private static final TraceStrategy mockTraceStrategy = mock(TraceStrategy.class);
+
+  static {
+    TraceTrampoline.setTraceStrategy(mockTraceStrategy);
+  }
+
+  @Rule public final ExpectedException exception = ExpectedException.none();
+
+  @Test
+  public void setTraceStrategy_already_initialized() {
+    exception.expect(IllegalStateException.class);
+
+    TraceTrampoline.setTraceStrategy(mockTraceStrategy);
+  }
+
+  @Test
+  @SuppressWarnings("MustBeClosedChecker")
+  public void startScopedSpan() {
+    Closeable mockCloseable = mock(Closeable.class);
+    Mockito.when(mockTraceStrategy.startScopedSpan("test")).thenReturn(mockCloseable);
+
+    Closeable closeable = TraceTrampoline.startScopedSpan("test");
+
+    Mockito.verify(mockTraceStrategy).startScopedSpan("test");
+    assertThat(closeable).isSameAs(mockCloseable);
+  }
+}