Add Wrapper for WriteListener. (#1984)

* Add Wrapper for WriteListener.

* add try block after attach context and detach context in finally block.

Author: rghetia

contrib/http_servlet/build.gradle

@@ -10,10 +10,11 @@ description = 'OpenCensus Http Servlet Plugin'
 }
 
 dependencies {
+    compile libraries.grpc_context
     compile project(':opencensus-api')
     compile project(':opencensus-contrib-http-util')
 
-    compile "javax.servlet:javax.servlet-api:3.0.1"
+    compile "javax.servlet:javax.servlet-api:3.1.0"
 
     signature "org.codehaus.mojo.signature:java17:1.0@signature"
     signature "net.sf.androidscents.signature:android-api-level-14:4.0_r4@signature"

contrib/http_servlet/src/main/java/io/opencensus/contrib/http/servlet/WriteListenerWrapper.java

@@ -0,0 +1,70 @@
+/*
+ * Copyright 2019, 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.http.servlet;
+
+import static com.google.common.base.Preconditions.checkNotNull;
+
+import io.grpc.Context;
+import io.opencensus.common.ExperimentalApi;
+import java.io.IOException;
+import javax.servlet.WriteListener;
+
+/**
+ * This class is a wrapper class for {@link WriteListener}. It facilitates executing asynchronous
+ * onWritePossible method in the original context at the time of creating the listener.
+ *
+ * @since 0.25.0
+ */
+@ExperimentalApi
+public class WriteListenerWrapper implements WriteListener {
+  private final io.grpc.Context context;
+  private final WriteListener writeListener;
+
+  /**
+   * Creates an instance of {@code WriteListenerWrapper}. It saves current {@link Context} at the
+   * time of creation.
+   *
+   * @param writeListener {@link WriteListener} object being wrapped.
+   * @since 0.25.0
+   */
+  public WriteListenerWrapper(WriteListener writeListener) {
+    checkNotNull(writeListener, "WriteListener is null");
+    context = io.grpc.Context.current();
+    this.writeListener = writeListener;
+  }
+
+  /**
+   * It executes onWritePossible() method of the object being wrapped in the saved context. It saves
+   * current context before executing the method and restores it after it is finished executing.
+   *
+   * @since 0.25.0
+   */
+  @Override
+  public void onWritePossible() throws IOException {
+    Context previousContext = context.attach();
+    try {
+      writeListener.onWritePossible();
+    } finally {
+      context.detach(previousContext);
+    }
+  }
+
+  @Override
+  public void onError(final Throwable t) {
+    writeListener.onError(t);
+  }
+}

contrib/http_servlet/src/test/java/io/opencensus/contrib/http/servlet/WriteListenerWrapperTest.java

@@ -0,0 +1,71 @@
+/*
+ * Copyright 2019, 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.http.servlet;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import io.grpc.Context;
+import io.grpc.Context.Key;
+import java.io.IOException;
+import javax.servlet.ServletException;
+import javax.servlet.WriteListener;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Unit tests for {@link WriteListenerWrapper}. */
+@RunWith(JUnit4.class)
+public class WriteListenerWrapperTest {
+
+  @Test
+  public void testWriteOnPossibleWithContext() throws IOException, ServletException {
+
+    Key<String> key = Context.<String>key("test-key");
+    Context curr = Context.current();
+    assertThat(curr).isNotNull();
+    final Context parentContext = curr.withValue(key, "parent");
+    assertThat(parentContext).isNotNull();
+
+    Context prev = parentContext.attach();
+    try {
+      WriteListenerWrapper writeListener =
+          new WriteListenerWrapper(
+              new WriteListener() {
+                @Override
+                public void onWritePossible() throws IOException {
+                  Context curr = Context.current();
+                  assertThat(curr).isNotNull();
+                  assertThat(curr).isEqualTo(parentContext);
+                }
+
+                @Override
+                public void onError(Throwable t) {}
+              });
+
+      Context childContext = parentContext.withValue(key, "child");
+      assertThat(childContext).isNotNull();
+      assertThat(childContext.attach()).isNotNull();
+      try {
+        writeListener.onWritePossible();
+      } finally {
+        childContext.detach(parentContext);
+      }
+    } finally {
+      parentContext.detach(prev);
+    }
+  }
+}