From 9f84ce4a96d6d250b8874ccdb463104d5522cf61 Mon Sep 17 00:00:00 2001 From: Sztanojev Peter Date: Mon, 27 Jul 2026 08:09:51 +0000 Subject: [PATCH] feat(conn): add SetWriteDeadline to bound writes without allocating MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bounding a write requires a context, so a caller that wants a write deadline must allocate one per frame — a context.WithTimeout on its side and, for any cancelable context, a context.AfterFunc on ours. A server fanning one message out to many connections pays that per recipient rather than per message, where it is the dominant cost of sending: profiling a chat server holding 500 connections in one room at ~49,600 message deliveries a second attributed about 700 bytes of garbage per delivered message to it, and 99% of everything the server allocated under that load. A deadline is an instant, and an instant needs no allocation to express. SetWriteDeadline stores one in an atomic and writeFrame arms a timer created once per connection, using only Reset and Stop — the approach netconn.go already takes for net.Conn deadlines, exposed on Conn so a caller that can name its bound as a deadline does not have to buy a context to carry it. Paired with a non-cancelable context on Write, which #566 already made free, the write path allocates nothing. Existing callers are unaffected: the deadline is unset by default, and with none set writeFrame does one atomic load more than before. On the same server and load, replacing the per-frame context with a deadline: allocation over a 20s window falls from 689MB to zero sampled, CPU samples for identical delivered work from 14.33s to 12.50s, write(2) from 63.1% to 73.1% of the profile, runtime.mallocgc from 4.82% to below the reporting threshold, and mean delivery latency measured at the client from 4.20ms to 3.53ms. Two details worth noting for review. The timer is armed by writeFrame rather than by SetWriteDeadline, so that it is armed under writeFrameMu together with the frame it bounds and cannot disturb a frame already in flight. And armWriteDeadline must not close the connection itself: close takes writeFrameMu for a client connection via msgWriter.close, so closing inline deadlocks, which is why an expired deadline refuses the frame with os.ErrDeadlineExceeded and leaves the closing to the timer's goroutine. It reports bools rather than returning a stop closure for the reason #565 gives on the read path — a closure over the connection is an allocation per frame. ws_js.go gets the same method so the API does not differ by platform. Writes there are handed to the browser and never block, so the deadline is checked when a write is attempted rather than interrupting one. Refs #565, #566 --- conn.go | 86 ++++++++++++++++++++++++++++ write.go | 9 +++ writedeadline_test.go | 129 ++++++++++++++++++++++++++++++++++++++++++ ws_js.go | 31 ++++++++++ 4 files changed, 255 insertions(+) create mode 100644 writedeadline_test.go diff --git a/conn.go b/conn.go index bfc78aaa..0ea17ab9 100644 --- a/conn.go +++ b/conn.go @@ -7,11 +7,13 @@ import ( "context" "fmt" "io" + "math" "net" "runtime" "strconv" "sync" "sync/atomic" + "time" ) // MessageType represents the type of a WebSocket message. @@ -54,6 +56,20 @@ type Conn struct { readTimeoutStop atomic.Pointer[func() bool] writeTimeoutStop atomic.Pointer[func() bool] + // writeDeadline is the deadline set by SetWriteDeadline as unix nanos, + // or 0 for none. writeDeadlineTimer is created once per connection and + // thereafter only Reset and Stopped, neither of which allocates — the + // same approach netconn.go takes for net.Conn deadlines. + // + // The timer is armed by writeFrame rather than by SetWriteDeadline, so + // that it is armed under writeFrameMu together with the frame it + // bounds. That keeps it safe when several goroutines write to one + // connection: setting a deadline cannot disturb a frame already in + // flight, and frames written under one deadline share it instead of + // restarting the clock per frame. + writeDeadline atomic.Int64 + writeDeadlineTimer *time.Timer + // Read state. readMu *mu readHeaderBuf [8]byte @@ -135,6 +151,16 @@ func newConn(cfg connConfig) *Conn { } } + // One timer for the life of the connection, so that arming a write + // deadline is a Reset and disarming it a Stop, neither of which + // allocates. The duration is irrelevant because it is stopped + // immediately: there is no constructor for a stopped timer, and this is + // how netconn.go makes one. + c.writeDeadlineTimer = time.AfterFunc(math.MaxInt64, func() { + c.close() + }) + c.writeDeadlineTimer.Stop() + runtime.SetFinalizer(c, func(c *Conn) { c.close() }) @@ -181,6 +207,66 @@ func (c *Conn) setupWriteTimeout(ctx context.Context) bool { return true } +// SetWriteDeadline sets a deadline for writes on the connection. Writes +// that pass it close the connection, as an expired context passed to Write +// does. A zero t clears the deadline. +// +// It exists so that a deadline can be expressed without allocating. +// Bounding a write with a context costs a context.WithTimeout in the caller +// and a context.AfterFunc inside this package, per frame; a deadline is an +// instant, and an instant is one atomic store against a timer that is +// reused for the life of the connection. Callers that pass a +// non-cancellable context to Write and set the deadline here allocate +// nothing on the write path, which matters for a server fanning one message +// out to many connections, where that per-frame cost is paid per recipient. +// +// The deadline belongs to the connection, not to a call, so where several +// goroutines write to one connection the last to set it wins. Frames +// written under one deadline all wait for that same instant rather than +// restarting the clock, and setting a deadline never disturbs a frame +// already in flight. +func (c *Conn) SetWriteDeadline(t time.Time) error { + if t.IsZero() { + c.writeDeadline.Store(0) + return nil + } + c.writeDeadline.Store(t.UnixNano()) + return nil +} + +// armWriteDeadline prepares the write deadline for one frame. It reports +// whether the deadline has already passed, and whether the timer was armed +// and so needs stopping once the frame is done. The caller must hold +// writeFrameMu. +// +// It must not close the connection itself, even though an expired deadline +// means the connection is going away: close takes writeFrameMu for a client +// connection, via msgWriter.close, and this runs with that lock held. That +// is why expiry is reported back to the caller and the closing is left to +// the timer's own goroutine, which holds nothing. +// +// It reports bools rather than returning the stop func it would rather +// return, because a closure over c is a heap allocation on every frame, +// which would defeat the point of the mechanism. +func (c *Conn) armWriteDeadline() (expired, armed bool) { + deadline := c.writeDeadline.Load() + if deadline == 0 { + return false, false + } + + d := time.Until(time.Unix(0, deadline)) + if d <= 0 { + // Already past: refuse the frame, and let the timer do the closing + // from where that is safe. Deliberately left running. + c.writeDeadlineTimer.Reset(1) + return true, false + } + c.writeDeadlineTimer.Reset(d) + return false, true +} + +func (c *Conn) stopWriteDeadline() { c.writeDeadlineTimer.Stop() } + func (c *Conn) clearWriteTimeout() { swapTimeoutStop(&c.writeTimeoutStop, nil) } diff --git a/write.go b/write.go index 81bab4fa..76d12306 100644 --- a/write.go +++ b/write.go @@ -12,6 +12,7 @@ import ( "fmt" "io" "net" + "os" "time" "github.com/coder/websocket/internal/bpool" @@ -322,6 +323,14 @@ func (c *Conn) writeFrame(ctx context.Context, fin bool, flate bool, opcode opco defer c.clearWriteTimeout() } + // Armed here rather than in SetWriteDeadline so that it is under + // writeFrameMu together with the frame it bounds. + if expired, armed := c.armWriteDeadline(); expired { + return 0, os.ErrDeadlineExceeded + } else if armed { + defer c.stopWriteDeadline() + } + c.writeHeader.fin = fin c.writeHeader.opcode = opcode c.writeHeader.payloadLength = int64(len(p)) diff --git a/writedeadline_test.go b/writedeadline_test.go new file mode 100644 index 00000000..f6374833 --- /dev/null +++ b/writedeadline_test.go @@ -0,0 +1,129 @@ +//go:build !js + +package websocket_test + +import ( + "context" + "errors" + "os" + "testing" + "time" + + "github.com/coder/websocket" + "github.com/coder/websocket/internal/test/assert" + "github.com/coder/websocket/internal/test/wstest" +) + +// drainWriteDeadline attaches a peer that consumes messages and returns a +// function that stops it. +// +// CloseRead is not usable for these tests: it closes the connection when a +// data message arrives, which is exactly what they send. +func drainWriteDeadline(c *websocket.Conn) func() { + done := make(chan struct{}) + go func() { + defer close(done) + for { + if _, _, err := c.Read(context.Background()); err != nil { + return + } + } + }() + return func() { + c.CloseNow() + <-done + } +} + +func TestSetWriteDeadline(t *testing.T) { + t.Parallel() + + msg := []byte("hello") + + // A deadline in the future does not get in the way, and clearing it with + // the zero time puts things back as they were. + t.Run("live and cleared", func(t *testing.T) { + t.Parallel() + + c1, c2 := wstest.Pipe(nil, nil) + defer c1.CloseNow() + defer c2.CloseNow() + defer drainWriteDeadline(c2)() + + assert.Success(t, c1.SetWriteDeadline(time.Now().Add(time.Minute))) + assert.Success(t, c1.Write(context.Background(), websocket.MessageText, msg)) + + assert.Success(t, c1.SetWriteDeadline(time.Time{})) + assert.Success(t, c1.Write(context.Background(), websocket.MessageText, msg)) + }) + + // A deadline that has already passed closes the connection rather than + // writing, which is what an expired context passed to Write does. + t.Run("expired", func(t *testing.T) { + t.Parallel() + + c1, c2 := wstest.Pipe(nil, nil) + defer c1.CloseNow() + defer c2.CloseNow() + defer drainWriteDeadline(c2)() + + assert.Success(t, c1.SetWriteDeadline(time.Now().Add(-time.Second))) + + err := c1.Write(context.Background(), websocket.MessageText, msg) + if err == nil { + t.Fatal("write succeeded past its deadline") + } + if !errors.Is(err, os.ErrDeadlineExceeded) { + t.Fatalf("write past its deadline failed with %v, want os.ErrDeadlineExceeded", err) + } + + // And the connection goes away, as it does when a context passed to + // Write expires. That happens on the deadline timer's goroutine + // rather than inline, so give it a moment. + deadline := time.Now().Add(5 * time.Second) + for { + if _, _, err := c1.Read(context.Background()); err != nil { + break + } + if time.Now().After(deadline) { + t.Fatal("connection still alive after its write deadline passed") + } + } + }) +} + +// TestSetWriteDeadlineAllocs is the point of the method: a deadline set this +// way costs nothing per frame, where bounding a write with a context +// allocates one in the caller and a context.AfterFunc in here. +// +// It is deliberately not parallel — AllocsPerRun panics inside a parallel +// test. +func TestSetWriteDeadlineAllocs(t *testing.T) { + msg := []byte("hello") + + c1, c2 := wstest.Pipe(nil, nil) + defer c1.CloseNow() + defer c2.CloseNow() + defer drainWriteDeadline(c2)() + + write := func() { + if err := c1.Write(context.Background(), websocket.MessageText, msg); err != nil { + t.Fatal(err) + } + } + + // Compared against the same write with no deadline set, so this asserts + // what the method claims — that a deadline is free — rather than that a + // write allocates nothing, which is not this method's business. + assert.Success(t, c1.SetWriteDeadline(time.Time{})) + write() // warm any one-off buffers before counting + without := testing.AllocsPerRun(200, write) + + assert.Success(t, c1.SetWriteDeadline(time.Now().Add(time.Hour))) + with := testing.AllocsPerRun(200, write) + + if with > without { + t.Fatalf("write with a deadline allocated %v, without a deadline %v; a deadline should be free", with, without) + } + t.Logf("allocations per write: %v without a deadline, %v with one", without, with) +} diff --git a/ws_js.go b/ws_js.go index 026b75fc..ea7dd405 100644 --- a/ws_js.go +++ b/ws_js.go @@ -8,12 +8,14 @@ import ( "io" "net" "net/http" + "os" "reflect" "runtime" "strings" "sync" "sync/atomic" "syscall/js" + "time" "github.com/coder/websocket/internal/bpool" "github.com/coder/websocket/internal/wsjs" @@ -64,6 +66,32 @@ type Conn struct { readSignal chan struct{} readBufMu sync.Mutex readBuf []wsjs.MessageEvent + + // writeDeadline is the deadline set by SetWriteDeadline as unix nanos, + // or 0 for none. + writeDeadline atomic.Int64 +} + +// SetWriteDeadline sets a deadline for writes on the connection. Writes +// that pass it close the connection. A zero t clears the deadline. +// +// Writes here are handed to the browser and never block, so unlike the +// native implementation there is nothing to interrupt: the deadline is +// checked when a write is attempted, and a write attempted after it has +// passed closes the connection instead of sending. +func (c *Conn) SetWriteDeadline(t time.Time) error { + if t.IsZero() { + c.writeDeadline.Store(0) + return nil + } + c.writeDeadline.Store(t.UnixNano()) + return nil +} + +// writeDeadlineExceeded reports whether a deadline is set and has passed. +func (c *Conn) writeDeadlineExceeded() bool { + deadline := c.writeDeadline.Load() + return deadline != 0 && !time.Now().Before(time.Unix(0, deadline)) } func (c *Conn) close(err error, wasClean bool) { @@ -214,6 +242,9 @@ func (c *Conn) write(typ MessageType, p []byte) error { if c.isClosed() { return net.ErrClosed } + if c.writeDeadlineExceeded() { + return os.ErrDeadlineExceeded + } switch typ { case MessageBinary: return c.ws.SendBytes(p)