From b76622afeeaaad8ac57451bd6364f259ac40050f Mon Sep 17 00:00:00 2001 From: kvizyx Date: Fri, 2 Oct 2026 03:55:47 +0300 Subject: [PATCH 1/6] feat: add PreparedMessage for writing a message to many connections Broadcasting the same message with compression enabled compressed it once per connection, which dominates CPU cost when fanning out to many peers. PreparedMessage compresses the payload once and shares it across connections that compress without context takeover while other connections use the regular write path. --- README.md | 8 +- conn_test.go | 142 ++++++++++++++++ example_test.go | 10 +- export_test.go | 5 + internal/examples/broadcast/README.md | 37 +++++ internal/examples/broadcast/main.go | 59 +++++++ internal/examples/broadcast/server.go | 181 +++++++++++++++++++++ internal/examples/broadcast/server_test.go | 69 ++++++++ prepared.go | 85 ++++++++++ prepared_test.go | 89 ++++++++++ write.go | 33 ++++ ws_js.go | 23 ++- 12 files changed, 734 insertions(+), 7 deletions(-) create mode 100644 internal/examples/broadcast/README.md create mode 100644 internal/examples/broadcast/main.go create mode 100644 internal/examples/broadcast/server.go create mode 100644 internal/examples/broadcast/server_test.go create mode 100644 prepared.go create mode 100644 prepared_test.go diff --git a/README.md b/README.md index 3f174d5b..2bab4eb8 100644 --- a/README.md +++ b/README.md @@ -48,10 +48,11 @@ See GitHub issues for minor issues but the major future enhancements are: ## Examples -For a production quality example that demonstrates the complete API, see the +- Production quality example that demonstrates the complete API, see the [echo example](./internal/examples/echo). - -For a full stack example, see the [chat example](./internal/examples/chat). +- Full stack example, see the [chat example](./internal/examples/chat). +- Broadcasting a message to many connections with `PreparedMessage`, see the +[broadcast example](./internal/examples/broadcast). ### Server @@ -107,7 +108,6 @@ c.Close(websocket.StatusNormalClosure, "") Advantages of [gorilla/websocket](https://github.com/gorilla/websocket): - Mature and widely used -- [Prepared writes](https://pkg.go.dev/github.com/gorilla/websocket#PreparedMessage) - Configurable [buffer sizes](https://pkg.go.dev/github.com/gorilla/websocket#hdr-Buffers) Advantages of github.com/coder/websocket: diff --git a/conn_test.go b/conn_test.go index 01af8cd4..b1095603 100644 --- a/conn_test.go +++ b/conn_test.go @@ -61,6 +61,69 @@ func TestConn(t *testing.T) { } }) + t.Run("writePreparedConcurrent", func(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithTimeout(context.Background(), time.Second*30) + defer cancel() + + exp := strings.Repeat("prepared", 128) + pm := websocket.NewPreparedMessage(websocket.MessageText, []byte(exp)) + + modes := []websocket.CompressionMode{ + websocket.CompressionDisabled, + websocket.CompressionContextTakeover, + websocket.CompressionNoContextTakeover, + } + + // Write the same message with every compression mode from both sides, + // so connections with different settings share it and client masking + // runs concurrently on the shared payload. + var errs []<-chan error + for i := range 12 { + mode := modes[i%len(modes)] + client, server := wstest.Pipe(&websocket.DialOptions{ + CompressionMode: mode, + }, &websocket.AcceptOptions{ + CompressionMode: mode, + }) + t.Cleanup(func() { + client.CloseNow() + server.CloseNow() + }) + + w, r := client, server + if i%2 == 0 { + w, r = server, client + } + errs = append(errs, xsync.Go(func() error { + for range 5 { + err := w.WritePrepared(ctx, pm) + if err != nil { + return err + } + } + return nil + })) + errs = append(errs, xsync.Go(func() error { + for range 5 { + _, p, err := r.Read(ctx) + if err != nil { + return err + } + if string(p) != exp { + return fmt.Errorf("unexpected msg: %q", p) + } + } + return nil + })) + } + + for _, errc := range errs { + assert.Success(t, <-errc) + } + }) + t.Run("badClose", func(t *testing.T) { tt, c1, c2 := newConnTest(t, nil, nil) @@ -667,6 +730,85 @@ func BenchmarkConn(b *testing.B) { } } +func BenchmarkBroadcast(b *testing.B) { + const conns = 16 + + modes := []struct { + name string + mode websocket.CompressionMode + }{ + {"disabledCompress", websocket.CompressionDisabled}, + {"compressContextTakeover", websocket.CompressionContextTakeover}, + {"compressNoContext", websocket.CompressionNoContextTakeover}, + } + writes := []struct { + name string + write func(ctx context.Context, cs []*websocket.Conn, msg []byte) error + }{ + {"write", func(ctx context.Context, cs []*websocket.Conn, msg []byte) error { + for _, c := range cs { + err := c.Write(ctx, websocket.MessageText, msg) + if err != nil { + return err + } + } + return nil + }}, + {"writePrepared", func(ctx context.Context, cs []*websocket.Conn, msg []byte) error { + pm := websocket.NewPreparedMessage(websocket.MessageText, msg) + for _, c := range cs { + err := c.WritePrepared(ctx, pm) + if err != nil { + return err + } + } + return nil + }}, + } + + msg := []byte(strings.Repeat("1234", 256)) + + // Clients mask every write, so they are benchmarked separately. + for _, side := range []string{"server", "client"} { + for _, m := range modes { + for _, w := range writes { + b.Run(side+"/"+m.name+"/"+w.name, func(b *testing.B) { + ctx, cancel := context.WithTimeout(context.Background(), time.Minute) + defer cancel() + + writers := make([]*websocket.Conn, conns) + for i := range writers { + client, server := wstest.Pipe(&websocket.DialOptions{ + CompressionMode: m.mode, + }, &websocket.AcceptOptions{ + CompressionMode: m.mode, + }) + b.Cleanup(func() { + client.CloseNow() + server.CloseNow() + }) + writers[i] = server + if side == "client" { + writers[i] = client + } + writers[i].DiscardWrites() + } + + b.SetBytes(int64(len(msg) * conns)) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + err := w.write(ctx, writers, msg) + if err != nil { + b.Fatal(err) + } + } + }) + } + } + } +} + func echoServer(w http.ResponseWriter, r *http.Request, opts *websocket.AcceptOptions) (err error) { defer errd.Wrap(&err, "echo server failed") diff --git a/example_test.go b/example_test.go index 7026e311..aa22762d 100644 --- a/example_test.go +++ b/example_test.go @@ -118,7 +118,7 @@ func Example_writeOnly() { } func Example_crossOrigin() { - // This handler demonstrates how to safely accept cross origin WebSockets + // This handler demonstrates how to safely accept cross-origin WebSockets // from the origin example.com. fn := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { c, err := websocket.Accept(w, r, &websocket.AcceptOptions{ @@ -165,7 +165,13 @@ func Example_fullStackChat() { // https://github.com/nhooyr/websocket/tree/master/internal/examples/chat } -// This example demonstrates a echo server. +// This example demonstrates an echo server. func Example_echo() { // https://github.com/nhooyr/websocket/tree/master/internal/examples/echo } + +// This example demonstrates broadcasting messages to many subscribers with +// a PreparedMessage. +func ExampleConn_writePrepared() { + // https://github.com/coder/websocket/tree/master/internal/examples/broadcast +} diff --git a/export_test.go b/export_test.go index e0071922..41ceb59c 100644 --- a/export_test.go +++ b/export_test.go @@ -3,6 +3,7 @@ package websocket import ( + "io" "net" "github.com/coder/websocket/internal/util" @@ -17,6 +18,10 @@ func (c *Conn) RecordBytesWritten() *int { return &bytesWritten } +func (c *Conn) DiscardWrites() { + c.bw.Reset(io.Discard) +} + func (c *Conn) RecordBytesRead() *int { var bytesRead int c.br.Reset(util.ReaderFunc(func(p []byte) (int, error) { diff --git a/internal/examples/broadcast/README.md b/internal/examples/broadcast/README.md new file mode 100644 index 00000000..b0ce57e4 --- /dev/null +++ b/internal/examples/broadcast/README.md @@ -0,0 +1,37 @@ +# Broadcast Example + +This directory contains a broadcast server example using github.com/coder/websocket. + +```bash +$ cd internal/examples/broadcast +$ go run . localhost:8080 +listening on http://127.0.0.1:8080 +``` + +Subscribe with a WebSocket client like [websocat](https://github.com/vi/websocat) and publish +messages with curl: + +```bash +$ websocat ws://127.0.0.1:8080/subscribe +$ curl --data-binary 'hello' http://127.0.0.1:8080/publish +``` + +Every published message is delivered to all subscribers. + +## Structure + +The server is in `server.go`. Subscribers connect to the WebSocket `/subscribe` endpoint and +messages are published via the HTTP POST `/publish` endpoint, so that you can easily publish +with curl. + +Each published message is wrapped in a single `PreparedMessage` and queued for every subscriber. +Connections are accepted with `CompressionNoContextTakeover`, so a message that reaches the +compression threshold is compressed once per broadcast instead of once per subscriber. +Subscribers that do not support compression receive the same message uncompressed. + +Publishing never blocks: a subscriber that cannot keep up with its queue is disconnected. + +`server_test.go` contains a test that subscribes clients with and without compression and +ensures every client receives every published message. + +`main.go` brings it all together so that you can run it and play around with it. diff --git a/internal/examples/broadcast/main.go b/internal/examples/broadcast/main.go new file mode 100644 index 00000000..1edb37a5 --- /dev/null +++ b/internal/examples/broadcast/main.go @@ -0,0 +1,59 @@ +package main + +import ( + "context" + "errors" + "log" + "net" + "net/http" + "os" + "os/signal" + "time" +) + +func main() { + log.SetFlags(0) + + err := run() + if err != nil { + log.Fatal(err) + } +} + +// run starts a http.Server for the passed in address +// with all requests handled by broadcastServer. +func run() error { + if len(os.Args) < 2 { + return errors.New("please provide an address to listen on as the first argument") + } + + l, err := net.Listen("tcp", os.Args[1]) + if err != nil { + return err + } + log.Printf("listening on http://%v", l.Addr()) + + s := &http.Server{ + Handler: newBroadcastServer(log.Printf), + ReadTimeout: time.Second * 10, + WriteTimeout: time.Second * 10, + } + errc := make(chan error, 1) + go func() { + errc <- s.Serve(l) + }() + + sigs := make(chan os.Signal, 1) + signal.Notify(sigs, os.Interrupt) + select { + case err := <-errc: + log.Printf("failed to serve: %v", err) + case sig := <-sigs: + log.Printf("terminating: %v", sig) + } + + ctx, cancel := context.WithTimeout(context.Background(), time.Second*10) + defer cancel() + + return s.Shutdown(ctx) +} diff --git a/internal/examples/broadcast/server.go b/internal/examples/broadcast/server.go new file mode 100644 index 00000000..dc7113f4 --- /dev/null +++ b/internal/examples/broadcast/server.go @@ -0,0 +1,181 @@ +package main + +import ( + "context" + "errors" + "io" + "net/http" + "sync" + "time" + + "github.com/coder/websocket" +) + +// broadcastServer broadcasts every message published to /publish to all +// WebSocket clients subscribed to /subscribe. +// +// Each message is wrapped in a single websocket.PreparedMessage, so it is +// compressed once per broadcast instead of once per subscriber. +type broadcastServer struct { + // subscriberMessageBuffer controls the max number + // of messages that can be queued for a subscriber + // before it is kicked. + // + // Defaults to 16. + subscriberMessageBuffer int + + // logf controls where logs are sent. + logf func(f string, v ...any) + + // serveMux routes the various endpoints to the appropriate handler. + serveMux http.ServeMux + + subscribersMu sync.Mutex + subscribers map[*subscriber]struct{} +} + +// newBroadcastServer constructs a broadcastServer with the defaults. +func newBroadcastServer(logf func(f string, v ...any)) *broadcastServer { + bs := &broadcastServer{ + subscriberMessageBuffer: 16, + logf: logf, + subscribers: make(map[*subscriber]struct{}), + } + bs.serveMux.HandleFunc("/subscribe", bs.subscribeHandler) + bs.serveMux.HandleFunc("/publish", bs.publishHandler) + + return bs +} + +// subscriber represents a subscriber. +// Messages are sent on the msgs channel and if the client +// cannot keep up with the messages, the slow channel is closed. +type subscriber struct { + msgs chan *websocket.PreparedMessage + slow chan struct{} + slowOnce sync.Once +} + +func (s *subscriber) closeSlow() { + s.slowOnce.Do(func() { + close(s.slow) + }) +} + +func (bs *broadcastServer) ServeHTTP(w http.ResponseWriter, r *http.Request) { + bs.serveMux.ServeHTTP(w, r) +} + +// subscribeHandler accepts the WebSocket connection and then subscribes +// it to all future messages. +func (bs *broadcastServer) subscribeHandler(w http.ResponseWriter, r *http.Request) { + err := bs.subscribe(w, r) + if errors.Is(err, context.Canceled) { + return + } + if websocket.CloseStatus(err) == websocket.StatusNormalClosure || + websocket.CloseStatus(err) == websocket.StatusGoingAway { + return + } + if err != nil { + bs.logf("%v", err) + return + } +} + +// publishHandler reads the request body with a limit of 65536 bytes and then +// publishes the received message. +func (bs *broadcastServer) publishHandler(w http.ResponseWriter, r *http.Request) { + if r.Method != "POST" { + http.Error(w, http.StatusText(http.StatusMethodNotAllowed), http.StatusMethodNotAllowed) + return + } + body := http.MaxBytesReader(w, r.Body, 65536) + msg, err := io.ReadAll(body) + if err != nil { + http.Error(w, http.StatusText(http.StatusRequestEntityTooLarge), http.StatusRequestEntityTooLarge) + return + } + + bs.publish(websocket.NewPreparedMessage(websocket.MessageText, msg)) + + w.WriteHeader(http.StatusAccepted) +} + +// subscribe subscribes the given WebSocket to all broadcast messages. +// The subscriber is registered before the WebSocket is accepted, so it +// receives every message published after the client's Dial returns. +// +// It uses CloseRead to keep reading from the connection to process control +// messages and cancel the context if the connection drops. +func (bs *broadcastServer) subscribe(w http.ResponseWriter, r *http.Request) error { + s := &subscriber{ + msgs: make(chan *websocket.PreparedMessage, bs.subscriberMessageBuffer), + slow: make(chan struct{}), + } + bs.addSubscriber(s) + defer bs.deleteSubscriber(s) + + // Without context takeover, connections compress messages independently + // of each other, which allows sharing the compressed payload. + c, err := websocket.Accept(w, r, &websocket.AcceptOptions{ + CompressionMode: websocket.CompressionNoContextTakeover, + }) + if err != nil { + return err + } + defer c.CloseNow() + + ctx := c.CloseRead(context.Background()) + + for { + select { + case pm := <-s.msgs: + err := writeTimeout(ctx, time.Second*5, c, pm) + if err != nil { + return err + } + case <-s.slow: + return c.Close(websocket.StatusPolicyViolation, "connection too slow to keep up with messages") + case <-ctx.Done(): + return ctx.Err() + } + } +} + +// publish publishes the pm to all subscribers. +// It never blocks and so messages to slow subscribers +// are dropped. +func (bs *broadcastServer) publish(pm *websocket.PreparedMessage) { + bs.subscribersMu.Lock() + defer bs.subscribersMu.Unlock() + + for s := range bs.subscribers { + select { + case s.msgs <- pm: + default: + s.closeSlow() + } + } +} + +// addSubscriber registers a subscriber. +func (bs *broadcastServer) addSubscriber(s *subscriber) { + bs.subscribersMu.Lock() + bs.subscribers[s] = struct{}{} + bs.subscribersMu.Unlock() +} + +// deleteSubscriber deletes the given subscriber. +func (bs *broadcastServer) deleteSubscriber(s *subscriber) { + bs.subscribersMu.Lock() + delete(bs.subscribers, s) + bs.subscribersMu.Unlock() +} + +func writeTimeout(ctx context.Context, timeout time.Duration, c *websocket.Conn, pm *websocket.PreparedMessage) error { + ctx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + return c.WritePrepared(ctx, pm) +} diff --git a/internal/examples/broadcast/server_test.go b/internal/examples/broadcast/server_test.go new file mode 100644 index 00000000..77b21f01 --- /dev/null +++ b/internal/examples/broadcast/server_test.go @@ -0,0 +1,69 @@ +package main + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/coder/websocket" +) + +// Test_broadcastServer subscribes clients with and without compression, +// publishes messages below and above the compression threshold and ensures +// every client receives all of them. +func Test_broadcastServer(t *testing.T) { + t.Parallel() + + s := httptest.NewServer(newBroadcastServer(t.Logf)) + defer s.Close() + + ctx, cancel := context.WithTimeout(context.Background(), time.Minute) + defer cancel() + + var clients []*websocket.Conn + for i := 0; i < 4; i++ { + mode := websocket.CompressionDisabled + if i%2 == 0 { + mode = websocket.CompressionNoContextTakeover + } + c, _, err := websocket.Dial(ctx, s.URL+"/subscribe", &websocket.DialOptions{ + CompressionMode: mode, + }) + if err != nil { + t.Fatal(err) + } + defer c.Close(websocket.StatusInternalError, "the sky is falling") + clients = append(clients, c) + } + + msgs := []string{ + "hello", + strings.Repeat("hello world ", 512), + } + for _, msg := range msgs { + resp, err := http.Post(s.URL+"/publish", "text/plain", strings.NewReader(msg)) + if err != nil { + t.Fatal(err) + } + resp.Body.Close() + if resp.StatusCode != http.StatusAccepted { + t.Fatalf("expected status %v but got %v", http.StatusAccepted, resp.StatusCode) + } + } + + for i, c := range clients { + for _, exp := range msgs { + _, b, err := c.Read(ctx) + if err != nil { + t.Fatal(err) + } + if string(b) != exp { + t.Fatalf("client %v: expected %q but got %q", i, exp, b) + } + } + c.Close(websocket.StatusNormalClosure, "") + } +} diff --git a/prepared.go b/prepared.go new file mode 100644 index 00000000..290dd7d2 --- /dev/null +++ b/prepared.go @@ -0,0 +1,85 @@ +//go:build !js + +package websocket + +import ( + "bytes" + "fmt" + "sync" +) + +// PreparedMessage is a message that can be written to many connections +// efficiently with Conn.WritePrepared, such as when broadcasting the same +// update to every subscriber. +// +// With compression enabled, most of the cost of writing a message is +// compressing it. Conn.Write compresses the message once per connection, while +// Conn.WritePrepared compresses it once and shares the compressed payload between +// connections. Sharing is only possible when all the following hold: +// +// - The connection compresses without context takeover. For a server, this +// is the case with CompressionNoContextTakeover, or with +// CompressionContextTakeover when the client negotiates +// server_no_context_takeover. For a client, client_no_context_takeover +// plays the same role. +// - The message is at least as long as the compression threshold, 512 bytes by +// default for CompressionNoContextTakeover. Smaller messages are not +// compressed. +// +// Otherwise, Conn.WritePrepared costs the same as Conn.Write. With context takeover, +// each connection compresses against the sliding window of its own previous +// messages, so the compressed payload differs per connection and cannot be +// shared. Without compression there is no expensive work to share. Clients +// generate a new masking key for every write in all cases. +// +// The compressed payload is computed on the first write that needs it and +// kept for the lifetime of the PreparedMessage, so a PreparedMessage can be +// written any number of times, for example a greeting sent to every new +// connection. It is safe for concurrent use. +// +// In Wasm, the browser handles framing and compression, so Conn.WritePrepared is +// equivalent to Conn.Write. +type PreparedMessage struct { + typ MessageType + data []byte + + compressOnce sync.Once + compressed []byte + compressErr error +} + +// NewPreparedMessage returns a PreparedMessage of the specified type with cloned data. +func NewPreparedMessage(typ MessageType, data []byte) *PreparedMessage { + return &PreparedMessage{ + typ: typ, + data: bytes.Clone(data), + } +} + +// compress returns the payload of PreparedMessage compressed without context takeover. +func (pm *PreparedMessage) compress() ([]byte, error) { + pm.compressOnce.Do(func() { + var buf bytes.Buffer + + tw := &trimLastFourBytesWriter{ + w: &buf, + } + fw := getFlateWriter(tw) + defer putFlateWriter(fw) + + _, err := fw.Write(pm.data) + if err != nil { + pm.compressErr = fmt.Errorf("failed to compress: %w", err) + return + } + + err = fw.Flush() + if err != nil { + pm.compressErr = fmt.Errorf("failed to flush compression: %w", err) + return + } + + pm.compressed = buf.Bytes() + }) + return pm.compressed, pm.compressErr +} diff --git a/prepared_test.go b/prepared_test.go new file mode 100644 index 00000000..526afac0 --- /dev/null +++ b/prepared_test.go @@ -0,0 +1,89 @@ +//go:build !js + +package websocket + +import ( + "bufio" + "bytes" + "context" + "strings" + "testing" + + "github.com/coder/websocket/internal/test/assert" +) + +type bufConn struct { + bytes.Buffer +} + +func (*bufConn) Close() error { + return nil +} + +// TestWritePreparedWire verifies that WritePrepared puts the same bytes on the +// wire as Write, including for messages written after it, and that the +// compressed payload is shared only when the connection compresses the +// message without context takeover. +func TestWritePreparedWire(t *testing.T) { + t.Parallel() + + largeMsg := strings.Repeat("prepared message ", 100) + smallMsg := "small" + nextMsg := []byte(strings.Repeat("next message ", 100)) + + testCases := []struct { + name string + copts *compressionOptions + msg string + shared bool // Whether the compressed payload of the message is shared. + }{ + {"Disabled", nil, largeMsg, false}, + {"ContextTakeover/AboveThreshold", CompressionContextTakeover.opts(), largeMsg, false}, + {"ContextTakeover/BelowThreshold", CompressionContextTakeover.opts(), smallMsg, false}, + {"NoContextTakeover/AboveThreshold", CompressionNoContextTakeover.opts(), largeMsg, true}, + {"NoContextTakeover/AtThreshold", CompressionNoContextTakeover.opts(), strings.Repeat("a", 512), true}, + {"NoContextTakeover/BelowThreshold", CompressionNoContextTakeover.opts(), strings.Repeat("a", 511), false}, + // Only the server side of the negotiated options applies to server writes. + {"ServerNoContextTakeover", &compressionOptions{serverNoContextTakeover: true}, largeMsg, true}, + {"ClientNoContextTakeover", &compressionOptions{clientNoContextTakeover: true}, largeMsg, false}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + newServerConn := func() (*Conn, *bufConn) { + rwc := &bufConn{} + c := newConn(connConfig{ + rwc: rwc, + copts: tc.copts, + br: bufio.NewReader(rwc), + bw: bufio.NewWriter(rwc), + }) + t.Cleanup(func() { + c.CloseNow() + }) + return c, rwc + } + + ctx := context.Background() + + data := []byte(tc.msg) + pm := NewPreparedMessage(MessageText, data) + copy(data, "modified") // NewPreparedMessage must copy data. + + want, wantBuf := newServerConn() + got, gotBuf := newServerConn() + for range 2 { + assert.Success(t, want.Write(ctx, MessageText, []byte(tc.msg))) + assert.Success(t, want.Write(ctx, MessageBinary, nextMsg)) + + assert.Success(t, got.WritePrepared(ctx, pm)) + assert.Success(t, got.Write(ctx, MessageBinary, nextMsg)) + } + + assert.Equal(t, "wire bytes", wantBuf.Bytes(), gotBuf.Bytes()) + assert.Equal(t, "shared", tc.shared, pm.compressed != nil) + }) + } +} diff --git a/write.go b/write.go index 81bab4fa..0fe1954e 100644 --- a/write.go +++ b/write.go @@ -48,6 +48,39 @@ func (c *Conn) Write(ctx context.Context, typ MessageType, p []byte) error { return nil } +// WritePrepared writes a prepared message to the connection. +// +// It behaves like Write but reuses the compressed payload of pm across +// connections where possible. See PreparedMessage. +func (c *Conn) WritePrepared(ctx context.Context, pm *PreparedMessage) error { + err := c.writePrepared(ctx, pm) + if err != nil { + return fmt.Errorf("failed to write prepared msg: %w", err) + } + return nil +} + +func (c *Conn) writePrepared(ctx context.Context, pm *PreparedMessage) error { + if !c.flate() || len(pm.data) < c.flateThreshold || c.msgWriter.flateContextTakeover() { + _, err := c.write(ctx, pm.typ, pm.data) + return err + } + + p, err := pm.compress() + if err != nil { + return err + } + + err = c.msgWriter.reset(ctx, pm.typ) + if err != nil { + return err + } + defer c.msgWriter.mu.unlock() + + _, err = c.writeFrame(ctx, true, true, c.msgWriter.opcode, p) + return err +} + type msgWriter struct { c *Conn diff --git a/ws_js.go b/ws_js.go index 026b75fc..cdf25a38 100644 --- a/ws_js.go +++ b/ws_js.go @@ -194,7 +194,7 @@ func (c *Conn) Ping(ctx context.Context) error { } // Write writes a message of the given type to the connection. -// Always non blocking. +// Always non-blocking. func (c *Conn) Write(ctx context.Context, typ MessageType, p []byte) error { err := c.write(typ, p) if err != nil { @@ -210,6 +210,27 @@ func (c *Conn) Write(ctx context.Context, typ MessageType, p []byte) error { return nil } +// PreparedMessage is a message that can be written to many connections with +// Conn.WritePrepared. +type PreparedMessage struct { + typ MessageType + data []byte +} + +// NewPreparedMessage returns a PreparedMessage of the specified type with cloned data. +func NewPreparedMessage(typ MessageType, data []byte) *PreparedMessage { + return &PreparedMessage{ + typ: typ, + data: bytes.Clone(data), + } +} + +// WritePrepared writes a prepared message to the connection. +// The browser handles framing and compression, so it is equivalent to Write. +func (c *Conn) WritePrepared(ctx context.Context, pm *PreparedMessage) error { + return c.Write(ctx, pm.typ, pm.data) +} + func (c *Conn) write(typ MessageType, p []byte) error { if c.isClosed() { return net.ErrClosed From 2d37f1cf0310e288736674bbc25a32109984dd30 Mon Sep 17 00:00:00 2001 From: kvizyx Date: Fri, 2 Oct 2026 16:36:49 +0300 Subject: [PATCH 2/6] docs: add prepared message to highlights --- README.md | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 2bab4eb8..f9f00241 100644 --- a/README.md +++ b/README.md @@ -29,6 +29,7 @@ go get github.com/coder/websocket - [net.Conn](https://pkg.go.dev/github.com/coder/websocket#NetConn) wrapper - [Ping pong](https://pkg.go.dev/github.com/coder/websocket#Conn.Ping) API - [RFC 7692](https://tools.ietf.org/html/rfc7692) permessage-deflate compression +- [Prepared message](https://pkg.go.dev/github.com/coder/websocket#PreparedMessage) for efficient broadcasting - [CloseRead](https://pkg.go.dev/github.com/coder/websocket#Conn.CloseRead) helper for write only connections - Compile to [Wasm](https://pkg.go.dev/github.com/coder/websocket#hdr-Wasm) @@ -49,10 +50,10 @@ See GitHub issues for minor issues but the major future enhancements are: ## Examples - Production quality example that demonstrates the complete API, see the -[echo example](./internal/examples/echo). + [echo example](./internal/examples/echo). - Full stack example, see the [chat example](./internal/examples/chat). -- Broadcasting a message to many connections with `PreparedMessage`, see the -[broadcast example](./internal/examples/broadcast). +- Broadcasting a message to many connections, see the + [broadcast example](./internal/examples/broadcast). ### Server From 1ac0d68db14c7c6d879f779140e7d02a69d9cea9 Mon Sep 17 00:00:00 2001 From: kvizyx Date: Fri, 2 Oct 2026 16:38:30 +0300 Subject: [PATCH 3/6] feat: skip compression for messages it does not make smaller --- compress.go | 5 ++++- compress_test.go | 9 +++++++++ prepared.go | 19 ++++++++++++++----- prepared_test.go | 7 ++++++- write.go | 18 +++++++++++++++--- 5 files changed, 48 insertions(+), 10 deletions(-) diff --git a/compress.go b/compress.go index 41bd5bdb..c8b67469 100644 --- a/compress.go +++ b/compress.go @@ -41,10 +41,13 @@ const ( // // This means less efficient compression as the sliding window from previous messages will not be used but the // memory overhead will be lower as there will be no fixed cost for the flate.Writer nor the 32 KB sliding window. - // Especially if the connections are long lived and seldom written to. + // Especially if the connections are long-lived and seldom written to. // // Thus, it uses less memory than CompressionContextTakeover but compresses less efficiently. // + // A message written with Conn.Write or Conn.WritePrepared is sent uncompressed if compression does not make it + // smaller, such as already compressed data. + // // If the peer does not support CompressionNoContextTakeover then we will fall back to CompressionDisabled. CompressionNoContextTakeover ) diff --git a/compress_test.go b/compress_test.go index 81adce9a..c09f61a8 100644 --- a/compress_test.go +++ b/compress_test.go @@ -76,6 +76,8 @@ func TestWriteSingleFrameCompressed(t *testing.T) { largeMsg = []byte(strings.Repeat("hello world ", 100)) smallMsg = []byte("small message") + // Random bytes do not compress, so compression makes them larger. + incompressibleMsg = xrand.Bytes(1024) ) testCases := []struct { @@ -88,6 +90,10 @@ func TestWriteSingleFrameCompressed(t *testing.T) { {"NoContextTakeover/AboveThreshold", CompressionNoContextTakeover, largeMsg, true}, {"ContextTakeover/BelowThreshold", CompressionContextTakeover, smallMsg, false}, {"NoContextTakeover/BelowThreshold", CompressionNoContextTakeover, smallMsg, false}, + // With context takeover, the message is in the sliding window, so it + // stays compressed to keep the peer's window in sync. + {"ContextTakeover/Incompressible", CompressionContextTakeover, incompressibleMsg, true}, + {"NoContextTakeover/Incompressible", CompressionNoContextTakeover, incompressibleMsg, false}, } for _, tc := range testCases { @@ -126,6 +132,9 @@ func TestWriteSingleFrameCompressed(t *testing.T) { assert.Equal(t, "opcode", opText, h.opcode) assert.Equal(t, "rsv1 (compressed)", tc.wantRsv1, h.rsv1) + if !tc.wantRsv1 { + assert.Equal(t, "payload length", int64(len(tc.msg)), h.payloadLength) + } assert.Equal(t, "fin", true, h.fin) err = <-writeDone diff --git a/prepared.go b/prepared.go index 290dd7d2..487ba17a 100644 --- a/prepared.go +++ b/prepared.go @@ -43,8 +43,10 @@ type PreparedMessage struct { typ MessageType data []byte + // Written to connections that compress without context takeover. compressOnce sync.Once - compressed []byte + deflated []byte // Compressed data, set only if smaller than data. + compressed bool // Whether deflated is written instead of data. compressErr error } @@ -56,8 +58,10 @@ func NewPreparedMessage(typ MessageType, data []byte) *PreparedMessage { } } -// compress returns the payload of PreparedMessage compressed without context takeover. -func (pm *PreparedMessage) compress() ([]byte, error) { +// compress returns the bytes to write on connections that compress without +// context takeover and whether they are compressed. They are the original +// data if compression does not make it smaller. +func (pm *PreparedMessage) compress() ([]byte, bool, error) { pm.compressOnce.Do(func() { var buf bytes.Buffer @@ -79,7 +83,12 @@ func (pm *PreparedMessage) compress() ([]byte, error) { return } - pm.compressed = buf.Bytes() + if buf.Len() < len(pm.data) { + pm.deflated, pm.compressed = buf.Bytes(), true + } }) - return pm.compressed, pm.compressErr + if !pm.compressed { + return pm.data, false, pm.compressErr + } + return pm.deflated, true, nil } diff --git a/prepared_test.go b/prepared_test.go index 526afac0..b57feee5 100644 --- a/prepared_test.go +++ b/prepared_test.go @@ -10,6 +10,7 @@ import ( "testing" "github.com/coder/websocket/internal/test/assert" + "github.com/coder/websocket/internal/test/xrand" ) type bufConn struct { @@ -30,6 +31,8 @@ func TestWritePreparedWire(t *testing.T) { largeMsg := strings.Repeat("prepared message ", 100) smallMsg := "small" nextMsg := []byte(strings.Repeat("next message ", 100)) + // Random bytes do not compress, so compression makes them larger. + incompressibleMsg := string(xrand.Bytes(1024)) testCases := []struct { name string @@ -43,6 +46,8 @@ func TestWritePreparedWire(t *testing.T) { {"NoContextTakeover/AboveThreshold", CompressionNoContextTakeover.opts(), largeMsg, true}, {"NoContextTakeover/AtThreshold", CompressionNoContextTakeover.opts(), strings.Repeat("a", 512), true}, {"NoContextTakeover/BelowThreshold", CompressionNoContextTakeover.opts(), strings.Repeat("a", 511), false}, + {"ContextTakeover/Incompressible", CompressionContextTakeover.opts(), incompressibleMsg, false}, + {"NoContextTakeover/Incompressible", CompressionNoContextTakeover.opts(), incompressibleMsg, false}, // Only the server side of the negotiated options applies to server writes. {"ServerNoContextTakeover", &compressionOptions{serverNoContextTakeover: true}, largeMsg, true}, {"ClientNoContextTakeover", &compressionOptions{clientNoContextTakeover: true}, largeMsg, false}, @@ -83,7 +88,7 @@ func TestWritePreparedWire(t *testing.T) { } assert.Equal(t, "wire bytes", wantBuf.Bytes(), gotBuf.Bytes()) - assert.Equal(t, "shared", tc.shared, pm.compressed != nil) + assert.Equal(t, "shared", tc.shared, pm.compressed) }) } } diff --git a/write.go b/write.go index 0fe1954e..f9b52563 100644 --- a/write.go +++ b/write.go @@ -66,7 +66,7 @@ func (c *Conn) writePrepared(ctx context.Context, pm *PreparedMessage) error { return err } - p, err := pm.compress() + p, compressed, err := pm.compress() if err != nil { return err } @@ -77,7 +77,7 @@ func (c *Conn) writePrepared(ctx context.Context, pm *PreparedMessage) error { } defer c.msgWriter.mu.unlock() - _, err = c.writeFrame(ctx, true, true, c.msgWriter.opcode, p) + _, err = c.writeFrame(ctx, true, compressed, c.msgWriter.opcode, p) return err } @@ -213,7 +213,19 @@ func (mw *msgWriter) writeCompressedFrame(ctx context.Context, p []byte) (int, e mw.closed = true - _, err = mw.c.writeFrame(ctx, true, true, mw.opcode, buf.Bytes()) + // Without context takeover, the peer keeps no history between messages, + // so a message that compression does not make smaller is sent as is. + // With context takeover, the message is already in our sliding window + // and must be sent compressed to keep the peer's window in sync, as + // uncompressed messages do not affect it. + // See https://tools.ietf.org/html/rfc7692#section-7.2.3.2 + compressed := mw.flateContextTakeover() || buf.Len() < len(p) + b := p + if compressed { + b = buf.Bytes() + } + + _, err = mw.c.writeFrame(ctx, true, compressed, mw.opcode, b) if err != nil { return 0, err } From 385007ab415c9bd1ce719dbbd9f09a721d00f63b Mon Sep 17 00:00:00 2001 From: kvizyx Date: Fri, 2 Oct 2026 17:14:51 +0300 Subject: [PATCH 4/6] test: cover prepared write in fuzz and Wasm tests --- conn_test.go | 5 +++++ internal/test/wstest/echo.go | 15 ++++++++++++++- ws_js_test.go | 22 ++++++++++++++++++++++ 3 files changed, 41 insertions(+), 1 deletion(-) diff --git a/conn_test.go b/conn_test.go index b1095603..8d5fd26c 100644 --- a/conn_test.go +++ b/conn_test.go @@ -50,9 +50,14 @@ func TestConn(t *testing.T) { c1.SetReadLimit(131072) + // Interleave prepared writes with regular writes, as they share + // the compression state of the connection. for range 5 { err := wstest.Echo(tt.ctx, c1, 131072) assert.Success(t, err) + + err = wstest.EchoPrepared(tt.ctx, c1, 131072) + assert.Success(t, err) } err := c1.Close(websocket.StatusNormalClosure, "") diff --git a/internal/test/wstest/echo.go b/internal/test/wstest/echo.go index c0c8dcd7..18078433 100644 --- a/internal/test/wstest/echo.go +++ b/internal/test/wstest/echo.go @@ -49,6 +49,19 @@ func EchoLoop(ctx context.Context, c *websocket.Conn) error { // Echo writes a message and ensures the same is sent back on c. func Echo(ctx context.Context, c *websocket.Conn, max int) error { + return echo(ctx, c, max, func(typ websocket.MessageType, msg []byte) error { + return c.Write(ctx, typ, msg) + }) +} + +// EchoPrepared is like Echo but writes the message with Conn.WritePrepared. +func EchoPrepared(ctx context.Context, c *websocket.Conn, max int) error { + return echo(ctx, c, max, func(typ websocket.MessageType, msg []byte) error { + return c.WritePrepared(ctx, websocket.NewPreparedMessage(typ, msg)) + }) +} + +func echo(ctx context.Context, c *websocket.Conn, max int, write func(websocket.MessageType, []byte) error) error { expType := websocket.MessageBinary if xrand.Bool() { expType = websocket.MessageText @@ -57,7 +70,7 @@ func Echo(ctx context.Context, c *websocket.Conn, max int) error { msg := randMessage(expType, xrand.Int(max)) writeErr := xsync.Go(func() error { - return c.Write(ctx, expType, msg) + return write(expType, msg) }) actType, act, err := c.Read(ctx) diff --git a/ws_js_test.go b/ws_js_test.go index 1fa242f4..f0e1e89b 100644 --- a/ws_js_test.go +++ b/ws_js_test.go @@ -37,6 +37,28 @@ func TestWasm(t *testing.T) { assert.Success(t, err) } +func TestWasmPrepared(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithTimeout(context.Background(), time.Minute) + defer cancel() + + c, _, err := websocket.Dial(ctx, os.Getenv("WS_ECHO_SERVER_URL"), &websocket.DialOptions{ + Subprotocols: []string{"echo"}, + }) + assert.Success(t, err) + defer c.Close(websocket.StatusInternalError, "") + + c.SetReadLimit(65536) + for range 10 { + err = wstest.EchoPrepared(ctx, c, 65536) + assert.Success(t, err) + } + + err = c.Close(websocket.StatusNormalClosure, "") + assert.Success(t, err) +} + func TestWasmDialTimeout(t *testing.T) { t.Parallel() From 6e4f00815ed7251c0b2fdf99299e27da13b735d6 Mon Sep 17 00:00:00 2001 From: kvizyx Date: Sat, 3 Oct 2026 00:29:53 +0300 Subject: [PATCH 5/6] test: clarify compressed expectation in prepared write test --- prepared_test.go | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/prepared_test.go b/prepared_test.go index b57feee5..b9c19e5c 100644 --- a/prepared_test.go +++ b/prepared_test.go @@ -23,8 +23,8 @@ func (*bufConn) Close() error { // TestWritePreparedWire verifies that WritePrepared puts the same bytes on the // wire as Write, including for messages written after it, and that the -// compressed payload is shared only when the connection compresses the -// message without context takeover. +// shared compressed payload is written only when the connection compresses +// without context takeover and compression makes the message smaller. func TestWritePreparedWire(t *testing.T) { t.Parallel() @@ -35,10 +35,13 @@ func TestWritePreparedWire(t *testing.T) { incompressibleMsg := string(xrand.Bytes(1024)) testCases := []struct { - name string - copts *compressionOptions - msg string - shared bool // Whether the compressed payload of the message is shared. + name string + copts *compressionOptions + msg string + // Whether the shared compressed payload is written. It is false both + // when the regular write path is used and when the message is written + // as is because compression does not make it smaller. + wantCompressed bool }{ {"Disabled", nil, largeMsg, false}, {"ContextTakeover/AboveThreshold", CompressionContextTakeover.opts(), largeMsg, false}, @@ -66,7 +69,7 @@ func TestWritePreparedWire(t *testing.T) { bw: bufio.NewWriter(rwc), }) t.Cleanup(func() { - c.CloseNow() + _ = c.CloseNow() }) return c, rwc } @@ -88,7 +91,7 @@ func TestWritePreparedWire(t *testing.T) { } assert.Equal(t, "wire bytes", wantBuf.Bytes(), gotBuf.Bytes()) - assert.Equal(t, "shared", tc.shared, pm.compressed) + assert.Equal(t, "compressed", tc.wantCompressed, pm.compressed) }) } } From 9bffa3c9a98ef8e5280b44b6a50571af1c471a45 Mon Sep 17 00:00:00 2001 From: kvizyx Date: Sat, 3 Oct 2026 00:44:24 +0300 Subject: [PATCH 6/6] fix(examples): reject invalid UTF-8 in broadcast example --- example_test.go | 2 +- internal/examples/broadcast/README.md | 5 +++-- internal/examples/broadcast/server.go | 10 +++++++++- internal/examples/broadcast/server_test.go | 15 +++++++++++---- prepared.go | 6 +++--- 5 files changed, 27 insertions(+), 11 deletions(-) diff --git a/example_test.go b/example_test.go index aa22762d..e43c4810 100644 --- a/example_test.go +++ b/example_test.go @@ -172,6 +172,6 @@ func Example_echo() { // This example demonstrates broadcasting messages to many subscribers with // a PreparedMessage. -func ExampleConn_writePrepared() { +func ExampleConn_WritePrepared() { // https://github.com/coder/websocket/tree/master/internal/examples/broadcast } diff --git a/internal/examples/broadcast/README.md b/internal/examples/broadcast/README.md index b0ce57e4..e1e5d1d3 100644 --- a/internal/examples/broadcast/README.md +++ b/internal/examples/broadcast/README.md @@ -16,7 +16,8 @@ $ websocat ws://127.0.0.1:8080/subscribe $ curl --data-binary 'hello' http://127.0.0.1:8080/publish ``` -Every published message is delivered to all subscribers. +Every published message is delivered to all subscribers. Messages are sent as text, so +`/publish` rejects bodies that are not valid UTF-8. ## Structure @@ -32,6 +33,6 @@ Subscribers that do not support compression receive the same message uncompresse Publishing never blocks: a subscriber that cannot keep up with its queue is disconnected. `server_test.go` contains a test that subscribes clients with and without compression and -ensures every client receives every published message. +ensures every client receives every published message, while invalid UTF-8 is rejected. `main.go` brings it all together so that you can run it and play around with it. diff --git a/internal/examples/broadcast/server.go b/internal/examples/broadcast/server.go index dc7113f4..91af8330 100644 --- a/internal/examples/broadcast/server.go +++ b/internal/examples/broadcast/server.go @@ -7,6 +7,7 @@ import ( "net/http" "sync" "time" + "unicode/utf8" "github.com/coder/websocket" ) @@ -85,8 +86,11 @@ func (bs *broadcastServer) subscribeHandler(w http.ResponseWriter, r *http.Reque // publishHandler reads the request body with a limit of 65536 bytes and then // publishes the received message. +// +// Messages are sent as text, which must be valid UTF-8, so other bodies are +// rejected instead of being sent as invalid text messages. func (bs *broadcastServer) publishHandler(w http.ResponseWriter, r *http.Request) { - if r.Method != "POST" { + if r.Method != http.MethodPost { http.Error(w, http.StatusText(http.StatusMethodNotAllowed), http.StatusMethodNotAllowed) return } @@ -96,6 +100,10 @@ func (bs *broadcastServer) publishHandler(w http.ResponseWriter, r *http.Request http.Error(w, http.StatusText(http.StatusRequestEntityTooLarge), http.StatusRequestEntityTooLarge) return } + if !utf8.Valid(msg) { + http.Error(w, "message must be valid UTF-8", http.StatusBadRequest) + return + } bs.publish(websocket.NewPreparedMessage(websocket.MessageText, msg)) diff --git a/internal/examples/broadcast/server_test.go b/internal/examples/broadcast/server_test.go index 77b21f01..82d07948 100644 --- a/internal/examples/broadcast/server_test.go +++ b/internal/examples/broadcast/server_test.go @@ -13,7 +13,7 @@ import ( // Test_broadcastServer subscribes clients with and without compression, // publishes messages below and above the compression threshold and ensures -// every client receives all of them. +// every client receives all of them, while invalid UTF-8 is rejected. func Test_broadcastServer(t *testing.T) { t.Parallel() @@ -43,17 +43,24 @@ func Test_broadcastServer(t *testing.T) { "hello", strings.Repeat("hello world ", 512), } - for _, msg := range msgs { + publish := func(msg string, expStatus int) { resp, err := http.Post(s.URL+"/publish", "text/plain", strings.NewReader(msg)) if err != nil { t.Fatal(err) } resp.Body.Close() - if resp.StatusCode != http.StatusAccepted { - t.Fatalf("expected status %v but got %v", http.StatusAccepted, resp.StatusCode) + if resp.StatusCode != expStatus { + t.Fatalf("expected status %v but got %v", expStatus, resp.StatusCode) } } + // Published first, so reading the valid messages below fails if it was + // delivered. + publish("\xff", http.StatusBadRequest) + for _, msg := range msgs { + publish(msg, http.StatusAccepted) + } + for i, c := range clients { for _, exp := range msgs { _, b, err := c.Read(ctx) diff --git a/prepared.go b/prepared.go index 487ba17a..8408b631 100644 --- a/prepared.go +++ b/prepared.go @@ -22,9 +22,9 @@ import ( // CompressionContextTakeover when the client negotiates // server_no_context_takeover. For a client, client_no_context_takeover // plays the same role. -// - The message is at least as long as the compression threshold, 512 bytes by -// default for CompressionNoContextTakeover. Smaller messages are not -// compressed. +// - The message is at least as long as the compression threshold, 512 bytes +// by default when compressing without context takeover. Smaller messages +// are not compressed. // // Otherwise, Conn.WritePrepared costs the same as Conn.Write. With context takeover, // each connection compresses against the sliding window of its own previous