Fix 9P response writer shutdown

cmd/olliesrv/server.go

@@ -206,15 +206,30 @@ func (s *Server) Start(conn net.Conn) {
 		}
 		s.mu.Unlock()
 	}()
-
 	responses := make(chan *plan9.Fcall, 16)
 	var wg sync.WaitGroup
+	var writerWG sync.WaitGroup
+	writerWG.Add(1)
 	go func() {
+		defer writerWG.Done()
 		for resp := range responses {
-			plan9.WriteFcall(conn, resp)
+			if err := plan9.WriteFcall(conn, resp); err != nil {
+				s.log.Error("write: %v", err)
+				connCancel()
+				conn.Close()
+				return
+			}
 		}
 	}()
 
+	stop := func() {
+		connCancel()
+		conn.Close()
+		wg.Wait()
+		close(responses)
+		writerWG.Wait()
+	}
+
 	for {
 		// Set read deadline so we can check for shutdown periodically
 		conn.SetReadDeadline(time.Now().Add(1 * time.Second))
@@ -263,11 +278,13 @@ func (s *Server) Start(conn net.Conn) {
 				cs.mu.Unlock()
 				wg.Done()
 			}()
-			responses <- resp
+			select {
+			case responses <- resp:
+			case <-connCtx.Done():
+			}
 		}(fc, reqCtx)
 	}
-	wg.Wait()
-	close(responses)
+	stop()
 }
 
 func flush(cs *connState, fc *plan9.Fcall) *plan9.Fcall {

cmd/toolsrv/p9.go

@@ -49,45 +49,59 @@ type connState struct {
 
 // serve9P handles a single 9P connection using the virtfs tree.
 func serve9P(ctx context.Context, conn net.Conn, tree *virtfs.Tree, srv *server.Server, noAuth bool) {
+	connCtx, connCancel := context.WithCancel(ctx)
+	defer connCancel()
 	defer conn.Close()
 
 	cs := &connState{
 		fids:     make(map[uint32]*fid),
 		authFids: make(map[uint32]*authFid),
-		ctx:      ctx,
+		ctx:      connCtx,
 		noAuth:   noAuth,
 		srv:      srv,
 	}
 
 	responses := make(chan *plan9.Fcall, 16)
 	var wg sync.WaitGroup
+	var writerWG sync.WaitGroup
 
 	// Writer goroutine
+	writerWG.Add(1)
 	go func() {
+		defer writerWG.Done()
 		for resp := range responses {
 			if err := plan9.WriteFcall(conn, resp); err != nil {
 				logger.Warn("write error: %v", err)
+				connCancel()
+				conn.Close()
 				return
 			}
 		}
 	}()
 
+	stop := func() {
+		connCancel()
+		// Closing the connection interrupts ReadFcall and any blocked writes.
+		conn.Close()
+		wg.Wait()
+		close(responses)
+		writerWG.Wait()
+	}
+
 	for {
 		select {
-		case <-ctx.Done():
-			wg.Wait()
-			close(responses)
+		case <-connCtx.Done():
+			stop()
 			return
 		default:
 		}
 
 		fc, err := plan9.ReadFcall(conn)
 		if err != nil {
-			if err != io.EOF {
+			if err != io.EOF && connCtx.Err() == nil {
 				logger.Warn("read error: %v", err)
 			}
-			wg.Wait()
-			close(responses)
+			stop()
 			return
 		}
 
@@ -96,7 +110,10 @@ func serve9P(ctx context.Context, conn net.Conn, tree *virtfs.Tree, srv *server.
 		go func(fc *plan9.Fcall) {
 			defer wg.Done()
 			resp := handleFcall(cs, fc, tree)
-			responses <- resp
+			select {
+			case responses <- resp:
+			case <-connCtx.Done():
+			}
 		}(fc)
 	}
 }