Fix test flake (#4347)

Author: stefanvanburen

.golangci.yml

@@ -328,6 +328,11 @@ linters:
         # This is safe (we aren't traversing filesystem boundaries).
         path: private/pkg/storage/storageos/bucket.go
         text: os.Rename
+      - linters:
+          - containedctx
+        # connCtx is cancelled when the connection is done; the context lifetime is
+        # tied to the struct, not passed per-call.
+        path: private/buf/buflsp/buflsp.go
       - linters:
           - containedctx
         # we actually want to embed a context here

private/buf/buflsp/buflsp.go

@@ -61,6 +61,17 @@ func Serve(
 	logger.Info("starting LSP server")
 
 	conn := jsonrpc2.NewConn(stream)
+	// connCtx is a context scoped to the connection's lifetime. It is cancelled
+	// when the connection is done (or when ctx is cancelled), so that background
+	// goroutines (e.g. RunChecks) do not outlive the connection.
+	connCtx, connCancel := context.WithCancel(context.Background())
+	go func() {
+		select {
+		case <-ctx.Done():
+		case <-conn.Done():
+		}
+		connCancel()
+	}()
 	lsp := &lsp{
 		conn: conn,
 		client: protocol.ClientDispatcher(
@@ -76,6 +87,8 @@ func Serve(
 		queryExecutor: queryExecutor,
 		opener:        source.NewMap(nil),
 		irSession:     new(ir.Session),
+		connCtx:       connCtx,
+		connCancel:    connCancel,
 	}
 	lsp.fileManager = newFileManager(lsp)
 	lsp.workspaceManager = newWorkspaceManager(lsp)
@@ -100,9 +113,11 @@ func Serve(
 // Its handler methods are not defined in buflsp.go; they are defined in other files, grouped
 // according to the groupings in
 type lsp struct {
-	conn      jsonrpc2.Conn
-	client    protocol.Client
-	container appext.Container
+	conn       jsonrpc2.Conn
+	client     protocol.Client
+	container  appext.Container
+	connCtx    context.Context    // cancelled when the connection is done
+	connCancel context.CancelFunc // cancels connCtx
 
 	logger           *slog.Logger
 	bufVersion       string // buf version, set at server creation
@@ -150,7 +165,10 @@ func (l *lsp) newHandler() (jsonrpc2.Handler, error) {
 		return nil, err
 	}
 	actual := protocol.ServerHandler(server, nil)
-	return jsonrpc2.AsyncHandler(func(ctx context.Context, reply jsonrpc2.Replier, req jsonrpc2.Request) error {
+	// [protocol.CancelHandler] intercepts $/cancelRequest notifications from the client and
+	// cancels the context of the matching in-flight request. It must wrap AsyncHandler so
+	// that the cancellable context is the one running inside each spawned goroutine.
+	return protocol.CancelHandler(jsonrpc2.AsyncHandler(func(ctx context.Context, reply jsonrpc2.Replier, req jsonrpc2.Request) error {
 		l.logger.Debug(
 			"handling request",
 			slog.String("method", req.Method()),
@@ -181,5 +199,5 @@ func (l *lsp) newHandler() (jsonrpc2.Handler, error) {
 			)
 		}
 		return nil
-	}), nil
+	})), nil
 }

private/buf/buflsp/file.go

@@ -1062,7 +1062,7 @@ func (f *file) RunChecks(ctx context.Context) {
 	}
 
 	const checkTimeout = 30 * time.Second
-	ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), checkTimeout)
+	ctx, cancel := context.WithTimeout(f.lsp.connCtx, checkTimeout)
 	f.cancelChecks = cancel
 
 	go func() {
@@ -1079,7 +1079,7 @@ func (f *file) RunChecks(ctx context.Context) {
 			); err != nil {
 				var fileAnnotationSet bufanalysis.FileAnnotationSet
 				if !errors.As(err, &fileAnnotationSet) {
-					if errors.Is(err, context.Canceled) {
+					if errors.Is(err, context.Canceled) || ctx.Err() != nil {
 						f.lsp.logger.DebugContext(ctx, "checks cancelled", slog.String("uri", f.uri.Filename()), xslog.ErrorAttr(err))
 					} else if errors.Is(err, context.DeadlineExceeded) {
 						f.lsp.logger.WarnContext(ctx, "checks deadline exceeded", slog.String("uri", f.uri.Filename()), xslog.ErrorAttr(err))

private/buf/buflsp/image.go

@@ -85,7 +85,9 @@ func buildImage(
 	var diagnostics []protocol.Diagnostic
 	compiled, err := compiler.Compile(ctx, path)
 	if err != nil {
-		logger.Warn("error building image", slog.String("path", path), xslog.ErrorAttr(err))
+		if !errors.Is(err, context.Canceled) {
+			logger.WarnContext(ctx, "error building image", slog.String("path", path), xslog.ErrorAttr(err))
+		}
 		var errorWithPos reporter.ErrorWithPos
 		if errors.As(err, &errorWithPos) {
 			diagnostics = []protocol.Diagnostic{newDiagnostic(errorWithPos, false, opener, logger)}

private/buf/buflsp/server.go

@@ -193,6 +193,12 @@ func (s *server) SetTrace(
 // The client will wait until Shutdown returns, and then call Exit.
 func (s *server) Shutdown(ctx context.Context) error {
 	s.lsp.shutdown = true
+	// Cancel any in-progress checks for all tracked files, then cancel the
+	// connection context to stop any remaining background goroutines.
+	for _, file := range s.lsp.fileManager.uriToFile.Range {
+		file.CancelChecks(ctx)
+	}
+	s.lsp.connCancel()
 	return nil
 }