Support access the tracer context in spring gateway filter (#539)

Co-authored-by: darknesstm <[email protected]>

.github/workflows/plugins-test.3.yaml

@@ -90,6 +90,7 @@ jobs:
           - dbcp-2.x-scenario
           - jsonrpc4j-1.x-scenario
           - gateway-3.x-scenario
+          - gateway-3.x-filter-context-scenario
           - neo4j-4.x-scenario
           - oracle-scenario
           - druid-1.x-scenario

CHANGES.md

@@ -133,6 +133,7 @@ Callable {
 
 * Support Jdk17 ZGC metric collect
 * Support Jetty 11.x plugin
+* Support access the sky-walking tracer context in spring gateway filter
 * Fix the scenario of using the HBase plugin with spring-data-hadoop.
 * Add RocketMQ 5.x plugin
 * Fix the conflict between the logging kernel and the JDK threadpool plugin.

apm-application-toolkit/apm-toolkit-webflux/src/main/java/org/apache/skywalking/apm/toolkit/webflux/WebFluxSkyWalkingTraceContext.java

@@ -0,0 +1,79 @@
+/*
+ * 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.toolkit.webflux;
+
+import org.springframework.web.server.ServerWebExchange;
+
+import java.util.Optional;
+
+/**
+ * TraceContext for WebFlux.
+ */
+public class WebFluxSkyWalkingTraceContext {
+     /**
+     * Try to get the traceId of current trace context.
+     *
+     * @param serverWebExchange - EnhancedInstance that contains the tracing context
+     * @return traceId, if it exists, or empty {@link String}.
+     */
+    public static String traceId(ServerWebExchange serverWebExchange) {
+        return "";
+    }
+
+    /**
+     * Try to get the segmentId of current trace context.
+     *
+     * @param serverWebExchange - EnhancedInstance that contains the tracing context
+     * @return segmentId, if it exists, or empty {@link String}.
+     */
+    public static String segmentId(ServerWebExchange serverWebExchange) {
+        return "";
+    }
+
+    /**
+     * Try to get the spanId of current trace context. The spanId is a negative number when the trace context is
+     * missing.
+     *
+     * @param serverWebExchange - EnhancedInstance that contains the tracing context
+     * @return spanId, if it exists, or empty {@link String}.
+     */
+    public static int spanId(ServerWebExchange serverWebExchange) {
+        return -1;
+    }
+
+    /**
+     * Try to get the custom value from trace context.
+     *
+     * @param serverWebExchange - EnhancedInstance that contains the tracing context
+     * @return custom data value.
+     */
+    public static Optional<String> getCorrelation(ServerWebExchange serverWebExchange, String key) {
+        return Optional.empty();
+    }
+
+    /**
+     * Put the custom key/value into trace context.
+     *
+     * @param serverWebExchange - EnhancedInstance that contains the tracing context
+     * @return previous value if it exists.
+     */
+    public static Optional<String> putCorrelation(ServerWebExchange serverWebExchange, String key, String value) {
+        return Optional.empty();
+    }
+}

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

@@ -19,11 +19,11 @@
 
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.stream.Collectors;
 import org.apache.skywalking.apm.agent.core.base64.Base64;
 import org.apache.skywalking.apm.agent.core.conf.Config;
@@ -49,7 +49,7 @@ public class CorrelationContext {
     }
 
     public CorrelationContext() {
-        this.data = new HashMap<>(Config.Correlation.ELEMENT_MAX_NUMBER);
+        this.data = new ConcurrentHashMap<>(Config.Correlation.ELEMENT_MAX_NUMBER);
     }
 
     /**

apm-sniffer/apm-toolkit-activation/apm-toolkit-webflux-activation/src/main/java/org/apache/skywalking/apm/toolkit/activation/webflux/WebFluxSkyWalkingCorrelationContextGetInterceptor.java

@@ -0,0 +1,49 @@
+/*
+ * 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.toolkit.activation.webflux;
+
+import org.apache.skywalking.apm.agent.core.context.ContextSnapshot;
+import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.MethodInterceptResult;
+
+import java.lang.reflect.Method;
+import java.util.Optional;
+
+public class WebFluxSkyWalkingCorrelationContextGetInterceptor extends WebFluxSkyWalkingStaticMethodsAroundInterceptor {
+
+    @Override
+    public void beforeMethod(Class clazz, Method method, Object[] allArguments, Class<?>[] parameterTypes, MethodInterceptResult result) {
+        final ContextSnapshot contextSnapshot = getContextSnapshot(allArguments[0]);
+        if (contextSnapshot == null || contextSnapshot.getCorrelationContext() == null) {
+            return;
+        }
+
+        final String key = (String) allArguments[1];
+        final Optional<String> data = contextSnapshot.getCorrelationContext().get(key);
+
+        result.defineReturnValue(data);
+    }
+
+    @Override
+    public Object afterMethod(Class clazz, Method method, Object[] allArguments, Class<?>[] parameterTypes, Object ret) {
+        return ret;
+    }
+
+    @Override
+    public void handleMethodException(Class clazz, Method method, Object[] allArguments, Class<?>[] parameterTypes, Throwable t) {
+    }
+}

apm-sniffer/apm-toolkit-activation/apm-toolkit-webflux-activation/src/main/java/org/apache/skywalking/apm/toolkit/activation/webflux/WebFluxSkyWalkingCorrelationContextPutInterceptor.java

@@ -0,0 +1,50 @@
+/*
+ * 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.toolkit.activation.webflux;
+
+import org.apache.skywalking.apm.agent.core.context.ContextSnapshot;
+import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.MethodInterceptResult;
+
+import java.lang.reflect.Method;
+import java.util.Optional;
+
+public class WebFluxSkyWalkingCorrelationContextPutInterceptor extends WebFluxSkyWalkingStaticMethodsAroundInterceptor {
+
+    @Override
+    public void beforeMethod(Class clazz, Method method, Object[] allArguments, Class<?>[] parameterTypes, MethodInterceptResult result) {
+        final ContextSnapshot contextSnapshot = getContextSnapshot(allArguments[0]);
+        if (contextSnapshot == null || contextSnapshot.getCorrelationContext() == null) {
+            return;
+        }
+
+        final String key = (String) allArguments[1];
+        final String value = (String) allArguments[2];
+        final Optional<String> previous = contextSnapshot.getCorrelationContext().put(key, value);
+
+        result.defineReturnValue(previous);
+    }
+
+    @Override
+    public Object afterMethod(Class clazz, Method method, Object[] allArguments, Class<?>[] parameterTypes, Object ret) {
+        return ret;
+    }
+
+    @Override
+    public void handleMethodException(Class clazz, Method method, Object[] allArguments, Class<?>[] parameterTypes, Throwable t) {
+    }
+}

apm-sniffer/apm-toolkit-activation/apm-toolkit-webflux-activation/src/main/java/org/apache/skywalking/apm/toolkit/activation/webflux/WebFluxSkyWalkingOperatorsInterceptor.java

@@ -22,18 +22,15 @@
 import org.apache.skywalking.apm.agent.core.context.ContextSnapshot;
 import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.EnhancedInstance;
 import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.MethodInterceptResult;
-import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.StaticMethodsAroundInterceptor;
 import org.apache.skywalking.apm.network.trace.component.ComponentsDefine;
 import org.springframework.web.server.ServerWebExchange;
-import org.springframework.web.server.ServerWebExchangeDecorator;
-import org.springframework.web.server.adapter.DefaultServerWebExchange;
 import reactor.util.context.Context;
 
 import java.lang.reflect.Method;
 
 /**
  */
-public class WebFluxSkyWalkingOperatorsInterceptor implements StaticMethodsAroundInterceptor {
+public class WebFluxSkyWalkingOperatorsInterceptor extends WebFluxSkyWalkingStaticMethodsAroundInterceptor {
 
     @Override
     public void beforeMethod(Class clazz, Method method, Object[] allArguments, Class<?>[] parameterTypes,
@@ -67,14 +64,4 @@ public void handleMethodException(Class clazz, Method method, Object[] allArgume
         ContextManager.activeSpan().log(t);
     }
 
-    private static EnhancedInstance getInstance(Object o) {
-        EnhancedInstance instance = null;
-        if (o instanceof DefaultServerWebExchange && o instanceof EnhancedInstance) {
-            instance = (EnhancedInstance) o;
-        } else if (o instanceof ServerWebExchangeDecorator) {
-            ServerWebExchange delegate = ((ServerWebExchangeDecorator) o).getDelegate();
-            return getInstance(delegate);
-        }
-        return instance;
-    }
 }

apm-sniffer/apm-toolkit-activation/apm-toolkit-webflux-activation/src/main/java/org/apache/skywalking/apm/toolkit/activation/webflux/WebFluxSkyWalkingSegmentIDInterceptor.java

@@ -0,0 +1,53 @@
+/*
+ * 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.toolkit.activation.webflux;
+
+import org.apache.skywalking.apm.agent.core.context.ContextSnapshot;
+import org.apache.skywalking.apm.agent.core.logging.api.ILog;
+import org.apache.skywalking.apm.agent.core.logging.api.LogManager;
+import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.MethodInterceptResult;
+
+import java.lang.reflect.Method;
+
+public class WebFluxSkyWalkingSegmentIDInterceptor extends WebFluxSkyWalkingStaticMethodsAroundInterceptor {
+
+    private static final ILog LOGGER = LogManager.getLogger(WebFluxSkyWalkingSegmentIDInterceptor.class);
+
+    @Override
+    public void beforeMethod(Class clazz, Method method, Object[] allArguments, Class<?>[] parameterTypes,
+                             MethodInterceptResult result) {
+        final ContextSnapshot contextSnapshot = getContextSnapshot(allArguments[0]);
+        if (contextSnapshot == null || contextSnapshot.getCorrelationContext() == null) {
+            return;
+        }
+        result.defineReturnValue(contextSnapshot.getTraceSegmentId());
+    }
+
+    @Override
+    public Object afterMethod(Class clazz, Method method, Object[] allArguments, Class<?>[] parameterTypes,
+                              Object ret) {
+        return ret;
+    }
+
+    @Override
+    public void handleMethodException(Class clazz, Method method, Object[] allArguments, Class<?>[] parameterTypes,
+                                      Throwable t) {
+        LOGGER.error("Failed to get segment Id.", t);
+    }
+}

apm-sniffer/apm-toolkit-activation/apm-toolkit-webflux-activation/src/main/java/org/apache/skywalking/apm/toolkit/activation/webflux/WebFluxSkyWalkingSpanIDInterceptor.java

@@ -0,0 +1,53 @@
+/*
+ * 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.toolkit.activation.webflux;
+
+import org.apache.skywalking.apm.agent.core.context.ContextSnapshot;
+import org.apache.skywalking.apm.agent.core.logging.api.ILog;
+import org.apache.skywalking.apm.agent.core.logging.api.LogManager;
+import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.MethodInterceptResult;
+
+import java.lang.reflect.Method;
+
+public class WebFluxSkyWalkingSpanIDInterceptor extends WebFluxSkyWalkingStaticMethodsAroundInterceptor {
+
+    private static final ILog LOGGER = LogManager.getLogger(WebFluxSkyWalkingSpanIDInterceptor.class);
+
+    @Override
+    public void beforeMethod(Class clazz, Method method, Object[] allArguments, Class<?>[] parameterTypes,
+                             MethodInterceptResult result) {
+        final ContextSnapshot contextSnapshot = getContextSnapshot(allArguments[0]);
+        if (contextSnapshot == null || contextSnapshot.getCorrelationContext() == null) {
+            return;
+        }
+        result.defineReturnValue(contextSnapshot.getSpanId());
+    }
+
+    @Override
+    public Object afterMethod(Class clazz, Method method, Object[] allArguments, Class<?>[] parameterTypes,
+                              Object ret) {
+        return ret;
+    }
+
+    @Override
+    public void handleMethodException(Class clazz, Method method, Object[] allArguments, Class<?>[] parameterTypes,
+                                      Throwable t) {
+        LOGGER.error("Failed to getDefault span Id.", t);
+    }
+}