Use sync.Gate in p9.connState.

sync.WaitGroup.Add(positive delta) is illegal if the WaitGroup counter is zero
and WaitGroup.Wait() may be called concurrently. This is problematic for
p9.connState.pendingWg, which counts inflight requests (so transitions from
zero are normal) and is waited-upon when receiving from the underlying Unix
domain socket returns an error, e.g. during connection shutdown. (Even if the
socket has been closed, new requests can still be concurrently received via
flipcall channels.)

PiperOrigin-RevId: 359416057

1 files changed, 13 insertions, 8 deletions

diff --git a/pkg/p9/server.go b/pkg/p9/server.go
index 8c5c434fd..290d5b9ce 100644
--- a/pkg/p9/server.go
+++ b/pkg/p9/server.go
@@ -81,8 +81,8 @@ type connState struct {
 	// version 0 implies 9P2000.L.
 	version uint32
 
-	// pendingWg counts requests that are still being handled.
-	pendingWg sync.WaitGroup
+	// reqGate counts requests that are still being handled.
+	reqGate sync.Gate
 
 	// -- below relates to the legacy handler --
 
@@ -481,9 +481,13 @@ func (cs *connState) lookupChannel(id uint32) *channel {
 
 // handle handles a single message.
 func (cs *connState) handle(m message) (r message) {
-	cs.pendingWg.Add(1)
+	if !cs.reqGate.Enter() {
+		// connState.stop() has been called; the connection is shutting down.
+		r = newErr(syscall.ECONNRESET)
+		return
+	}
 	defer func() {
-		cs.pendingWg.Done()
+		cs.reqGate.Leave()
 		if r == nil {
 			// Don't allow a panic to propagate.
 			err := recover()
@@ -594,10 +598,11 @@ func (cs *connState) handleRequests() {
 }
 
 func (cs *connState) stop() {
-	// Wait for completion of all inflight requests. This is mostly so that if
-	// a request is stuck, the sandbox supervisor has the opportunity to kill
-	// us with SIGABRT to get a stack dump of the offending handler.
-	cs.pendingWg.Wait()
+	// Stop new requests from proceeding, and wait for completion of all
+	// inflight requests. This is mostly so that if a request is stuck, the
+	// sandbox supervisor has the opportunity to kill us with SIGABRT to get a
+	// stack dump of the offending handler.
+	cs.reqGate.Close()
 
 	// Free the channels.
 	cs.channelMu.Lock()