Add unsetCurrentContext to ContextPropagator interface (#898)

* Add unsetCurrentContext to ContextPropagator interface

* add unit test for unset

* fix order in unset context in workflow thread

Author: shijiesheng

src/main/java/com/uber/cadence/context/ContextPropagator.java

@@ -136,4 +136,7 @@ public interface ContextPropagator {
 
   /** Sets the current context */
   void setCurrentContext(Object context);
+
+  /** Unsets the current context. This is called when the context is no longer needed */
+  void unsetCurrentContext();
 }

src/main/java/com/uber/cadence/internal/context/ContextThreadLocal.java

@@ -67,4 +67,10 @@ public static void propagateContextToCurrentThread(Map<String, Object> contextDa
       }
     }
   }
+
+  public static void unsetCurrentContext() {
+    for (ContextPropagator propagator : contextPropagators.get()) {
+      propagator.unsetCurrentContext();
+    }
+  }
 }

src/main/java/com/uber/cadence/internal/sync/WorkflowThreadImpl.java

@@ -132,6 +132,7 @@ public void run() {
         }
         threadContext.setUnhandledException(e);
       } finally {
+        ContextThreadLocal.unsetCurrentContext();
         DeterministicRunnerImpl.setCurrentThreadInternal(null);
         threadContext.setStatus(Status.DONE);
         thread.setName(originalName);

src/main/java/com/uber/cadence/internal/worker/ActivityWorker.java

@@ -174,6 +174,7 @@ public void handle(PollForActivityTaskResponse task) throws Exception {
         MDC.remove(LoggerTag.ACTIVITY_TYPE);
         MDC.remove(LoggerTag.WORKFLOW_ID);
         MDC.remove(LoggerTag.RUN_ID);
+        unsetCurrentContext();
       }
     }
 
@@ -200,6 +201,12 @@ void propagateContext(PollForActivityTaskResponse response) {
       }
     }
 
+    void unsetCurrentContext() {
+      for (ContextPropagator propagator : options.getContextPropagators()) {
+        propagator.unsetCurrentContext();
+      }
+    }
+
     @Override
     public Throwable wrapFailure(PollForActivityTaskResponse task, Throwable failure) {
       WorkflowExecution execution = task.getWorkflowExecution();

src/main/java/com/uber/cadence/internal/worker/LocalActivityWorker.java

@@ -22,6 +22,7 @@
 import com.uber.cadence.MarkerRecordedEventAttributes;
 import com.uber.cadence.PollForActivityTaskResponse;
 import com.uber.cadence.common.RetryOptions;
+import com.uber.cadence.context.ContextPropagator;
 import com.uber.cadence.internal.common.LocalActivityMarkerData;
 import com.uber.cadence.internal.metrics.MetricsTag;
 import com.uber.cadence.internal.metrics.MetricsType;
@@ -168,6 +169,7 @@ public void handle(Task task) throws Exception {
         task.eventConsumer.accept(event);
       } finally {
         span.finish();
+        unsetCurrentContext();
       }
     }
 
@@ -241,6 +243,10 @@ private void propagateContext(ExecuteLocalActivityParameters params) {
         .ifPresent(this::restoreContext);
   }
 
+  private void unsetCurrentContext() {
+    options.getContextPropagators().forEach(ContextPropagator::unsetCurrentContext);
+  }
+
   private void restoreContext(Map<String, byte[]> context) {
     options
         .getContextPropagators()

src/test/java/com/uber/cadence/activity/LocalActivityContextPropagationTest.java

@@ -218,6 +218,9 @@ public void setCurrentContext(Object context) {
       String propagatedContextName = ((Context) context).getContextName();
       wrapperContext.newContext(propagatedContextName);
     }
+
+    @Override
+    public void unsetCurrentContext() {}
   }
 
   private class WrapperContext {

src/test/java/com/uber/cadence/internal/testing/WorkflowTestingTest.java

@@ -17,9 +17,7 @@
 
 package com.uber.cadence.internal.testing;
 
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertTrue;
-import static org.junit.Assert.fail;
+import static org.junit.Assert.*;
 import static org.mockito.Matchers.anyString;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
@@ -45,21 +43,10 @@
 import com.uber.cadence.testing.TestEnvironmentOptions;
 import com.uber.cadence.testing.TestWorkflowEnvironment;
 import com.uber.cadence.worker.Worker;
-import com.uber.cadence.workflow.ActivityTimeoutException;
-import com.uber.cadence.workflow.Async;
-import com.uber.cadence.workflow.ChildWorkflowOptions;
-import com.uber.cadence.workflow.ChildWorkflowTimedOutException;
-import com.uber.cadence.workflow.Promise;
-import com.uber.cadence.workflow.SignalMethod;
-import com.uber.cadence.workflow.Workflow;
-import com.uber.cadence.workflow.WorkflowMethod;
+import com.uber.cadence.workflow.*;
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
-import java.util.Collections;
-import java.util.List;
-import java.util.Map;
-import java.util.Optional;
-import java.util.UUID;
+import java.util.*;
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
@@ -742,6 +729,37 @@ public void testMockedChildSimulatedTimeout() {
 
   public static class TestContextPropagator implements ContextPropagator {
 
+    //    public static class TestContext {
+    //        private final String testKey;
+    //        private final Integer level;
+    //
+    //        public TestContext(String testKey, Integer level) {
+    //            this.testKey = testKey;
+    //            this.level = level;
+    //        }
+    //
+    //        public String getTestKey() {
+    //            return testKey;
+    //        }
+    //        public Integer getLevel() {
+    //            return level;
+    //        }
+    //
+    //        public Map<String, byte[]> ToMap() {
+    //            return "TestContext{testKey=" + testKey + ", level=" + level + "}";
+    //        }
+    //    }
+
+    private int level;
+
+    TestContextPropagator() {
+      level = 0;
+    }
+
+    public int getLevel() {
+      return level;
+    }
+
     @Override
     public String getName() {
       return this.getClass().getName();
@@ -774,6 +792,13 @@ public Object getCurrentContext() {
     @Override
     public void setCurrentContext(Object context) {
       MDC.put("test", String.valueOf(context));
+      level++;
+    }
+
+    @Override
+    public void unsetCurrentContext() {
+      MDC.remove("test");
+      level--;
     }
   }
 
@@ -793,13 +818,15 @@ public void testWorkflowContextPropagation() {
     testEnvironment.start();
     MDC.put("test", "testing123");
     WorkflowClient client = testEnvironment.newWorkflowClient();
+    TestContextPropagator propagator = new TestContextPropagator();
     WorkflowOptions options =
         new WorkflowOptions.Builder()
-            .setContextPropagators(Collections.singletonList(new TestContextPropagator()))
+            .setContextPropagators(Collections.singletonList(propagator))
             .build();
     TestWorkflow workflow = client.newWorkflowStub(TestWorkflow.class, options);
     String result = workflow.workflow1("input1");
     assertEquals("testing123", result);
+    assertEquals(0, propagator.getLevel());
   }
 
   public static class ContextPropagationParentWorkflowImpl implements ParentWorkflow {