internal/jsonrpc2_v2: eliminate a temporary connection leak in (*Server).run

Prior to this CL, (*Server).run only filters out inactive connections
when it accepts a new connection. If existing connections complete,
their associated resources can't be garbage-collected until either the
next connection is accepted or the Listener is closed.

This change moves the open-connection accounting to an explicit hook
passed to newConnection, eliminating the need to call Wait entirely.

For golang/go#46047

Change-Id: I3732cb463fcea0c142f17f2b1510fdfd2dbc81da
Reviewed-on: https://go-review.googlesource.com/c/tools/+/388774
Auto-Submit: Bryan Mills <bcmills@google.com>
Reviewed-by: Ian Cottrell <iancottrell@google.com>
gopls-CI: kokoro <noreply+kokoro@google.com>
Run-TryBot: Bryan Mills <bcmills@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>

internal/jsonrpc2/conn.go

@@ -74,6 +74,7 @@ type Connection struct {
 	handler Handler
 
 	onInternalError func(error)
+	onDone          func()
 }
 
 // inFlightState records the state of the incoming and outgoing calls on a
@@ -113,6 +114,9 @@ func (c *Connection) updateInFlight(f func(*inFlightState)) {
 	idle := s.incoming == 0 && len(s.outgoing) == 0 && !s.handlerRunning
 	if idle && (s.closing || s.readErr != nil) && !s.closed {
 		c.closeErr <- c.closer.Close()
+		if c.onDone != nil {
+			c.onDone()
+		}
 		s.closed = true
 	}
 }
@@ -131,8 +135,13 @@ func (o ConnectionOptions) Bind(context.Context, *Connection) (ConnectionOptions
 }
 
 // newConnection creates a new connection and runs it.
+//
 // This is used by the Dial and Serve functions to build the actual connection.
-func newConnection(bindCtx context.Context, rwc io.ReadWriteCloser, binder Binder) (*Connection, error) {
+//
+// The connection is closed automatically (and its resources cleaned up) when
+// the last request has completed after the underlying ReadWriteCloser breaks,
+// but it may be stopped earlier by calling Close (for a clean shutdown).
+func newConnection(bindCtx context.Context, rwc io.ReadWriteCloser, binder Binder, onDone func()) (*Connection, error) {
 	// TODO: Should we create a new event span here?
 	// This will propagate cancellation from ctx; should it?
 	ctx := notDone{bindCtx}
@@ -141,6 +150,7 @@ func newConnection(bindCtx context.Context, rwc io.ReadWriteCloser, binder Binde
 		closer:   rwc,
 		closeErr: make(chan error, 1),
 		writer:   make(chan Writer, 1),
+		onDone:   onDone,
 	}
 
 	options, err := binder.Bind(bindCtx, c)

internal/jsonrpc2/serve.go

@@ -54,7 +54,7 @@ func Dial(ctx context.Context, dialer Dialer, binder Binder) (*Connection, error
 	if err != nil {
 		return nil, err
 	}
-	return newConnection(ctx, rwc, binder)
+	return newConnection(ctx, rwc, binder, nil)
 }
 
 // Serve starts a new server listening for incoming connections and returns
@@ -84,25 +84,25 @@ func (s *Server) Wait() error {
 // duration, otherwise it exits only on error.
 func (s *Server) run(ctx context.Context) {
 	defer s.async.done()
-	var activeConns []*Connection
+
+	var activeConns sync.WaitGroup
 	for {
-		// we never close the accepted connection, we rely on the other end
-		// closing or the socket closing itself naturally
+		// We never close the accepted connection — we rely on the other end
+		// closing or the socket closing itself naturally.
 		rwc, err := s.listener.Accept(ctx)
 		if err != nil {
 			if !isClosingError(err) {
 				s.async.setError(err)
 			}
-			// we are done generating new connections for good
+			// We are done generating new connections for good.
 			break
 		}
 
-		// see if any connections were closed while we were waiting
-		activeConns = onlyActive(activeConns)
-
-		// a new inbound connection,
-		conn, err := newConnection(ctx, rwc, s.binder)
+		// A new inbound connection.
+		activeConns.Add(1)
+		_, err = newConnection(ctx, rwc, s.binder, activeConns.Done)
 		if err != nil {
+			activeConns.Done()
 			if !isClosingError(err) {
 				s.async.setError(err)
 				s.listener.Close()
@@ -110,27 +110,8 @@ func (s *Server) run(ctx context.Context) {
 			}
 			continue
 		}
-		activeConns = append(activeConns, conn)
 	}
-
-	// wait for all active conns to finish
-	for _, c := range activeConns {
-		c.Wait()
-	}
-}
-
-func onlyActive(conns []*Connection) []*Connection {
-	i := 0
-	for _, c := range conns {
-		c.updateInFlight(func(s *inFlightState) {
-			if !s.closed {
-				conns[i] = c
-				i++
-			}
-		})
-	}
-	// trim the slice down
-	return conns[:i]
+	activeConns.Wait()
 }
 
 // isClosingError reports if the error occurs normally during the process of