Fix RuntimeContext cannot store anything when the trace is not creat (#…

…1480)

* Fix RuntimeContext cannot store anything when the trace is not create

* Add RuntimecontextSnapshot and remove debug info

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/conf/RuntimeContextConfiguration.java

@@ -0,0 +1,28 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.skywalking.apm.agent.core.conf;
+
+public class RuntimeContextConfiguration {
+
+    public static String[] NEED_PROPAGATE_CONTEXT_KEY = new String[] {
+        "SW_REQUEST",
+        "SW_RESPONSE"
+    };
+}

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/AbstractTracerContext.java

@@ -105,8 +105,4 @@ public interface AbstractTracerContext {
      */
     void stopSpan(AbstractSpan span);
 
-    /**
-     * @return the runtime context from current tracing context.
-     */
-    RuntimeContext getRuntimeContext();
 }

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/ContextManager.java

@@ -42,6 +42,7 @@
 public class ContextManager implements TracingContextListener, BootService, IgnoreTracerContextListener {
     private static final ILog logger = LogManager.getLogger(ContextManager.class);
     private static ThreadLocal<AbstractTracerContext> CONTEXT = new ThreadLocal<AbstractTracerContext>();
+    private static ThreadLocal<RuntimeContext> RUNTIME_CONTEXT = new ThreadLocal<RuntimeContext>();
     private static ContextManagerExtendService EXTEND_SERVICE;
 
     private static AbstractTracerContext getOrCreate(String operationName, boolean forceSampling) {
@@ -198,11 +199,12 @@ public static boolean isActive() {
     }
 
     public static RuntimeContext getRuntimeContext() {
-        if (isActive()) {
-            return get().getRuntimeContext();
-        } else {
-            throw new IllegalStateException("No active context");
+        RuntimeContext runtimeContext = RUNTIME_CONTEXT.get();
+        if (runtimeContext == null) {
+            runtimeContext = new RuntimeContext(RUNTIME_CONTEXT);
+            RUNTIME_CONTEXT.set(runtimeContext);
         }
-    }
 
+        return runtimeContext;
+    }
 }

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/IgnoredTracerContext.java

@@ -37,13 +37,6 @@ public class IgnoredTracerContext implements AbstractTracerContext {
 
     private int stackDepth;
 
-    /**
-     * Runtime context of the ignored context
-     *
-     * The context should work even no trace, in order to avoid the unexpected status.
-     */
-    private RuntimeContext runtimeContext;
-
     public IgnoredTracerContext() {
         this.stackDepth = 0;
     }
@@ -102,13 +95,6 @@ public void stopSpan(AbstractSpan span) {
         }
     }
 
-    @Override public RuntimeContext getRuntimeContext() {
-        if (runtimeContext == null) {
-            runtimeContext = new RuntimeContext();
-        }
-        return runtimeContext;
-    }
-
     public static class ListenerManager {
         private static List<IgnoreTracerContextListener> LISTENERS = new LinkedList<IgnoreTracerContextListener>();
 

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/RuntimeContext.java

@@ -18,20 +18,28 @@
 
 package org.apache.skywalking.apm.agent.core.context;
 
+import java.util.HashMap;
+import java.util.Iterator;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
+import org.apache.skywalking.apm.agent.core.conf.RuntimeContextConfiguration;
 
 /**
  * RuntimeContext is alive during the tracing context.
  * It will not be serialized to the collector, and always stays in the same context only.
  *
  * In most cases, it means it only stays in a single thread for context propagation.
  *
- * @author wusheng
+ * @author wusheng, ascrutae
  */
 public class RuntimeContext {
+    private final ThreadLocal<RuntimeContext> contextThreadLocal;
     private Map context = new ConcurrentHashMap(0);
 
+    public RuntimeContext(ThreadLocal<RuntimeContext> contextThreadLocal) {
+        this.contextThreadLocal = contextThreadLocal;
+    }
+
     public void put(Object key, Object value) {
         context.put(key, value);
     }
@@ -43,4 +51,32 @@ public Object get(Object key) {
     public <T> T get(Object key, Class<T> type) {
         return (T)context.get(key);
     }
+
+    public void remove(Object key) {
+        context.remove(key);
+
+        if (context.isEmpty()) {
+            contextThreadLocal.remove();
+        }
+    }
+
+    public RuntimeContextSnapshot capture() {
+        Map runtimeContextMap = new HashMap();
+        for (String key : RuntimeContextConfiguration.NEED_PROPAGATE_CONTEXT_KEY) {
+            Object value = this.get(key);
+            if (value != null) {
+                runtimeContextMap.put(key, value);
+            }
+        }
+
+        return new RuntimeContextSnapshot(runtimeContextMap);
+    }
+
+    public void accept(RuntimeContextSnapshot snapshot) {
+        Iterator<Map.Entry> iterator = snapshot.iterator();
+        while (iterator.hasNext()) {
+            Map.Entry runtimeContextItem = iterator.next();
+            ContextManager.getRuntimeContext().put(runtimeContextItem.getKey(), runtimeContextItem.getValue());
+        }
+    }
 }

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/RuntimeContextSnapshot.java

@@ -0,0 +1,35 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.skywalking.apm.agent.core.context;
+
+import java.util.Iterator;
+import java.util.Map;
+
+public class RuntimeContextSnapshot {
+    private final Map map;
+
+    public RuntimeContextSnapshot(Map map) {
+        this.map = map;
+    }
+
+    public Iterator<Map.Entry> iterator() {
+        return map.entrySet().iterator();
+    }
+}

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/TracingContext.java

@@ -79,11 +79,6 @@ public class TracingContext implements AbstractTracerContext {
      */
     private int spanIdGenerator;
 
-    /**
-     * Runtime context of the tracing context
-     */
-    private RuntimeContext runtimeContext;
-
     /**
      * Initialize all fields with default value.
      */
@@ -424,14 +419,6 @@ public void stopSpan(AbstractSpan span) {
         }
     }
 
-    @Override
-    public RuntimeContext getRuntimeContext() {
-        if (runtimeContext == null) {
-            runtimeContext = new RuntimeContext();
-        }
-        return runtimeContext;
-    }
-
     /**
      * Finish this context, and notify all {@link TracingContextListener}s, managed by {@link
      * TracingContext.ListenerManager}

apm-sniffer/apm-sdk-plugin/hystrix-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/hystrix/v1/HystrixCommandRunInterceptor.java

@@ -18,7 +18,6 @@
 
 package org.apache.skywalking.apm.plugin.hystrix.v1;
 
-import com.netflix.hystrix.HystrixCommand;
 import java.lang.reflect.Method;
 import org.apache.skywalking.apm.agent.core.context.ContextManager;
 import org.apache.skywalking.apm.agent.core.context.ContextSnapshot;
@@ -28,8 +27,6 @@
 import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.MethodInterceptResult;
 import org.apache.skywalking.apm.network.trace.component.ComponentsDefine;
 
-import static org.apache.skywalking.apm.plugin.hystrix.v1.Constants.ISOLATE_STRATEGY_KEY_IN_RUNNING_CONTEXT;
-
 public class HystrixCommandRunInterceptor implements InstanceMethodsAroundInterceptor {
     @Override
     public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
@@ -46,8 +43,6 @@ public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allAr
         }
         // Because of `fall back` method running in other thread. so we need capture concurrent span for tracing.
         enhanceRequireObjectCache.setContextSnapshot(ContextManager.capture());
-
-        ContextManager.getRuntimeContext().put(ISOLATE_STRATEGY_KEY_IN_RUNNING_CONTEXT, ((HystrixCommand)objInst).getProperties().executionIsolationStrategy().get().name().toUpperCase());
     }
 
     @Override

apm-sniffer/apm-sdk-plugin/hystrix-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/hystrix/v1/HystrixConcurrencyStrategyInterceptor.java

@@ -0,0 +1,45 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.skywalking.apm.plugin.hystrix.v1;
+
+import com.netflix.hystrix.strategy.concurrency.HystrixConcurrencyStrategy;
+import java.lang.reflect.Method;
+import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.EnhancedInstance;
+import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.InstanceMethodsAroundInterceptor;
+import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.MethodInterceptResult;
+
+public class HystrixConcurrencyStrategyInterceptor implements InstanceMethodsAroundInterceptor {
+    @Override
+    public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
+        MethodInterceptResult result) throws Throwable {
+
+    }
+
+    @Override
+    public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
+        Object ret) throws Throwable {
+        return new SWHystrixConcurrencyStrategyWrapper((HystrixConcurrencyStrategy)ret);
+    }
+
+    @Override public void handleMethodException(EnhancedInstance objInst, Method method, Object[] allArguments,
+        Class<?>[] argumentsTypes, Throwable t) {
+
+    }
+}

apm-sniffer/apm-sdk-plugin/hystrix-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/hystrix/v1/SWHystrixConcurrencyStrategyWrapper.java

@@ -0,0 +1,64 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.skywalking.apm.plugin.hystrix.v1;
+
+import com.netflix.hystrix.strategy.concurrency.HystrixConcurrencyStrategy;
+import java.util.concurrent.Callable;
+import org.apache.skywalking.apm.agent.core.conf.RuntimeContextConfiguration;
+import org.apache.skywalking.apm.agent.core.context.ContextManager;
+import org.apache.skywalking.apm.agent.core.context.RuntimeContextSnapshot;
+
+public class SWHystrixConcurrencyStrategyWrapper extends HystrixConcurrencyStrategy {
+
+    private final HystrixConcurrencyStrategy delegate;
+
+    public SWHystrixConcurrencyStrategyWrapper(
+        HystrixConcurrencyStrategy delegate) {
+        this.delegate = delegate;
+    }
+
+    @Override
+    public <T> Callable<T> wrapCallable(Callable<T> callable) {
+        return new WrappedCallable<T>(ContextManager.getRuntimeContext().capture(), super.wrapCallable(callable));
+    }
+
+    static class WrappedCallable<T> implements Callable<T> {
+
+        private final RuntimeContextSnapshot contextSnapshot;
+        private final Callable<T> target;
+
+        WrappedCallable(RuntimeContextSnapshot contextSnapshot, Callable<T> target) {
+            this.contextSnapshot = contextSnapshot;
+            this.target = target;
+        }
+
+        @Override public T call() throws Exception {
+            try {
+                ContextManager.getRuntimeContext().accept(contextSnapshot);
+                return target.call();
+            } finally {
+                for (String key : RuntimeContextConfiguration.NEED_PROPAGATE_CONTEXT_KEY) {
+                    ContextManager.getRuntimeContext().remove(key);
+                }
+            }
+        }
+    }
+
+}

apm-sniffer/apm-sdk-plugin/hystrix-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/hystrix/v1/define/HystrixPluginsInstrumentation.java

@@ -32,6 +32,8 @@ public class HystrixPluginsInstrumentation extends ClassInstanceMethodsEnhancePl
 
     public static final String INTERCEPT_CLASS = "org.apache.skywalking.apm.plugin.hystrix.v1.HystrixPluginsInterceptor";
     public static final String ENHANCE_METHOD = "getCommandExecutionHook";
+    public static final String GET_CONCURRENCY_STRATEGY_METHOD = "getConcurrencyStrategy";
+    public static final String GET_CONCURRENCY_STRATEGY_INTERCEPT_CLASS = "org.apache.skywalking.apm.plugin.hystrix.v1.HystrixConcurrencyStrategyInterceptor";
     public static final String ENHANCE_CLASS = "com.netflix.hystrix.strategy.HystrixPlugins";
 
     @Override
@@ -51,6 +53,19 @@ protected InstanceMethodsInterceptPoint[] getInstanceMethodsInterceptPoints() {
                     return INTERCEPT_CLASS;
                 }
 
+                @Override public boolean isOverrideArgs() {
+                    return false;
+                }
+            },
+            new InstanceMethodsInterceptPoint() {
+                @Override public ElementMatcher<MethodDescription> getMethodsMatcher() {
+                    return named(GET_CONCURRENCY_STRATEGY_METHOD);
+                }
+
+                @Override public String getMethodsInterceptor() {
+                    return GET_CONCURRENCY_STRATEGY_INTERCEPT_CLASS;
+                }
+
                 @Override public boolean isOverrideArgs() {
                     return false;
                 }