diff --git a/flash/src/main/java/dev/relism/flash/HttpServer.java b/flash/src/main/java/dev/relism/flash/HttpServer.java index c354ee5..a00853d 100644 --- a/flash/src/main/java/dev/relism/flash/HttpServer.java +++ b/flash/src/main/java/dev/relism/flash/HttpServer.java @@ -34,9 +34,14 @@ import java.util.concurrent.atomic.AtomicReference; * Pure I/O transport layer. Owns the {@link ServerSocket}, the virtual-thread * executor, and the keep-alive accept loop. Routing is delegated to HTTP and WS routers. * - *

Allocation model (unchanged)

+ *

Allocation model

* @@ -122,6 +127,19 @@ class HttpServer implements ServerHandle { private static final ThreadLocal LONG_BUF = ThreadLocal.withInitial(() -> new byte[20]); + /** + * Relay buffer for copying a streaming {@link Response} body to the client — shared by + * {@link #writeStreamingBody}'s non-chunked path and {@link #writeChunked}, so both draw + * from the same reused array instead of each allocating its own {@code byte[8192]} (the + * non-chunked path previously relied on {@link InputStream#transferTo}, which allocates + * internally on every call). Sized to match the pre-existing behavior this replaces, not + * newly tuned — not exposed as a {@link FlashConfiguration} tunable since nothing here + * needed one before. + */ + private static final int STREAM_RELAY_BUFFER_SIZE = 8192; + private static final ThreadLocal STREAM_RELAY_BUFFER = + ThreadLocal.withInitial(() -> new byte[STREAM_RELAY_BUFFER_SIZE]); + private static final int SHA1_LEN = 20; private static final int WS_ACCEPT_LEN = 28; @@ -248,7 +266,7 @@ class HttpServer implements ServerHandle { out.flush(); request.drain(); WebSocketSession session = new WebSocketSession( - in, rawOut, configuration.getWsFrameBufferSize()); + in, rawOut, configuration.getWsFrameBufferSize(), request, false); runWsLoop(session, wsHandler); return; } @@ -420,7 +438,7 @@ class HttpServer implements ServerHandle { out.write(CRLF); out.write(keepAlive ? CONNECTION_KEEPALIVE : CONNECTION_CLOSE); out.write(CRLF); - response.getStream().transferTo(out); + relay(response.getStream(), out); } else { out.write(TRANSFER_CHUNKED); out.write(keepAlive ? CONNECTION_KEEPALIVE : CONNECTION_CLOSE); @@ -429,6 +447,18 @@ class HttpServer implements ServerHandle { } } + /** + * Copies {@code in} to {@code out} until EOF, same contract as {@link InputStream#transferTo} + * — but via {@link #STREAM_RELAY_BUFFER} instead of a fresh {@code byte[]} per call, which is + * what {@code transferTo}'s own (JDK-internal) implementation would otherwise allocate on + * every streamed response. + */ + private static void relay(InputStream in, OutputStream out) throws IOException { + byte[] buf = STREAM_RELAY_BUFFER.get(); + int n; + while ((n = in.read(buf)) > 0) out.write(buf, 0, n); + } + private static void writeStatusPhrase(OutputStream out, int statusCode) throws IOException { byte[] phrase = HttpStatus.bytesForCode(statusCode); if (phrase != null) out.write(phrase); @@ -447,7 +477,7 @@ class HttpServer implements ServerHandle { } private static void writeChunked(OutputStream out, InputStream stream) throws IOException { - byte[] buf = new byte[8192]; + byte[] buf = STREAM_RELAY_BUFFER.get(); int n; while ((n = stream.read(buf)) > 0) { writeHex(out, n); diff --git a/flash/src/main/java/dev/relism/flash/models/HeaderMap.java b/flash/src/main/java/dev/relism/flash/models/HeaderMap.java index c241eb2..090d6ec 100644 --- a/flash/src/main/java/dev/relism/flash/models/HeaderMap.java +++ b/flash/src/main/java/dev/relism/flash/models/HeaderMap.java @@ -36,6 +36,13 @@ public class HeaderMap { private int sectionStart; private int sectionEnd; + // Lazily created, then reused for the life of this HeaderMap (i.e. the connection — + // see the class javadoc) across every #forEach call and every header within a call. + // Same idiom as #view's per-call anonymous ByteView, just amortized to zero allocations + // instead of two per header: the slices are repositioned in place, not reallocated. + private Slice nameSlice; + private Slice valueSlice; + /** Resets this map to the header section {@code buffer[sectionStart, sectionEnd)}. */ public void reset(byte[] buffer, int sectionStart, int sectionEnd) { this.buffer = buffer; @@ -43,6 +50,62 @@ public class HeaderMap { this.sectionEnd = sectionEnd; } + /** + * Visits every header in declaration order without allocating — no per-header {@code + * String}/{@link ByteView}/list-entry object, unlike {@link #all()}. {@code name}/{@code + * value} are the same two {@link ByteView} instances on every call, repositioned in place; + * they are valid only for the duration of that single {@link HeaderConsumer#accept} call — + * same "do not retain past the handler" rule as {@link #view}, just per-invocation instead + * of per-request. Prefer a non-capturing or field-reusing {@link HeaderConsumer} (see its + * javadoc) if the call site itself needs to stay allocation-free too. + * + *

Exists for callers that must handle an open-ended set of header names — e.g. a reverse + * proxy forwarding whatever the client sent — where {@link #first}/{@link #all}'s per-name + * lookup isn't usable because the set of names isn't known upfront. + */ + public void forEach(HeaderConsumer consumer) { + if (buffer == null) return; + if (nameSlice == null) { + nameSlice = new Slice(); + valueSlice = new Slice(); + } + int i = sectionStart; + while (i < sectionEnd) { + int lineEnd = findCR(i); + int colon = findColon(i, lineEnd); + if (colon != -1) { + int vs = skipSpaces(colon + 1, lineEnd); + nameSlice.start = i; + nameSlice.len = colon - i; + valueSlice.start = vs; + valueSlice.len = lineEnd - vs; + consumer.accept(nameSlice, valueSlice); + } + i = lineEnd + 2; + } + } + + /** + * Callback for {@link #forEach}. Implement with a reusable, field-holding instance (reset + * before each {@code forEach} call) rather than a capturing lambda if the call site itself + * needs to be allocation-free too — a capturing lambda is its own per-call allocation, same + * as anywhere else on a hot path (see {@code docs/CODE-STYLE.md} in the Pathway project for + * the idiom this mirrors). + */ + @FunctionalInterface + public interface HeaderConsumer { + void accept(ByteView name, ByteView value); + } + + /** Mutable zero-copy slice into {@link #buffer} — see {@link #forEach}. */ + private final class Slice implements ByteView { + int start; + int len; + + @Override public int length() { return len; } + @Override public byte byteAt(int i) { return buffer[start + i]; } + } + /** Returns the first value of header {@code name} (case-insensitive), or {@code null}. */ public String first(String name) { long r = findFirst(name); diff --git a/flash/src/main/java/dev/relism/flash/websocket/WebSocketSession.java b/flash/src/main/java/dev/relism/flash/websocket/WebSocketSession.java index 9e2256a..1f10cef 100644 --- a/flash/src/main/java/dev/relism/flash/websocket/WebSocketSession.java +++ b/flash/src/main/java/dev/relism/flash/websocket/WebSocketSession.java @@ -1,9 +1,12 @@ package dev.relism.flash.websocket; +import dev.relism.flash.models.Request; + import java.io.EOFException; import java.io.IOException; import java.io.InputStream; import java.io.OutputStream; +import java.util.concurrent.ThreadLocalRandom; import java.util.concurrent.atomic.AtomicBoolean; /** @@ -39,21 +42,42 @@ public final class WebSocketSession { private final InputStream in; private final OutputStream out; private final byte[] readBuf; + private final Request request; + private final boolean maskOutgoing; private final AtomicBoolean open = new AtomicBoolean(true); private int closeCode = 1000; - private final byte[] hdrScratch = new byte[10]; + /** 1 opcode byte + up to 8 extended-length bytes + up to 4 mask-key bytes (masked mode only). */ + private final byte[] hdrScratch = new byte[14]; public WebSocketSession(InputStream in, OutputStream out, int bufferSize) { - this.in = in; - this.out = out; - this.readBuf = new byte[bufferSize]; + this(in, out, bufferSize, null, false); + } + + /** + * @param request the HTTP request that upgraded to this session, or {@code null} if the + * caller has no use for it (e.g. a session opened as a WS client rather + * than accepted as a WS server). Stored as-is, no copy. + * @param maskOutgoing {@code true} if this session is acting as a WS client — RFC 6455 + * requires client-to-server frames to be masked, unlike the server-to-client + * direction {@link #writeFrame} originally only supported. See {@link + * #writeFrame} for how masking is applied without allocating. + */ + public WebSocketSession(InputStream in, OutputStream out, int bufferSize, Request request, boolean maskOutgoing) { + this.in = in; + this.out = out; + this.readBuf = new byte[bufferSize]; + this.request = request; + this.maskOutgoing = maskOutgoing; } public boolean isOpen() { return open.get(); } public int closeCode() { return closeCode; } + /** The request that upgraded this connection, or {@code null} — see the 4-arg constructor. */ + public Request request() { return request; } + // ── Public send API ──────────────────────────────────────────────────── public void sendText(byte[] utf8, int off, int len) throws IOException { @@ -146,8 +170,9 @@ public final class WebSocketSession { // ── Private ──────────────────────────────────────────────────────────── /** - * Encodes the WS frame header into {@link #hdrScratch} (at most 10 bytes), - * then writes header + payload in two bulk calls to the raw socket stream. + * Encodes the WS frame header into {@link #hdrScratch} (at most 14 bytes: 1 opcode + up to 8 + * extended-length + up to 4 mask-key), then writes header + payload in two bulk calls to the + * raw socket stream. * *

No {@code flush()} — {@code out} is the unbuffered socket {@link OutputStream} * (see {@code HttpServer#process}). Each {@code write()} lands directly in the @@ -156,19 +181,28 @@ public final class WebSocketSession { * (header then payload) will be merged into a single TCP segment by the kernel * because they arrive faster than the ACK from the peer — exactly the coalescing * we want, at zero cost. + * + *

{@link #maskOutgoing} (client mode): RFC 6455 requires every client-to-server frame + * to be masked. The mask key is generated into {@link #hdrScratch} (no new allocation — same + * fixed field every frame reuses) and the payload is masked in place via {@link + * #unmaskInPlace} — XOR is its own inverse, so the exact routine {@link #readFrame} already + * uses to unmask an inbound payload masks an outbound one too, with no separate code path and + * no copy. This mutates the caller's {@code payload} array as a side effect: callers using + * masked mode must not reuse that buffer expecting it unchanged after the call. */ private void writeFrame(byte opcode, byte[] payload, int off, int len) throws IOException { synchronized (out) { int hlen = 0; hdrScratch[hlen++] = (byte) (0x80 | opcode); + int maskBit = maskOutgoing ? 0x80 : 0x00; if (len <= 125) { - hdrScratch[hlen++] = (byte) len; + hdrScratch[hlen++] = (byte) (maskBit | len); } else if (len <= 0xFFFF) { - hdrScratch[hlen++] = 126; + hdrScratch[hlen++] = (byte) (maskBit | 126); hdrScratch[hlen++] = (byte) ((len >> 8) & 0xFF); hdrScratch[hlen++] = (byte) (len & 0xFF); } else { - hdrScratch[hlen++] = 127; + hdrScratch[hlen++] = (byte) (maskBit | 127); hdrScratch[hlen++] = 0; hdrScratch[hlen++] = 0; hdrScratch[hlen++] = 0; hdrScratch[hlen++] = 0; hdrScratch[hlen++] = (byte) ((len >> 24) & 0xFF); @@ -176,6 +210,17 @@ public final class WebSocketSession { hdrScratch[hlen++] = (byte) ((len >> 8) & 0xFF); hdrScratch[hlen++] = (byte) (len & 0xFF); } + if (maskOutgoing) { + // ThreadLocalRandom needs no seeding/allocation per call; the four mask bytes are + // carved out of one int, never boxed. + int mask = ThreadLocalRandom.current().nextInt(); + byte m0 = (byte) (mask >>> 24), m1 = (byte) (mask >>> 16), m2 = (byte) (mask >>> 8), m3 = (byte) mask; + hdrScratch[hlen++] = m0; + hdrScratch[hlen++] = m1; + hdrScratch[hlen++] = m2; + hdrScratch[hlen++] = m3; + unmaskInPlace(payload, off, len, m0, m1, m2, m3); + } out.write(hdrScratch, 0, hlen); out.write(payload, off, len); // No flush — TCP_NODELAY handles delivery. See Javadoc above. diff --git a/flash/src/test/java/dev/relism/flash/HttpServerTest.java b/flash/src/test/java/dev/relism/flash/HttpServerTest.java index 1d32d6e..ff4705f 100644 --- a/flash/src/test/java/dev/relism/flash/HttpServerTest.java +++ b/flash/src/test/java/dev/relism/flash/HttpServerTest.java @@ -228,4 +228,37 @@ class HttpServerTest { assertTrue(second.contains("Connection: keep-alive")); } } + + /** + * Regression guard for the shared {@code STREAM_RELAY_BUFFER}: both the non-chunked + * ({@code /api/stream}) and chunked ({@code /api/chunked-out}) streaming paths reuse the same + * per-connection buffer now — sending one of each back to back on one connection must not + * leave either response corrupted by the other reusing the array mid-transfer. + */ + @Test + void testKeepAlive_streamedAndChunkedResponsesOnSameConnectionDontCorruptEachOther() throws Exception { + String streamReq = "GET /api/stream HTTP/1.1\r\nHost: localhost\r\n\r\n"; + String chunkedReq = "GET /api/chunked-out HTTP/1.1\r\nHost: localhost\r\n\r\n"; + + try (Socket socket = new Socket("127.0.0.1", port); + OutputStream out = socket.getOutputStream(); + InputStream in = socket.getInputStream()) { + socket.setSoTimeout(SOCKET_TIMEOUT_MS); + out.write(streamReq.getBytes(StandardCharsets.UTF_8)); + out.write(chunkedReq.getBytes(StandardCharsets.UTF_8)); + out.write(streamReq.getBytes(StandardCharsets.UTF_8)); + out.flush(); + + String first = readOneResponse(in); + String second = readOneResponse(in); + String third = readOneResponse(in); + + assertTrue(first.contains("Content-Length: 23")); + assertTrue(first.endsWith("streaming response body")); + assertTrue(second.contains("Transfer-Encoding: chunked")); + assertTrue(second.endsWith("streaming response body")); + assertTrue(third.contains("Content-Length: 23")); + assertTrue(third.endsWith("streaming response body")); + } + } } diff --git a/flash/src/test/java/dev/relism/flash/models/HeaderMapTest.java b/flash/src/test/java/dev/relism/flash/models/HeaderMapTest.java index dc9d14b..08a6ad0 100644 --- a/flash/src/test/java/dev/relism/flash/models/HeaderMapTest.java +++ b/flash/src/test/java/dev/relism/flash/models/HeaderMapTest.java @@ -90,4 +90,41 @@ class HeaderMapTest { assertTrue(map.all("Host").isEmpty()); assertTrue(map.all().isEmpty()); } + + // --- forEach --- + + @Test + void forEach_visitsEveryHeaderInDeclarationOrder() { + HeaderMap map = parse("Host: localhost", "Accept: text/plain", "Cookie: a=1"); + List seen = new java.util.ArrayList<>(); + map.forEach((name, value) -> seen.add(toStr(name) + "=" + toStr(value))); + assertEquals(List.of("Host=localhost", "Accept=text/plain", "Cookie=a=1"), seen); + } + + @Test + void forEach_emptyMap_neverInvokesConsumer() { + HeaderMap map = new HeaderMap(); + map.forEach((name, value) -> fail("must not be called on an empty map")); + } + + @Test + void forEach_reusesTheSameTwoViewInstancesAcrossEveryHeader() { + // The zero-allocation contract: forEach must reposition two ByteViews in place, not + // allocate a fresh pair per header — same instances across all three calls here. + HeaderMap map = parse("A: 1", "B: 2", "C: 3"); + List names = new java.util.ArrayList<>(); + List values = new java.util.ArrayList<>(); + map.forEach((name, value) -> { names.add(name); values.add(value); }); + + assertSame(names.get(0), names.get(1)); + assertSame(names.get(1), names.get(2)); + assertSame(values.get(0), values.get(1)); + assertSame(values.get(1), values.get(2)); + } + + private static String toStr(ByteView v) { + byte[] b = new byte[v.length()]; + for (int i = 0; i < b.length; i++) b[i] = v.byteAt(i); + return new String(b, StandardCharsets.UTF_8); + } } diff --git a/flash/src/test/java/dev/relism/flash/websocket/WebSocketSessionTest.java b/flash/src/test/java/dev/relism/flash/websocket/WebSocketSessionTest.java index dac5f12..21f202a 100644 --- a/flash/src/test/java/dev/relism/flash/websocket/WebSocketSessionTest.java +++ b/flash/src/test/java/dev/relism/flash/websocket/WebSocketSessionTest.java @@ -1,14 +1,79 @@ package dev.relism.flash.websocket; +import dev.relism.flash.http.HttpMethod; +import dev.relism.flash.models.HeaderMap; +import dev.relism.flash.models.Request; +import dev.relism.flash.models.RequestLine; +import dev.relism.fpr.core.ByteView; import org.junit.jupiter.api.Test; import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; +import java.nio.charset.StandardCharsets; import static org.junit.jupiter.api.Assertions.*; class WebSocketSessionTest { + private static ByteView viewOf(String s) { + byte[] bytes = s.getBytes(StandardCharsets.UTF_8); + return new ByteView() { + public int length() { return bytes.length; } + public byte byteAt(int idx) { return bytes[idx]; } + }; + } + + @Test + void request_returnsWhatWasPassedToConstructor() { + RequestLine line = new RequestLine(HttpMethod.GET, viewOf("/chat"), null, viewOf("HTTP/1.1"), new HeaderMap()); + Request req = new Request(line, new byte[0]); + WebSocketSession session = new WebSocketSession( + new ByteArrayInputStream(new byte[0]), new ByteArrayOutputStream(), 64, req, false); + + assertSame(req, session.request()); + } + + @Test + void request_defaultsToNullOnThreeArgConstructor() { + WebSocketSession session = new WebSocketSession(new ByteArrayInputStream(new byte[0]), new ByteArrayOutputStream(), 64); + + assertNull(session.request()); + } + + @Test + void sendText_masksWhenActingAsClient() throws Exception { + ByteArrayOutputStream out = new ByteArrayOutputStream(); + WebSocketSession session = new WebSocketSession( + new ByteArrayInputStream(new byte[0]), out, 64, null, true); + + byte[] payload = "hi".getBytes(); + session.sendText(payload, 0, payload.length); + + byte[] bytes = out.toByteArray(); + assertEquals((byte) 0x81, bytes[0]); // FIN + TEXT + assertEquals((byte) (0x80 | 2), bytes[1]); // masked bit + length 2 + byte m0 = bytes[2], m1 = bytes[3], m2 = bytes[4], m3 = bytes[5]; + assertEquals((byte) ('h' ^ m0), bytes[6]); + assertEquals((byte) ('i' ^ m1), bytes[7]); + // The caller's buffer is mutated in place by the mask (documented, zero-copy tradeoff). + assertEquals((byte) ('h' ^ m0), payload[0]); + } + + @Test + void sendText_doesNotMaskWhenActingAsServer() throws Exception { + ByteArrayOutputStream out = new ByteArrayOutputStream(); + WebSocketSession session = new WebSocketSession(new ByteArrayInputStream(new byte[0]), out, 64); + + byte[] payload = "hi".getBytes(); + session.sendText(payload, 0, payload.length); + + byte[] bytes = out.toByteArray(); + assertEquals((byte) 0x81, bytes[0]); + assertEquals((byte) 2, bytes[1]); // no masked bit + assertEquals('h', bytes[2]); + assertEquals('i', bytes[3]); + } + @Test void close_setsClosedAndWritesFrame() throws Exception { ByteArrayOutputStream out = new ByteArrayOutputStream();