diff --git a/cmd/olliesrv/server.go b/cmd/olliesrv/server.go index fbf16dd..b957df2 100644 --- a/cmd/olliesrv/server.go +++ b/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 { diff --git a/cmd/toolsrv/p9.go b/cmd/toolsrv/p9.go index cf84196..0896fdc 100644 --- a/cmd/toolsrv/p9.go +++ b/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) } }