Provide creating scope for async request http-servlet. (#1714)

* Provide creating scope for async request http-servlet.

* create listener and set it as attribute earlier in the request processing.

* fix format issue.

* change since 0.19.1 to 0.20.0

3 files changed, 72 insertions, 4 deletions

diff --git a/contrib/http_servlet/src/main/java/io/opencensus/contrib/http/servlet/OcHttpServletFilter.java b/contrib/http_servlet/src/main/java/io/opencensus/contrib/http/servlet/OcHttpServletFilter.java
index 1e73c501..e8b82428 100644
--- a/contrib/http_servlet/src/main/java/io/opencensus/contrib/http/servlet/OcHttpServletFilter.java
+++ b/contrib/http_servlet/src/main/java/io/opencensus/contrib/http/servlet/OcHttpServletFilter.java
@@ -92,6 +92,9 @@ public class OcHttpServletFilter implements Filter {
       HttpServletResponse httpResp = (HttpServletResponse) response;
 
       HttpRequestContext context = handler.handleStart(httpReq, httpReq);
+      OcHttpServletListener listener = new OcHttpServletListener(handler, context);
+      httpReq.setAttribute(OcHttpServletUtil.OPENCENSUS_SERVLET_LISTENER, listener);
+
       int length = httpReq.getContentLength();
       if (length > 0) {
         handler.handleMessageReceived(context, length);
@@ -105,7 +108,6 @@ public class OcHttpServletFilter implements Filter {
       }
 
       if (httpReq.isAsyncStarted()) {
-        OcHttpServletListener listener = new OcHttpServletListener(handler, context);
         AsyncContext async = httpReq.getAsyncContext();
         async.addListener(listener, httpReq, httpResp);
       } else {
diff --git a/contrib/http_servlet/src/main/java/io/opencensus/contrib/http/servlet/OcHttpServletListener.java b/contrib/http_servlet/src/main/java/io/opencensus/contrib/http/servlet/OcHttpServletListener.java
index 7684bc50..0c128b9a 100644
--- a/contrib/http_servlet/src/main/java/io/opencensus/contrib/http/servlet/OcHttpServletListener.java
+++ b/contrib/http_servlet/src/main/java/io/opencensus/contrib/http/servlet/OcHttpServletListener.java
@@ -18,9 +18,12 @@ package io.opencensus.contrib.http.servlet;
 
 import static com.google.common.base.Preconditions.checkNotNull;
 
+import com.google.errorprone.annotations.MustBeClosed;
 import io.opencensus.common.ExperimentalApi;
+import io.opencensus.common.Scope;
 import io.opencensus.contrib.http.HttpRequestContext;
 import io.opencensus.contrib.http.HttpServerHandler;
+import io.opencensus.trace.Tracing;
 import java.io.Closeable;
 import javax.servlet.AsyncContext;
 import javax.servlet.AsyncEvent;
@@ -74,8 +77,6 @@ public final class OcHttpServletListener implements Closeable, AsyncListener {
   }
 
   @Override
-  @SuppressWarnings(
-      "MustBeClosedChecker") // Close will happen in onTimeout or onError or onComplete method
   public void onStartAsync(AsyncEvent event) {
     AsyncContext eventAsyncContext = event.getAsyncContext();
     if (eventAsyncContext != null) {
@@ -91,4 +92,9 @@ public final class OcHttpServletListener implements Closeable, AsyncListener {
         (HttpServletResponse) event.getSuppliedResponse(),
         null);
   }
+
+  @MustBeClosed
+  Scope withSpan() {
+    return Tracing.getTracer().withSpan(handler.getSpanFromContext(context));
+  }
 }
diff --git a/contrib/http_servlet/src/main/java/io/opencensus/contrib/http/servlet/OcHttpServletUtil.java b/contrib/http_servlet/src/main/java/io/opencensus/contrib/http/servlet/OcHttpServletUtil.java
index 2a1f359b..616b9c27 100644
--- a/contrib/http_servlet/src/main/java/io/opencensus/contrib/http/servlet/OcHttpServletUtil.java
+++ b/contrib/http_servlet/src/main/java/io/opencensus/contrib/http/servlet/OcHttpServletUtil.java
@@ -16,13 +16,21 @@
 
 package io.opencensus.contrib.http.servlet;
 
+import static com.google.common.base.Preconditions.checkNotNull;
+
+import com.google.errorprone.annotations.MustBeClosed;
+import io.opencensus.common.Scope;
 import io.opencensus.contrib.http.HttpRequestContext;
 import io.opencensus.contrib.http.HttpServerHandler;
+import io.opencensus.trace.BlankSpan;
+import io.opencensus.trace.Tracing;
+import javax.servlet.ServletRequest;
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
 
-class OcHttpServletUtil {
+public class OcHttpServletUtil {
   static final String CONTENT_LENGTH = "Content-Length";
+  static final String OPENCENSUS_SERVLET_LISTENER = "opencensus.servlet.listener";
 
   private OcHttpServletUtil() {}
 
@@ -41,4 +49,56 @@ class OcHttpServletUtil {
       }
     }
   }
+
+  /**
+   * Enters the scope of code where the given {@link ServletRequest} will be processed and returns
+   * an object that represents the scope. The scope is exited when the returned object is closed. A
+   * span created for the {@link ServletRequest} is set to the current Context.
+   *
+   * <p>Supports try-with-resource idiom.
+   *
+   * <p>Example of usage:
+   *
+   * <pre>{@code
+   * void AsyncRequestProcessor(AsyncContext asyncCtx) {
+   *   try (Scope ws = OcHttpServletUtil.withScope(asyncCtx.getRequest) {
+   *     tracer.getCurrentSpan().addAnnotation("my annotation");
+   *     doSomeOtherWork();  // Here "span" is the current Span.
+   *   }
+   * }
+   * }</pre>
+   *
+   * <p>Prior to Java SE 7, you can use a finally block to ensure that a resource is closed
+   * regardless of whether the try statement completes normally or abruptly.
+   *
+   * <p>Example of usage prior to Java SE7:
+   *
+   * <pre>{@code
+   * void AsyncRequestProcessor(AsyncContext asyncCtx) {
+   *   Scope ws = OcHttpServletUtil.withScope(asyncCtx.getRequest)
+   *   try {
+   *     tracer.getCurrentSpan().addAnnotation("my annotation");
+   *     doSomeOtherWork();  // Here "span" is the current Span.
+   *   } finally {
+   *     ws.close();
+   *   }
+   * }
+   * }</pre>
+   *
+   * @param request The {@link ServletRequest} request that is about to be processed.
+   * @return an object that defines a scope where the span associated with the given {@link
+   *     ServletRequest} will be set to the current Context.
+   * @throws NullPointerException if {@code request} is {@code null}.
+   * @since 0.20.0
+   */
+  @MustBeClosed
+  public static Scope withScope(ServletRequest request) throws NullPointerException {
+    checkNotNull(request, "request");
+    OcHttpServletListener listener =
+        (OcHttpServletListener) request.getAttribute(OPENCENSUS_SERVLET_LISTENER);
+    if (listener != null) {
+      return listener.withSpan();
+    }
+    return Tracing.getTracer().withSpan(BlankSpan.INSTANCE);
+  }
 }