Skip to content

s3: add blob fetch rpc #1063

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
bootjp merged 1 commit into from
Jul 19, 2026
Merged

s3: add blob fetch rpc #1063

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
26 changes: 23 additions & 3 deletions adapter/admin_grpc.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -478,11 +478,31 @@ func sortedGroupIDs(m map[uint64]AdminGroup) []uint64 {
// package-qualify the service name) does not silently bypass the auth gate.
var adminMethodPrefix = "/" + pb.Admin_ServiceDesc.ServiceName + "/"

func adminTokenProtectedMethod(fullMethod string) bool {
return strings.HasPrefix(fullMethod, adminMethodPrefix)
}

// AdminTokenAuth builds a gRPC unary+stream interceptor pair enforcing
// "authorization: Bearer <token>" metadata against the supplied token. An
// empty token disables enforcement; callers should pair that mode with a
// --adminInsecureNoAuth flag so operators knowingly opt in.
func AdminTokenAuth(token string) (grpc.UnaryServerInterceptor, grpc.StreamServerInterceptor) {
return bearerTokenAuth(token, adminTokenProtectedMethod, "admin")
}

// S3BlobPeerTokenAuth protects the peer-only blob transport with a credential
// distinct from the read-only Admin token.
func S3BlobPeerTokenAuth(token string) (grpc.UnaryServerInterceptor, grpc.StreamServerInterceptor) {
return bearerTokenAuth(token, func(fullMethod string) bool {
return strings.HasPrefix(fullMethod, "/"+pb.S3BlobFetch_ServiceDesc.ServiceName+"/")
}, "S3 blob peer")
}

func bearerTokenAuth(
token string,
protected func(string) bool,
credentialName string,
) (grpc.UnaryServerInterceptor, grpc.StreamServerInterceptor) {
if token == "" {
return nil, nil
}
Expand All @@ -501,7 +521,7 @@ func AdminTokenAuth(token string) (grpc.UnaryServerInterceptor, grpc.StreamServe
return status.Error(codes.Unauthenticated, "authorization is not a bearer token")
}
if subtle.ConstantTimeCompare([]byte(got), expected) != 1 {
return status.Error(codes.Unauthenticated, "invalid admin token")
return status.Errorf(codes.Unauthenticated, "invalid %s token", credentialName)
}
return nil
}
Expand All @@ -511,7 +531,7 @@ func AdminTokenAuth(token string) (grpc.UnaryServerInterceptor, grpc.StreamServe
info *grpc.UnaryServerInfo,
handler grpc.UnaryHandler,
) (any, error) {
if !strings.HasPrefix(info.FullMethod, adminMethodPrefix) {
if !protected(info.FullMethod) {
return handler(ctx, req)
}
if err := check(ctx); err != nil {
Expand All @@ -525,7 +545,7 @@ func AdminTokenAuth(token string) (grpc.UnaryServerInterceptor, grpc.StreamServe
info *grpc.StreamServerInfo,
handler grpc.StreamHandler,
) error {
if !strings.HasPrefix(info.FullMethod, adminMethodPrefix) {
if !protected(info.FullMethod) {
return handler(srv, ss)
}
if err := check(ss.Context()); err != nil {
Expand Down
53 changes: 53 additions & 0 deletions adapter/admin_grpc_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -597,6 +597,59 @@ func TestAdminTokenAuthSkipsOtherServices(t *testing.T) {
}
}

func TestAdminTokenAuthDoesNotAuthorizeS3BlobFetch(t *testing.T) {
t.Parallel()
unary, stream := AdminTokenAuth("read-only-admin")
handler := func(_ context.Context, _ any) (any, error) { return "ok", nil }
info := &grpc.UnaryServerInfo{FullMethod: pb.S3BlobFetch_FetchChunkBlob_FullMethodName}
if _, err := unary(context.Background(), nil, info, handler); err != nil {
t.Fatalf("admin gate must skip peer method: %v", err)
}

streamInfo := &grpc.StreamServerInfo{FullMethod: pb.S3BlobFetch_PushChunkBlob_FullMethodName}
called := false
err := stream(nil, &adminAuthTestServerStream{ctx: context.Background()}, streamInfo, func(any, grpc.ServerStream) error {
called = true
return nil
})
if err != nil || !called {
t.Fatalf("admin gate must skip peer stream: err = %v, called = %v", err, called)
}
}

func TestS3BlobPeerTokenAuthRejectsAdminToken(t *testing.T) {
t.Parallel()
unary, stream := S3BlobPeerTokenAuth("peer-secret")
handler := func(_ context.Context, _ any) (any, error) { return "ok", nil }
info := &grpc.UnaryServerInfo{FullMethod: pb.S3BlobFetch_FetchChunkBlob_FullMethodName}
adminCtx := metadata.NewIncomingContext(context.Background(), metadata.Pairs("authorization", "Bearer read-only-admin"))
if _, err := unary(adminCtx, nil, info, handler); status.Code(err) != codes.Unauthenticated {
t.Fatalf("admin token code = %v, want %v", status.Code(err), codes.Unauthenticated)
}

peerCtx := metadata.NewIncomingContext(context.Background(), metadata.Pairs("authorization", "Bearer peer-secret"))
streamInfo := &grpc.StreamServerInfo{FullMethod: pb.S3BlobFetch_PushChunkBlob_FullMethodName}
called := false
err := stream(nil, &adminAuthTestServerStream{ctx: peerCtx}, streamInfo, func(any, grpc.ServerStream) error {
called = true
return nil
})
if err != nil || !called {
t.Fatalf("peer-authenticated stream err = %v, called = %v", err, called)
}
}

type adminAuthTestServerStream struct {
ctx context.Context
}

func (s *adminAuthTestServerStream) Context() context.Context { return s.ctx }
func (*adminAuthTestServerStream) SetHeader(metadata.MD) error { return nil }
func (*adminAuthTestServerStream) SendHeader(metadata.MD) error { return nil }
func (*adminAuthTestServerStream) SetTrailer(metadata.MD) {}
func (*adminAuthTestServerStream) SendMsg(any) error { return nil }
func (*adminAuthTestServerStream) RecvMsg(any) error { return nil }

func TestAdminTokenAuthEmptyTokenDisabled(t *testing.T) {
t.Parallel()
unary, stream := AdminTokenAuth("")
Expand Down
99 changes: 73 additions & 26 deletions adapter/s3.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -112,8 +112,13 @@ type S3Server struct {
putAdmission *s3PutAdmission
putAdmissionObserver S3PutAdmissionObserver
blobOffloadEnabled bool
blobOffloadGCReady bool
blobOffloadChecker S3BlobOffloadCapabilityChecker
blobOffloadObserver S3BlobOffloadObserver
blobCluster S3BlobCluster
blobLocalStores S3BlobLocalStoreResolver
blobMinReplicas int
blobPushBlocked func() bool
}

type s3BucketMeta struct {
Expand All @@ -139,12 +144,14 @@ type s3ObjectManifest struct {
}

type s3ObjectPart struct {
PartNo uint64 `json:"part_no"`
ETag string `json:"etag"`
SizeBytes int64 `json:"size_bytes"`
ChunkCount uint64 `json:"chunk_count"`
ChunkSizes []uint64 `json:"chunk_sizes,omitempty"`
PartVersion uint64 `json:"part_version,omitempty"`
PartNo uint64 `json:"part_no"`
ETag string `json:"etag"`
SizeBytes int64 `json:"size_bytes"`
ChunkCount uint64 `json:"chunk_count"`
ChunkSizes []uint64 `json:"chunk_sizes,omitempty"`
PartVersion uint64 `json:"part_version,omitempty"`
ChunkRefVersion uint64 `json:"chunk_ref_version,omitempty"`
Offloaded bool `json:"offloaded,omitempty"`
}

type s3ContinuationToken struct {
Expand Down Expand Up @@ -282,12 +289,14 @@ type s3UploadMeta struct {
}

type s3PartDescriptor struct {
PartNo uint64 `json:"part_no"`
ETag string `json:"etag"`
SizeBytes int64 `json:"size_bytes"`
ChunkCount uint64 `json:"chunk_count"`
ChunkSizes []uint64 `json:"chunk_sizes,omitempty"`
PartVersion uint64 `json:"part_version,omitempty"`
PartNo uint64 `json:"part_no"`
ETag string `json:"etag"`
SizeBytes int64 `json:"size_bytes"`
ChunkCount uint64 `json:"chunk_count"`
ChunkSizes []uint64 `json:"chunk_sizes,omitempty"`
PartVersion uint64 `json:"part_version,omitempty"`
ChunkRefVersion uint64 `json:"chunk_ref_version,omitempty"`
Offloaded bool `json:"offloaded,omitempty"`
}

type s3InitiateMultipartUploadResult struct {
Expand Down Expand Up @@ -347,7 +356,10 @@ func NewS3Server(listen net.Listener, s3Addr string, st store.MVCCStore, coordin
leaderS3: cloneLeaderAddrMap(leaderS3),
cleanupSem: make(chan struct{}, s3ManifestCleanupWorkers),
putAdmission: newS3PutAdmissionFromEnv(),
blobOffloadEnabled: newS3BlobOffloadEnabledFromEnv(),
blobOffloadEnabled: S3BlobOffloadEnabledFromEnv(),
}
if localStores, ok := st.(S3BlobLocalStoreResolver); ok {
s.blobLocalStores = localStores
}
for _, opt := range opts {
if opt != nil {
Expand Down Expand Up @@ -390,6 +402,9 @@ func (s *S3Server) Run() error {
}

func (s *S3Server) Stop() {
if s != nil && s.blobCluster != nil {
_ = s.blobCluster.Close()
}
if s != nil && s.httpServer != nil {
_ = s.httpServer.Shutdown(context.Background())
}
Expand Down Expand Up @@ -881,9 +896,9 @@ func (s *S3Server) putObject(w http.ResponseWriter, r *http.Request, bucket stri
if !s.admitS3PutRequest(w, r, bucket, objectKey, s3MaxObjectSizeBytes, "object exceeds maximum allowed size") {
return
}
s.observeS3BlobOffloadDecision(r.Context())
offload := s.observeS3BlobOffloadDecision(r.Context()).mode == s3BlobOffloadModeOffload
upload, uploadBodyErr, uploadErr := s.uploadS3ObjectData(
r.Context(), r, streamBody, state, bucket, objectKey, expectedPayloadSHA,
r.Context(), r, streamBody, state, bucket, objectKey, expectedPayloadSHA, offload,
)
if uploadErr != nil || uploadBodyErr != nil {
s.cleanupS3PutObjectChunks(r.Context(), state, upload, bucket, objectKey)
Expand Down Expand Up @@ -956,6 +971,12 @@ func (s *S3Server) getObject(w http.ResponseWriter, r *http.Request, bucket stri

// GET without Range: stream the full object.
if rangeHeader == "" {
if s3ManifestHasOffloadedParts(manifest) {
if err := s.ensureS3ObjectRangeLocal(r.Context(), bucket, meta.Generation, objectKey, manifest, readTS, 0, manifest.SizeBytes); err != nil {
writeS3InternalError(w, err)
return
}
}
writeS3ObjectHeaders(w.Header(), manifest)
w.WriteHeader(http.StatusOK)
s.streamObjectChunks(w, r, bucket, meta.Generation, objectKey, manifest, readTS, 0, manifest.SizeBytes)
Expand All @@ -971,6 +992,12 @@ func (s *S3Server) getObject(w http.ResponseWriter, r *http.Request, bucket stri
}

contentLength := rangeEnd - rangeStart + 1
if s3ManifestHasOffloadedParts(manifest) {
if err := s.ensureS3ObjectRangeLocal(r.Context(), bucket, meta.Generation, objectKey, manifest, readTS, rangeStart, contentLength); err != nil {
writeS3InternalError(w, err)
return
}
}
writeS3ObjectHeaders(w.Header(), manifest)
w.Header().Set("Content-Length", strconv.FormatInt(contentLength, 10))
w.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", rangeStart, rangeEnd, manifest.SizeBytes))
Expand Down Expand Up @@ -1006,13 +1033,11 @@ func (s *S3Server) streamObjectChunks(w http.ResponseWriter, r *http.Request, bu
)
return
}
chunkKey := s3keys.VersionedBlobKey(bucket, generation, objectKey, manifest.UploadID, part.PartNo, chunkIndex, part.PartVersion)
chunk, err := s.store.GetAt(r.Context(), chunkKey, readTS)
chunk, err := s.readS3ObjectChunk(r.Context(), bucket, generation, objectKey, manifest.UploadID, part, chunkIndex, chunkSize, readTS)
if err != nil {
slog.ErrorContext(r.Context(), "streamObjectChunks: GetAt failed",
slog.ErrorContext(r.Context(), "streamObjectChunks: read chunk failed",
"bucket", bucket,
"object_key", objectKey,
"chunk_key", string(chunkKey),
"err", err,
)
return
Expand Down Expand Up @@ -1233,9 +1258,9 @@ func (s *S3Server) uploadPart(w http.ResponseWriter, r *http.Request, bucket str
if !s.admitS3PutRequest(w, r, bucket, objectKey, s3MaxPartSizeBytes, "part exceeds maximum allowed size") {
return
}
s.observeS3BlobOffloadDecision(r.Context())
offload := s.observeS3BlobOffloadDecision(r.Context()).mode == s3BlobOffloadModeOffload
upload, previous, uploadBodyErr, uploadErr := s.storeS3UploadPart(
r.Context(), r, streamBody, state, bucket, objectKey, uploadID, admissionProtocol,
r.Context(), r, streamBody, state, bucket, objectKey, uploadID, admissionProtocol, offload,
)
if uploadErr != nil || uploadBodyErr != nil {
s.writeS3ChunkUploadError(w, uploadBodyErr, uploadErr, bucket, objectKey)
Expand All @@ -1246,6 +1271,7 @@ func (s *S3Server) uploadPart(w http.ResponseWriter, r *http.Request, bucket str
s.cleanupPartBlobsAsync(
bucket, state.meta.Generation, objectKey, uploadID,
previous.PartNo, previous.ChunkCount, previous.PartVersion,
previous.ChunkRefVersion, previous.Offloaded,
)
}
w.Header().Set("ETag", quoteS3ETag(upload.ETag))
Expand Down Expand Up @@ -1441,8 +1467,18 @@ func (s *S3Server) cleanupUploadParts(ctx context.Context, bucket string, genera
// cleanupPartBlobsAsync asynchronously deletes the blob chunk keys for a single
// upload part. It is used to garbage-collect orphaned chunks when a part
// descriptor write fails after the chunks have already been committed.
// partVersion must match the value used when writing the chunk keys.
func (s *S3Server) cleanupPartBlobsAsync(bucket string, generation uint64, objectKey string, uploadID string, partNo uint64, chunkCount uint64, partVersion uint64) {
// The applicable version must match the legacy blob or offloaded chunkref keys.
func (s *S3Server) cleanupPartBlobsAsync(
bucket string,
generation uint64,
objectKey string,
uploadID string,
partNo uint64,
chunkCount uint64,
partVersion uint64,
chunkRefVersion uint64,
offloaded bool,
) {
select {
case s.cleanupSem <- struct{}{}:
default:
Expand Down Expand Up @@ -1470,9 +1506,13 @@ func (s *S3Server) cleanupPartBlobsAsync(bucket string, generation uint64, objec
pending = pending[:0]
}
for i := uint64(0); i < chunkCount; i++ {
key := s3keys.VersionedBlobKey(bucket, generation, objectKey, uploadID, partNo, i, partVersion)
if offloaded {
key = s3keys.VersionedChunkRefKey(bucket, generation, objectKey, uploadID, partNo, i, chunkRefVersion)
}
pending = append(pending, &kv.Elem[kv.OP]{
Op: kv.Del,
Key: s3keys.VersionedBlobKey(bucket, generation, objectKey, uploadID, partNo, i, partVersion),
Key: key,
})
if len(pending) >= s3MetaBatchOps {
flush()
Expand All @@ -1495,9 +1535,12 @@ func (s *S3Server) cleanupUploadDataAsync(bucket string, generation uint64, obje
defer cancel()
// Delete part descriptors.
s.cleanupUploadParts(ctx, bucket, generation, objectKey, uploadID)
// Delete blob chunks.
// Delete legacy blob chunks and offloaded chunk references. Content-
// addressed chunkblobs are reclaimed by the reference-counted GC.
blobPrefix := s3keys.BlobPrefixForUpload(bucket, generation, objectKey, uploadID)
s.deleteByPrefix(ctx, blobPrefix, bucket, generation, objectKey, uploadID)
chunkRefPrefix := s3keys.ChunkRefPrefixForUpload(bucket, generation, objectKey, uploadID)
s.deleteByPrefix(ctx, chunkRefPrefix, bucket, generation, objectKey, uploadID)
}()
}

Expand Down Expand Up @@ -1807,9 +1850,13 @@ func (s *S3Server) appendPartBlobKeys(pending *[]*kv.Elem[kv.OP], bucket string,
if err != nil {
return false
}
key := s3keys.VersionedBlobKey(bucket, generation, objectKey, uploadID, part.PartNo, chunkIndex, part.PartVersion)
if part.Offloaded {
key = s3keys.VersionedChunkRefKey(bucket, generation, objectKey, uploadID, part.PartNo, chunkIndex, part.ChunkRefVersion)
}
*pending = append(*pending, &kv.Elem[kv.OP]{
Op: kv.Del,
Key: s3keys.VersionedBlobKey(bucket, generation, objectKey, uploadID, part.PartNo, chunkIndex, part.PartVersion),
Key: key,
})
if len(*pending) >= s3MetaBatchOps {
flush()
Expand Down
15 changes: 12 additions & 3 deletions adapter/s3_admin_objects.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -538,9 +538,18 @@ func (r *s3AdminObjectReader) ensureBuffered() error {
r.chunkIdx = 0
continue
}
chunkKey := s3keys.VersionedBlobKey(r.bucket, r.generation, r.key,
r.manifest.UploadID, part.PartNo, r.chunkIdx, part.PartVersion)
chunk, err := r.server.store.GetAt(r.ctx, chunkKey, r.readTS)
expectedSize := part.ChunkSizes[r.chunkIdx]
chunk, err := r.server.readS3ObjectChunk(
r.ctx,
r.bucket,
r.generation,
r.key,
r.manifest.UploadID,
part,
r.chunkIdx,
expectedSize,
r.readTS,
)
if err != nil {
return errors.WithStack(err)
}
Expand Down
Loading
Loading