From 8ac22fc2fd5e212cf9d9242165980217ddc0e39f Mon Sep 17 00:00:00 2001
From: Madhura Bhave <mbhave@pivotal.io>
Date: Fri, 17 Mar 2017 12:04:55 -0700
Subject: [PATCH] Add elapsed time to the Trace Actuator output

Closes gh-8654
---
 .../boot/actuate/trace/TraceProperties.java         |  6 ++++++
 .../boot/actuate/trace/WebRequestTraceFilter.java   | 10 ++++++++++
 .../actuate/trace/WebRequestTraceFilterTests.java   | 13 +++++++++++++
 3 files changed, 29 insertions(+)

diff --git a/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/trace/TraceProperties.java b/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/trace/TraceProperties.java
index 9129372700..3c01ac6879 100644
--- a/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/trace/TraceProperties.java
+++ b/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/trace/TraceProperties.java
@@ -43,6 +43,7 @@ public class TraceProperties {
 		defaultIncludes.add(Include.RESPONSE_HEADERS);
 		defaultIncludes.add(Include.COOKIES);
 		defaultIncludes.add(Include.ERRORS);
+		defaultIncludes.add(Include.TIME_TAKEN);
 		DEFAULT_INCLUDES = Collections.unmodifiableSet(defaultIncludes);
 	}
 
@@ -140,6 +141,11 @@ public class TraceProperties {
 		 */
 		REMOTE_USER,
 
+		/**
+		 * Include the time taken to service the request.
+		 */
+		TIME_TAKEN
+
 	}
 
 }
diff --git a/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/trace/WebRequestTraceFilter.java b/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/trace/WebRequestTraceFilter.java
index 853135680e..0ee7bfba89 100644
--- a/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/trace/WebRequestTraceFilter.java
+++ b/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/trace/WebRequestTraceFilter.java
@@ -97,10 +97,13 @@ public class WebRequestTraceFilter extends OncePerRequestFilter implements Order
 		this.order = order;
 	}
 
+
+
 	@Override
 	protected void doFilterInternal(HttpServletRequest request,
 			HttpServletResponse response, FilterChain filterChain)
 					throws ServletException, IOException {
+		long startTime = System.currentTimeMillis();
 		Map<String, Object> trace = getTrace(request);
 		logTrace(request, trace);
 		int status = HttpStatus.INTERNAL_SERVER_ERROR.value();
@@ -109,6 +112,8 @@ public class WebRequestTraceFilter extends OncePerRequestFilter implements Order
 			status = response.getStatus();
 		}
 		finally {
+			long endTime = System.currentTimeMillis();
+			addTimeTaken(startTime, endTime, trace);
 			enhanceTrace(trace, status == response.getStatus() ? response
 					: new CustomStatusResponseWrapper(response, status));
 			this.repository.add(trace);
@@ -195,6 +200,11 @@ public class WebRequestTraceFilter extends OncePerRequestFilter implements Order
 	protected void postProcessRequestHeaders(Map<String, Object> headers) {
 	}
 
+	private void addTimeTaken(long startTime, long endTime, Map<String, Object> trace) {
+		long timeTaken = endTime - startTime;
+		add(trace, Include.TIME_TAKEN, "timeTaken", "" + timeTaken);
+	}
+
 	@SuppressWarnings("unchecked")
 	protected void enhanceTrace(Map<String, Object> trace, HttpServletResponse response) {
 		if (isIncluded(Include.RESPONSE_HEADERS)) {
diff --git a/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/trace/WebRequestTraceFilterTests.java b/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/trace/WebRequestTraceFilterTests.java
index 829dafb9a3..b8c7698e6c 100644
--- a/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/trace/WebRequestTraceFilterTests.java
+++ b/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/trace/WebRequestTraceFilterTests.java
@@ -33,6 +33,7 @@ import org.junit.Test;
 
 import org.springframework.boot.actuate.trace.TraceProperties.Include;
 import org.springframework.boot.autoconfigure.web.DefaultErrorAttributes;
+import org.springframework.mock.web.MockFilterChain;
 import org.springframework.mock.web.MockHttpServletRequest;
 import org.springframework.mock.web.MockHttpServletResponse;
 
@@ -242,6 +243,18 @@ public class WebRequestTraceFilterTests {
 		assertThat(map.get("status").toString()).isEqualTo("404");
 	}
 
+	@Test
+	@SuppressWarnings("unchecked")
+	public void filterAddsTimeTaken() throws Exception {
+		MockHttpServletRequest request = spy(new MockHttpServletRequest("GET", "/foo"));
+		MockHttpServletResponse response = new MockHttpServletResponse();
+		MockFilterChain chain = new MockFilterChain();
+		this.filter.doFilter(request, response, chain);
+		String timeTaken = (String) this.repository.findAll()
+				.iterator().next().getInfo().get("timeTaken");
+		assertThat(timeTaken).isNotNull();
+	}
+
 	@Test
 	public void filterHasError() {
 		this.filter.setErrorAttributes(new DefaultErrorAttributes());