blob: removes race condition between ctx cancellation and stream close (#167489)

trunk-io[bot] · web-flow · commit 5bb249bae668 · 2026-04-06T20:29:25.000Z
blob: removes race condition between ctx cancellation and stream close
diff --git a/pkg/blobs/client.go b/pkg/blobs/client.go
@@ -124,11 +124,15 @@ func (c *remoteClient) readFileWithFlowControl(
 }
 
 type streamWriter struct {
+	ctx context.Context
 	s   blobspb.RPCBlob_PutStreamClient
 	buf blobspb.StreamChunk
 }
 
 func (w *streamWriter) Write(p []byte) (int, error) {
+	if err := w.ctx.Err(); err != nil {
+		return 0, err
+	}
 	n := 0
 	for len(p) > 0 {
 		l := copy(w.buf.Payload[:cap(w.buf.Payload)], p)
@@ -145,6 +149,9 @@ func (w *streamWriter) Write(p []byte) (int, error) {
 }
 
 func (w *streamWriter) Close() error {
+	if err := w.ctx.Err(); err != nil {
+		return err
+	}
 	_, err := w.s.CloseAndRecv()
 	return err
 }
@@ -156,7 +163,7 @@ func (c *remoteClient) Writer(ctx context.Context, file string) (io.WriteCloser,
 		return nil, err
 	}
 	buf := make([]byte, 0, ChunkSize)
-	return &streamWriter{s: stream, buf: blobspb.StreamChunk{Payload: buf}}, nil
+	return &streamWriter{ctx: ctx, s: stream, buf: blobspb.StreamChunk{Payload: buf}}, nil
 }
 
 func (c *remoteClient) List(ctx context.Context, pattern string) ([]string, error) {

Original file line number	Diff line number	Diff line change
`@@ -124,11 +124,15 @@ func (c *remoteClient) readFileWithFlowControl(`
`124`	`124`	`}`
`125`	`125`
`126`	`126`	`type streamWriter struct {`
	`127`	`+ ctx context.Context`
`127`	`128`	`s blobspb.RPCBlob_PutStreamClient`
`128`	`129`	`buf blobspb.StreamChunk`
`129`	`130`	`}`
`130`	`131`
`131`	`132`	`func (w *streamWriter) Write(p []byte) (int, error) {`
	`133`	`+ if err := w.ctx.Err(); err != nil {`
	`134`	`+ return 0, err`
	`135`	`+ }`
`132`	`136`	`n := 0`
`133`	`137`	`for len(p) > 0 {`
`134`	`138`	`l := copy(w.buf.Payload[:cap(w.buf.Payload)], p)`
`@@ -145,6 +149,9 @@ func (w *streamWriter) Write(p []byte) (int, error) {`
`145`	`149`	`}`
`146`	`150`
`147`	`151`	`func (w *streamWriter) Close() error {`
	`152`	`+ if err := w.ctx.Err(); err != nil {`
	`153`	`+ return err`
	`154`	`+ }`
`148`	`155`	`_, err := w.s.CloseAndRecv()`
`149`	`156`	`return err`
`150`	`157`	`}`
`@@ -156,7 +163,7 @@ func (c *remoteClient) Writer(ctx context.Context, file string) (io.WriteCloser,`
`156`	`163`	`return nil, err`
`157`	`164`	`}`
`158`	`165`	`buf := make([]byte, 0, ChunkSize)`
`159`		`- return &streamWriter{s: stream, buf: blobspb.StreamChunk{Payload: buf}}, nil`
	`166`	`+ return &streamWriter{ctx: ctx, s: stream, buf: blobspb.StreamChunk{Payload: buf}}, nil`
`160`	`167`	`}`
`161`	`168`
`162`	`169`	`func (c *remoteClient) List(ctx context.Context, pattern string) ([]string, error) {`