internal/jsonrpc2_v2: move the idle timeout handling out of the server

Change-Id: Ia7dedd84171f4fc4818fb00b154c21b04887391f
Reviewed-on: https://go-review.googlesource.com/c/tools/+/305129
Trust: Ian Cottrell <iancottrell@google.com>
Run-TryBot: Ian Cottrell <iancottrell@google.com>
gopls-CI: kokoro <noreply+kokoro@google.com>
TryBot-Result: Go Bot <gobot@golang.org>
Reviewed-by: Robert Findley <rfindley@google.com>

internal/jsonrpc2/jsonrpc2.go

@@ -68,6 +68,15 @@ func (a *async) done() {
 	close(a.ready)
 }
 
+func (a *async) isDone() bool {
+	select {
+	case <-a.ready:
+		return true
+	default:
+		return false
+	}
+}
+
 func (a *async) wait() error {
 	<-a.ready
 	err := <-a.errBox

internal/jsonrpc2/jsonrpc2_test.go

@@ -130,7 +130,7 @@ func testConnection(t *testing.T, framer jsonrpc2.Framer) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	server, err := jsonrpc2.Serve(ctx, listener, binder{framer, nil}, jsonrpc2.ServeOptions{})
+	server, err := jsonrpc2.Serve(ctx, listener, binder{framer, nil})
 	if err != nil {
 		t.Fatal(err)
 	}

internal/jsonrpc2/serve.go

@@ -7,9 +7,9 @@ package jsonrpc2
 import (
 	"context"
 	"io"
+	"sync"
 	"time"
 
-	"golang.org/x/tools/internal/event"
 	errors "golang.org/x/xerrors"
 )
 
@@ -39,17 +39,9 @@ type Dialer interface {
 type Server struct {
 	listener Listener
 	binder   Binder
-	options  ServeOptions // a copy of the config that started this server
 	async    async
 }
 
-// ServeOptions holds the options to the Serve function.
-//TODO: kill ServeOptions and push timeout into the listener
-type ServeOptions struct {
-	// IdleTimeout is the maximum amount of time to remain idle and running.
-	IdleTimeout time.Duration
-}
-
 // Dial uses the dialer to make a new connection, wraps the returned
 // reader and writer using the framer to make a stream, and then builds
 // a connection on top of that stream using the binder.
@@ -69,11 +61,10 @@ func Dial(ctx context.Context, dialer Dialer, binder Binder) (*Connection, error
 // You can call Wait to block on the server, or Shutdown to get the sever to
 // terminate gracefully.
 // To notice incoming connections, use an intercepting Binder.
-func Serve(ctx context.Context, listener Listener, binder Binder, options ServeOptions) (*Server, error) {
+func Serve(ctx context.Context, listener Listener, binder Binder) (*Server, error) {
 	server := &Server{
 		listener: listener,
 		binder:   binder,
-		options:  options,
 	}
 	server.async.init()
 	go server.run(ctx)
@@ -90,94 +81,49 @@ func (s *Server) Wait() error {
 // duration, otherwise it exits only on error.
 func (s *Server) run(ctx context.Context) {
 	defer s.async.done()
-	// Max duration: ~290 years; surely that's long enough.
+	var activeConns []*Connection
-	const forever = 1<<63 - 1
-	idleTimeout := s.options.IdleTimeout
-	if idleTimeout <= 0 {
-		idleTimeout = forever
-	}
-	idleTimer := time.NewTimer(idleTimeout)
-
-	// run a goroutine that listens for incoming connections and posts them
-	// back to the worker
-	newStreams := make(chan io.ReadWriteCloser)
-	go func() {
-		for {
-			// 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) {
-					event.Error(ctx, "Accept", err)
-				}
-				// signal we are done generating new connections for good
-				close(newStreams)
-				return
-			}
-			newStreams <- rwc
-		}
-	}()
-
-	closedConns := make(chan struct{})
-	activeConns := 0
-	lnClosed := false
 	for {
-		select {
+		// we never close the accepted connection, we rely on the other end
-		case rwc := <-newStreams:
+		// closing or the socket closing itself naturally
-			// whatever happes we are not idle anymore
+		rwc, err := s.listener.Accept(ctx)
-			idleTimer.Stop()
+		if err != nil {
-			if rwc == nil {
+			if !isClosingError(err) {
-				// the net listener has been closed
+				s.async.setError(err)
-				lnClosed = true
-				if activeConns == 0 {
-					// accept is done and there are no active connections, so just stop now
-					return
-				}
-				// replace the channel with one that will never trigger
-				// this is save because the only writer has already quit
-				newStreams = nil
-				// and then wait for all active connections to stop
-				continue
 			}
-			// a new inbound connection,
+			// we are done generating new connections for good
-			conn, err := newConnection(ctx, rwc, s.binder)
+			break
-			if err != nil {
+		}
-				if !isClosingError(err) {
+
-					event.Error(ctx, "NewConn", err)
+		// see if any connections were closed while we were waiting
-				}
+		activeConns = onlyActive(activeConns)
-				continue
+
+		// a new inbound connection,
+		conn, err := newConnection(ctx, rwc, s.binder)
+		if err != nil {
+			if !isClosingError(err) {
+				s.async.setError(err)
 			}
-			// register the new conn as active
+			continue
-			activeConns++
+		}
-			// wrap the conn in a close monitor
+		activeConns = append(activeConns, conn)
-			//TODO: we do this to maintain our active count correctly, is there a better way?
+	}
-			go func() {
+
-				err := conn.Wait()
+	// wait for all active conns to finish
-				if err != nil && !isClosingError(err) {
+	for _, c := range activeConns {
-					event.Error(ctx, "closed a connection", err)
+		c.Wait()
-				}
+	}
-				closedConns <- struct{}{}
+}
-			}()
+
-		case <-closedConns:
+func onlyActive(conns []*Connection) []*Connection {
-			activeConns--
+	i := 0
-			if activeConns == 0 {
+	for _, c := range conns {
-				// no more active connections, restart the idle timer
+		if !c.async.isDone() {
-				if lnClosed {
+			conns[i] = c
-					// we can never get a new connection, so we are done
+			i++
-					return
-				}
-				// we are idle, but might get a new connection still
-				idleTimer.Reset(idleTimeout)
-			}
-		case <-idleTimer.C:
-			// no activity for a while, time to stop serving
-			s.async.setError(ErrIdleTimeout)
-			return
-		case <-ctx.Done():
-			s.async.setError(ctx.Err())
-			return
 		}
 	}
+	// trim the slice down
+	return conns[:i]
 }
 
 // isClosingError reports if the error occurs normally during the process of
@@ -206,3 +152,114 @@ func isClosingError(err error) bool {
 
 	return false
 }
+
+// NewIdleListener wraps a listener with an idle timeout.
+// When there are no active connections for at least the timeout duration a
+// call to accept will fail with ErrIdleTimeout.
+func NewIdleListener(timeout time.Duration, wrap Listener) Listener {
+	l := &idleListener{
+		timeout:    timeout,
+		wrapped:    wrap,
+		newConns:   make(chan *idleCloser),
+		closed:     make(chan struct{}),
+		wasTimeout: make(chan struct{}),
+	}
+	go l.run()
+	return l
+}
+
+type idleListener struct {
+	wrapped    Listener
+	timeout    time.Duration
+	newConns   chan *idleCloser
+	closed     chan struct{}
+	wasTimeout chan struct{}
+	closeOnce  sync.Once
+}
+
+type idleCloser struct {
+	wrapped   io.ReadWriteCloser
+	closed    chan struct{}
+	closeOnce sync.Once
+}
+
+func (c *idleCloser) Read(p []byte) (int, error) {
+	n, err := c.wrapped.Read(p)
+	if err != nil && isClosingError(err) {
+		c.closeOnce.Do(func() { close(c.closed) })
+	}
+	return n, err
+}
+
+func (c *idleCloser) Write(p []byte) (int, error) {
+	// we do not close on write failure, we rely on the wrapped writer to do that
+	// if it is appropriate, which we will detect in the next read.
+	return c.wrapped.Write(p)
+}
+
+func (c *idleCloser) Close() error {
+	// we rely on closing the wrapped stream to signal to the next read that we
+	// are closed, rather than triggering the closed signal directly
+	return c.wrapped.Close()
+}
+
+func (l *idleListener) Accept(ctx context.Context) (io.ReadWriteCloser, error) {
+	rwc, err := l.wrapped.Accept(ctx)
+	if err != nil {
+		if isClosingError(err) {
+			// underlying listener was closed
+			l.closeOnce.Do(func() { close(l.closed) })
+			// was it closed because of the idle timeout?
+			select {
+			case <-l.wasTimeout:
+				err = ErrIdleTimeout
+			default:
+			}
+		}
+		return nil, err
+	}
+	conn := &idleCloser{
+		wrapped: rwc,
+		closed:  make(chan struct{}),
+	}
+	l.newConns <- conn
+	return conn, err
+}
+
+func (l *idleListener) Close() error {
+	defer l.closeOnce.Do(func() { close(l.closed) })
+	return l.wrapped.Close()
+}
+
+func (l *idleListener) Dialer() Dialer {
+	return l.wrapped.Dialer()
+}
+
+func (l *idleListener) run() {
+	var conns []*idleCloser
+	for {
+		var firstClosed chan struct{} // left at nil if there are no active conns
+		var timeout <-chan time.Time  // left at nil if there are  active conns
+		if len(conns) > 0 {
+			firstClosed = conns[0].closed
+		} else {
+			timeout = time.After(l.timeout)
+		}
+		select {
+		case <-l.closed:
+			// the main listener closed, no need to keep going
+			return
+		case conn := <-l.newConns:
+			// a new conn arrived, add it to the list
+			conns = append(conns, conn)
+		case <-timeout:
+			// we timed out, only happens when there are no active conns
+			// close the underlying listener, and allow the normal closing process to happen
+			close(l.wasTimeout)
+			l.wrapped.Close()
+		case <-firstClosed:
+			// a conn closed, remove it from the active list
+			conns = conns[:copy(conns, conns[1:])]
+		}
+	}
+}

internal/jsonrpc2/serve_test.go

@@ -23,12 +23,9 @@ func TestIdleTimeout(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+	listener = jsonrpc2.NewIdleListener(100*time.Millisecond, listener)
 	defer listener.Close()
-
+	server, err := jsonrpc2.Serve(ctx, listener, jsonrpc2.ConnectionOptions{})
-	server, err := jsonrpc2.Serve(ctx, listener, jsonrpc2.ConnectionOptions{},
-		jsonrpc2.ServeOptions{
-			IdleTimeout: 100 * time.Millisecond,
-		})
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -119,10 +116,9 @@ func TestServe(t *testing.T) {
 }
 
 func newFake(ctx context.Context, l jsonrpc2.Listener) (*jsonrpc2.Connection, func(context.Context), error) {
+	l = jsonrpc2.NewIdleListener(100*time.Millisecond, l)
 	server, err := jsonrpc2.Serve(ctx, l, jsonrpc2.ConnectionOptions{
 		Handler: fakeHandler{},
-	}, jsonrpc2.ServeOptions{
-		IdleTimeout: 100 * time.Millisecond,
 	})
 	if err != nil {
 		return nil, nil, err