Skip to content

Commit 1759671

Browse files
itssimonitssimon
committed
Capture chunked response size
1 parent 2dcf32b commit 1759671

File tree

3 files changed

+108
-10
lines changed

3 files changed

+108
-10
lines changed

‎src/main/java/io/apitally/spring/ApitallyFilter.java‎

Lines changed: 84 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,11 @@
2222
import io.apitally.common.dto.Response;
2323
import jakarta.servlet.FilterChain;
2424
import jakarta.servlet.ServletException;
25+
import jakarta.servlet.ServletOutputStream;
26+
import jakarta.servlet.WriteListener;
2527
import jakarta.servlet.http.HttpServletRequest;
2628
import jakarta.servlet.http.HttpServletResponse;
29+
import jakarta.servlet.http.HttpServletResponseWrapper;
2730
import jakarta.validation.ConstraintViolation;
2831
import jakarta.validation.ConstraintViolationException;
2932

@@ -58,14 +61,17 @@ protected void doFilterInternal(@NonNull HttpServletRequest request, @NonNull Ht
5861
ContentCachingResponseWrapper cachingResponse = shouldCacheResponse
5962
? new ContentCachingResponseWrapper(response)
6063
: null;
64+
CountingResponseWrapper countingResponse = cachingResponse == null
65+
? new CountingResponseWrapper(response)
66+
: null;
6167

6268
Exception exception = null;
6369
final long startTime = System.currentTimeMillis();
6470

6571
try {
6672
filterChain.doFilter(
6773
cachingRequest != null ? cachingRequest : request,
68-
cachingResponse != null ? cachingResponse : response);
74+
cachingResponse != null ? cachingResponse : countingResponse);
6975
} catch (Exception e) {
7076
exception = e;
7177
throw e;
@@ -92,9 +98,11 @@ protected void doFilterInternal(@NonNull HttpServletRequest request, @NonNull Ht
9298
final long requestSize = requestContentLength >= 0 ? requestContentLength
9399
: cachingRequest != null ? requestBody.length : -1;
94100
final long responseContentLength = getResponseContentLength(response);
95-
final long responseSize = responseContentLength >= 0
96-
? responseContentLength
97-
: (cachingResponse != null ? responseBody.length : -1);
101+
final long responseSize = responseContentLength >= 0 ? responseContentLength
102+
: (cachingResponse != null ? responseBody.length
103+
: countingResponse != null
104+
? countingResponse.getByteCount()
105+
: -1);
98106
client.requestCounter
99107
.addRequest(consumerIdentifier, request.getMethod(), path, response.getStatus(),
100108
responseTimeInMillis,
@@ -154,7 +162,6 @@ protected void doFilterInternal(@NonNull HttpServletRequest request, @NonNull Ht
154162
logger.error("Error in Apitally filter", e);
155163
}
156164
}
157-
158165
}
159166

160167
private static long getResponseContentLength(HttpServletResponse response) {
@@ -167,4 +174,76 @@ private static long getResponseContentLength(HttpServletResponse response) {
167174
}
168175
return -1L;
169176
}
177+
178+
private static class CountingResponseWrapper extends HttpServletResponseWrapper {
179+
private CountingServletOutputStream countingStream;
180+
181+
public CountingResponseWrapper(HttpServletResponse response) {
182+
super(response);
183+
}
184+
185+
@Override
186+
public ServletOutputStream getOutputStream() throws IOException {
187+
if (countingStream == null) {
188+
countingStream = new CountingServletOutputStream(super.getOutputStream());
189+
}
190+
return countingStream;
191+
}
192+
193+
public long getByteCount() {
194+
return countingStream != null ? countingStream.getByteCount() : 0;
195+
}
196+
}
197+
198+
private static class CountingServletOutputStream extends ServletOutputStream {
199+
private final ServletOutputStream outputStream;
200+
private long byteCount;
201+
202+
public CountingServletOutputStream(ServletOutputStream outputStream) {
203+
this.outputStream = outputStream;
204+
this.byteCount = 0;
205+
}
206+
207+
@Override
208+
public boolean isReady() {
209+
return outputStream.isReady();
210+
}
211+
212+
@Override
213+
public void setWriteListener(WriteListener writeListener) {
214+
outputStream.setWriteListener(writeListener);
215+
}
216+
217+
@Override
218+
public void write(int b) throws IOException {
219+
outputStream.write(b);
220+
byteCount++;
221+
}
222+
223+
@Override
224+
public void write(byte[] b) throws IOException {
225+
outputStream.write(b);
226+
byteCount += b.length;
227+
}
228+
229+
@Override
230+
public void write(byte[] b, int off, int len) throws IOException {
231+
outputStream.write(b, off, len);
232+
byteCount += len;
233+
}
234+
235+
@Override
236+
public void flush() throws IOException {
237+
outputStream.flush();
238+
}
239+
240+
@Override
241+
public void close() throws IOException {
242+
outputStream.close();
243+
}
244+
245+
public long getByteCount() {
246+
return byteCount;
247+
}
248+
}
170249
}

‎src/test/java/io/apitally/spring/ApitallyFilterTest.java‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -77,10 +77,12 @@ void setUp() {
7777
apitallyClient.validationErrorCounter.getAndResetValidationErrors();
7878
apitallyClient.serverErrorCounter.getAndResetServerErrors();
7979
apitallyClient.consumerRegistry.getAndResetConsumers();
80+
apitallyClient.requestLogger.getConfig().setEnabled(true);
8081
}
8182

8283
@Test
8384
void testRequestCounter() {
85+
apitallyClient.requestLogger.getConfig().setEnabled(false);
8486
ResponseEntity<String> response;
8587

8688
response = restTemplate.getForEntity("/items", String.class);
@@ -95,13 +97,16 @@ void testRequestCounter() {
9597
response = restTemplate.getForEntity("/items/2", String.class);
9698
assertTrue(response.getStatusCode().is2xxSuccessful());
9799

100+
response = restTemplate.getForEntity("/stream", String.class);
101+
assertTrue(response.getStatusCode().is2xxSuccessful());
102+
98103
response = restTemplate.getForEntity("/throw", String.class);
99104
assertTrue(response.getStatusCode().is5xxServerError());
100105

101106
delay(100);
102107

103108
List<Requests> requests = apitallyClient.requestCounter.getAndResetRequests();
104-
assertEquals(4, requests.size());
109+
assertEquals(5, requests.size());
105110
assertTrue(requests.stream()
106111
.anyMatch(r -> r.getMethod().equals("GET")
107112
&& r.getPath().equals("/items")
@@ -117,6 +122,12 @@ void testRequestCounter() {
117122
&& r.getPath().equals("/items/{id}")
118123
&& r.getStatusCode() == 400
119124
&& r.getRequestCount() == 1));
125+
assertTrue(requests.stream().anyMatch(
126+
r -> r.getMethod().equals("GET")
127+
&& r.getPath().equals("/stream")
128+
&& r.getStatusCode() == 200
129+
&& r.getRequestCount() == 1
130+
&& r.getResponseSizeSum() == 14));
120131
assertTrue(requests.stream().anyMatch(
121132
r -> r.getMethod().equals("GET")
122133
&& r.getPath().equals("/throw")

‎src/test/java/io/apitally/spring/app/TestController.java‎

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,9 @@
1717
import org.springframework.web.bind.annotation.RequestParam;
1818
import org.springframework.web.bind.annotation.ResponseStatus;
1919
import org.springframework.web.bind.annotation.RestController;
20+
import org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody;
2021

2122
import io.apitally.spring.ApitallyConsumer;
22-
2323
import jakarta.servlet.http.HttpServletRequest;
2424
import jakarta.validation.ConstraintViolationException;
2525
import jakarta.validation.Valid;
@@ -61,9 +61,17 @@ public void updateItem(@Valid @RequestBody TestItem newItem, @PathVariable @Min(
6161
public void deleteItem(@PathVariable @Min(1) Integer id) {
6262
}
6363

64-
@GetMapping(value = "/healthz", produces = "application/json; charset=utf-8")
65-
public String getHealthCheck() {
66-
return "{ \"healthy\" : true }";
64+
@GetMapping("/stream")
65+
public ResponseEntity<StreamingResponseBody> getItemsStream(HttpServletRequest request) {
66+
return ResponseEntity
67+
.ok()
68+
.header("Transfer-Encoding", "chunked")
69+
.header("Content-Type", "text/plain")
70+
.body(out -> {
71+
out.write(("Item 1" + "\n").getBytes());
72+
out.write(("Item 2" + "\n").getBytes());
73+
out.flush();
74+
});
6775
}
6876

6977
@GetMapping(value = "/throw", produces = "application/json; charset=utf-8")

0 commit comments

Comments
 (0)