Skip to content

Commit fd2c59e

Browse files
authored
fix(cas): handle client cancellation (#378)
Signed-off-by: Miguel Martinez Trivino <miguel@chainloop.dev>
1 parent af920e3 commit fd2c59e

File tree

1 file changed

+12
-5
lines changed

1 file changed

+12
-5
lines changed

app/artifact-cas/internal/service/bytestream.go

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,8 @@ import (
3434
kerrors "github.com/go-kratos/kratos/v2/errors"
3535
"github.com/go-kratos/kratos/v2/log"
3636
"google.golang.org/genproto/googleapis/bytestream"
37+
"google.golang.org/grpc/codes"
38+
"google.golang.org/grpc/status"
3739
)
3840

3941
// Implements the bytestream interface
@@ -109,7 +111,7 @@ func (s *ByteStreamService) Write(stream bytestream.ByteStream_WriteServer) erro
109111

110112
// Now it's time to check if the data provider has sent an error
111113
if err != nil {
112-
if errors.Is(err, context.Canceled) {
114+
if errors.Is(err, context.Canceled) || status.Code(err) == codes.Canceled {
113115
s.log.Infow("msg", "upload canceled", "digest", req.resource.Digest, "name", req.resource.FileName)
114116
return nil
115117
}
@@ -192,11 +194,16 @@ func bufferStream(ctx context.Context, stream bytestream.ByteStream_WriteServer,
192194
default:
193195
// Extract the next chunk of data from the stream request
194196
req, err := getWriteRequest(stream)
195-
// There is nothing else to read or the client has marked the upload as finished
196-
if errors.Is(err, io.EOF) || req.GetFinishWrite() {
197+
if err != nil {
198+
// If we have finished reading the stream we don't consider it a real error
199+
if !errors.Is(err, io.EOF) {
200+
bufferErr = err
201+
}
197202
return
198-
} else if err != nil { // If there is any other error we return it
199-
bufferErr = err
203+
}
204+
205+
// Check if the client has finished sending data
206+
if req.GetFinishWrite() {
200207
return
201208
}
202209

0 commit comments

Comments
 (0)