Skip to content

Add ZSTD ReadOffset support #362

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
snakethatlovesstaticlibs wants to merge 2 commits into
base: main
Choose a base branch
from
+102 −4
Open

Add ZSTD ReadOffset support #362

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
29a8df5
Add ZSTD `ReadOffset` support
snakethatlovesstaticlibs Aug 20, 2026
65a0cb0
Don't test old behavior
snakethatlovesstaticlibs Aug 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 24 additions & 4 deletions pkg/blobstore/grpcservers/byte_stream_server.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,13 +35,19 @@ func NewByteStreamServer(blobAccess blobstore.BlobAccess, readChunkSize int, zst
}

func (s *byteStreamServer) Read(in *bytestream.ReadRequest, out bytestream.ByteStream_ReadServer) error {
if in.ReadLimit != 0 {

@EdSchouten EdSchouten Aug 22, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What's up with this change? I don't think we have any code for respecting this, even for the uncompressed case.

@snakethatlovesstaticlibs snakethatlovesstaticlibs Aug 24, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This change leaves existing behavior as is for the uncompressed case, but changes the return value to InvalidArgument for the compressed case, because the spec wants it that way

// When downloading compressed blobs:
// * `ReadRequest.read_offset` refers to the offset in the uncompressed form
//   of the blob.
// * Servers MUST return `INVALID_ARGUMENT` if `ReadRequest.read_limit` is
//   non-zero.

from https://github.com/bazelbuild/remote-apis/blob/becdd8f9ff811df88a22d3eadd6341753d51d167/build/bazel/remote/execution/v2/remote_execution.proto#L313-L317

It probably doesn't matter since clients likely won't send a read limit for the compressed case, but I figured I'd fix it anyway. Please let me know if you want it reverted and I can do that too

return status.Error(codes.Unimplemented, "This service does not support downloading partial files")
}
digest, compressor, err := digest.NewDigestFromByteStreamReadPath(in.ResourceName)
if err != nil {
return err
}
if in.ReadLimit != 0 {
if compressor != remoteexecution.Compressor_IDENTITY {
// REAPI requires non-zero read limits on compressed ByteStream
// reads to be rejected with INVALID_ARGUMENT.
// https://github.com/bazelbuild/remote-apis/blob/becdd8f9ff811df88a22d3eadd6341753d51d167/build/bazel/remote/execution/v2/remote_execution.proto#L313-L317
return status.Error(codes.InvalidArgument, "Read limits are not permitted for compressed blobs")
}
return status.Error(codes.Unimplemented, "This service does not support downloading partial files")
}
ctx := out.Context()
switch compressor {
case remoteexecution.Compressor_IDENTITY:
Expand Down Expand Up @@ -69,7 +75,21 @@ func (s *byteStreamServer) Read(in *bytestream.ReadRequest, out bytestream.ByteS
return status.Errorf(codes.ResourceExhausted, "Failed to acquire ZSTD encoder: %v", err)
}
defer encoder.Close()
return b.IntoWriter(encoder)

r := b.ToChunkReader(in.ReadOffset, s.readChunkSize)
defer r.Close()
for {
chunk, err := r.Read()
if err == io.EOF {
return nil
}
if err != nil {
return err
}
if _, err := encoder.Write(chunk); err != nil {
return err
}
}
default:
return status.Errorf(codes.Unimplemented, "This service does not support downloading compression type: %s", compressor)
}
Expand Down
78 changes: 78 additions & 0 deletions pkg/blobstore/grpcservers/byte_stream_server_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,32 @@ import (
"go.uber.org/mock/gomock"
)

func readAndDecompressZSTD(t *testing.T, ctx context.Context, client bytestream.ByteStreamClient, resourceName string, readOffset int64) []byte {
t.Helper()
req, err := client.Read(ctx, &bytestream.ReadRequest{
ResourceName: resourceName,
ReadOffset: readOffset,
})
require.NoError(t, err)

var compressedData []byte
for {
response, err := req.Recv()
if err == io.EOF {
break
}
require.NoError(t, err)
compressedData = append(compressedData, response.Data...)
}

decoder, err := zstd.NewReader(nil)
require.NoError(t, err)
decompressedData, err := decoder.DecodeAll(compressedData, nil)
decoder.Close()
require.NoError(t, err)
return decompressedData
}

func TestByteStreamServer(t *testing.T) {
ctrl, ctx := gomock.WithContext(context.Background(), t)

Expand Down Expand Up @@ -202,6 +228,58 @@ func TestByteStreamServer(t *testing.T) {
require.Less(t, len(compressedData), len(originalData))
})

t.Run("ReadZSTDCompressionWithOffset", func(t *testing.T) {
originalData := []byte("This is a test message that should be compressed with ZSTD")
blobAccess.EXPECT().Get(
gomock.Any(),
digest.MustNewDigest("", remoteexecution.DigestFunction_SHA256, "8b2c3f8a9d0e1f2a3b4c5d6e7f8a9b0c1d2e3f4a5b6c7d8e9f0a1b2c3d4e5f61", 58),
).Return(buffer.NewValidatedBufferFromByteSlice(originalData))

decompressedData := readAndDecompressZSTD(t, ctx, client, "compressed-blobs/zstd/8b2c3f8a9d0e1f2a3b4c5d6e7f8a9b0c1d2e3f4a5b6c7d8e9f0a1b2c3d4e5f61/58", 17)
require.Equal(t, originalData[17:], decompressedData)
})

t.Run("ReadZSTDCompressionWithReadLimit", func(t *testing.T) {
req, err := client.Read(ctx, &bytestream.ReadRequest{
ResourceName: "compressed-blobs/zstd/8b2c3f8a9d0e1f2a3b4c5d6e7f8a9b0c1d2e3f4a5b6c7d8e9f0a1b2c3d4e5f61/58",
ReadLimit: 10,
})
require.NoError(t, err)
_, err = req.Recv()
testutil.RequireEqualStatus(t, status.Error(codes.InvalidArgument, "Read limits are not permitted for compressed blobs"), err)
})

t.Run("ReadZSTDCompressionResumedMultipleTimes", func(t *testing.T) {
originalData := []byte("This is a test message that should survive multiple interrupted and resumed downloads")
digestFunction := digest.MustNewFunction("", remoteexecution.DigestFunction_SHA256)
generator := digestFunction.NewGenerator(int64(len(originalData)))
_, err := generator.Write(originalData)
require.NoError(t, err)
blobDigest := generator.Sum()
resourceName := fmt.Sprintf("compressed-blobs/zstd/%s/%d", blobDigest.GetHashString(), len(originalData))
offsets := []int64{0, 13, 47, int64(len(originalData))}

var downloadedData []byte
for i := 0; i < len(offsets)-1; i++ {
blobAccess.EXPECT().Get(
gomock.Any(),
blobDigest,
).Return(buffer.NewValidatedBufferFromByteSlice(originalData))

decompressedData := readAndDecompressZSTD(t, ctx, client, resourceName, offsets[i])

bytesConsumedBeforeInterruption := offsets[i+1] - offsets[i]
require.GreaterOrEqual(t, int64(len(decompressedData)), bytesConsumedBeforeInterruption)
downloadedData = append(downloadedData, decompressedData[:bytesConsumedBeforeInterruption]...)
}
require.Len(t, downloadedData, len(originalData))
downloadedDataGenerator := digestFunction.NewGenerator(int64(len(downloadedData)))
_, err = downloadedDataGenerator.Write(downloadedData)
require.NoError(t, err)
require.Equal(t, blobDigest, downloadedDataGenerator.Sum())
require.Equal(t, originalData, downloadedData)
})

t.Run("ReadUnsupportedCompression", func(t *testing.T) {
// Test reading with unsupported compression type.
req, err := client.Read(ctx, &bytestream.ReadRequest{
Expand Down