diff --git a/go.mod b/go.mod index c3377efee..041944c21 100644 --- a/go.mod +++ b/go.mod @@ -4,6 +4,7 @@ go 1.25.0 require ( github.com/containerd/log v0.1.0 + github.com/moby/locker v1.0.1 golang.org/x/sys v0.46.0 google.golang.org/genproto/googleapis/rpc v0.0.0-20260226221140-a57be14db171 google.golang.org/grpc v1.81.1 diff --git a/go.sum b/go.sum index c565d3a54..cb6cdeed1 100644 --- a/go.sum +++ b/go.sum @@ -6,6 +6,8 @@ github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps= github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= +github.com/moby/locker v1.0.1 h1:fOXqR41zeveg4fFODix+1Ch4mj/gT0NE1XJbp/epuBg= +github.com/moby/locker v1.0.1/go.mod h1:S7SDdo5zpBK84bzzVlKr2V0hz+7x9hWbYC/kq7oQppc= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/sirupsen/logrus v1.9.4 h1:TsZE7l11zFCLZnZ+teH4Umoq5BhEIfIzfRDZ1Uzql2w= diff --git a/server.go b/server.go index b8c0689cc..b954c201c 100644 --- a/server.go +++ b/server.go @@ -22,12 +22,14 @@ import ( "io" "math/rand" "net" + "strconv" "sync" "sync/atomic" "syscall" "time" "github.com/containerd/log" + "github.com/moby/locker" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" ) @@ -289,6 +291,7 @@ func (s *Server) newConn(conn net.Conn, handshake any) (*serverConn, error) { conn: conn, handshake: handshake, shutdown: make(chan struct{}), + locker: locker.New(), } c.setState(connStateIdle) if err := s.addConnection(c); err != nil { @@ -306,6 +309,7 @@ type serverConn struct { shutdownOnce sync.Once shutdown chan struct{} // forced shutdown, used by close + locker *locker.Locker } func (c *serverConn) getState() (connState, bool) { @@ -478,17 +482,22 @@ func (c *serverConn) run(sctx context.Context) { } return nil } + lockID := strconv.FormatUint(uint64(id), 10) + c.locker.Lock(lockID) sh, err := c.server.services.handle(ctx, &req, respond) if err != nil { status, _ := status.FromError(err) if !sendStatus(mh.StreamID, status) { + c.locker.Unlock(lockID) return } + c.locker.Unlock(lockID) continue } streams.Store(id, sh) atomic.AddInt32(&active, 1) + c.locker.Unlock(lockID) } // TODO: else we must ignore this for future compat. log this? } @@ -547,8 +556,11 @@ func (c *serverConn) run(sctx context.Context) { // The ttrpc protocol currently does not support the case where // the server is localClosed but not remoteClosed. Once the server // is closing, the whole stream may be considered finished + lockID := strconv.FormatUint(uint64(response.id), 10) + c.locker.Lock(lockID) streams.Delete(response.id) atomic.AddInt32(&active, -1) + c.locker.Unlock(lockID) } case err := <-recvErr: // TODO(stevvooe): Not wildly clear what we should do in this