modelcontextprotocol/go-sdk
Watch
1
0
Fork
You've already forked go-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

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

Browse source 
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>
This commit is contained in:
Ian Cottrell 2021-03-25 22:16:30 -04:00
commit 76e263b8c9
4 changed files with 163 additions and 101 deletions
Download patch file Download diff file Expand all files Collapse all files

9
internal/jsonrpc2/jsonrpc2.go
View file

@ -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

2
internal/jsonrpc2/jsonrpc2_test.go
View file

@ -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)
}

245
internal/jsonrpc2/serve.go
View file

@ -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.
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
var activeConns []*Connection
for {
select {
case rwc := <-newStreams:
// whatever happes we are not idle anymore
idleTimer.Stop()
if rwc == nil {
// the net listener has been closed
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
// 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)
}
// a new inbound connection,
conn, err := newConnection(ctx, rwc, s.binder)
if err != nil {
if !isClosingError(err) {
event.Error(ctx, "NewConn", err)
}
continue
// 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)
if err != nil {
if !isClosingError(err) {
s.async.setError(err)
}
// register the new conn as active
activeConns++
// wrap the conn in a close monitor
//TODO: we do this to maintain our active count correctly, is there a better way?
go func() {
err := conn.Wait()
if err != nil && !isClosingError(err) {
event.Error(ctx, "closed a connection", err)
}
closedConns <- struct{}{}
}()
case <-closedConns:
activeConns--
if activeConns == 0 {
// no more active connections, restart the idle timer
if lnClosed {
// we can never get a new connection, so we are done
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
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 {
if !c.async.isDone() {
conns[i] = c
i++
}
}
// 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:])]
}
}
}

10
internal/jsonrpc2/serve_test.go
View file

@ -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{},
jsonrpc2.ServeOptions{
IdleTimeout: 100 * time.Millisecond,
})
server, err := jsonrpc2.Serve(ctx, listener, jsonrpc2.ConnectionOptions{})
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