diff --git a/core/streaming/proxy/streaming.go b/core/streaming/proxy/streaming.go index 6bd91fbb04..9c582c4861 100644 --- a/core/streaming/proxy/streaming.go +++ b/core/streaming/proxy/streaming.go @@ -21,6 +21,7 @@ import ( "errors" "fmt" "io" + "sync" streamingapi "github.com/containerd/containerd/api/services/streaming/v1" "github.com/containerd/errdefs" @@ -105,10 +106,13 @@ func (sc *streamCreator) Create(ctx context.Context, id string) (streaming.Strea } type clientStream struct { - s streamingapi.TTRPCStreaming_StreamClient + sendMu sync.Mutex + s streamingapi.TTRPCStreaming_StreamClient } func (cs *clientStream) Send(a typeurl.Any) (err error) { + cs.sendMu.Lock() + defer cs.sendMu.Unlock() err = cs.s.Send(typeurl.MarshalProto(a)) if !errors.Is(err, io.EOF) { err = errgrpc.ToNative(err) @@ -125,5 +129,7 @@ func (cs *clientStream) Recv() (a typeurl.Any, err error) { } func (cs *clientStream) Close() error { + cs.sendMu.Lock() + defer cs.sendMu.Unlock() return cs.s.CloseSend() } diff --git a/core/streaming/proxy/streaming_test.go b/core/streaming/proxy/streaming_test.go new file mode 100644 index 0000000000..6a2dec4eff --- /dev/null +++ b/core/streaming/proxy/streaming_test.go @@ -0,0 +1,61 @@ +/* + Copyright The containerd Authors. + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. +*/ + +package proxy + +import ( + "io" + "testing" + + "github.com/containerd/typeurl/v2" + "google.golang.org/protobuf/types/known/anypb" +) + +// unsafeStreamClient writes a flag in CloseSend that Send reads, as ttrpc does. +type unsafeStreamClient struct { + closed bool +} + +func (c *unsafeStreamClient) Send(*anypb.Any) error { + if c.closed { + return io.EOF + } + return nil +} + +func (c *unsafeStreamClient) Recv() (*anypb.Any, error) { return &anypb.Any{}, nil } +func (c *unsafeStreamClient) CloseSend() error { c.closed = true; return nil } +func (c *unsafeStreamClient) SendMsg(m any) error { return nil } +func (c *unsafeStreamClient) RecvMsg(m any) error { return nil } + +// TestClientStreamSendClose races under -race unless Send and Close are serialized. +func TestClientStreamSendClose(t *testing.T) { + cs := &clientStream{s: &unsafeStreamClient{}} + var a typeurl.Any = &anypb.Any{} + + done := make(chan struct{}) + go func() { + defer close(done) + for range 100 { + cs.Send(a) + } + }() + + if err := cs.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + <-done +} diff --git a/core/transfer/streaming/reader.go b/core/transfer/streaming/reader.go index e989b0d732..312c34d762 100644 --- a/core/transfer/streaming/reader.go +++ b/core/transfer/streaming/reader.go @@ -18,9 +18,9 @@ package streaming import ( "context" - "errors" "fmt" "io" + "sync/atomic" transferapi "github.com/containerd/containerd/api/types/transfer" "github.com/containerd/containerd/v2/core/streaming" @@ -30,7 +30,7 @@ import ( type readByteStream struct { ctx context.Context stream streaming.Stream - window int32 + window atomic.Int32 updated chan struct{} errCh chan error remaining []byte @@ -40,13 +40,12 @@ func ReadByteStream(ctx context.Context, stream streaming.Stream) io.ReadCloser rbs := &readByteStream{ ctx: ctx, stream: stream, - window: 0, - errCh: make(chan error), + errCh: make(chan error, 1), updated: make(chan struct{}, 1), } go func() { for { - if rbs.window >= windowSize { + if rbs.window.Load() >= windowSize { select { case <-ctx.Done(): return @@ -62,11 +61,11 @@ func ReadByteStream(ctx context.Context, stream streaming.Stream) io.ReadCloser rbs.errCh <- err return } - if err := stream.Send(anyType); err == nil { - rbs.window += windowSize - } else if !errors.Is(err, io.EOF) { - rbs.errCh <- err + if err := stream.Send(anyType); err != nil { + // Read reports stream errors. + return } + rbs.window.Add(windowSize) } }() @@ -105,9 +104,12 @@ func (r *readByteStream) Read(p []byte) (n int, err error) { if len(v.Data) > plen { r.remaining = v.Data[plen:] } - r.window = r.window - int32(n) - if r.window < windowSize { - r.updated <- struct{}{} + if r.window.Add(-int32(len(v.Data))) < windowSize { + select { + case r.updated <- struct{}{}: + default: + // Don't block if the window goroutine has ended. + } } return n, nil default: diff --git a/core/transfer/streaming/reader_test.go b/core/transfer/streaming/reader_test.go new file mode 100644 index 0000000000..68ac1cfac3 --- /dev/null +++ b/core/transfer/streaming/reader_test.go @@ -0,0 +1,45 @@ +/* + Copyright The containerd Authors. + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. +*/ + +package streaming + +import ( + "bytes" + "io" + "testing" +) + +// TestReadByteStream races under -race unless the window counter is atomic. +func TestReadByteStream(t *testing.T) { + expected := bytes.Repeat([]byte("containerd"), windowSize/4) + + rs, ws := pipeStream() + w := WriteByteStream(t.Context(), ws) + r := ReadByteStream(t.Context(), rs) + + go func() { + io.Copy(w, bytes.NewReader(expected)) + w.Close() + }() + + actual, err := io.ReadAll(r) + if err != nil { + t.Fatal(err) + } + if !bytes.Equal(expected, actual) { + t.Fatalf("read %d bytes, want %d", len(actual), len(expected)) + } +}