|
| 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 | +} |
0 commit comments