Skip to content

feat: add bounded ZSTD encoder/decoder pool with memory management #327

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

Merged
EdSchouten merged 18 commits into from
Mar 12, 2026
Merged

feat: add bounded ZSTD encoder/decoder pool with memory management #327

Show file tree
Hide file tree
Changes from 15 commits
Commits
Show all changes
18 commits
Select commit Hold shift + click to select a range
d171a61
Add bounded ZSTD encoder/decoder pool with memory management
aron-muon Jan 5, 2026
4c0ebeb
fix: remove duplicate package declaration in test file
aron-muon Jan 5, 2026
bb3536b
fix: added test
aron-muon Jan 5, 2026
53f9407
Integrate BoundedZstdPool into CAS blob access
aron-muon Jan 5, 2026
92b8cd9
ci: add better tests for literal testing of releasing pools
aron-muon Feb 27, 2026
399afb4
fix: address review findings in ZSTD pool implementation
aron-muon Mar 2, 2026
e3e08bd
refactor: address maintainer review feedback
aron-muon Mar 2, 2026
4d90aed
Merge pull request #1 from aron-muon/zstd-memory-management
aron-muon Mar 2, 2026
99f0de4
ci: fix lint
aron-muon Mar 2, 2026
2003438
Merge branch 'main' into zstd-memory-management
aron-muon Mar 3, 2026
ee28769
feat: make ZSTD pool parameters configurable via protobuf
aron-muon Mar 5, 2026
46ed930
refactor: move ZSTD pool config to top-level ApplicationConfiguration
aron-muon Mar 10, 2026
582533f
refactor: create pkg/zstd package with Pool/Encoder/Decoder interfaces
aron-muon Mar 10, 2026
bdf8749
refactor: address review feedback on zstd pool interfaces
aron-muon Mar 11, 2026
cf9d61a
feat: add MetricsPool wrapper for ZSTD pool observability
aron-muon Mar 11, 2026
3811e52
refactor: address review feedback on zstd pool interfaces
aron-muon Mar 11, 2026
d930e97
fix: gate CAS client compression on per-backend enable_compression flag
aron-muon Mar 11, 2026
cb55ac7
chore: fix lint
aron-muon Mar 11, 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
1 change: 1 addition & 0 deletions cmd/bb_copy/BUILD.bazel
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ go_library(
"//pkg/program",
"//pkg/proto/configuration/bb_copy",
"//pkg/util",
"//pkg/zstd",
"@org_golang_google_grpc//codes",
"@org_golang_google_grpc//status",
],
Expand Down
4 changes: 3 additions & 1 deletion cmd/bb_copy/main.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ import (
"github.com/buildbarn/bb-storage/pkg/program"
"github.com/buildbarn/bb-storage/pkg/proto/configuration/bb_copy"
"github.com/buildbarn/bb-storage/pkg/util"
bb_zstd "github.com/buildbarn/bb-storage/pkg/zstd"

"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
Expand Down Expand Up @@ -43,7 +44,8 @@ func main() {

blobAccessCreator := blobstore_configuration.NewCASBlobAccessCreator(
grpcClientFactory,
int(configuration.MaximumMessageSizeBytes))
int(configuration.MaximumMessageSizeBytes),
bb_zstd.NewPoolFromConfiguration(nil))
source, err := blobstore_configuration.NewBlobAccessFromConfiguration(
dependenciesGroup,
configuration.Source,
Expand Down
1 change: 1 addition & 0 deletions cmd/bb_replicator/BUILD.bazel
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ go_library(
"//pkg/proto/configuration/bb_replicator",
"//pkg/proto/replicator",
"//pkg/util",
"//pkg/zstd",
"@org_golang_google_grpc//:grpc",
"@org_golang_google_grpc//codes",
"@org_golang_google_grpc//status",
Expand Down
4 changes: 3 additions & 1 deletion cmd/bb_replicator/main.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ import (
"github.com/buildbarn/bb-storage/pkg/proto/configuration/bb_replicator"
replicator_pb "github.com/buildbarn/bb-storage/pkg/proto/replicator"
"github.com/buildbarn/bb-storage/pkg/util"
bb_zstd "github.com/buildbarn/bb-storage/pkg/zstd"

"google.golang.org/grpc"
"google.golang.org/grpc/codes"
Expand All @@ -34,7 +35,8 @@ func main() {

blobAccessCreator := blobstore_configuration.NewCASBlobAccessCreator(
grpcClientFactory,
int(configuration.MaximumMessageSizeBytes))
int(configuration.MaximumMessageSizeBytes),
bb_zstd.NewPoolFromConfiguration(nil))
source, err := blobstore_configuration.NewBlobAccessFromConfiguration(
dependenciesGroup,
configuration.Source,
Expand Down
1 change: 1 addition & 0 deletions cmd/bb_storage/BUILD.bazel
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ go_library(
"//pkg/proto/icas",
"//pkg/proto/iscc",
"//pkg/util",
"//pkg/zstd",
"@bazel_remote_apis//build/bazel/remote/execution/v2:remote_execution_go_proto",
"@bazel_remote_apis//build/bazel/semver:semver_go_proto",
"@org_golang_google_genproto_googleapis_bytestream//:bytestream",
Expand Down
10 changes: 8 additions & 2 deletions cmd/bb_storage/main.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ import (
"github.com/buildbarn/bb-storage/pkg/proto/icas"
"github.com/buildbarn/bb-storage/pkg/proto/iscc"
"github.com/buildbarn/bb-storage/pkg/util"
bb_zstd "github.com/buildbarn/bb-storage/pkg/zstd"

"google.golang.org/genproto/googleapis/bytestream"
"google.golang.org/grpc"
Expand All @@ -42,6 +43,9 @@ func main() {
return util.StatusWrap(err, "Failed to apply global configuration options")
}

// Create a process-wide ZSTD compression pool.
zstdPool := bb_zstd.NewPoolFromConfiguration(configuration.ZstdCompression)

// Providers for data returned by ServerCapabilities.cache_capabilities
// as part of the GetCapabilities() call. We permit these calls
// if the client is permitted to at least one method against one
Expand All @@ -58,7 +62,8 @@ func main() {
configuration.ContentAddressableStorage,
blobstore_configuration.NewCASBlobAccessCreator(
grpcClientFactory,
int(configuration.MaximumMessageSizeBytes)),
int(configuration.MaximumMessageSizeBytes),
zstdPool),
grpcClientFactory)
if err != nil {
return util.StatusWrap(err, "Failed to create Content Addressable Storage")
Expand Down Expand Up @@ -185,7 +190,8 @@ func main() {
s,
grpcservers.NewByteStreamServer(
contentAddressableStorage,
1<<16))
1<<16,
zstdPool))
}
if actionCache != nil {
remoteexecution.RegisterActionCacheServer(
Expand Down
4 changes: 3 additions & 1 deletion pkg/blobstore/BUILD.bazel
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,10 +43,10 @@ go_library(
"//pkg/proto/icas",
"//pkg/proto/iscc",
"//pkg/util",
"//pkg/zstd",
"@bazel_remote_apis//build/bazel/remote/execution/v2:remote_execution_go_proto",
"@com_github_aws_aws_sdk_go_v2//aws",
"@com_github_aws_aws_sdk_go_v2_service_s3//:s3",
"@com_github_klauspost_compress//zstd",
"@com_github_prometheus_client_golang//prometheus",
"@org_golang_google_grpc//codes",
"@org_golang_google_grpc//status",
Expand Down Expand Up @@ -82,11 +82,13 @@ go_test(
"//pkg/proto/icas",
"//pkg/testutil",
"//pkg/util",
"//pkg/zstd",
"@bazel_remote_apis//build/bazel/remote/execution/v2:remote_execution_go_proto",
"@bazel_remote_apis//build/bazel/semver:semver_go_proto",
"@com_github_aws_aws_sdk_go_v2//aws",
"@com_github_aws_aws_sdk_go_v2_service_s3//:s3",
"@com_github_aws_aws_sdk_go_v2_service_s3//types",
"@com_github_klauspost_compress//zstd",
"@com_github_stretchr_testify//require",
"@org_golang_google_grpc//codes",
"@org_golang_google_grpc//status",
Expand Down
1 change: 1 addition & 0 deletions pkg/blobstore/configuration/BUILD.bazel
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@ go_library(
"//pkg/proto/configuration/digest",
"//pkg/random",
"//pkg/util",
"//pkg/zstd",
"@bazel_remote_apis//build/bazel/remote/execution/v2:remote_execution_go_proto",
"@com_github_aws_aws_sdk_go_v2_service_s3//:s3",
"@com_github_fxtlabs_primes//:primes",
Expand Down
10 changes: 7 additions & 3 deletions pkg/blobstore/configuration/cas_blob_access_creator.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import (
"github.com/buildbarn/bb-storage/pkg/blobstore"
"github.com/buildbarn/bb-storage/pkg/blobstore/grpcclients"
"github.com/buildbarn/bb-storage/pkg/blobstore/local"
bb_zstd "github.com/buildbarn/bb-storage/pkg/zstd"
"github.com/buildbarn/bb-storage/pkg/capabilities"
"github.com/buildbarn/bb-storage/pkg/cloud/aws"
"github.com/buildbarn/bb-storage/pkg/cloud/gcp"
Expand Down Expand Up @@ -38,18 +39,20 @@ type casBlobAccessCreator struct {
casBlobReplicatorCreator

maximumMessageSizeBytes int
zstdPool bb_zstd.Pool
}

// NewCASBlobAccessCreator creates a BlobAccessCreator that can be
// provided to NewBlobAccessFromConfiguration() to construct a
// BlobAccess that is suitable for accessing the Content Addressable
// Storage.
func NewCASBlobAccessCreator(grpcClientFactory grpc.ClientFactory, maximumMessageSizeBytes int) BlobAccessCreator {
func NewCASBlobAccessCreator(grpcClientFactory grpc.ClientFactory, maximumMessageSizeBytes int, zstdPool bb_zstd.Pool) BlobAccessCreator {
return &casBlobAccessCreator{
casBlobReplicatorCreator: casBlobReplicatorCreator{
grpcClientFactory: grpcClientFactory,
},
maximumMessageSizeBytes: maximumMessageSizeBytes,
zstdPool: zstdPool,
}
}

Expand Down Expand Up @@ -96,7 +99,7 @@ func (bac *casBlobAccessCreator) NewCustomBlobAccess(terminationGroup program.Gr
// TODO: Should we provide a configuration option, so
// that digest.KeyWithoutInstance can be used?
return BlobAccessInfo{
BlobAccess: grpcclients.NewCASBlobAccess(client, uuid.NewRandom, 64<<10, backend.Grpc.EnableCompression),
BlobAccess: grpcclients.NewCASBlobAccess(client, uuid.NewRandom, 64<<10, backend.Grpc.EnableZstdCompression, bac.zstdPool),
DigestKeyFormat: digest.KeyWithInstance,
}, "grpc", nil
case *pb.BlobAccessConfiguration_ReferenceExpanding:
Expand Down Expand Up @@ -158,7 +161,8 @@ func (bac *casBlobAccessCreator) NewCustomBlobAccess(terminationGroup program.Gr
},
s3.NewFromConfig(awsConfig),
gcsClient,
bac.maximumMessageSizeBytes),
bac.maximumMessageSizeBytes,
bac.zstdPool),
DigestKeyFormat: indirectContentAddressableStorage.DigestKeyFormat,
}, "reference_expanding", nil
default:
Expand Down
3 changes: 2 additions & 1 deletion pkg/blobstore/configuration/new_blob_access.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ import (
digest_pb "github.com/buildbarn/bb-storage/pkg/proto/configuration/digest"
"github.com/buildbarn/bb-storage/pkg/random"
"github.com/buildbarn/bb-storage/pkg/util"
bb_zstd "github.com/buildbarn/bb-storage/pkg/zstd"
"github.com/fxtlabs/primes"

"google.golang.org/grpc/codes"
Expand Down Expand Up @@ -606,7 +607,7 @@ func NewCASAndACBlobAccessFromConfiguration(terminationGroup program.Group, conf
contentAddressableStorage, err := NewBlobAccessFromConfiguration(
terminationGroup,
configuration.GetContentAddressableStorage(),
NewCASBlobAccessCreator(grpcClientFactory, maximumMessageSizeBytes))
NewCASBlobAccessCreator(grpcClientFactory, maximumMessageSizeBytes, bb_zstd.NewPoolFromConfiguration(nil)))
if err != nil {
return nil, nil, util.StatusWrap(err, "Failed to create Content Addressable Storage")
}
Expand Down
7 changes: 5 additions & 2 deletions pkg/blobstore/grpcclients/BUILD.bazel
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,9 +20,9 @@ go_library(
"//pkg/proto/icas",
"//pkg/proto/iscc",
"//pkg/util",
"//pkg/zstd",
"@bazel_remote_apis//build/bazel/remote/execution/v2:remote_execution_go_proto",
"@com_github_google_uuid//:uuid",
"@com_github_klauspost_compress//zstd",
"@org_golang_google_genproto_googleapis_bytestream//:bytestream",
"@org_golang_google_grpc//:grpc",
"@org_golang_google_grpc//codes",
Expand All @@ -33,14 +33,17 @@ go_library(

go_test(
name = "grpcclients_test",
srcs = ["cas_blob_access_test.go"],
srcs = [
"cas_blob_access_test.go",
],
deps = [
":grpcclients",
"//internal/mock",
"//pkg/blobstore/buffer",
"//pkg/digest",
"//pkg/testutil",
"//pkg/util",
"//pkg/zstd",
"@bazel_remote_apis//build/bazel/remote/execution/v2:remote_execution_go_proto",
"@bazel_remote_apis//build/bazel/semver:semver_go_proto",
"@com_github_google_uuid//:uuid",
Expand Down
Loading
Loading