diff --git a/client_test.go b/client_test.go index 758eea6cd..7816435f7 100644 --- a/client_test.go +++ b/client_test.go @@ -796,3 +796,156 @@ func (cli clientReadOnEOF) OnTraffic(c Conn) (action Action) { }{data: data, err: err} return None } + +// TestClientSafeContext verifies that a client-side Conn's SafeContext()/ +// SetSafeContext() behave correctly on connections created via +// Client.Dial/DialContext/EnrollContext: DialContext's ctx argument is +// reflected by both Context() and SafeContext(), switching the stored +// concrete type (including nil) across calls never panics, and +// SafeContext() is safe to call concurrently from a goroutine other than +// the event-loop goroutine. +func TestClientSafeContext(t *testing.T) { + ln, err := net.Listen("tcp", "127.0.0.1:9978") + assert.NoError(t, err) + defer ln.Close() //nolint:errcheck + + err = goPool.DefaultWorkerPool.Submit(func() { + for i := 0; i < 5; i++ { + conn, err := ln.Accept() + if err != nil { + return + } + err = goPool.DefaultWorkerPool.Submit(func() { + buf := make([]byte, 4) + for { + _, err := io.ReadFull(conn, buf) + if err != nil { + conn.Close() //nolint:errcheck + return + } + _, err = conn.Write([]byte("pong")) + if err != nil { + conn.Close() //nolint:errcheck + return + } + } + }) + assert.NoError(t, err) + } + }) + assert.NoError(t, err) + + ev := &clientSafeContextEvents{tester: t, done: make(chan struct{})} + cli, err := NewClient(ev) + assert.NoError(t, err) + defer cli.Stop() //nolint:errcheck + + err = cli.Start() + assert.NoError(t, err) + + initialCtx := &safeCtxPayloadA{n: 42} + ev.initialCtx = initialCtx + c, err := cli.DialContext("tcp", "127.0.0.1:9978", initialCtx) + assert.NoError(t, err) + + select { + case <-ev.done: + case <-time.After(5 * time.Second): + t.Fatal("timeout waiting for test completion") + } + + // release() (which nils out safeCtx) runs on the event-loop goroutine + // immediately after OnClose returns, racing with this goroutine waking + // up from ev.done, so poll briefly instead of asserting once. + assert.Eventually(t, func() bool { + return c.SafeContext() == nil + }, time.Second, time.Millisecond, "SafeContext() must be nil once the connection is released") +} + +type clientSafeContextEvents struct { + BuiltinEventEngine + tester *testing.T + done chan struct{} + + initialCtx any + + trafficCount int32 + backgroundHits int32 + stopBackground chan struct{} +} + +func (ev *clientSafeContextEvents) OnOpen(c Conn) (out []byte, action Action) { + // OnOpen always runs before any OnTraffic event for this connection, so + // asserting here is guaranteed to observe the dial-time ctx before it can + // be overwritten by OnTraffic's SetSafeContext calls. + assert.Equal(ev.tester, ev.initialCtx, c.Context()) + assert.Equal(ev.tester, ev.initialCtx, c.SafeContext()) + + ev.stopBackground = make(chan struct{}) + // Exercise SafeContext() concurrently from a goroutine other than the + // event-loop goroutine. + err := goPool.DefaultWorkerPool.Submit(func() { + for { + select { + case <-ev.stopBackground: + return + default: + } + switch v := c.SafeContext().(type) { + case nil: + case *safeCtxPayloadA: + _ = v.n + case *safeCtxPayloadB: + _ = v.s + case int: + default: + assert.Failf(ev.tester, "unexpected SafeContext type", "%T", v) + } + atomic.AddInt32(&ev.backgroundHits, 1) + } + }) + assert.NoError(ev.tester, err) + + out = []byte("ping") + return +} + +func (ev *clientSafeContextEvents) OnTraffic(c Conn) (action Action) { + buf, err := c.Next(-1) + assert.NoError(ev.tester, err) + assert.Equal(ev.tester, "pong", string(buf)) + + n := atomic.AddInt32(&ev.trafficCount, 1) + switch n % 4 { + case 0: + c.SetSafeContext(&safeCtxPayloadA{n: n}) + case 1: + c.SetSafeContext(&safeCtxPayloadB{s: "hi"}) + case 2: + c.SetSafeContext(int(n)) + case 3: + c.SetSafeContext(nil) + } + + if n >= 5 { + action = Close + return + } + + _, err = c.Write([]byte("ping")) + assert.NoError(ev.tester, err) + return +} + +func (ev *clientSafeContextEvents) OnClose(c Conn, _ error) (action Action) { + // SafeContext() must still be safely callable (no panic) right up + // until release(), which happens after OnClose returns. + _ = c.SafeContext() + + close(ev.stopBackground) + assert.Greater(ev.tester, atomic.LoadInt32(&ev.backgroundHits), int32(0), + "expected background goroutine to observe SafeContext() at least once") + + close(ev.done) + return +} diff --git a/client_unix.go b/client_unix.go index c14b1cc4a..1a50b1d25 100644 --- a/client_unix.go +++ b/client_unix.go @@ -274,7 +274,8 @@ func (cli *Client) EnrollContext(c net.Conn, ctx any) (Conn, error) { default: return nil, errorx.ErrUnsupportedProtocol } - gc.ctx = ctx + gc.SetContext(ctx) + gc.SetSafeContext(ctx) connOpened := make(chan struct{}) ccb := &connWithCallback{c: gc, cb: func() { diff --git a/connection_unix.go b/connection_unix.go index a982bffca..943f5eca4 100644 --- a/connection_unix.go +++ b/connection_unix.go @@ -20,6 +20,7 @@ import ( "io" "net" "os" + "sync/atomic" "time" "golang.org/x/sys/unix" @@ -39,6 +40,7 @@ type conn struct { fd int // file descriptor gfd gfd.GFD // gnet file descriptor ctx any // user-defined context + safeCtx atomic.Pointer[any] // safe user-defined context remote unix.Sockaddr // remote socket address proto string // protocol name: "tcp", "udp", or "unix". localAddr net.Addr // local addr @@ -91,6 +93,7 @@ func (c *conn) release() { c.opened = false c.isEOF = false c.ctx = nil + c.safeCtx.Store(nil) c.buffer = nil if addr, ok := c.localAddr.(*net.TCPAddr); ok && len(c.loop.listeners) == 0 && len(addr.Zone) > 0 { bsPool.Put(bs.StringToBytes(addr.Zone)) @@ -559,3 +562,14 @@ func (*conn) SetReadDeadline(_ time.Time) error { func (*conn) SetWriteDeadline(_ time.Time) error { return errorx.ErrUnsupportedOp } + +func (c *conn) SafeContext() (ctx any) { + if p := c.safeCtx.Load(); p != nil { + return *p + } + return nil +} + +func (c *conn) SetSafeContext(ctx any) { + c.safeCtx.Store(&ctx) +} diff --git a/connection_windows.go b/connection_windows.go index 618150a99..d20e8595a 100644 --- a/connection_windows.go +++ b/connection_windows.go @@ -19,6 +19,7 @@ import ( "io" "net" "os" + "sync/atomic" "syscall" "time" @@ -52,14 +53,15 @@ type openConn struct { type conn struct { pc net.PacketConn - ctx any // user-defined context - loop *eventloop // owner event-loop - buffer *bbPool.ByteBuffer // reuse memory of inbound data as a temporary buffer - cache []byte // temporary cache for the inbound data - rawConn net.Conn // original connection - localAddr net.Addr // local server addr - remoteAddr net.Addr // remote addr - inboundBuffer elastic.RingBuffer // buffer for data from the remote + ctx any // user-defined context + safeCtx atomic.Pointer[any] // safe user-defined context + loop *eventloop // owner event-loop + buffer *bbPool.ByteBuffer // reuse memory of inbound data as a temporary buffer + cache []byte // temporary cache for the inbound data + rawConn net.Conn // original connection + localAddr net.Addr // local server addr + remoteAddr net.Addr // remote addr + inboundBuffer elastic.RingBuffer // buffer for data from the remote } func packTCPConn(c *conn, buf []byte) *tcpConn { @@ -84,7 +86,7 @@ func packUDPConn(c *conn, buf []byte) *udpConn { } func newStreamConn(el *eventloop, nc net.Conn, ctx any) (c *conn) { - return &conn{ + c = &conn{ ctx: ctx, loop: el, buffer: bbPool.Get(), @@ -92,10 +94,13 @@ func newStreamConn(el *eventloop, nc net.Conn, ctx any) (c *conn) { localAddr: nc.LocalAddr(), remoteAddr: nc.RemoteAddr(), } + c.SetSafeContext(ctx) + return c } func (c *conn) release() { c.ctx = nil + c.safeCtx.Store(nil) c.localAddr = nil if c.rawConn != nil { c.rawConn = nil @@ -107,7 +112,7 @@ func (c *conn) release() { } func newUDPConn(el *eventloop, pc net.PacketConn, rc net.Conn, localAddr, remoteAddr net.Addr, ctx any) *conn { - return &conn{ + c := &conn{ ctx: ctx, pc: pc, rawConn: rc, @@ -116,6 +121,8 @@ func newUDPConn(el *eventloop, pc net.PacketConn, rc net.Conn, localAddr, remote localAddr: localAddr, remoteAddr: remoteAddr, } + c.SetSafeContext(ctx) + return c } func (c *conn) resetBuffer() { @@ -573,3 +580,14 @@ func (*conn) SetReadDeadline(_ time.Time) error { func (*conn) SetWriteDeadline(_ time.Time) error { return errorx.ErrUnsupportedOp } + +func (c *conn) SafeContext() (ctx any) { + if p := c.safeCtx.Load(); p != nil { + return *p + } + return nil +} + +func (c *conn) SetSafeContext(ctx any) { + c.safeCtx.Store(&ctx) +} diff --git a/eventloop_unix.go b/eventloop_unix.go index 4e089cf3b..5553db475 100644 --- a/eventloop_unix.go +++ b/eventloop_unix.go @@ -183,7 +183,8 @@ func (el *eventloop) enroll(c net.Conn, addr net.Addr, ctx any) (resCh chan Regi return } - gc.ctx = ctx + gc.SetContext(ctx) + gc.SetSafeContext(ctx) connOpened := make(chan struct{}) ccb := &connWithCallback{c: gc, cb: func() { diff --git a/gnet.go b/gnet.go index d87953787..36d9c31c9 100644 --- a/gnet.go +++ b/gnet.go @@ -430,6 +430,10 @@ type Conn interface { // you must invoke it within any method in EventHandler. Context() (ctx any) + // SafeContext is the concurrency-safe version of Context that can + // be invoked on any goroutine. + SafeContext() (ctx any) + // EventLoop returns the event-loop that the connection belongs to. // The returned EventLoop is concurrency-safe. EventLoop() EventLoop @@ -438,12 +442,20 @@ type Conn interface { // you must invoke it within any method in EventHandler. SetContext(ctx any) + // SetSafeContext is the concurrency-safe version of SetContext that can + // be invoked on any goroutine. + SetSafeContext(ctx any) + // LocalAddr is the connection's local socket address, it's not concurrency-safe, - // you must invoke it within any method in EventHandler. + // you must invoke it and use the returned value within any method in EventHandler. + // If you must use the returned value outside of EventHandler, you should make + // a copy of it by calling net.Addr.String(). LocalAddr() net.Addr // RemoteAddr is the connection's remote address, it's not concurrency-safe, - // you must invoke it within any method in EventHandler. + // you must invoke it and use the returned value within any method in EventHandler. + // If you must use the returned value outside of EventHandler, you should make + // a copy of it by calling net.Addr.String(). RemoteAddr() net.Addr // Wake triggers an OnTraffic event for the current connection, it's concurrency-safe. diff --git a/gnet_test.go b/gnet_test.go index df62fa0aa..0cf498d3f 100644 --- a/gnet_test.go +++ b/gnet_test.go @@ -1724,8 +1724,8 @@ func (t *testDisconnectedAsyncWriteServer) OnTick() (delay time.Duration, action func TestDisconnectedAsyncWrite(t *testing.T) { t.Run("async-write", func(t *testing.T) { - events := &testDisconnectedAsyncWriteServer{tester: t, addr: ":10000"} - err := Run(events, "tcp://:10000", WithTicker(true)) + events := &testDisconnectedAsyncWriteServer{tester: t, addr: ":19999"} + err := Run(events, "tcp://:19999", WithTicker(true)) assert.NoError(t, err) }) t.Run("async-writev", func(t *testing.T) { @@ -2100,8 +2100,8 @@ func TestUDPSendtoServer(t *testing.T) { // while sendto fails on macOS with EINVAL (invalid argument) that might // also occur on more BSD systems: FreeBSD, OpenBSD, NetBSD, and DragonflyBSD. // To pass this test on all platforms, specify an explicit IP address here. - // addr := ":10000" - addr := "127.0.0.1:10000" + // addr := ":19999" + addr := "127.0.0.1:19999" events := &testUDPSendtoServer{tester: t, addr: addr} err := Run(events, "udp://"+addr, WithTicker(true)) assert.NoError(t, err) @@ -2323,7 +2323,7 @@ func (p *streamProxyServer) OnTick() (time.Duration, Action) { func TestStreamProxyServer(t *testing.T) { t.Run("tcp-proxy-server", func(t *testing.T) { - addr := "tcp://127.0.0.1:10000" + addr := "tcp://127.0.0.1:19999" backendServers := []string{ "tcp://127.0.0.1:10001", "tcp://127.0.0.1:10002", @@ -2618,7 +2618,7 @@ func (p *udpProxyServer) OnTick() (delay time.Duration, action Action) { func TestUDPProxyServer(t *testing.T) { t.Run("backend-udp-proxy-server", func(t *testing.T) { - addr := "udp://127.0.0.1:10000" + addr := "udp://127.0.0.1:19999" backendServers := []string{ "udp://127.0.0.1:10001", "udp://127.0.0.1:10002", @@ -2731,3 +2731,156 @@ func testUDPProxyServer(t *testing.T, addr string, backendServers []string, mult backends.Wait() //nolint:errcheck } + +// TestSafeContext verifies that Conn.SafeContext()/SetSafeContext() behave +// correctly: the value set is always the value retrieved, mixing different +// concrete types (including nil) across calls never panics, and calling +// SafeContext() concurrently from goroutines other than the event-loop +// goroutine is race-free, per the "concurrency-safe" contract documented +// on the Conn interface. +func TestSafeContext(t *testing.T) { + testSafeContext(t, "tcp", ":9979") +} + +type safeCtxPayloadA struct{ n int32 } +type safeCtxPayloadB struct{ s string } + +type testSafeContextServer struct { + BuiltinEventEngine + tester *testing.T + network, addr string + + trafficCount int32 + backgroundHits int32 + stopBackground chan struct{} + done chan struct{} + closedConn Conn +} + +func (s *testSafeContextServer) OnBoot(Engine) (action Action) { + err := goPool.DefaultWorkerPool.Submit(func() { + c, err := net.Dial(s.network, s.addr) + assert.NoError(s.tester, err) + defer c.Close() //nolint:errcheck + + for i := 0; i < 5; i++ { + _, err = c.Write([]byte("ping")) + assert.NoError(s.tester, err) + buf := make([]byte, 4) + _, err = io.ReadFull(c, buf) + assert.NoError(s.tester, err) + assert.Equal(s.tester, "pong", string(buf)) + } + }) + assert.NoError(s.tester, err) + return None +} + +func (s *testSafeContextServer) OnOpen(c Conn) (out []byte, action Action) { + // SafeContext() must return nil before SetSafeContext() is ever called. + assert.Nil(s.tester, c.SafeContext(), "expected nil SafeContext before first SetSafeContext") + + c.SetContext(c) + c.SetSafeContext(&safeCtxPayloadA{n: 0}) + + s.stopBackground = make(chan struct{}) + // Exercise SafeContext() concurrently from a goroutine other than the + // event-loop goroutine, which is the entire point of "Safe" in its name. + err := goPool.DefaultWorkerPool.Submit(func() { + for { + select { + case <-s.stopBackground: + return + default: + } + switch v := c.SafeContext().(type) { + case nil: + case *safeCtxPayloadA: + _ = v.n + case *safeCtxPayloadB: + _ = v.s + case int: + default: + assert.Failf(s.tester, "unexpected SafeContext type", "%T", v) + } + atomic.AddInt32(&s.backgroundHits, 1) + } + }) + assert.NoError(s.tester, err) + + return +} + +func (s *testSafeContextServer) OnTraffic(c Conn) (action Action) { + buf, err := c.Next(-1) + assert.NoError(s.tester, err) + assert.Equal(s.tester, "ping", string(buf)) + + n := atomic.AddInt32(&s.trafficCount, 1) + // Rotate through different concrete types, including nil, on every + // call - this must never panic. + switch n % 4 { + case 0: + c.SetSafeContext(&safeCtxPayloadA{n: n}) + got, ok := c.SafeContext().(*safeCtxPayloadA) + assert.True(s.tester, ok, "expected *safeCtxPayloadA") + assert.EqualValues(s.tester, n, got.n) + case 1: + c.SetSafeContext(&safeCtxPayloadB{s: "hello"}) + got, ok := c.SafeContext().(*safeCtxPayloadB) + assert.True(s.tester, ok, "expected *safeCtxPayloadB") + assert.Equal(s.tester, "hello", got.s) + case 2: + c.SetSafeContext(int(n)) + got, ok := c.SafeContext().(int) + assert.True(s.tester, ok, "expected int") + assert.EqualValues(s.tester, n, got) + case 3: + c.SetSafeContext(nil) + assert.Nil(s.tester, c.SafeContext(), "expected nil after SetSafeContext(nil)") + } + + // Context() (the non-safe variant) must be untouched by SafeContext calls. + assert.Equal(s.tester, c, c.Context(), "invalid context") + + _, err = c.Write([]byte("pong")) + assert.NoError(s.tester, err) + + if n >= 5 { + action = Close + } + return +} + +func (s *testSafeContextServer) OnClose(c Conn, _ error) (action Action) { + // SafeContext() must still be safely callable (no panic) right up + // until release(), which happens after OnClose returns. + _ = c.SafeContext() + + close(s.stopBackground) + assert.Greater(s.tester, atomic.LoadInt32(&s.backgroundHits), int32(0), + "expected background goroutine to observe SafeContext() at least once") + + s.closedConn = c + close(s.done) + action = Shutdown + return +} + +func testSafeContext(t *testing.T, network, addr string) { + events := &testSafeContextServer{tester: t, network: network, addr: addr, done: make(chan struct{})} + err := Run(events, network+"://"+addr) + assert.NoError(t, err) + select { + case <-events.done: + case <-time.After(5 * time.Second): + t.Fatal("timeout waiting for test completion") + } + + // release() (which nils out safeCtx) runs on the event-loop goroutine + // immediately after OnClose returns, racing with this goroutine waking + // up from events.done, so poll briefly instead of asserting once. + assert.Eventually(t, func() bool { + return events.closedConn.SafeContext() == nil + }, time.Second, time.Millisecond, "SafeContext() must be nil once the connection is released") +}