Fix aynsc servlet (#63)

(cherry picked from commit b9ef652)

sofa-tracer-plugins/sofa-tracer-springmvc-plugin/src/main/java/com/alipay/sofa/tracer/plugins/springmvc/SpringMvcSofaTracerFilter.java

@@ -77,19 +77,15 @@ public void doFilter(ServletRequest servletRequest, ServletResponse servletRespo
             springMvcSpan.setTag(CommonSpanTags.REQUEST_URL, request.getRequestURL().toString());
             springMvcSpan.setTag(CommonSpanTags.METHOD, request.getMethod());
             springMvcSpan.setTag(CommonSpanTags.REQ_SIZE, request.getContentLength());
-            OutputStreamResponseWrapper outputStreamResponseWrapper = new OutputStreamResponseWrapper(
-                response);
+            //wrapper
+            ResponseWrapper responseWrapper = new ResponseWrapper(response);
 
             //filter begin
-            filterChain.doFilter(servletRequest, outputStreamResponseWrapper);
+            filterChain.doFilter(servletRequest, responseWrapper);
             //filter end
 
-            httpStatus = outputStreamResponseWrapper.getStatus();
-            ServletOutputStream servletOutputStream = outputStreamResponseWrapper.getOutputStream();
-            if (servletOutputStream instanceof ServletOutputStreamWrapper) {
-                ServletOutputStreamWrapper servletOutputStreamWrapper = (ServletOutputStreamWrapper) servletOutputStream;
-                responseSize = servletOutputStreamWrapper.getContentLengthWrapper();
-            }
+            httpStatus = responseWrapper.getStatus();
+            responseSize = responseWrapper.getContentLength();
         } catch (Exception e) {
             SelfLog.error("Spring MVC Tracer error occurs in SpringMvcSofaTracerFilter.doFilter.",
                 e);
@@ -137,71 +133,28 @@ public SofaTracerSpanContext getSpanContextFromRequest(HttpServletRequest reques
         return spanContext;
     }
 
-    class ServletOutputStreamWrapper extends ServletOutputStream {
-
-        private long                contentLengthWrapper = 0;
+    class ResponseWrapper extends HttpServletResponseWrapper {
 
-        private ServletOutputStream servletOutputStream;
-
-        public ServletOutputStreamWrapper(ServletOutputStream servletOutputStream) {
-            this.servletOutputStream = servletOutputStream;
-        }
-
-        @Override
-        public boolean isReady() {
-            return this.servletOutputStream.isReady();
-        }
-
-        @Override
-        public void setWriteListener(WriteListener writeListener) {
-            this.servletOutputStream.setWriteListener(writeListener);
-        }
-
-        @Override
-        public void write(int i) throws IOException {
-            this.servletOutputStream.write(i);
-        }
-
-        @Override
-        public void write(byte[] b, int off, int len) throws IOException {
-            this.servletOutputStream.write(b, off, len);
-            this.contentLengthWrapper += len;
-        }
-
-        public long getContentLengthWrapper() {
-            return contentLengthWrapper;
-        }
-    }
-
-    class OutputStreamResponseWrapper extends HttpServletResponseWrapper {
-
-        private HttpServletResponse        httpServletResponse;
-
-        private ServletOutputStreamWrapper servletOutputStreamWrapper;
+        int contentLength = 0;
 
         /**
          * @param httpServletResponse httpServletResponse
          */
-        public OutputStreamResponseWrapper(HttpServletResponse httpServletResponse)
-                                                                                   throws IOException {
+        public ResponseWrapper(HttpServletResponse httpServletResponse) throws IOException {
             super(httpServletResponse);
-            this.httpServletResponse = httpServletResponse;
         }
 
+        /**
+         * @see javax.servlet.ServletResponseWrapper#setContentLength(int)
+         */
         @Override
-        public ServletOutputStream getOutputStream() throws IOException {
-            //first init and second return
-            if (this.servletOutputStreamWrapper == null) {
-                ServletOutputStream servletOutputStream = this.httpServletResponse
-                    .getOutputStream();
-                this.servletOutputStreamWrapper = new ServletOutputStreamWrapper(
-                    servletOutputStream);
-            }
-            return this.servletOutputStreamWrapper;
+        public void setContentLength(int len) {
+            contentLength = len;
+            super.setContentLength(len);
         }
 
-        public HttpServletResponse getHttpServletResponse() {
-            return httpServletResponse;
+        public int getContentLength() {
+            return contentLength;
         }
     }
 

tracer-sofa-boot-starter/src/test/java/com/alipay/sofa/tracer/boot/base/controller/SampleRestController.java

@@ -20,6 +20,10 @@
 import org.springframework.web.bind.annotation.RequestParam;
 import org.springframework.web.bind.annotation.RestController;
 
+import javax.servlet.AsyncContext;
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
+import java.io.IOException;
 import java.util.concurrent.atomic.AtomicLong;
 
 /**
@@ -31,9 +35,11 @@
 @RestController
 public class SampleRestController {
 
-    private static final String template = "Hello, %s!";
+    public static String        ASYNC_RESP = "Hello World!";
 
-    private final AtomicLong    counter  = new AtomicLong();
+    private static final String template   = "Hello, %s!";
+
+    private final AtomicLong    counter    = new AtomicLong();
 
     @RequestMapping("/greeting")
     public Greeting greeting(@RequestParam(value = "name", defaultValue = "World") String name) {
@@ -53,6 +59,14 @@ public Greeting noDigestLog() {
         return greeting;
     }
 
+    @RequestMapping("/asyncServlet")
+    public void asyncServlet(HttpServletRequest request, HttpServletResponse response)
+                                                                                      throws IOException {
+        AsyncContext asyncContext = request.startAsync();
+        asyncContext.getResponse().getWriter().write(ASYNC_RESP);
+        asyncContext.complete();
+    }
+
     public static class Greeting {
 
         private boolean success = false;

tracer-sofa-boot-starter/src/test/java/com/alipay/sofa/tracer/boot/springmvc/SpringMvcAsyncTest.java

@@ -0,0 +1,82 @@
+/*
+ * 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 com.alipay.sofa.tracer.boot.springmvc;
+
+import com.alipay.common.tracer.core.configuration.SofaTracerConfiguration;
+import com.alipay.common.tracer.core.reporter.digest.manager.SofaTracerDigestReporterAsyncManager;
+import com.alipay.sofa.tracer.boot.base.AbstractTestBase;
+import com.alipay.sofa.tracer.boot.base.controller.SampleRestController;
+import com.alipay.sofa.tracer.plugins.springmvc.SpringMvcLogEnum;
+import com.alipay.sofa.tracer.plugins.springmvc.SpringMvcTracer;
+import org.apache.commons.io.FileUtils;
+import org.junit.Test;
+import org.springframework.http.HttpStatus;
+import org.springframework.http.ResponseEntity;
+import org.springframework.test.context.ActiveProfiles;
+
+import java.io.File;
+import java.lang.reflect.Field;
+import java.util.HashMap;
+import java.util.List;
+
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+/**
+ * SpringMvcFilterTest
+ *
+ * @author yangguanchao
+ * @since 2018/08/31
+ */
+@ActiveProfiles("async")
+public class SpringMvcAsyncTest extends AbstractTestBase {
+
+    @Test
+    public void testAsyncServlet() throws Exception {
+        //clear log file handler and reoutput
+        reflectSpringMVCClear();
+        //avoid close digest print
+        SofaTracerConfiguration.setProperty(SofaTracerConfiguration.DISABLE_DIGEST_LOG_KEY,
+            new HashMap<String, String>());
+        assertNotNull(testRestTemplate);
+        String restUrl = urlHttpPrefix + "/asyncServlet";
+
+        ResponseEntity<String> response = testRestTemplate.getForEntity(restUrl, String.class);
+        String greetingResponse = response.getBody();
+        assertTrue(greetingResponse.equals(SampleRestController.ASYNC_RESP));
+        assertTrue(response.getStatusCode() == HttpStatus.OK);
+        // http://docs.spring.io/spring-boot/docs/1.4.2.RELEASE/reference/htmlsingle/#boot-features-testing
+        Thread.sleep(3000);
+        //wait for async output
+        List<String> contents = FileUtils.readLines(new File(logDirectoryPath
+                                                             + File.separator
+                                                             + SpringMvcLogEnum.SPRING_MVC_DIGEST
+                                                                 .getDefaultLogName()));
+        assertTrue(contents.size() == 1);
+    }
+
+    private static void reflectSpringMVCClear() throws NoSuchFieldException, IllegalAccessException {
+        Field field = SpringMvcTracer.class.getDeclaredField("springMvcTracer");
+        field.setAccessible(true);
+        field.set(null, null);
+        //clear
+        Field fieldAsync = SofaTracerDigestReporterAsyncManager.class
+            .getDeclaredField("asyncCommonDigestAppenderManager");
+        fieldAsync.setAccessible(true);
+        fieldAsync.set(null, null);
+    }
+}

tracer-sofa-boot-starter/src/test/resources/application-async.properties

@@ -0,0 +1,21 @@
+#
+# 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.
+#
+
+spring.application.name=SpringMvcAsyncServlet
+
+#default 60s and now reset to 1
+com.alipay.sofa.tracer.statLogInterval=1