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: rework concurrency in idleListener

Browse source 
This eliminates a race between a successful Accept call and a
concurrent Close call, which previously could have shut down the
'run' goroutine before Accept sent to the newConns channel, causing
Accept to deadlock.

In fact, it eliminates the long-running 'run' goroutine entirely
(replacing it with a time.Timer), and in the process avoids leaking
O(N) closed connections when only the very first one is long-lived.

It also eliminates a potential double-Close bug: if the run method had
called l.wrapped.Close due to an idle timeout, a subsequent call to
Close would invoke l.wrapped.Close again. The io.Closer method
explicitly documents doubled Close calls as undefined behavior, and
many Closer implementations (especially test fakes) panic or deadlock
in that case.

It also eliminates a timer leak if the Listener rapidly oscillates
between active and idle: previously the implementation used
time.After, but it now uses an explicit time.Timer which can be
stopped (and garbage-collected) when the listener becomes active.

Idleness is now tracked based on the connection's Close method rather
than Read: we have no guarantee in general that a caller will ever
actually invoke Read (if, for example, they Close the connection as
soon as it is dialed), but we can reasonably expect a caller to at
least try to ensure that Close is always called.

We now also verify, using a finalizer on a best-effort basis, that the
Close method on each connection is called. We use the finalizer to
verify the Close call — rather than to close the connection implicitly
— because closing the connection in a finalizer would delay the start
of the idle timer by an arbitrary and unbounded duration after the
last connection is actually no longer in use.

Fixes golang/go#46047.
Fixes golang/go#51435.
For golang/go#46520.
For golang/go#49387.

Change-Id: If173a3ed7a44aff14734b72c8340122e8d020f26
Reviewed-on: https://go-review.googlesource.com/c/tools/+/388597
Run-TryBot: Bryan Mills <bcmills@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>
Reviewed-by: Ian Cottrell <iancottrell@google.com>
Auto-Submit: Bryan Mills <bcmills@google.com>
gopls-CI: kokoro <noreply+kokoro@google.com>
This commit is contained in:
Bryan C. Mills 2022-02-28 17:19:32 -05:00 • committed by Gopher Robot
commit 9821e82f8e
3 changed files with 334 additions and 112 deletions
Download patch file Download diff file Expand all files Collapse all files

252
internal/jsonrpc2/serve.go
View file

@ -7,6 +7,7 @@ package jsonrpc2
import (
"context"
"errors"
"fmt"
"io"
"runtime"
"strings"
@ -171,80 +172,121 @@ func isClosingError(err error) bool {
}
// 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.
//
// When there are no active connections for at least the timeout duration,
// calls to Accept will fail with ErrIdleTimeout.
//
// A connection is considered inactive as soon as its Close method is called.
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{}),
wrapped: wrap,
timeout: timeout,
active: make(chan int, 1),
timedOut: make(chan struct{}),
idleTimer: make(chan *time.Timer, 1),
}
go l.run()
l.idleTimer <- time.AfterFunc(l.timeout, l.timerExpired)
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()
wrapped Listener
timeout time.Duration
// Only one of these channels is receivable at any given time.
active chan int // count of active connections; closed when Close is called if not timed out
timedOut chan struct{} // closed when the idle timer expires
idleTimer chan *time.Timer // holds the timer only when idle
}
// Accept accepts an incoming connection.
//
// If an incoming connection is accepted concurrent to the listener being closed
// due to idleness, the new connection is immediately closed.
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:
}
}
if err != nil && !isClosingError(err) {
return nil, err
}
conn := &idleCloser{
wrapped: rwc,
closed: make(chan struct{}),
select {
case n, ok := <-l.active:
if err != nil {
if ok {
l.active <- n
}
return nil, err
}
if ok {
l.active <- n + 1
} else {
// l.wrapped.Close Close has been called, but Accept returned a
// connection. This race can occur with concurrent Accept and Close calls
// with any net.Listener, and it is benign: since the listener was closed
// explicitly, it can't have also timed out.
}
return l.newConn(rwc), nil
case <-l.timedOut:
if err == nil {
// Keeping the connection open would leave the listener simultaneously
// active and closed due to idleness, which would be contradictory and
// confusing. Close the connection and pretend that it never happened.
rwc.Close()
}
return nil, ErrIdleTimeout
case timer := <-l.idleTimer:
if err != nil {
// The idle timer hasn't run yet, so err can't be ErrIdleTimeout.
// Leave the idle timer as it was and return whatever error we got.
l.idleTimer <- timer
return nil, err
}
if !timer.Stop() {
// Failed to stop the timer — the timer goroutine is in the process of
// firing. Send the timer back to the timer goroutine so that it can
// safely close the timedOut channel, and then wait for the listener to
// actually be closed before we return ErrIdleTimeout.
l.idleTimer <- timer
rwc.Close()
<-l.timedOut
return nil, ErrIdleTimeout
}
l.active <- 1
return l.newConn(rwc), nil
}
l.newConns <- conn
return conn, err
}
func (l *idleListener) Close() error {
defer l.closeOnce.Do(func() { close(l.closed) })
select {
case _, ok := <-l.active:
if ok {
close(l.active)
}
case <-l.timedOut:
// Already closed by the timer; take care not to double-close if the caller
// only explicitly invokes this Close method once, since the io.Closer
// interface explicitly leaves doubled Close calls undefined.
return ErrIdleTimeout
case timer := <-l.idleTimer:
if !timer.Stop() {
// Couldn't stop the timer. It shouldn't take long to run, so just wait
// (so that the Listener is guaranteed to be closed before we return)
// and pretend that this call happened afterward.
// That way we won't leak any timers or goroutines when Close returns.
l.idleTimer <- timer
<-l.timedOut
return ErrIdleTimeout
}
close(l.active)
}
return l.wrapped.Close()
}
@ -252,31 +294,83 @@ 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
func (l *idleListener) timerExpired() {
select {
case n, ok := <-l.active:
if ok {
panic(fmt.Sprintf("jsonrpc2: idleListener idle timer fired with %d connections still active", n))
} else {
timeout = time.After(l.timeout)
panic("jsonrpc2: Close finished with idle timer still running")
}
select {
case <-l.closed:
// the main listener closed, no need to keep going
case <-l.timedOut:
panic("jsonrpc2: idleListener idle timer fired more than once")
case <-l.idleTimer:
// The timer for this very call!
}
// Close the Listener with all channels still blocked to ensure that this call
// to l.wrapped.Close doesn't race with the one in l.Close.
defer close(l.timedOut)
l.wrapped.Close()
}
func (l *idleListener) connClosed() {
select {
case n, ok := <-l.active:
if !ok {
// l is already closed, so it can't close due to idleness,
// and we don't need to track the number of active connections any more.
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:])]
}
n--
if n == 0 {
l.idleTimer <- time.AfterFunc(l.timeout, l.timerExpired)
} else {
l.active <- n
}
case <-l.timedOut:
panic("jsonrpc2: idleListener idle timer fired before last active connection was closed")
case <-l.idleTimer:
panic("jsonrpc2: idleListener idle timer active before last active connection was closed")
}
}
type idleListenerConn struct {
wrapped io.ReadWriteCloser
l *idleListener
closeOnce sync.Once
}
func (l *idleListener) newConn(rwc io.ReadWriteCloser) *idleListenerConn {
c := &idleListenerConn{
wrapped: rwc,
l: l,
}
// A caller that forgets to call Close may disrupt the idleListener's
// accounting, even though the file descriptor for the underlying connection
// may eventually be garbage-collected anyway.
//
// Set a (best-effort) finalizer to verify that a Close call always occurs.
// (We will clear the finalizer explicitly in Close.)
runtime.SetFinalizer(c, func(c *idleListenerConn) {
panic("jsonrpc2: IdleListener connection became unreachable without a call to Close")
})
return c
}
func (c *idleListenerConn) Read(p []byte) (int, error) { return c.wrapped.Read(p) }
func (c *idleListenerConn) Write(p []byte) (int, error) { return c.wrapped.Write(p) }
func (c *idleListenerConn) Close() error {
defer c.closeOnce.Do(func() {
c.l.connClosed()
runtime.SetFinalizer(c, nil)
})
return c.wrapped.Close()
}

190
internal/jsonrpc2/serve_test.go
View file

@ -7,6 +7,8 @@ package jsonrpc2_test
import (
"context"
"errors"
"fmt"
"runtime/debug"
"testing"
"time"
@ -16,48 +18,117 @@ import (
func TestIdleTimeout(t *testing.T) {
stacktest.NoLeak(t)
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
listener, err := jsonrpc2.NetListener(ctx, "tcp", "localhost:0", jsonrpc2.NetListenOptions{})
if err != nil {
t.Fatal(err)
}
listener = jsonrpc2.NewIdleListener(100*time.Millisecond, listener)
defer listener.Close()
server, err := jsonrpc2.Serve(ctx, listener, jsonrpc2.ConnectionOptions{})
if err != nil {
t.Fatal(err)
}
// Use a panicking time.AfterFunc instead of context.WithTimeout so that we
// get a goroutine dump on failure. We expect the test to take on the order of
// a few tens of milliseconds at most, so 10s should be several orders of
// magnitude of headroom.
timer := time.AfterFunc(10*time.Second, func() {
debug.SetTraceback("all")
panic("TestIdleTimeout deadlocked")
})
defer timer.Stop()
ctx := context.Background()
connect := func() *jsonrpc2.Connection {
client, err := jsonrpc2.Dial(ctx,
listener.Dialer(),
jsonrpc2.ConnectionOptions{})
try := func(d time.Duration) (longEnough bool) {
listener, err := jsonrpc2.NetListener(ctx, "tcp", "localhost:0", jsonrpc2.NetListenOptions{})
if err != nil {
t.Fatal(err)
}
return client
}
// Exercise some connection/disconnection patterns, and then assert that when
// our timer fires, the server exits.
conn1 := connect()
conn2 := connect()
if err := conn1.Close(); err != nil {
t.Fatalf("conn1.Close failed with error: %v", err)
}
if err := conn2.Close(); err != nil {
t.Fatalf("conn2.Close failed with error: %v", err)
}
conn3 := connect()
if err := conn3.Close(); err != nil {
t.Fatalf("conn3.Close failed with error: %v", err)
idleStart := time.Now()
listener = jsonrpc2.NewIdleListener(d, listener)
defer listener.Close()
server, err := jsonrpc2.Serve(ctx, listener, jsonrpc2.ConnectionOptions{})
if err != nil {
t.Fatal(err)
}
// Exercise some connection/disconnection patterns, and then assert that when
// our timer fires, the server exits.
conn1, err := jsonrpc2.Dial(ctx, listener.Dialer(), jsonrpc2.ConnectionOptions{})
if err != nil {
if since := time.Since(idleStart); since < d {
t.Fatalf("conn1 failed to connect after %v: %v", since, err)
}
t.Log("jsonrpc2.Dial:", err)
return false // Took to long to dial, so the failure could have been due to the idle timeout.
}
// On the server side, Accept can race with the connection timing out.
// Send a call and wait for the response to ensure that the connection was
// actually fully accepted.
ac := conn1.Call(ctx, "ping", nil)
if err := ac.Await(ctx, nil); !errors.Is(err, jsonrpc2.ErrMethodNotFound) {
if since := time.Since(idleStart); since < d {
t.Fatalf("conn1 broken after %v: %v", since, err)
}
t.Log(`conn1.Call(ctx, "ping", nil):`, err)
conn1.Close()
return false
}
conn2, err := jsonrpc2.Dial(ctx, listener.Dialer(), jsonrpc2.ConnectionOptions{})
if err != nil {
conn1.Close()
if since := time.Since(idleStart); since < d {
t.Fatalf("conn2 failed to connect while non-idle: %v", err)
}
t.Log("jsonrpc2.Dial:", err)
return false
}
if err := conn1.Close(); err != nil {
t.Fatalf("conn1.Close failed with error: %v", err)
}
idleStart = time.Now()
if err := conn2.Close(); err != nil {
t.Fatalf("conn2.Close failed with error: %v", err)
}
conn3, err := jsonrpc2.Dial(ctx, listener.Dialer(), jsonrpc2.ConnectionOptions{})
if err != nil {
if since := time.Since(idleStart); since < d {
t.Fatalf("conn3 failed to connect after %v: %v", since, err)
}
t.Log("jsonrpc2.Dial:", err)
return false // Took to long to dial, so the failure could have been due to the idle timeout.
}
ac = conn3.Call(ctx, "ping", nil)
if err := ac.Await(ctx, nil); !errors.Is(err, jsonrpc2.ErrMethodNotFound) {
if since := time.Since(idleStart); since < d {
t.Fatalf("conn3 broken after %v: %v", since, err)
}
t.Log(`conn3.Call(ctx, "ping", nil):`, err)
conn3.Close()
return false
}
idleStart = time.Now()
if err := conn3.Close(); err != nil {
t.Fatalf("conn3.Close failed with error: %v", err)
}
serverError := server.Wait()
if !errors.Is(serverError, jsonrpc2.ErrIdleTimeout) {
t.Errorf("run() returned error %v, want %v", serverError, jsonrpc2.ErrIdleTimeout)
}
if since := time.Since(idleStart); since < d {
t.Errorf("server shut down after %v idle; want at least %v", since, d)
}
return true
}
serverError := server.Wait()
if !errors.Is(serverError, jsonrpc2.ErrIdleTimeout) {
t.Errorf("run() returned error %v, want %v", serverError, jsonrpc2.ErrIdleTimeout)
d := 1 * time.Millisecond
for {
t.Logf("testing with idle timout %v", d)
if !try(d) {
d *= 2
continue
}
break
}
}
@ -142,3 +213,52 @@ func newFake(t *testing.T, ctx context.Context, l jsonrpc2.Listener) (*jsonrpc2.
server.Wait()
}, nil
}
// TestIdleListenerAcceptCloseRace checks for the Accept/Close race fixed in CL 388597.
//
// (A bug in the idleListener implementation caused a successful Accept to block
// on sending to a background goroutine that could have already exited.)
func TestIdleListenerAcceptCloseRace(t *testing.T) {
ctx := context.Background()
n := 10
// Each iteration of the loop appears to take around a millisecond, so to
// avoid spurious failures we'll set the watchdog for three orders of
// magnitude longer. When the bug was present, this reproduced the deadlock
// reliably on a Linux workstation when run with -count=100, which should be
// frequent enough to show up on the Go build dashboard if it regresses.
watchdog := time.Duration(n) * 1000 * time.Millisecond
timer := time.AfterFunc(watchdog, func() {
debug.SetTraceback("all")
panic(fmt.Sprintf("TestAcceptCloseRace deadlocked after %v", watchdog))
})
defer timer.Stop()
for ; n > 0; n-- {
listener, err := jsonrpc2.NetPipeListener(ctx)
if err != nil {
t.Fatal(err)
}
listener = jsonrpc2.NewIdleListener(24*time.Hour, listener)
done := make(chan struct{})
go func() {
conn, err := jsonrpc2.Dial(ctx, listener.Dialer(), jsonrpc2.ConnectionOptions{})
listener.Close()
if err == nil {
conn.Close()
}
close(done)
}()
// Accept may return a non-nil error if Close closes the underlying network
// connection before the wrapped Accept call unblocks. However, it must not
// deadlock!
c, err := listener.Accept(ctx)
if err == nil {
c.Close()
}
<-done
}
}

8
internal/jsonrpc2/wire.go
View file

@ -76,3 +76,11 @@ func NewError(code int64, message string) error {
func (err *wireError) Error() string {
return err.Message
}
func (err *wireError) Is(other error) bool {
w, ok := other.(*wireError)
if !ok {
return false
}
return err.Code == w.Code
}