a request context on the ipc server, cancelled by Close (V-638)

serveConn dispatched under context.Background(), so a dispatch in flight during shutdown could not be told to stop and closeGrace could only abandon it. The server now carries a context, Close cancels it, and each conn derives its own so nothing outlives the connection. A zero-value Server built outside Listen falls back to Background; two wiring tests do that. Client.Close read c.conn with no lock while roundtrip re-dialed and dropped it, which -race caught on the new test. The conn field now has its own mutex, held only across a read or an assignment, so Close and the watchdog reach the connection without queueing behind the call they are interrupting.
2026-08-06 22:59:02 +04:00
parent 123b9aa961
commit cbd8077d2c
3 changed files with 117 additions and 13 deletions
@@ -117,3 +117,54 @@ func TestClientDeadlineBoundsACall(t *testing.T) {
 		t.Errorf("Ping took %v, want the context deadline to bound it", elapsed)
 	}
 }
+
+// blockingAPI parks Presence until its context is cancelled and records what
+// cancelled it. Every other method is the unimplemented floor.
+type blockingAPI struct {
+	UnimplementedCoreAPI
+	entered chan struct{}
+	err     chan error
+}
+
+func (b *blockingAPI) Presence(ctx context.Context) (Presence, error) {
+	close(b.entered)
+	<-ctx.Done()
+	b.err <- ctx.Err()
+	return Presence{}, ctx.Err()
+}
+
+// serveConn dispatched under context.Background() until V-638, so Close could
+// only abandon a dispatch in flight and never tell it to stop.
+func TestServerCloseCancelsADispatchInFlight(t *testing.T) {
+	api := &blockingAPI{entered: make(chan struct{}), err: make(chan error, 1)}
+	srv, err := Listen(filepath.Join(t.TempDir(), "core.sock"), api)
+	if err != nil {
+		t.Fatalf("listen: %v", err)
+	}
+	served := make(chan struct{})
+	go func() { _ = srv.Serve(); close(served) }()
+
+	cli, err := Dial(srv.Path())
+	if err != nil {
+		t.Fatalf("dial: %v", err)
+	}
+	defer cli.Close()
+	go func() { _, _ = cli.Presence(context.Background()) }()
+
+	select {
+	case <-api.entered:
+	case <-time.After(5 * time.Second):
+		t.Fatal("the handler was never dispatched")
+	}
+
+	_ = srv.Close()
+	<-served
+	select {
+	case got := <-api.err:
+		if !errors.Is(got, context.Canceled) {
+			t.Errorf("handler saw %v, want context.Canceled", got)
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("Close did not cancel the dispatch")
+	}
+}
@@ -26,7 +26,12 @@ type Client struct {
 	conn net.Conn
 	path string       // the address as configured, kept for errors and logs
 	addr netaddr.Addr // parsed, so a dropped conn can be re-dialed (core restart)
-	mu   sync.Mutex
+	mu   sync.Mutex   // one request at a time, so a frame and its reply pair up
+
+	// connMu guards the conn field alone, and is held only across an assignment
+	// or a read. It exists so Close and the cancellation watchdog can reach the
+	// connection without waiting for the call that is holding c.mu (V-638).
+	connMu sync.Mutex
 }

 // defaultCallTimeout bounds a call whose context carries no deadline. It is
@@ -109,11 +114,18 @@ func Dial(path string) (*Client, error) {
 	return &Client{conn: c, path: path, addr: addr}, nil
 }

+// Close closes the connection out from under a call in flight, on purpose: a
+// shutdown must not wait out a parked read. It takes connMu and never c.mu, so
+// it cannot block behind the call it is interrupting.
 func (c *Client) Close() error {
+	c.connMu.Lock()
+	defer c.connMu.Unlock()
 	if c.conn == nil {
 		return nil
 	}
-	return c.conn.Close()
+	err := c.conn.Close()
+	c.conn = nil
+	return err
 }

 // DialWait is Dial with patience: it retries with capped backoff until the
@@ -223,14 +235,15 @@ func (c *Client) call(ctx context.Context, m Method, params, result any) error {
 // calling drop, because drop wants c.mu and the caller is holding it — the
 // closed socket fails the read, and roundtrip drops it on the way out.
 func (c *Client) roundtrip(ctx context.Context, m Method, raw json.RawMessage, resp *Response) error {
-	if c.conn == nil {
-		conn, err := netaddr.Dial(c.addr)
+	conn := c.currentConn()
+	if conn == nil {
+		dialed, err := netaddr.Dial(c.addr)
 		if err != nil {
 			return fmt.Errorf("%w: dial %s: %v", errWriteLost, c.addr, err)
 		}
-		c.conn = conn
+		c.setConn(dialed)
+		conn = dialed
 	}
-	conn := c.conn
 	if dl, ok := ctx.Deadline(); ok {
 		_ = conn.SetDeadline(dl)
 	} else {
@@ -248,11 +261,11 @@ func (c *Client) roundtrip(ctx context.Context, m Method, raw json.RawMessage, r
 		}
 	}()

-	if err := writeFrame(c.conn, Request{Method: m, Params: raw}); err != nil {
+	if err := writeFrame(conn, Request{Method: m, Params: raw}); err != nil {
 		c.drop()
 		return fmt.Errorf("%w: %v", errWriteLost, err)
 	}
-	if err := readFrame(c.conn, resp); err != nil {
+	if err := readFrame(conn, resp); err != nil {
 		c.drop()
 		return fmt.Errorf("%w: %v", errReadLost, err)
 	}
@@ -261,12 +274,26 @@ func (c *Client) roundtrip(ctx context.Context, m Method, raw json.RawMessage, r

 // drop closes and forgets the current conn so the next call re-dials.
 func (c *Client) drop() {
+	c.connMu.Lock()
+	defer c.connMu.Unlock()
 	if c.conn != nil {
 		_ = c.conn.Close()
 		c.conn = nil
 	}
 }

+func (c *Client) currentConn() net.Conn {
+	c.connMu.Lock()
+	defer c.connMu.Unlock()
+	return c.conn
+}
+
+func (c *Client) setConn(conn net.Conn) {
+	c.connMu.Lock()
+	defer c.connMu.Unlock()
+	c.conn = conn
+}
+
 // hydrate rehydrates a wire RpcError into the matching package sentinel. The
 // code↔sentinel table is the only place the wire "knows" about errors; keep it
 // in sync with codeOf in wire.go.
@@ -30,6 +30,14 @@ type Server struct {
 	done   chan struct{}
 	accept sync.Mutex // guards wg.Add vs Close's wg.Wait sequence

+	// ctx — server-scoped, cancelled by Close, and the parent of every request
+	// context. serveConn dispatched under context.Background() until V-638, so
+	// a dispatch in flight during shutdown could not be told to stop and the
+	// closeGrace below could only abandon it. Cancelling gives a handler that
+	// respects its context the chance to return instead.
+	ctx    context.Context
+	cancel context.CancelFunc
+
 	// conns — every accepted connection still being served. Close needs these
 	// because closing the listener does nothing to a connection already
 	// accepted: serveConn is parked in readFrame waiting for a peer that may
@@ -208,11 +216,14 @@ func Listen(path string, api CoreAPI) (*Server, error) {
 	if err != nil {
 		return nil, err
 	}
+	ctx, cancel := context.WithCancel(context.Background())
 	s := &Server{
-		path: path,
-		addr: addr,
-		ln:   ln,
-		done: make(chan struct{}),
+		path:   path,
+		addr:   addr,
+		ln:     ln,
+		done:   make(chan struct{}),
+		ctx:    ctx,
+		cancel: cancel,
 	}
 	s.api.Store(api)
 	return s, nil
@@ -250,7 +261,10 @@ func (s *Server) Serve() error {

 func (s *Server) serveConn(c net.Conn) {
 	caller, callerOK := peerCaller(c)
-	ctx := context.Background()
+	// Derived from the server's, so Close cancels a dispatch in flight, and
+	// cancelled when this conn ends so nothing a handler spawned outlives it.
+	ctx, cancel := context.WithCancel(s.serverContext())
+	defer cancel()
 	if callerOK {
 		ctx = WithCaller(ctx, caller)
 	}
@@ -274,6 +288,15 @@ func (s *Server) serveConn(c net.Conn) {
 	}
 }

+// serverContext is s.ctx, or Background for a Server built as a zero value
+// rather than by Listen (the wiring tests do that).
+func (s *Server) serverContext() context.Context {
+	if s.ctx == nil {
+		return context.Background()
+	}
+	return s.ctx
+}
+
 func (s *Server) safeDispatch(ctx context.Context, req Request) (result json.RawMessage, err error) {
 	defer func() {
 		if r := recover(); r != nil {
@@ -720,6 +743,9 @@ func (s *Server) Close() error {
 	default:
 		close(s.done)
 	}
+	if s.cancel != nil {
+		s.cancel()
+	}
 	err := s.ln.Close()
 	// Closing the listener stops new connections; it does nothing to the ones
 	// already accepted. Close those too, or every serveConn parked in readFrame