From 29dbf7ecd9d3dda16c77fc604847642a6a3f8fe4 Mon Sep 17 00:00:00 2001 From: wuwen Date: Wed, 23 Jun 2021 11:18:16 +0800 Subject: [PATCH] Spring MVC customize traced requests handler add postHandle method. --- .../instrument/web/mvc/HandlerParser.java | 16 ++ ...panCustomizingAsyncHandlerInterceptor.java | 10 + .../SpanCustomizingHandlerInterceptor.java | 4 + ...erParserPostHandleWebIntegrationTests.java | 218 ++++++++++++++++++ 4 files changed, 248 insertions(+) create mode 100644 tests/brave/spring-cloud-sleuth-instrumentation-mvc-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/web/TraceCustomHandlerParserPostHandleWebIntegrationTests.java diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/mvc/HandlerParser.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/mvc/HandlerParser.java index 59e3cdbd3..dc2623359 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/mvc/HandlerParser.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/mvc/HandlerParser.java @@ -23,6 +23,7 @@ import org.springframework.cloud.sleuth.SpanCustomizer; import org.springframework.cloud.sleuth.docs.AssertingSpanCustomizer; import org.springframework.web.method.HandlerMethod; import org.springframework.web.servlet.HandlerInterceptor; +import org.springframework.web.servlet.ModelAndView; /** * Spring MVC specific type used to customize traced requests based on the handler. @@ -41,6 +42,10 @@ public class HandlerParser { @Override protected void preHandle(HttpServletRequest request, Object handler, SpanCustomizer customizer) { } + + @Override + protected void postHandle(HttpServletRequest request, Object handler, ModelAndView modelAndView, SpanCustomizer customizer) { + } }; /** Simple class name that processed the request. ex BookController */ @@ -74,6 +79,17 @@ public class HandlerParser { } } + /** + * Invoked posterior to request invocation during + * {@link HandlerInterceptor#postHandle(HttpServletRequest, HttpServletResponse, Object, ModelAndView)}. + * + * @param request request + * @param handler handler + * @param customizer span customizer + */ + protected void postHandle(HttpServletRequest request, Object handler, ModelAndView modelAndView, SpanCustomizer customizer) { + } + /* * Intentionally public for @Autowired to work without explicit binding */ diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/mvc/SpanCustomizingAsyncHandlerInterceptor.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/mvc/SpanCustomizingAsyncHandlerInterceptor.java index 466cbc894..f7ee3e9a8 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/mvc/SpanCustomizingAsyncHandlerInterceptor.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/mvc/SpanCustomizingAsyncHandlerInterceptor.java @@ -24,6 +24,7 @@ import org.springframework.cloud.sleuth.SpanCustomizer; import org.springframework.cloud.sleuth.instrument.web.servlet.TracingFilter; import org.springframework.web.servlet.AsyncHandlerInterceptor; import org.springframework.web.servlet.HandlerInterceptor; +import org.springframework.web.servlet.ModelAndView; import org.springframework.web.servlet.handler.HandlerInterceptorAdapter; import static org.springframework.cloud.sleuth.instrument.web.mvc.SpanCustomizingHandlerInterceptor.setErrorAttribute; @@ -51,6 +52,15 @@ public final class SpanCustomizingAsyncHandlerInterceptor extends HandlerInterce return true; } + @Override + public void postHandle(HttpServletRequest request, HttpServletResponse response, Object handler, + ModelAndView modelAndView) { + Object span = request.getAttribute(SpanCustomizer.class.getName()); + if (span instanceof SpanCustomizer) { + handlerParser.postHandle(request, handler, modelAndView, (SpanCustomizer) span); + } + } + /** * Sets the "error" and "http.route" attributes so that the {@link TracingFilter} can * read them. diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/mvc/SpanCustomizingHandlerInterceptor.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/mvc/SpanCustomizingHandlerInterceptor.java index 6174413ce..de66fab11 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/mvc/SpanCustomizingHandlerInterceptor.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/mvc/SpanCustomizingHandlerInterceptor.java @@ -65,6 +65,10 @@ public final class SpanCustomizingHandlerInterceptor implements HandlerIntercept @Override public void postHandle(HttpServletRequest request, HttpServletResponse response, Object handler, ModelAndView modelAndView) { + Object span = request.getAttribute(SpanCustomizer.class.getName()); + if (span instanceof SpanCustomizer) { + handlerParser.postHandle(request, handler, modelAndView, (SpanCustomizer) span); + } } /** Sets the "error" attribute so that the {@link TracingFilter} can read it. */ diff --git a/tests/brave/spring-cloud-sleuth-instrumentation-mvc-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/web/TraceCustomHandlerParserPostHandleWebIntegrationTests.java b/tests/brave/spring-cloud-sleuth-instrumentation-mvc-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/web/TraceCustomHandlerParserPostHandleWebIntegrationTests.java new file mode 100644 index 000000000..5ab37edc6 --- /dev/null +++ b/tests/brave/spring-cloud-sleuth-instrumentation-mvc-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/web/TraceCustomHandlerParserPostHandleWebIntegrationTests.java @@ -0,0 +1,218 @@ +/* + * Copyright 2013-2021 the original author or 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 + * + * https://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.springframework.cloud.sleuth.brave.instrument.web; + +import java.io.UnsupportedEncodingException; + +import javax.servlet.http.HttpServletRequest; + +import brave.Span.Kind; +import brave.handler.MutableSpan; +import brave.handler.SpanHandler; +import brave.propagation.CurrentTraceContext; +import brave.sampler.Sampler; +import brave.test.IntegrationTestSpanHandler; +import org.apache.commons.logging.Log; +import org.apache.commons.logging.LogFactory; +import org.junit.ClassRule; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; + +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.boot.autoconfigure.EnableAutoConfiguration; +import org.springframework.boot.test.context.SpringBootTest; +import org.springframework.boot.test.system.OutputCaptureExtension; +import org.springframework.boot.web.servlet.FilterRegistrationBean; +import org.springframework.cloud.sleuth.SpanCustomizer; +import org.springframework.cloud.sleuth.instrument.web.mvc.HandlerParser; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; +import org.springframework.core.env.Environment; +import org.springframework.http.client.ClientHttpResponse; +import org.springframework.web.bind.annotation.PostMapping; +import org.springframework.web.bind.annotation.RequestBody; +import org.springframework.web.bind.annotation.RestController; +import org.springframework.web.client.DefaultResponseErrorHandler; +import org.springframework.web.client.RestTemplate; +import org.springframework.web.filter.AbstractRequestLoggingFilter; +import org.springframework.web.filter.CommonsRequestLoggingFilter; +import org.springframework.web.servlet.ModelAndView; +import org.springframework.web.util.ContentCachingRequestWrapper; +import org.springframework.web.util.WebUtils; + +import static org.assertj.core.api.BDDAssertions.then; + +/** + * @author Wu Wen + */ +@SpringBootTest(classes = TraceCustomHandlerParserPostHandleWebIntegrationTests.Config.class, + webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT) +@ExtendWith(OutputCaptureExtension.class) +public class TraceCustomHandlerParserPostHandleWebIntegrationTests { + + @ClassRule + public static IntegrationTestSpanHandler spanHandler = new IntegrationTestSpanHandler(); + + @Autowired + CurrentTraceContext currentTraceContext; + + @Autowired + Environment environment; + + @Test + public void should_tag_payload() { + + new RestTemplate().postForEntity("http://localhost:" + port() + "/user", new UserController.User("WuWen"), + UserController.User.class); + + then(this.currentTraceContext.get()).isNull(); + MutableSpan span = spanHandler.takeRemoteSpan(Kind.SERVER); + + then(span.tags()).containsEntry("preHandle.payload.isNull", "null"); + then(span.tags()).containsEntry("http.payload", "{\"name\":\"WuWen\"}"); + } + + private int port() { + return this.environment.getProperty("local.server.port", Integer.class); + } + + @EnableAutoConfiguration( + // spring boot test will otherwise instrument the client and server with the + // same bean factory which isn't expected + excludeName = "org.springframework.cloud.sleuth.brave.instrument.web.client.TraceWebClientAutoConfiguration") + @Configuration(proxyBeanMethods = false) + public static class Config { + + @Bean + UserController goodController() { + return new UserController(); + } + + @Bean + SpanHandler testSpanHandler() { + return spanHandler; + } + + @Bean + Sampler alwaysSampler() { + return Sampler.ALWAYS_SAMPLE; + } + + @Bean + CommonsRequestLoggingFilter commonsRequestLoggingFilter() { + CommonsRequestLoggingFilter filter = new CommonsRequestLoggingFilter(); + filter.setIncludePayload(true); + filter.setMaxPayloadLength(512); + return filter; + } + + @Bean + public FilterRegistrationBean requestLoggingFilter(CommonsRequestLoggingFilter filter) { + FilterRegistrationBean registration = new FilterRegistrationBean<>(); + registration.addUrlPatterns("/*"); + registration.setFilter(filter); + return registration; + } + + @Bean + HandlerParser handlerParser() { + return new HandlerParser() { + + @Override + protected void preHandle(HttpServletRequest request, Object handler, SpanCustomizer customizer) { + customizer.tag("preHandle.payload.isNull", getMessagePayload(request)); + } + + @Override + protected void postHandle(HttpServletRequest request, Object handler, ModelAndView modelAndView, + SpanCustomizer customizer) { + customizer.tag("http.payload", getMessagePayload(request)); + } + + protected String getMessagePayload(HttpServletRequest request) { + ContentCachingRequestWrapper wrapper = WebUtils.getNativeRequest(request, + ContentCachingRequestWrapper.class); + if (wrapper != null) { + byte[] buf = wrapper.getContentAsByteArray(); + if (buf.length > 0) { + int length = Math.min(buf.length, 512); + try { + return new String(buf, 0, length, wrapper.getCharacterEncoding()); + } + catch (UnsupportedEncodingException ex) { + return "[unknown]"; + } + } + } + return "null"; + } + }; + } + + @Bean + RestTemplate restTemplate() { + RestTemplate restTemplate = new RestTemplate(); + restTemplate.setErrorHandler(new DefaultResponseErrorHandler() { + @Override + public void handleError(ClientHttpResponse response) { + } + }); + return restTemplate; + } + + } + + @RestController + public static class UserController { + + private static final Log log = LogFactory.getLog(UserController.class); + + @PostMapping("/user") + public User user(@RequestBody User user) { + log.info("Test request body."); + return user; + } + + public static class User { + + private String name; + + public User() { + } + + public User(String name) { + this.name = name; + } + + public String getName() { + return name; + } + + public void setName(String name) { + this.name = name; + } + + @Override + public String toString() { + return "User{" + "name='" + name + '\'' + '}'; + } + + } + + } + +}