client: force-close abandoned client streams via runtime cleanup

Attach a runtime.AddCleanup callback to every clientStream returned by
NewStream so that callers who drop the stream without consuming it have
the underlying *stream force-closed and removed from the connection's
stream map. Without this safety net a leaked stream's recv buffer fills
and the connection's read loop only recovers via the 1-second
ErrStreamFull fallback, leaving streamID slots and goroutines pinned
until that timeout fires.

The cleanup sets recvErr to errStreamAbandoned (unexported - the cleanup
runs after the clientStream is unreachable, so no caller is left to
match on it as a sentinel) and the abandon surfaces in the receive
loop's "failed to handle message" log via the existing error path.

Also adds NewClientWithContext so callers - notably tests - can supply a
parent context whose attached logger is used for the client's internal
goroutines, and a unit test that drives GC, asserts errStreamAbandoned
on the closed stream, verifies the connection is not deadlocked, and
captures the abandon log through the supplied context's logger.

Bumps the Go minimum to 1.24 for runtime.AddCleanup.

Signed-off-by: Derek McGowan <derek@mcg.dev>

Author: dmcgowan

client.go

@@ -22,6 +22,7 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"runtime"
 	"strings"
 	"sync"
 	"syscall"
@@ -107,9 +108,18 @@ func chainUnaryInterceptors(interceptors []UnaryClientInterceptor, final Invoker
 	}
 }
 
-// NewClient creates a new ttrpc client using the given connection
+// NewClient creates a new ttrpc client using the given connection.
+// It is equivalent to [NewClientWithContext] with [context.Background] as
+// the parent context.
 func NewClient(conn net.Conn, opts ...ClientOpts) *Client {
-	ctx, cancel := context.WithCancel(context.Background())
+	return NewClientWithContext(context.Background(), conn, opts...)
+}
+
+// NewClientWithContext creates a new ttrpc client using the given connection,
+// deriving the client's internal context from ctx. Cancellation of ctx
+// shuts the client down; the client's own Close does not cancel ctx.
+func NewClientWithContext(ctx context.Context, conn net.Conn, opts ...ClientOpts) *Client {
+	ctx, cancel := context.WithCancel(ctx)
 	channel := newChannel(conn)
 	c := &Client{
 		codec:           codec{},
@@ -434,10 +444,22 @@ func (c *Client) createStream(flags uint8, b []byte, recvBuf int) (*stream, erro
 }
 
 func (c *Client) deleteStream(s *stream) {
+	c.deleteStreamWithError(s, nil)
+}
+
+// deleteStreamWithError removes the stream from the client and closes it,
+// propagating the supplied error to anyone still observing the stream via
+// receive (the connection read loop) or RecvMsg. A nil error closes the
+// stream with the default ErrClosed.
+//
+// The stream is closed before being removed from the map so that any
+// in-flight message dispatch in the read loop observes recvErr through the
+// normal receive path rather than falling through to "inactive stream".
+func (c *Client) deleteStreamWithError(s *stream, err error) {
+	s.closeWithError(err)
 	c.streamLock.Lock()
 	delete(c.streams, s.id)
 	c.streamLock.Unlock()
-	s.closeWithError(nil)
 }
 
 func (c *Client) getStream(sid streamID) *stream {
@@ -522,12 +544,46 @@ func (c *Client) NewStream(ctx context.Context, desc *StreamDesc, service, metho
 		return nil, err
 	}
 
-	return &clientStream{
+	cs := &clientStream{
 		ctx:  ctx,
 		s:    s,
 		c:    c,
 		desc: desc,
-	}, nil
+	}
+	// Attach a cleanup as a safety net for callers that drop the stream
+	// without consuming it to completion. In the common case the stream is
+	// already closed by the time GC reaches it and the cleanup is a no-op.
+	// If it is still open, the caller leaked the stream; force-close it
+	// with errStreamAbandoned so the connection's receive loop is not
+	// blocked by a buffer that will never drain, and the abandon surfaces
+	// in logs with a specific cause.
+	runtime.AddCleanup(cs, finalizeClientStream, clientStreamCleanupArgs{c: c, s: s})
+	return cs, nil
+}
+
+// clientStreamCleanupArgs carries the state needed by finalizeClientStream.
+// It must not reference the *clientStream that the cleanup is attached to,
+// otherwise the cleanup would never fire.
+type clientStreamCleanupArgs struct {
+	c *Client
+	s *stream
+}
+
+// finalizeClientStream is the runtime.AddCleanup callback registered for each
+// clientStream returned by NewStream. The fast path is the common case: the
+// stream has already been closed (recvClose is closed) and there is nothing
+// to do. The slow path indicates the caller dropped the stream without
+// closing it; force-close the stream with errStreamAbandoned so the abandon
+// surfaces through the connection read loop's "failed to handle message"
+// log when a frame is in flight, and so the receive loop is not blocked by
+// a buffer that will never drain.
+func finalizeClientStream(args clientStreamCleanupArgs) {
+	select {
+	case <-args.s.recvClose:
+		return
+	default:
+	}
+	args.c.deleteStreamWithError(args.s, errStreamAbandoned)
 }
 
 func (c *Client) dispatch(ctx context.Context, req *Request, resp *Response) error {

errors.go

@@ -44,6 +44,14 @@ var (
 	ErrStreamFull = errors.New("ttrpc: stream buffer full")
 )
 
+// errStreamAbandoned is set on a stream's recvErr by the runtime cleanup
+// when the caller dropped a clientStream without closing it. It is not
+// exported because it cannot reach external callers: by the time the
+// cleanup runs every reference to the clientStream is gone, so no RecvMsg
+// or dispatch is left to observe it. Its purpose is to differentiate the
+// abandon case in the connection read loop's error log from a normal close.
+var errStreamAbandoned = errors.New("ttrpc: stream abandoned by caller")
+
 // OversizedMessageErr is used to indicate refusal to send an oversized message.
 // It wraps a ResourceExhausted grpc Status together with the offending message
 // length.

go.mod

@@ -1,6 +1,6 @@
 module github.com/containerd/ttrpc
 
-go 1.23
+go 1.24
 
 require (
 	github.com/containerd/log v0.1.0
@@ -13,4 +13,4 @@ require (
 	google.golang.org/protobuf v1.36.0
 )
 
-require github.com/sirupsen/logrus v1.9.3 // indirect
+require github.com/sirupsen/logrus v1.9.3

services.go

@@ -163,11 +163,17 @@ func (s *serviceSet) handle(ctx context.Context, req *Request, respond func(*sta
 	return nil, status.Errorf(codes.Unimplemented, "method %v", req.Method)
 }
 
-// streamRecvBufferSize is the buffer size for stream recv channels. It
-// should be large enough to absorb normal bursts without hitting the
-// 1-second timeout fallback in receive/data, but small enough that
-// per-stream memory overhead stays trivial.
-const streamRecvBufferSize = 64
+// streamRecvBufferSize is the buffer size for stream recv channels.
+//
+// Consumers are expected to either process incoming messages immediately
+// or hand them off to a goroutine, so the buffer only needs to absorb
+// brief scheduling jitter (a GC pause, a blocked syscall, the time to
+// pass a value to a worker channel). 16 is comfortably above realistic
+// jitter without inflating per-stream memory. Sustained slowness will
+// hit the streamFullTimeout fallback in stream.receive / streamHandler.data,
+// at which point either runtime.AddCleanup (for abandoned streams) or
+// ErrStreamFull (for buggy held-but-unconsumed consumers) recovers.
+const streamRecvBufferSize = 16
 
 type streamHandler struct {
 	ctx     context.Context
@@ -203,7 +209,7 @@ func (s *streamHandler) data(unmarshal Unmarshaler) error {
 			return nil
 		case <-s.ctx.Done():
 			return s.ctx.Err()
-		case <-time.After(time.Second):
+		case <-time.After(streamFullTimeout):
 			return ErrStreamFull
 		}
 	}

stream.go

@@ -29,6 +29,23 @@ type streamMessage struct {
 	payload []byte
 }
 
+// streamFullTimeout bounds how long the receive loop will wait for a stream's
+// recv buffer to drain before giving up. The fallback prevents a single
+// unconsumed stream from indefinitely blocking the connection-level receive
+// loop.
+//
+// Most buffer fillups in practice are abandoned streams (the caller dropped
+// the clientStream without consuming it), and those are handled faster by
+// the runtime.AddCleanup attached in NewStream than by waiting out this
+// timeout. Five seconds therefore primarily bounds head-of-line blocking
+// for the rarer "held but not consumed" case (a goroutine leak in the
+// caller), where neither the cleanup nor the contract that consumers
+// process immediately or hand off to a goroutine can help.
+//
+// Exposed as a var (rather than const) so tests can extend it to observe
+// the abandon-via-cleanup unblock path without racing the timeout.
+var streamFullTimeout = 5 * time.Second
+
 type stream struct {
 	id     streamID
 	sender sender
@@ -92,7 +109,7 @@ func (s *stream) receive(ctx context.Context, msg *streamMessage) error {
 			return nil
 		case <-ctx.Done():
 			return ctx.Err()
-		case <-time.After(time.Second):
+		case <-time.After(streamFullTimeout):
 			s.closeWithError(ErrStreamFull)
 			return ErrStreamFull
 		}