fix: better resource accounting

client.go

@@ -35,6 +35,10 @@ var (
 
 	log = logging.Logger("rpc")
 
+	// Normal channel closure must not send cancellation for a request ID that
+	// a retry may already have reused on a new connection.
+	errChannelClosed = xerrors.New("RPC channel closed")
+
 	_defaultHTTPClient = &http.Client{
 		Transport: &http.Transport{
 			Proxy: http.ProxyFromEnvironment,
@@ -367,6 +371,14 @@ func (c *client) setupRequestChan() chan clientRequest {
 			select {
 			case resp = <-cr.ready:
 				break loop
+			case <-c.exiting:
+				// A response may have been queued immediately before shutdown.
+				select {
+				case resp = <-cr.ready:
+					break loop
+				default:
+					return clientResponse{}, xerrors.New("websocket routine exiting")
+				}
 			case <-ctxDone: // send cancel request
 				ctxDone = nil
 
@@ -432,16 +444,21 @@ func (c *client) makeOutChan(ctx context.Context, ftyp reflect.Type, valOut int)
 	retVal := reflect.Zero(ftyp.Out(valOut))
 
 	chCtor := func() (context.Context, func([]byte, bool)) {
+		chanCtx, cancel := context.WithCancelCause(ctx)
 		// unpack chan type to make sure it's reflect.BothDir
 		ctyp := reflect.ChanOf(reflect.BothDir, ftyp.Out(valOut).Elem())
 		ch := reflect.MakeChan(ctyp, 0) // todo: buffer?
 		retVal = ch.Convert(ftyp.Out(valOut))
 
 		incoming := make(chan reflect.Value, 32)
 
-		// gorotuine to handle buffering of items
+		// goroutine to handle buffering of items
 		go func() {
+			defer cancel(errChannelClosed)
+			defer ch.Close()
+
 			buf := (&list.List{}).Init()
+			incomingCh := incoming
 
 			for {
 				front := buf.Front()
@@ -450,15 +467,15 @@ func (c *client) makeOutChan(ctx context.Context, ftyp reflect.Type, valOut int)
 				ncases := 1
 				cases[0] = reflect.SelectCase{
 					Dir:  reflect.SelectRecv,
-					Chan: reflect.ValueOf(ctx.Done()),
+					Chan: reflect.ValueOf(chanCtx.Done()),
 				}
 
 				incomingCase := -1
-				if incoming != nil {
+				if incomingCh != nil {
 					incomingCase = ncases
 					cases[ncases] = reflect.SelectCase{
 						Dir:  reflect.SelectRecv,
-						Chan: reflect.ValueOf(incoming),
+						Chan: reflect.ValueOf(incomingCh),
 					}
 					ncases++
 				}
@@ -478,7 +495,6 @@ func (c *client) makeOutChan(ctx context.Context, ftyp reflect.Type, valOut int)
 
 				switch chosen {
 				case 0:
-					ch.Close()
 					return
 				case incomingCase:
 					if ok {
@@ -492,21 +508,20 @@ func (c *client) makeOutChan(ctx context.Context, ftyp reflect.Type, valOut int)
 							}
 						}
 					} else {
-						incoming = nil
+						incomingCh = nil
 					}
 
 				case sendCase:
 					buf.Remove(front)
 				}
 
-				if incoming == nil && buf.Len() == 0 {
-					ch.Close()
+				if incomingCh == nil && buf.Len() == 0 {
 					return
 				}
 			}
 		}()
 
-		return ctx, func(result []byte, ok bool) {
+		return chanCtx, func(result []byte, ok bool) {
 			if !ok {
 				close(incoming)
 				return
@@ -518,14 +533,14 @@ func (c *client) makeOutChan(ctx context.Context, ftyp reflect.Type, valOut int)
 				return
 			}
 
-			if ctx.Err() != nil {
-				log.Errorf("got rpc message with cancelled context: %s", ctx.Err())
+			if chanCtx.Err() != nil {
+				log.Errorf("got rpc message with cancelled context: %s", chanCtx.Err())
 				return
 			}
 
 			select {
 			case incoming <- val:
-			case <-ctx.Done():
+			case <-chanCtx.Done():
 			}
 		}
 	}

client_test.go

@@ -0,0 +1,155 @@
+package jsonrpc
+
+import (
+	"context"
+	"errors"
+	"reflect"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+)
+
+type requestWaitContext struct {
+	context.Context
+	resume <-chan struct{}
+}
+
+func (c requestWaitContext) Done() <-chan struct{} {
+	<-c.resume
+	return c.Context.Done()
+}
+
+func TestClientPendingRequestUnblocksOnExit(t *testing.T) {
+	for _, queuedReply := range []bool{false, true} {
+		name := "no reply"
+		if queuedReply {
+			name = "queued reply"
+		}
+		t.Run(name, func(t *testing.T) {
+			attempts := 1
+			if queuedReply {
+				// Both receive cases are ready on every attempt. Select chooses
+				// between them randomly, so exercise the synchronized race repeatedly.
+				attempts = 64
+			}
+			for range attempts {
+				checkPendingRequestExit(t, queuedReply)
+			}
+		})
+	}
+}
+
+func checkPendingRequestExit(t *testing.T, queuedReply bool) {
+	t.Helper()
+	exiting := make(chan struct{})
+	stop := sync.OnceFunc(func() { close(exiting) })
+	defer stop()
+	resume := make(chan struct{})
+	unblock := sync.OnceFunc(func() { close(resume) })
+	defer unblock()
+	c := &client{exiting: exiting}
+	requests := c.setupRequestChan()
+	cr := clientRequest{
+		req:   request{Jsonrpc: "2.0", ID: float64(1)},
+		ready: make(chan clientResponse, 1),
+	}
+	type outcome struct {
+		response clientResponse
+		err      error
+	}
+	result := make(chan outcome, 1)
+	go func() {
+		// Done is called after submission and before waiting for a response.
+		ctx := requestWaitContext{Context: context.Background(), resume: resume}
+		response, err := c.doRequest(ctx, cr)
+		result <- outcome{response, err}
+	}()
+
+	select {
+	case <-requests:
+	case <-time.After(5 * time.Second):
+		t.Fatal("request was not submitted")
+	}
+	if queuedReply {
+		cr.ready <- clientResponse{Jsonrpc: "2.0", ID: cr.req.ID}
+	}
+	stop()
+	// The reply cannot be consumed until the connection has also exited.
+	unblock()
+
+	select {
+	case got := <-result:
+		if queuedReply {
+			if got.err != nil || got.response.ID != cr.req.ID {
+				t.Fatalf("queued response lost on exit: response=%+v error=%v", got.response, got.err)
+			}
+		} else if got.err == nil || !strings.Contains(got.err.Error(), "exiting") {
+			t.Fatalf("pending request error = %v, want connection exit error", got.err)
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("pending request did not unblock when connection exited")
+	}
+}
+
+func TestClientChannelCompletionDrainsBufferedValues(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	value, constructor := (&client{}).makeOutChan(ctx, reflect.TypeFor[func() <-chan int](), 0)
+	lifetime, sink := constructor()
+	output := value().Interface().(<-chan int)
+
+	for _, encoded := range []string{"1", "2", "3"} {
+		sink([]byte(encoded), true)
+	}
+	sink(nil, false)
+
+	for want := 1; want <= 3; want++ {
+		select {
+		case got, ok := <-output:
+			if !ok || got != want {
+				t.Fatalf("channel value = %d, open=%v; want %d", got, ok, want)
+			}
+		case <-time.After(5 * time.Second):
+			t.Fatal("buffered channel values did not drain")
+		}
+	}
+	select {
+	case _, ok := <-output:
+		if ok {
+			t.Fatal("channel stayed open after all buffered values drained")
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("completed output channel did not close")
+	}
+	select {
+	case <-lifetime.Done():
+	case <-time.After(5 * time.Second):
+		t.Fatal("completed channel context was retained")
+	}
+	if context.Cause(lifetime) != errChannelClosed || ctx.Err() != nil {
+		t.Fatalf("normal completion cause=%v, parent error=%v", context.Cause(lifetime), ctx.Err())
+	}
+}
+
+func TestClientChannelCancellationClosesOutput(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	value, constructor := (&client{}).makeOutChan(ctx, reflect.TypeFor[func() <-chan int](), 0)
+	lifetime, sink := constructor()
+	output := value().Interface().(<-chan int)
+	cancel()
+
+	select {
+	case _, ok := <-output:
+		if ok {
+			t.Fatal("canceled output channel stayed open")
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("canceling the caller did not stop the channel worker")
+	}
+	if !errors.Is(context.Cause(lifetime), context.Canceled) {
+		t.Fatalf("caller cancellation cause = %v, want context.Canceled", context.Cause(lifetime))
+	}
+	sink(nil, false)
+}

handler.go

@@ -149,7 +149,10 @@ func (s *handler) register(namespace string, r any) {
 // Handle
 
 type rpcErrFunc func(w func(func(io.Writer)), req *request, code ErrorCode, err error)
-type chanOut func(reflect.Value, any) error
+type chanOut struct {
+	reserve  func() error
+	register func(reflect.Value, any) error
+}
 
 func (s *handler) handleReader(ctx context.Context, r io.Reader, w io.Writer, rpcError rpcErrFunc) {
 	wf := func(cb func(io.Writer)) {
@@ -337,7 +340,9 @@ func (s *handler) createError(err error) *JSONRPCError {
 	return out
 }
 
-func (s *handler) handle(ctx context.Context, req request, w func(func(io.Writer)), rpcError rpcErrFunc, done func(keepCtx bool), chOut chanOut) {
+func (s *handler) handle(ctx context.Context, req request, w func(func(io.Writer)), rpcError rpcErrFunc, done func(keepCtx bool), chOut *chanOut) {
+	keepCtx := false
+	defer func() { done(keepCtx) }()
 	// Not sure if we need to sanitize the incoming req.Method or not.
 	ctx, span := s.getSpan(ctx, req)
 	ctx, _ = tag.New(ctx, tag.Insert(metrics.RPCMethod, req.Method))
@@ -352,20 +357,26 @@ func (s *handler) handle(ctx context.Context, req request, w func(func(io.Writer
 		if !ok {
 			rpcError(w, &req, rpcMethodNotFound, fmt.Errorf("method '%s' not found", req.Method))
 			stats.Record(ctx, metrics.RPCInvalidMethod.M(1))
-			done(false)
 			return
 		}
 	}
 
 	outCh := handler.valOut != -1 && handler.handlerFunc.Type().Out(handler.valOut).Kind() == reflect.Chan
-	defer done(outCh)
 
 	if chOut == nil && outCh {
 		rpcError(w, &req, rpcMethodNotFound, fmt.Errorf("method '%s' not supported in this mode (no out channel support)", req.Method))
 		stats.Record(ctx, metrics.RPCRequestError.M(1))
 		return
 	}
 
+	if outCh {
+		if err := chOut.reserve(); err != nil {
+			rpcError(w, &req, 1, err)
+			stats.Record(ctx, metrics.RPCRequestError.M(1))
+			return
+		}
+	}
+
 	callParams := make([]reflect.Value, 1+handler.hasCtx+handler.nParams)
 	callParams[0] = handler.receiver
 	if handler.hasCtx == 1 {
@@ -393,7 +404,6 @@ func (s *handler) handle(ctx context.Context, req request, w func(func(io.Writer
 		if len(ps) != handler.nParams {
 			rpcError(w, &req, rpcInvalidParams, fmt.Errorf("wrong param count (method '%s'): %d != %d", req.Method, len(ps), handler.nParams))
 			stats.Record(ctx, metrics.RPCRequestError.M(1))
-			done(false)
 			return
 		}
 
@@ -471,13 +481,20 @@ func (s *handler) handle(ctx context.Context, req request, w func(func(io.Writer
 	// check error as JSON-RPC spec prohibits error and value at the same time
 	if resp.Error == nil {
 		if res != nil && kind == reflect.Chan {
+			if callResult[handler.valOut].IsNil() {
+				rpcError(w, &req, 1, fmt.Errorf("method '%s' returned a nil channel", req.Method))
+				stats.Record(ctx, metrics.RPCResponseError.M(1))
+				return
+			}
+
 			// Channel responses are sent from channel control goroutine.
 			// Sending responses here could cause deadlocks on writeLk, or allow
 			// sending channel messages before this rpc call returns
 
 			//noinspection GoNilness // already checked above
-			err = chOut(callResult[handler.valOut], req.ID)
+			err = chOut.register(callResult[handler.valOut], req.ID)
 			if err == nil {
+				keepCtx = true
 				return // channel goroutine handles responding
 			}
 

options_server.go

@@ -14,8 +14,9 @@ type jsonrpcReverseClient struct{ reflect.Type }
 type ParamDecoder func(ctx context.Context, json []byte) (reflect.Value, error)
 
 type ServerConfig struct {
-	maxRequestSize int64
-	pingInterval   time.Duration
+	maxRequestSize   int64
+	pingInterval     time.Duration
+	maxSubscriptions int
 
 	paramDecoders map[reflect.Type]ParamDecoder
 	errors        *Errors
@@ -29,8 +30,9 @@ type ServerOption func(c *ServerConfig)
 
 func defaultServerConfig() ServerConfig {
 	return ServerConfig{
-		paramDecoders:  map[reflect.Type]ParamDecoder{},
-		maxRequestSize: DEFAULT_MAX_REQUEST_SIZE,
+		paramDecoders:    map[reflect.Type]ParamDecoder{},
+		maxRequestSize:   DEFAULT_MAX_REQUEST_SIZE,
+		maxSubscriptions: defaultMaxSubscriptions,
 
 		pingInterval:        5 * time.Second,
 		methodNameFormatter: DefaultMethodNameFormatter,
@@ -49,6 +51,15 @@ func WithMaxRequestSize(max int64) ServerOption {
 	}
 }
 
+// WithMaxSubscriptions limits pending and active subscriptions per WebSocket connection.
+// Values <= 0 use the default of 16384. Values above 65534 are capped to leave
+// room for the two internal reflect.Select cases.
+func WithMaxSubscriptions(max int) ServerOption {
+	return func(c *ServerConfig) {
+		c.maxSubscriptions = max
+	}
+}
+
 func WithServerErrors(es Errors) ServerOption {
 	return func(c *ServerConfig) {
 		c.errors = &es