Skip to content

Commit 5b40c01

Browse files
committed
Improve ResponseSubscribers performance
Fixes #1042 Signed-off-by: Daniel Garnier-Moiroux <[email protected]>
1 parent b6a46d8 commit 5b40c01

4 files changed

Lines changed: 561 additions & 1 deletion

File tree

‎mcp-core/src/main/java/io/modelcontextprotocol/client/transport/ResponseSubscribers.java‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -191,6 +191,17 @@ static final class Utf8LineDecoder {
191191

192192
private final StringBuilder leftover = new StringBuilder();
193193

194+
/**
195+
* How many leading characters of {@link #leftover} are already known to hold no
196+
* line terminator, so that the search for one resumes where the previous search
197+
* ended instead of restarting at the beginning of the buffer. Without it, a long
198+
* line is searched again in full for every chunk that arrives, which makes
199+
* reading an event cost time proportional to the square of its length.
200+
* @see <a href=
201+
* "https://github.com/modelcontextprotocol/java-sdk/issues/1042">#1042</a>
202+
*/
203+
private int scannedForLineTerminator = 0;
204+
194205
// Holds partial UTF-8 sequences left over from a previous chunk (max 3 bytes
195206
// for a BMP code point; 4 bytes for a supplementary one).
196207
private ByteBuffer pendingBytes = ByteBuffer.allocate(0);
@@ -263,6 +274,7 @@ List<String> flush() {
263274
last = last.substring(0, last.length() - 1);
264275
}
265276
leftover.setLength(0);
277+
this.scannedForLineTerminator = 0;
266278
lines.add(last);
267279
}
268280
return lines;
@@ -276,14 +288,18 @@ private void drainCharBuffer() {
276288

277289
private void extractCompletedLines(List<String> out) {
278290
int newlineIdx;
279-
while ((newlineIdx = leftover.indexOf("\n")) != -1) {
291+
while ((newlineIdx = leftover.indexOf("\n", this.scannedForLineTerminator)) != -1) {
280292
String line = leftover.substring(0, newlineIdx);
281293
if (line.endsWith("\r")) {
282294
line = line.substring(0, line.length() - 1);
283295
}
284296
out.add(line);
285297
leftover.delete(0, newlineIdx + 1);
298+
// What is left starts after the terminator, so none of it has been
299+
// searched yet.
300+
this.scannedForLineTerminator = 0;
286301
}
302+
this.scannedForLineTerminator = leftover.length();
287303
}
288304

289305
}
Lines changed: 219 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,219 @@
1+
/*
2+
* Copyright 2024-2026 the original author or authors.
3+
*/
4+
5+
package io.modelcontextprotocol.client.transport;
6+
7+
import java.nio.ByteBuffer;
8+
import java.nio.charset.StandardCharsets;
9+
import java.util.ArrayList;
10+
import java.util.List;
11+
import java.util.concurrent.Flow;
12+
13+
import org.junit.jupiter.api.Test;
14+
import org.junit.jupiter.api.Timeout;
15+
import org.slf4j.Logger;
16+
import org.slf4j.LoggerFactory;
17+
import reactor.adapter.JdkFlowAdapter;
18+
import reactor.core.publisher.Flux;
19+
20+
import io.modelcontextprotocol.client.transport.ResponseSubscribers.SseEvent;
21+
22+
import static org.assertj.core.api.Assertions.assertThat;
23+
24+
/**
25+
* Reproducer for the client-side SSE reading bottleneck reported in
26+
* <a href= "https://github.com/modelcontextprotocol/java-sdk/issues/1042">#1042</a>: a
27+
* tool response arriving as a single multi-megabyte {@code data:} line, which is what
28+
* compact JSON looks like on the wire, took ~5s to read where the same bytes read with
29+
* {@link java.net.http.HttpResponse.BodyHandlers#ofString()} took ~0.4s.
30+
*
31+
* <p>
32+
* What cost the time was the length of the line rather than the number of bytes, because
33+
* the buffered characters were gone over again every time a chunk arrived.
34+
* {@link #shouldReadOneLargeEventWithinBudgetOfManySmallOnes()} therefore measures the
35+
* same payload twice, once as one long line and once split over short ones, and asserts a
36+
* ratio rather than a duration, so that it keeps its meaning on a machine of any speed.
37+
*
38+
* <p>
39+
* Reading a single-line event in 16KiB chunks, best of three runs on the same machine:
40+
*
41+
* <pre>
42+
* payload 2.0.0 (fromLineSubscriber) rescanning the line scanning each line once
43+
* 1MiB 210ms 12ms 2ms
44+
* 2MiB 810ms 32ms 6ms
45+
* 4MiB 3314ms 138ms 10ms
46+
* 8MiB 13140ms 561ms 18ms
47+
* </pre>
48+
*
49+
* <p>
50+
* The middle column read each chunk incrementally, which is ~25x quicker than what 2.0.0
51+
* shipped, but {@link ResponseSubscribers.Utf8LineDecoder} still searched its buffered
52+
* characters for a line terminator from the start of the buffer on every chunk, so eight
53+
* times the payload cost ~45x the time. Resuming that search where the previous one ended
54+
* gives the third column, which scales with the payload rather than with its square and
55+
* brings the ratio this test measures from ~12 to ~1.6.
56+
*
57+
* <p>
58+
* See {@code HttpClientStreamableHttpTransportLargeResponseTests} in {@code mcp-test} for
59+
* the same comparison end to end, over a real connection.
60+
*
61+
* @author Daniel Garnier-Moiroux
62+
*/
63+
class LargeSseEventDecodingTests {
64+
65+
private static final Logger logger = LoggerFactory.getLogger(LargeSseEventDecodingTests.class);
66+
67+
/**
68+
* Roughly what {@link java.net.http.HttpClient} hands to a body subscriber at a time.
69+
* The cost the report is about was paid per chunk, so the chunking is part of the
70+
* reproducer.
71+
*/
72+
private static final int CHUNK_SIZE = 16 * 1024;
73+
74+
private static final int MIB = 1024 * 1024;
75+
76+
/**
77+
* The payload size in the report: ~4MiB of compact JSON, and therefore ~4MiB with no
78+
* line terminator in it.
79+
*/
80+
private static final int PAYLOAD_SIZE = 4 * MIB;
81+
82+
/**
83+
* How much of {@link #PAYLOAD_SIZE} each event carries when the same total is split
84+
* over many events.
85+
*/
86+
private static final int SMALL_EVENT_SIZE = 64 * 1024;
87+
88+
/**
89+
* How much longer decoding the payload as one long line may take than decoding the
90+
* same bytes as short ones. A reader that goes over each line once is indifferent to
91+
* how long the lines are, which measures ~1.6 here; the bound leaves headroom over
92+
* that, and is far below what rescanning the line measured (~12x) or what 2.0.0
93+
* measured (~220x).
94+
*/
95+
private static final double MAX_SINGLE_EVENT_PENALTY = 4.0;
96+
97+
private static final int MAX_SIZE = 64 * MIB;
98+
99+
@Test
100+
@Timeout(60)
101+
void shouldDecodeMultiMegabyteSingleLineEventIntact() {
102+
String payload = payloadOfSize(PAYLOAD_SIZE);
103+
104+
List<SseEvent> events = decode(oneLargeEvent(payload));
105+
106+
assertThat(events).hasSize(1);
107+
assertThat(events.get(0).event()).isEqualTo("message");
108+
assertThat(events.get(0).data()).isEqualTo(payload);
109+
}
110+
111+
@Test
112+
@Timeout(300)
113+
void shouldReadOneLargeEventWithinBudgetOfManySmallOnes() {
114+
byte[] oneEvent = oneLargeEvent(payloadOfSize(PAYLOAD_SIZE));
115+
byte[] manyEvents = manySmallEvents(PAYLOAD_SIZE, SMALL_EVENT_SIZE);
116+
int smallEventCount = PAYLOAD_SIZE / SMALL_EVENT_SIZE;
117+
118+
// The reporter measured a JIT effect at this payload size: the first few large
119+
// reads spike before the hot loop settles. Warm up, then interleave the two
120+
// shapes and take the best of each, so the comparison reflects steady state.
121+
for (int i = 0; i < 5; i++) {
122+
decode(manyEvents);
123+
}
124+
long single = Long.MAX_VALUE;
125+
long split = Long.MAX_VALUE;
126+
for (int i = 0; i < 3; i++) {
127+
single = Math.min(single, timeDecode(oneEvent, 1));
128+
split = Math.min(split, timeDecode(manyEvents, smallEventCount));
129+
}
130+
131+
double penalty = (double) single / Math.max(split, 1);
132+
logger.info("decoded {}KiB as one event in {}ms and as {} events in {}ms: ratio {}", PAYLOAD_SIZE / 1024,
133+
single / 1_000_000, smallEventCount, split / 1_000_000, String.format("%.1f", penalty));
134+
logScaling();
135+
136+
assertThat(penalty)
137+
.as("decoding %dKiB as a single SSE event took %.1fx as long as decoding the same number of bytes as "
138+
+ "%dKiB events, so the cost of an event grows with the length of its line", PAYLOAD_SIZE / 1024,
139+
penalty, SMALL_EVENT_SIZE / 1024)
140+
.isLessThan(MAX_SINGLE_EVENT_PENALTY);
141+
}
142+
143+
/**
144+
* Logs how decoding one long line scales with its length, which is the shape the
145+
* report is about: doubling the payload should cost about twice the time, not four
146+
* times it. Not asserted, because the ratio above covers the same ground with a
147+
* baseline measured on the same machine.
148+
*/
149+
private void logScaling() {
150+
for (int payloadSize : new int[] { MIB, 2 * MIB, 4 * MIB, 8 * MIB }) {
151+
byte[] body = oneLargeEvent(payloadOfSize(payloadSize));
152+
long best = Math.min(timeDecode(body, 1), timeDecode(body, 1));
153+
logger.info("decoded a single-line event of {}KiB in {}ms", payloadSize / 1024, best / 1_000_000);
154+
}
155+
}
156+
157+
/**
158+
* Decodes the body once and returns how long it took, in nanoseconds.
159+
*/
160+
private static long timeDecode(byte[] body, int expectedEvents) {
161+
long start = System.nanoTime();
162+
List<SseEvent> events = decode(body);
163+
long elapsed = System.nanoTime() - start;
164+
assertThat(events).hasSize(expectedEvents);
165+
return elapsed;
166+
}
167+
168+
/**
169+
* Runs the body through the transport's SSE reading path, as
170+
* {@code HttpClientStreamableHttpTransport} does, chunked the way the HTTP client
171+
* chunks a response body.
172+
*/
173+
private static List<SseEvent> decode(byte[] body) {
174+
Flow.Publisher<List<ByteBuffer>> publisher = JdkFlowAdapter
175+
.publisherToFlowPublisher(Flux.fromIterable(chunk(body)));
176+
Flux<String> lines = ResponseSubscribers.decodeLines(publisher);
177+
return ResponseSubscribers.decodeSseResponse(lines, MAX_SIZE).collectList().block();
178+
}
179+
180+
private static List<List<ByteBuffer>> chunk(byte[] body) {
181+
List<List<ByteBuffer>> chunks = new ArrayList<>();
182+
for (int offset = 0; offset < body.length; offset += CHUNK_SIZE) {
183+
int length = Math.min(CHUNK_SIZE, body.length - offset);
184+
chunks.add(List.of(ByteBuffer.wrap(body, offset, length).asReadOnlyBuffer()));
185+
}
186+
return chunks;
187+
}
188+
189+
/**
190+
* An SSE {@code message} event carrying the whole payload on a single {@code data:}
191+
* line.
192+
*/
193+
private static byte[] oneLargeEvent(String payload) {
194+
return ("event: message\ndata: " + payload + "\n\n").getBytes(StandardCharsets.UTF_8);
195+
}
196+
197+
/**
198+
* The same {@code total} number of payload bytes, spread over events of
199+
* {@code eachSize} each.
200+
*/
201+
private static byte[] manySmallEvents(int total, int eachSize) {
202+
StringBuilder body = new StringBuilder(total + 4096);
203+
for (int i = 0; i < total / eachSize; i++) {
204+
body.append("event: message\ndata: ").append(payloadOfSize(eachSize)).append("\n\n");
205+
}
206+
return body.toString().getBytes(StandardCharsets.UTF_8);
207+
}
208+
209+
/**
210+
* A single-line JSON-RPC response of exactly {@code size} characters, none of them a
211+
* line terminator.
212+
*/
213+
private static String payloadOfSize(int size) {
214+
String prefix = "{\"jsonrpc\":\"2.0\",\"id\":\"test-id\",\"result\":{\"content\":\"";
215+
String suffix = "\"}}";
216+
return prefix + "a".repeat(size - prefix.length() - suffix.length()) + suffix;
217+
}
218+
219+
}

‎mcp-core/src/test/java/io/modelcontextprotocol/client/transport/Utf8LineDecoderTests.java‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,36 @@ void invalidUtf8Throws() {
116116
.isInstanceOf(RuntimeException.class);
117117
}
118118

119+
@Test
120+
void resumesSearchAfterTerminatorWhenLineWasSplitAcrossManyChunks() {
121+
// The decoder remembers how far it has searched for a terminator, so the chunk
122+
// that finally terminates a long line must not leave that mark behind and hide
123+
// the lines that follow it.
124+
Utf8LineDecoder dec = new Utf8LineDecoder();
125+
for (int i = 0; i < 10; i++) {
126+
assertThat(dec.decode(chunk("aaaa"))).isEmpty();
127+
}
128+
assertThat(dec.decode(chunk("\nsecond\nthird\n"))).containsExactly("a".repeat(40), "second", "third");
129+
assertThat(dec.flush()).isEmpty();
130+
}
131+
132+
@Test
133+
void resumesSearchAcrossChunkWhenTerminatorFollowsUnterminatedPrefix() {
134+
Utf8LineDecoder dec = new Utf8LineDecoder();
135+
assertThat(dec.decode(chunk("unterminated"))).isEmpty();
136+
assertThat(dec.decode(chunk("-still-going"))).isEmpty();
137+
assertThat(dec.decode(chunk("\n"))).containsExactly("unterminated-still-going");
138+
}
139+
140+
@Test
141+
void crLfSplitAcrossChunks() {
142+
// The CR ends one chunk and the LF opens the next, so the terminator sits exactly
143+
// at the point the previous search stopped.
144+
Utf8LineDecoder dec = new Utf8LineDecoder();
145+
assertThat(dec.decode(chunk("hello\r"))).isEmpty();
146+
assertThat(dec.decode(chunk("\nworld\r\n"))).containsExactly("hello", "world");
147+
}
148+
119149
@Test
120150
void linesLongerThanInternalCharBuffer() {
121151
// 4096 is the internal CharBuffer size; send a single line ~10k chars to force

0 commit comments

Comments
 (0)