README.md

@@ -17,59 +17,69 @@ It ensures compile-time guarantees for message structure, reduces runtime errors
 
 - Fully typed WebSocket sessions
 - Type-safe request/response mechanism
-- Optional typed notifications (Todo)
+- Typed notifications
 - Lightweight, minimal runtime overhead
 
 ## Installation
 
-#### Add the repository
+`session-java` is published to the Gitea package registry at `git.selimaj.dev`.
 
 ```groovy
 repositories {
-    maven { url 'https://jitpack.io' }
+    maven { url 'https://git.selimaj.dev/api/packages/selimaj-dev/maven' }
 }
-```
 
-#### Add the dependency
-```groovy
 dependencies {
-    implementation 'com.github.selimaj-dev:session-java:0.1.3'
+    implementation 'dev.selimaj:session-java:0.2.0'
 }
 ```
 
----
+Versions up to 0.1.6 were published through JitPack as `com.github.selimaj-dev:session-java`.
 
-### **Basic Example (client)**
+## Usage
 
 ```java
-Session session = Session.connect("ws://localhost:8080/");
+static final Method<String, String, String> DATA =
+        new Method<>("data", String.class, String.class, String.class);
 
-TextNode response = session.request(Methods.Data, TextNode.valueOf("Hello from client")).get();
+Session session = Session.connect("ws://localhost:8080/", 10, TimeUnit.SECONDS);
 
-System.out.println(response);
+// Handle requests and notifications from the server.
+session.onRequest(DATA, (id, data) -> SessionResult.ok("echo: " + data));
+session.onNotification(DATA, data -> System.out.println("notified: " + data));
+
+// Find out when the connection ends, from either side.
+session.onClose(() -> System.out.println("disconnected"));
+
+// Send a request, with or without a deadline.
+String reply = session.request(DATA, "Hello", Duration.ofSeconds(10)).get();
 
 session.close();
 ```
 
-### **Basic Example (server)**
+### Semantics
 
-#### ⚠️ Server functionality not done yet this is only for future examples
+- A request completes exceptionally with:
+  - `SessionErrorException` when the peer answers with an error response (`getError()` has the payload)
+  - `SessionClosedException` when the session closes before the response arrives
+  - `TimeoutException` when the timeout passes, if one was given
+- `onClose` handlers run exactly once, whether the session was closed locally, by the peer, or by a network failure. A handler registered after close runs immediately.
+- Sends are serialized, so the session is safe to use from multiple threads.
+- A request for an unknown method is answered with an error response instead of no reply.
 
-```java
-SessionServer server = SessionServer.bind("ws://localhost:8080/");
+## Publishing
 
-server.onClient((session, addr) -> {
-    session.onRequest(Methods.Data, (id, data) -> {
-        return SessionResult.ok(TextNode.valueOf("Response from server"));
-    });
-});
+```sh
+GITEA_TOKEN=<token with package:write> ./gradlew publish
 ```
 
 ## Protocol
 
+Every message is a JSON text frame with a `type` tag.
+
 #### Request
 
-The request `id` is separated from the peer, and will increment only on it's requests.
+The `id` is chosen by the sender and increments per peer.
 
 ```json
 { "type": "request", "id": 1, "method": "data", "data": "Hello from client" }
@@ -77,16 +87,22 @@ The request `id` is separated from the peer, and will increment only on it's req
 
 #### Response
 
-The response `id` **must** remain the same as the request.
+A response **must** carry the id of the request it answers.
 
 ```json
 { "type": "response", "id": 1, "result": "Hello from server" }
 ```
 
-#### Notifications
-
-A notification is a method that doesn't need validation or output, it simply notifies a peer for a specific information
+#### Error response
 
 ```json
-{ "type": "notification", "result": "Hello from server" }
+{ "type": "errorresponse", "id": 1, "error": "Invalid data" }
+```
+
+#### Notification
+
+A notification is fire-and-forget and gets no response.
+
+```json
+{ "type": "notification", "method": "data", "data": "Hello from server" }
 ```

build.gradle

@@ -3,13 +3,14 @@ plugins {
     id 'maven-publish'
 }
 
-group = "dev.selimaj.session"
-version = "0.1.3"
+group = "dev.selimaj"
+version = "0.2.0"
 
 java {
     toolchain {
         languageVersion = JavaLanguageVersion.of(17)
     }
+    withSourcesJar()
 }
 
 repositories {
@@ -17,18 +18,15 @@ repositories {
 }
 
 dependencies {
-    // JSON parsing and serialization
-    implementation 'com.fasterxml.jackson.core:jackson-databind:2.17.0'
-    implementation 'com.fasterxml.jackson.core:jackson-annotations:2.17.0'
-    implementation 'com.fasterxml.jackson.core:jackson-core:2.17.0'
+    // JSON parsing and serialization; JsonNode is part of the public API.
+    api 'com.fasterxml.jackson.core:jackson-databind:2.17.0'
+    api 'com.fasterxml.jackson.core:jackson-annotations:2.17.0'
+    api 'com.fasterxml.jackson.core:jackson-core:2.17.0'
 
-    // WebSocket API
-    implementation 'jakarta.websocket:jakarta.websocket-api:2.1.1'
-    // WebSocket client runtime
-    implementation 'org.glassfish.tyrus:tyrus-client:2.1.2'
-
-    // For testing
+    // For testing: a real WebSocket server for the client to talk to.
     testImplementation 'org.junit.jupiter:junit-jupiter:5.10.0'
+    testRuntimeOnly 'org.junit.platform:junit-platform-launcher'
+    testImplementation 'org.java-websocket:Java-WebSocket:1.6.0'
 }
 
 test {
@@ -38,7 +36,24 @@ test {
 publishing {
     publications {
         mavenJava(MavenPublication) {
+            artifactId = 'session-java'
             from components.java
         }
     }
+
+    repositories {
+        // Gitea package registry. Publish with a token that has package:write:
+        //   GITEA_TOKEN=... ./gradlew publish
+        maven {
+            name = 'Gitea'
+            url = uri('https://git.selimaj.dev/api/packages/selimaj-dev/maven')
+            credentials(HttpHeaderCredentials) {
+                name = 'Authorization'
+                value = "token ${findProperty('giteaToken') ?: System.getenv('GITEA_TOKEN') ?: ''}"
+            }
+            authentication {
+                header(HttpHeaderAuthentication)
+            }
+        }
+    }
 }

src/main/java/dev/selimaj/session/Session.java

@@ -3,6 +3,7 @@ package dev.selimaj.session;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.node.TextNode;
 
+import dev.selimaj.session.types.SessionClosedException;
 import dev.selimaj.session.types.SessionResult;
 import dev.selimaj.session.types.Message;
 import dev.selimaj.session.types.Method;
@@ -12,6 +13,7 @@ import dev.selimaj.session.types.NotificationHandler;
 import java.net.URI;
 import java.net.http.HttpClient;
 import java.net.http.WebSocket;
+import java.time.Duration;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 
@@ -48,27 +50,42 @@ public final class Session {
         return new Session(ws, listener);
     }
 
-    public void send(Message msg) throws Exception {
-        listener.send(ws, msg);
+    public CompletableFuture<Void> send(Message msg) {
+        return listener.send(ws, msg);
     }
 
-    public void respond(int to, JsonNode val) throws Exception {
-        listener.respond(ws, to, val);
+    public CompletableFuture<Void> respond(int to, JsonNode val) {
+        return listener.respond(ws, to, val);
     }
 
-    public void respondError(int to, JsonNode val) throws Exception {
-        listener.respondError(ws, to, val);
+    public CompletableFuture<Void> respondError(int to, JsonNode val) {
+        return listener.respondError(ws, to, val);
     }
 
-    public void notify(String method, Object data) throws Exception {
-        listener.notify(ws, method, data);
+    public CompletableFuture<Void> notify(String method, Object data) {
+        return listener.notify(ws, method, data);
     }
 
+    /**
+     * Sends a request and completes with the peer's response. Completes
+     * exceptionally with {@link dev.selimaj.session.types.SessionErrorException}
+     * on an error response, or {@link SessionClosedException} if the session
+     * closes first.
+     */
     public <Req, Res, Err> CompletableFuture<Res> request(
-            Method<Req, Res, Err> method, Req req) throws Exception {
+            Method<Req, Res, Err> method, Req req) {
         return listener.request(ws, method.getName(), req, method.getResClass());
     }
 
+    /**
+     * {@link #request(Method, Object)} that completes exceptionally with a
+     * {@link java.util.concurrent.TimeoutException} after {@code timeout}.
+     */
+    public <Req, Res, Err> CompletableFuture<Res> request(
+            Method<Req, Res, Err> method, Req req, Duration timeout) {
+        return request(method, req).orTimeout(timeout.toMillis(), TimeUnit.MILLISECONDS);
+    }
+
     public <Req, Res, Err> void onRequest(Method<Req, Res, Err> method,
             MethodHandler<Req, Res, Err> handler) {
         MethodHandler<JsonNode, JsonNode, JsonNode> wrapper = (id, value) -> {
@@ -76,7 +93,7 @@ public final class Session {
                 return handler.handle(id, listener.mapper.treeToValue(value, method.getReqClass()))
                         .intoJSON(listener.mapper);
             } catch (Exception e) {
-                return SessionResult.error(TextNode.valueOf(e.getMessage()));
+                return SessionResult.error(TextNode.valueOf(String.valueOf(e.getMessage())));
             }
         };
 
@@ -97,7 +114,28 @@ public final class Session {
         this.listener.notificationHandlers.put(method.getName(), wrapper);
     }
 
-    public void close() throws Exception {
-        ws.sendClose(0, "Session closed");
+    /**
+     * Runs {@code handler} once when the session closes, whether closed locally,
+     * by the peer, or by a connection failure. Runs immediately if the session
+     * is already closed.
+     */
+    public void onClose(Runnable handler) {
+        listener.closeHandlers.add(handler);
+
+        if (listener.isClosed()) {
+            listener.runCloseHandler(handler);
+        }
+    }
+
+    public boolean isOpen() {
+        return !listener.isClosed() && !ws.isInputClosed() && !ws.isOutputClosed();
+    }
+
+    public void close() {
+        listener.markClosed(new SessionClosedException("Session closed"));
+
+        if (!ws.isOutputClosed()) {
+            ws.sendClose(WebSocket.NORMAL_CLOSURE, "Session closed");
+        }
     }
 }

src/main/java/dev/selimaj/session/SessionListener.java

@@ -2,42 +2,81 @@ package dev.selimaj.session;
 
 import java.net.http.WebSocket;
 import java.nio.ByteBuffer;
+import java.util.List;
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.CompletionException;
 import java.util.concurrent.CompletionStage;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.TextNode;
 
 import dev.selimaj.session.types.Message;
 import dev.selimaj.session.types.MethodHandler;
 import dev.selimaj.session.types.NotificationHandler;
+import dev.selimaj.session.types.SessionClosedException;
+import dev.selimaj.session.types.SessionErrorException;
 
 public class SessionListener implements WebSocket.Listener {
     final ConcurrentHashMap<String, MethodHandler<JsonNode, JsonNode, JsonNode>> methods = new ConcurrentHashMap<>();
     final ConcurrentHashMap<String, NotificationHandler<JsonNode>> notificationHandlers = new ConcurrentHashMap<>();
+    final List<Runnable> closeHandlers = new CopyOnWriteArrayList<>();
 
     final ObjectMapper mapper = new ObjectMapper();
     private final ConcurrentHashMap<Integer, CompletableFuture<JsonNode>> pending = new ConcurrentHashMap<>();
     private final AtomicInteger id = new AtomicInteger();
+    private final AtomicBoolean closed = new AtomicBoolean();
+
+    /** Text frames of a message still being received. */
+    private final StringBuilder partial = new StringBuilder();
+
+    /**
+     * java.net.http.WebSocket allows only one outstanding text send, so every
+     * send is chained after the previous one.
+     */
+    private final Object sendLock = new Object();
+    private CompletableFuture<?> lastSend = CompletableFuture.completedFuture(null);
+
+    @Override
+    public void onOpen(WebSocket ws) {
+        ws.request(1);
+    }
 
     @Override
     public CompletionStage<?> onText(WebSocket ws, CharSequence data, boolean last) {
-        Message msg = null;
+        partial.append(data);
+
+        if (last) {
+            String text = partial.toString();
+            partial.setLength(0);
+            handleMessage(ws, text);
+        }
+
+        ws.request(1);
+        return null;
+    }
+
+    private void handleMessage(WebSocket ws, String text) {
+        Message msg;
 
         try {
-            msg = mapper.readValue(data.toString(), Message.class);
+            msg = mapper.readValue(text, Message.class);
         } catch (Exception e) {
             e.printStackTrace();
-            return null;
+            return;
         }
 
         if (msg instanceof Message.Request r) {
             MethodHandler<JsonNode, JsonNode, JsonNode> handler = methods.get(r.method());
 
-            if (handler != null) {
+            if (handler == null) {
+                respondError(ws, r.id(), TextNode.valueOf("Unknown method: " + r.method()));
+                return;
+            }
+
             try {
                 var res = handler.handle(r.id(), r.data());
 
@@ -46,8 +85,8 @@ public class SessionListener implements WebSocket.Listener {
                 } else {
                     respond(ws, r.id(), res.getJsonNode(mapper));
                 }
-                } catch (Exception ignored) {
-                }
+            } catch (Exception e) {
+                respondError(ws, r.id(), TextNode.valueOf(String.valueOf(e.getMessage())));
             }
         } else if (msg instanceof Message.Response r) {
             CompletableFuture<JsonNode> fut = pending.remove(r.id());
@@ -56,8 +95,7 @@ public class SessionListener implements WebSocket.Listener {
         } else if (msg instanceof Message.ErrorResponse r) {
             CompletableFuture<JsonNode> fut = pending.remove(r.id());
             if (fut != null)
-                fut.completeExceptionally(
-                        new RuntimeException(r.error().toString()));
+                fut.completeExceptionally(new SessionErrorException(r.error()));
         } else if (msg instanceof Message.Notification r) {
             NotificationHandler<JsonNode> handler = notificationHandlers.get(r.method());
 
@@ -69,14 +107,6 @@ public class SessionListener implements WebSocket.Listener {
                 }
             }
         }
-
-        ws.request(1);
-        return null;
-    }
-
-    @Override
-    public void onOpen(WebSocket ws) {
-        ws.request(1);
     }
 
     @Override
@@ -93,24 +123,80 @@ public class SessionListener implements WebSocket.Listener {
 
     @Override
     public CompletionStage<?> onClose(WebSocket ws, int statusCode, String reason) {
-        ws.request(1);
+        markClosed(new SessionClosedException("Closed by peer: " + statusCode + " " + reason));
         return null;
     }
 
-    void send(WebSocket ws, Message msg) throws Exception {
-        ws.sendText(mapper.writeValueAsString(msg), true);
+    @Override
+    public void onError(WebSocket ws, Throwable error) {
+        markClosed(new SessionClosedException("Connection failed", error));
     }
 
-    void respond(WebSocket ws, int to, JsonNode val) throws Exception {
-        send(ws, new Message.Response(to, val));
+    boolean isClosed() {
+        return closed.get();
     }
 
-    void respondError(WebSocket ws, int to, JsonNode val) throws Exception {
-        send(ws, new Message.ErrorResponse(to, val));
+    /**
+     * Fails every pending request and runs the close handlers, exactly once.
+     */
+    void markClosed(SessionClosedException cause) {
+        if (!closed.compareAndSet(false, true)) {
+            return;
         }
 
-    void notify(WebSocket ws, String method, Object data) throws Exception {
-        send(ws, new Message.Notification(method, mapper.valueToTree(data)));
+        pending.values().forEach(fut -> fut.completeExceptionally(cause));
+        pending.clear();
+
+        for (Runnable handler : closeHandlers) {
+            runCloseHandler(handler);
+        }
+    }
+
+    /** Runs a close handler unless another thread already claimed it. */
+    void runCloseHandler(Runnable handler) {
+        if (!closeHandlers.remove(handler)) {
+            return;
+        }
+
+        try {
+            handler.run();
+        } catch (Exception e) {
+            e.printStackTrace();
+        }
+    }
+
+    CompletableFuture<Void> send(WebSocket ws, Message msg) {
+        if (closed.get()) {
+            return CompletableFuture.failedFuture(new SessionClosedException("Session is closed"));
+        }
+
+        String text;
+        try {
+            text = mapper.writeValueAsString(msg);
+        } catch (Exception e) {
+            return CompletableFuture.failedFuture(e);
+        }
+
+        synchronized (sendLock) {
+            CompletableFuture<Void> sent = lastSend
+                    .handle((r, e) -> null)
+                    .thenCompose(v -> ws.sendText(text, true))
+                    .thenApply(w -> null);
+            lastSend = sent;
+            return sent;
+        }
+    }
+
+    CompletableFuture<Void> respond(WebSocket ws, int to, JsonNode val) {
+        return send(ws, new Message.Response(to, val));
+    }
+
+    CompletableFuture<Void> respondError(WebSocket ws, int to, JsonNode val) {
+        return send(ws, new Message.ErrorResponse(to, val));
+    }
+
+    CompletableFuture<Void> notify(WebSocket ws, String method, Object data) {
+        return send(ws, new Message.Notification(method, mapper.valueToTree(data)));
     }
 
     public <Req, Res> CompletableFuture<Res> request(
@@ -120,29 +206,33 @@ public class SessionListener implements WebSocket.Listener {
             Class<Res> resType) {
 
         int id = this.id.incrementAndGet();
+        CompletableFuture<JsonNode> response = new CompletableFuture<>();
 
-        return CompletableFuture
-                .runAsync(() -> {
-                    try {
-                        send(ws, new Message.Request(
-                                id,
-                                method,
-                                mapper.valueToTree(data)));
-                    } catch (Exception e) {
-                        throw new CompletionException(e);
+        // Registered before sending so a fast response can't be missed.
+        pending.put(id, response);
+
+        if (closed.get()) {
+            pending.remove(id);
+            return CompletableFuture.failedFuture(new SessionClosedException("Session is closed"));
         }
-                })
-                .thenCompose(v -> {
-                    CompletableFuture<JsonNode> fut = new CompletableFuture<>();
-                    pending.put(id, fut);
-                    return fut;
-                })
-                .thenApply(json -> {
+
+        send(ws, new Message.Request(id, method, mapper.valueToTree(data))).whenComplete((v, e) -> {
+            if (e != null) {
+                pending.remove(id);
+                response.completeExceptionally(e);
+            }
+        });
+
+        // Drop the entry however the caller's future ends (e.g. orTimeout).
+        CompletableFuture<Res> result = response.thenApply(json -> {
             try {
                 return mapper.treeToValue(json, resType);
             } catch (Exception e) {
-                        throw new CompletionException(e);
+                throw new java.util.concurrent.CompletionException(e);
             }
         });
+        result.whenComplete((r, e) -> pending.remove(id));
+
+        return result;
     }
 }

src/main/java/dev/selimaj/session/types/SessionClosedException.java

@@ -0,0 +1,12 @@
+package dev.selimaj.session.types;
+
+/** The session closed (or failed) before the operation could complete. */
+public class SessionClosedException extends RuntimeException {
+    public SessionClosedException(String message) {
+        super(message);
+    }
+
+    public SessionClosedException(String message, Throwable cause) {
+        super(message, cause);
+    }
+}

src/main/java/dev/selimaj/session/types/SessionErrorException.java

@@ -0,0 +1,18 @@
+package dev.selimaj.session.types;
+
+import com.fasterxml.jackson.databind.JsonNode;
+
+/** The peer answered a request with an error response. */
+public class SessionErrorException extends RuntimeException {
+    private final JsonNode error;
+
+    public SessionErrorException(JsonNode error) {
+        super(error == null ? "null" : error.isTextual() ? error.asText() : error.toString());
+        this.error = error;
+    }
+
+    /** The error payload exactly as the peer sent it. */
+    public JsonNode getError() {
+        return error;
+    }
+}

src/main/java/dev/selimaj/session/types/SessionResult.java

@@ -9,7 +9,7 @@ public record SessionResult<T, E>(boolean isError, Object value) {
     }
 
     public static <T, E> SessionResult<T, E> error(E err) {
-        return new SessionResult<>(false, err);
+        return new SessionResult<>(true, err);
     }
 
     public SessionResult<JsonNode, JsonNode> intoJSON(ObjectMapper mapper) {

src/test/java/dev/selimaj/session/Methods.java

@@ -1,10 +0,0 @@
-package dev.selimaj.session;
-
-import com.fasterxml.jackson.databind.node.TextNode;
-
-import dev.selimaj.session.types.Method;
-
-public class Methods {
-    public static final Method<TextNode, TextNode, TextNode> Data = new Method<>("data",
-            TextNode.class, TextNode.class, TextNode.class);
-}

src/test/java/dev/selimaj/session/SessionTest.java

@@ -1,34 +1,195 @@
 package dev.selimaj.session;
 
+import static org.junit.jupiter.api.Assertions.*;
+
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import org.java_websocket.enums.Opcode;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Test;
 
+import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.node.TextNode;
+
+import dev.selimaj.session.types.Method;
+import dev.selimaj.session.types.SessionClosedException;
+import dev.selimaj.session.types.SessionErrorException;
 import dev.selimaj.session.types.SessionResult;
 
-public class SessionTest {
-    static void testRequest() throws Exception {
-        Session session = Session.connect("ws://localhost:8080/");
+class SessionTest {
+    static final Method<String, String, String> ECHO = new Method<>("echo", String.class, String.class,
+            String.class);
+    static final Method<String, String, String> FAIL = new Method<>("fail", String.class, String.class,
+            String.class);
 
-        session.onRequest(Methods.Data, (id, data) -> {
-            return SessionResult.ok(TextNode.valueOf("Response from client"));
-        });
-
-        TextNode response = session.request(Methods.Data, TextNode.valueOf("Hello from client")).get();
-
-        System.out.println(response);
+    private final TestServer server = new TestServer();
+    private Session session;
 
+    @AfterEach
+    void tearDown() throws Exception {
+        if (session != null) {
             session.close();
         }
+        server.stop(1000);
+    }
+
+    private Session connect(String uri) throws Exception {
+        session = Session.connect(uri, 5, TimeUnit.SECONDS);
+        return session;
+    }
+
+    private static Throwable cause(CompletableFuture<?> future) {
+        ExecutionException e = assertThrows(ExecutionException.class, () -> future.get(5, TimeUnit.SECONDS));
+        return e.getCause();
+    }
 
     @Test
-    void requestTest() throws Exception {
+    void requestAndResponse() throws Exception {
+        Session s = connect(server.start(TestServer::echo));
+
+        assertEquals("hi", s.request(ECHO, "hi").get(5, TimeUnit.SECONDS));
     }
 
-    public static void main(String[] args) {
-        try {
-            testRequest();
-        } catch (Exception e) {
-            e.printStackTrace();
+    @Test
+    void errorResponseCompletesWithSessionErrorException() throws Exception {
+        Session s = connect(server.start((conn, msg) -> TestServer.send(conn,
+                "{\"type\":\"errorresponse\",\"id\":" + msg.get("id").asInt() + ",\"error\":\"nope\"}")));
+
+        Throwable cause = cause(s.request(FAIL, "x"));
+        SessionErrorException error = assertInstanceOf(SessionErrorException.class, cause);
+        assertEquals("nope", error.getMessage());
+        assertEquals(TextNode.valueOf("nope"), error.getError());
     }
+
+    @Test
+    void manyConcurrentRequestsAllComplete() throws Exception {
+        // The server answers synchronously, so responses race the client's
+        // bookkeeping, and requests are sent from many threads at once.
+        Session s = connect(server.start(TestServer::echo));
+
+        List<CompletableFuture<String>> futures = new ArrayList<>();
+        List<Thread> threads = new ArrayList<>();
+        for (int t = 0; t < 8; t++) {
+            int thread = t;
+            Thread th = new Thread(() -> {
+                for (int i = 0; i < 100; i++) {
+                    CompletableFuture<String> f = s.request(ECHO, thread + ":" + i);
+                    synchronized (futures) {
+                        futures.add(f);
+                    }
+                }
+            });
+            threads.add(th);
+            th.start();
+        }
+        for (Thread th : threads) {
+            th.join();
+        }
+
+        CompletableFuture.allOf(futures.toArray(CompletableFuture[]::new)).get(10, TimeUnit.SECONDS);
+        assertEquals(800, futures.size());
+    }
+
+    @Test
+    void pendingRequestsFailAndOnCloseRunsOnceWhenServerCloses() throws Exception {
+        Session s = connect(server.start((conn, msg) -> {
+        }));
+        AtomicInteger closes = new AtomicInteger();
+        s.onClose(closes::incrementAndGet);
+
+        CompletableFuture<String> pending = s.request(ECHO, "never answered");
+        server.received.poll(5, TimeUnit.SECONDS);
+        server.client.close(1001, "going away");
+
+        assertInstanceOf(SessionClosedException.class, cause(pending));
+        assertEquals(1, closes.get());
+        assertFalse(s.isOpen());
+        assertInstanceOf(SessionClosedException.class, cause(s.request(ECHO, "late")));
+    }
+
+    @Test
+    void requestTimeout() throws Exception {
+        Session s = connect(server.start((conn, msg) -> {
+        }));
+
+        assertInstanceOf(TimeoutException.class, cause(s.request(ECHO, "x", Duration.ofMillis(100))));
+    }
+
+    @Test
+    void fragmentedMessageIsReassembled() throws Exception {
+        Session s = connect(server.start((conn, msg) -> {
+            byte[] json = ("{\"type\":\"response\",\"id\":" + msg.get("id").asInt()
+                    + ",\"result\":\"reassembled\"}").getBytes(StandardCharsets.UTF_8);
+            int third = json.length / 3;
+            conn.sendFragmentedFrame(Opcode.TEXT, ByteBuffer.wrap(json, 0, third), false);
+            conn.sendFragmentedFrame(Opcode.TEXT, ByteBuffer.wrap(json, third, third), false);
+            conn.sendFragmentedFrame(Opcode.TEXT, ByteBuffer.wrap(json, 2 * third, json.length - 2 * third),
+                    true);
+        }));
+
+        assertEquals("reassembled", s.request(ECHO, "x").get(5, TimeUnit.SECONDS));
+    }
+
+    @Test
+    void handlerResultsAndUnknownMethodsAreAnswered() throws Exception {
+        Session s = connect(server.start((conn, msg) -> {
+        }));
+        s.onRequest(ECHO, (id, data) -> SessionResult.ok("echo:" + data));
+        s.onRequest(FAIL, (id, data) -> SessionResult.error("failed:" + data));
+
+        TestServer.send(server.client, "{\"type\":\"request\",\"id\":1,\"method\":\"echo\",\"data\":\"a\"}");
+        TestServer.send(server.client, "{\"type\":\"request\",\"id\":2,\"method\":\"fail\",\"data\":\"b\"}");
+        TestServer.send(server.client, "{\"type\":\"request\",\"id\":3,\"method\":\"nope\",\"data\":null}");
+
+        JsonNode ok = server.received.poll(5, TimeUnit.SECONDS);
+        JsonNode err = server.received.poll(5, TimeUnit.SECONDS);
+        JsonNode unknown = server.received.poll(5, TimeUnit.SECONDS);
+
+        assertEquals("response", ok.get("type").asText());
+        assertEquals("echo:a", ok.get("result").asText());
+        assertEquals("errorresponse", err.get("type").asText());
+        assertEquals("failed:b", err.get("error").asText());
+        assertEquals("errorresponse", unknown.get("type").asText());
+        assertEquals("Unknown method: nope", unknown.get("error").asText());
+    }
+
+    @Test
+    void notificationsAreDelivered() throws Exception {
+        Session s = connect(server.start((conn, msg) -> {
+        }));
+        CompletableFuture<String> received = new CompletableFuture<>();
+        s.onNotification(ECHO, received::complete);
+
+        TestServer.send(server.client, "{\"type\":\"notification\",\"method\":\"echo\",\"data\":\"hello\"}");
+
+        assertEquals("hello", received.get(5, TimeUnit.SECONDS));
+    }
+
+    @Test
+    void closeSendsNormalClosureAndFailsLaterRequests() throws Exception {
+        Session s = connect(server.start(TestServer::echo));
+        CountDownLatch closed = new CountDownLatch(1);
+        s.onClose(closed::countDown);
+
+        s.close();
+
+        assertTrue(closed.await(5, TimeUnit.SECONDS));
+        assertEquals(1000, server.closeCodes.poll(5, TimeUnit.SECONDS));
+        assertInstanceOf(SessionClosedException.class, cause(s.request(ECHO, "late")));
+
+        // Registered after close: runs immediately, once.
+        AtomicInteger late = new AtomicInteger();
+        s.onClose(late::incrementAndGet);
+        assertEquals(1, late.get());
     }
 }

src/test/java/dev/selimaj/session/TestServer.java

@@ -0,0 +1,84 @@
+package dev.selimaj.session;
+
+import java.net.InetSocketAddress;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.TimeUnit;
+import java.util.function.BiConsumer;
+
+import org.java_websocket.WebSocket;
+import org.java_websocket.handshake.ClientHandshake;
+import org.java_websocket.server.WebSocketServer;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+
+/** An in-process protocol peer. Each received JSON message goes to {@link #handler}. */
+class TestServer extends WebSocketServer {
+    static final ObjectMapper MAPPER = new ObjectMapper();
+
+    private final CountDownLatch started = new CountDownLatch(1);
+    final BlockingQueue<JsonNode> received = new LinkedBlockingQueue<>();
+    final BlockingQueue<Integer> closeCodes = new LinkedBlockingQueue<>();
+    volatile WebSocket client;
+    volatile BiConsumer<WebSocket, JsonNode> handler = (conn, msg) -> {
+    };
+
+    TestServer() {
+        super(new InetSocketAddress("127.0.0.1", 0));
+        setReuseAddr(true);
+    }
+
+    String start(BiConsumer<WebSocket, JsonNode> handler) throws InterruptedException {
+        this.handler = handler;
+        start();
+        if (!started.await(5, TimeUnit.SECONDS)) {
+            throw new IllegalStateException("test server did not start");
+        }
+        return "ws://127.0.0.1:" + getPort() + "/";
+    }
+
+    /** Echoes request data back as the result. */
+    static void echo(WebSocket conn, JsonNode msg) {
+        if ("request".equals(msg.path("type").asText())) {
+            send(conn, "{\"type\":\"response\",\"id\":" + msg.get("id").asInt()
+                    + ",\"result\":" + msg.get("data") + "}");
+        }
+    }
+
+    static void send(WebSocket conn, String json) {
+        conn.send(json);
+    }
+
+    @Override
+    public void onOpen(WebSocket conn, ClientHandshake handshake) {
+        client = conn;
+    }
+
+    @Override
+    public void onClose(WebSocket conn, int code, String reason, boolean remote) {
+        closeCodes.add(code);
+    }
+
+    @Override
+    public void onMessage(WebSocket conn, String message) {
+        try {
+            JsonNode msg = MAPPER.readTree(message);
+            received.add(msg);
+            handler.accept(conn, msg);
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public void onError(WebSocket conn, Exception ex) {
+        ex.printStackTrace();
+    }
+
+    @Override
+    public void onStart() {
+        started.countDown();
+    }
+}