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: eliminate isClosingErr

Browse source 
Also implement and use the Shutdown method, which was mentioned in a
doc comment in CL 292169 but not actually present at that time.

With proper synchronization, we don't need heuristics to determine
whether an error is due to a connection or listener being closed.
We know whether we have called Close (and why), and we can assume
that if we have called Close then that is probably the reason for
any returned error.

Fixes golang/go#56281.

Change-Id: I5e0ed7db0f736ca8df8cd8cf556b674e7a863a69
Reviewed-on: https://go-review.googlesource.com/c/tools/+/443675
gopls-CI: kokoro <noreply+kokoro@google.com>
Reviewed-by: Alan Donovan <adonovan@google.com>
Auto-Submit: Bryan Mills <bcmills@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>
Run-TryBot: Bryan Mills <bcmills@google.com>
This commit is contained in:
Bryan C. Mills 2022-10-18 10:13:17 -04:00 • committed by Gopher Robot
commit f882967b35
1 changed files with 39 additions and 59 deletions
Download patch file Download diff file Expand all files Collapse all files

98
internal/jsonrpc2/serve.go
View file

@ -6,13 +6,11 @@ package jsonrpc2
import (
"context"
"errors"
"fmt"
"io"
"runtime"
"strings"
"sync"
"syscall"
"sync/atomic"
"time"
)
@ -43,18 +41,31 @@ type Server struct {
listener Listener
binder Binder
async *async
shutdownOnce sync.Once
closing int32 // atomic: set to nonzero when Shutdown is called
}
// 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.
//
// The returned Connection will operate independently using the Preempter and/or
// Handler provided by the Binder, and will release its own resources when the
// connection is broken, but the caller may Close it earlier to stop accepting
// (or sending) new requests.
func Dial(ctx context.Context, dialer Dialer, binder Binder) (*Connection, error) {
// dial a server
rwc, err := dialer.Dial(ctx)
if err != nil {
return nil, err
}
return newConnection(ctx, rwc, binder, nil)
conn, err := newConnection(ctx, rwc, binder, nil)
if err != nil {
rwc.Close()
return nil, err
}
return conn, nil
}
// Serve starts a new server listening for incoming connections and returns
@ -79,6 +90,14 @@ func (s *Server) Wait() error {
return s.async.wait()
}
// Shutdown informs the server to stop accepting new connections.
func (s *Server) Shutdown() {
s.shutdownOnce.Do(func() {
atomic.StoreInt32(&s.closing, 1)
s.listener.Close()
})
}
// run accepts incoming connections from the listener,
// If IdleTimeout is non-zero, run exits after there are no clients for this
// duration, otherwise it exits only on error.
@ -87,11 +106,12 @@ func (s *Server) run(ctx context.Context) {
var activeConns sync.WaitGroup
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) {
// Only Shutdown closes the listener. If we get an error after Shutdown is
// called, assume that that was the cause and don't report the error;
// otherwise, report the error in case it is unexpected.
if atomic.LoadInt32(&s.closing) == 0 {
s.async.setError(err)
}
// We are done generating new connections for good.
@ -100,59 +120,16 @@ func (s *Server) run(ctx context.Context) {
// A new inbound connection.
activeConns.Add(1)
_, err = newConnection(ctx, rwc, s.binder, activeConns.Done)
_, err = newConnection(ctx, rwc, s.binder, activeConns.Done) // unregisters itself when done
if err != nil {
rwc.Close()
activeConns.Done()
if !isClosingError(err) {
s.async.setError(err)
s.listener.Close()
break
}
continue
s.async.setError(err)
}
}
activeConns.Wait()
}
// isClosingError reports if the error occurs normally during the process of
// closing a network connection. It uses imperfect heuristics that err on the
// side of false negatives, and should not be used for anything critical.
func isClosingError(err error) bool {
if err == nil {
return false
}
// Fully unwrap the error, so the following tests work.
for wrapped := err; wrapped != nil; wrapped = errors.Unwrap(err) {
err = wrapped
}
// Was it based on an EOF error?
if err == io.EOF {
return true
}
// Was it based on a closed pipe?
if err == io.ErrClosedPipe {
return true
}
if isErrClosed(err) {
return true
}
if runtime.GOOS == "plan9" {
// Error reading from a closed connection.
if err == syscall.EINVAL {
return true
}
// Error trying to accept a new connection from a closed listener.
if strings.HasSuffix(err.Error(), " listen hungup") {
return true
}
}
return false
}
// NewIdleListener wraps a listener with an idle timeout.
//
// When there are no active connections for at least the timeout duration,
@ -188,10 +165,6 @@ type idleListener struct {
func (l *idleListener) Accept(ctx context.Context) (io.ReadWriteCloser, error) {
rwc, err := l.wrapped.Accept(ctx)
if err != nil && !isClosingError(err) {
return nil, err
}
select {
case n, ok := <-l.active:
if err != nil {
@ -216,13 +189,20 @@ func (l *idleListener) Accept(ctx context.Context) (io.ReadWriteCloser, error) {
// active and closed due to idleness, which would be contradictory and
// confusing. Close the connection and pretend that it never happened.
rwc.Close()
} else {
// In theory the timeout could have raced with an unrelated error return
// from Accept. However, ErrIdleTimeout is arguably still valid (since we
// would have closed due to the timeout independent of the error), and the
// harm from returning a spurious ErrIdleTimeout is negliglible anyway.
}
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.
// The idle timer doesn't run until it receives itself from the idleTimer
// channel, so it can't have called l.wrapped.Close yet and thus err can't
// be ErrIdleTimeout. Leave the idle timer as it was and return whatever
// error we got.
l.idleTimer <- timer
return nil, err
}