diff --git a/adapter/admin_grpc.go b/adapter/admin_grpc.go index 1e147330b..247b37c3e 100644 --- a/adapter/admin_grpc.go +++ b/adapter/admin_grpc.go @@ -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 " 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 } @@ -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 } @@ -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 { @@ -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 { diff --git a/adapter/admin_grpc_test.go b/adapter/admin_grpc_test.go index 75fe0b9a2..6031679da 100644 --- a/adapter/admin_grpc_test.go +++ b/adapter/admin_grpc_test.go @@ -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("") diff --git a/adapter/s3.go b/adapter/s3.go index 6c4874157..86099e517 100644 --- a/adapter/s3.go +++ b/adapter/s3.go @@ -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 { @@ -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 { @@ -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 { @@ -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 { @@ -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()) } @@ -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) @@ -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) @@ -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)) @@ -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 @@ -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) @@ -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)) @@ -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: @@ -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() @@ -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) }() } @@ -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() diff --git a/adapter/s3_admin_objects.go b/adapter/s3_admin_objects.go index ecac803be..d01db3a0a 100644 --- a/adapter/s3_admin_objects.go +++ b/adapter/s3_admin_objects.go @@ -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) } diff --git a/adapter/s3_blob_cluster.go b/adapter/s3_blob_cluster.go new file mode 100644 index 000000000..c7f80233e --- /dev/null +++ b/adapter/s3_blob_cluster.go @@ -0,0 +1,400 @@ +package adapter + +import ( + "bytes" + "context" + "crypto/sha256" + "io" + "sort" + "strings" + "sync" + "time" + + "github.com/bootjp/elastickv/internal" + "github.com/bootjp/elastickv/kv" + pb "github.com/bootjp/elastickv/proto" + "github.com/bootjp/elastickv/store" + "github.com/cockroachdb/errors" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" +) + +const ( + s3BlobCapabilityRefreshInterval = 30 * time.Second + s3BlobCapabilityPollTimeout = 3 * time.Second + s3BlobPeerRPCTimeout = 30 * time.Second +) + +// S3BlobReplica is a peer endpoint in the Raft group that owns a chunkblob. +type S3BlobReplica struct { + NodeID string + Address string + Suffrage string +} + +// S3BlobCluster combines capability negotiation and peer-local blob transfer. +// Implementations return current owning-group membership for every operation; +// callers calculate durability from that authoritative set. +type S3BlobCluster interface { + S3BlobOffloadCapabilityChecker + SelfNodeID() string + ReplicasForChunk(ctx context.Context, chunkKey []byte) ([]S3BlobReplica, error) + PushChunkBlob(ctx context.Context, replica S3BlobReplica, digest [s3ChunkBlobSHA256Bytes]byte, payload []byte, commitTS uint64) error + FetchChunkBlob(ctx context.Context, replica S3BlobReplica, digest [s3ChunkBlobSHA256Bytes]byte) ([]byte, error) + Close() error +} + +// S3BlobLocalStoreResolver exposes this process's store for a chunkblob key. +// This bypasses leader routing because chunkblob rows are peer-local auxiliary +// state and never enter Raft. +type S3BlobLocalStoreResolver interface { + LocalStoreForKey(key []byte) (store.MVCCStore, bool) +} + +type grpcS3BlobCluster struct { + selfNodeID string + members kv.RaftMembershipCoordinator + adminToken string + peerToken string + now func() time.Time + capTimeout time.Duration + rpcTimeout time.Duration + + mu sync.Mutex + conns map[string]*grpc.ClientConn + capability s3BlobCapabilityCache + closeCalled bool +} + +type s3BlobCapabilityCache struct { + fingerprint string + expiresAt time.Time + allSupport bool +} + +// NewGRPCS3BlobCluster constructs the production peer client. An empty token +// always fails capability negotiation because S3BlobFetch is peer-only and is +// never exposed by the insecure Admin mode. +func NewGRPCS3BlobCluster(selfNodeID string, members kv.RaftMembershipCoordinator, adminToken, peerToken string) S3BlobCluster { + return &grpcS3BlobCluster{ + selfNodeID: strings.TrimSpace(selfNodeID), + members: members, + adminToken: adminToken, + peerToken: peerToken, + now: time.Now, + capTimeout: s3BlobCapabilityPollTimeout, + rpcTimeout: s3BlobPeerRPCTimeout, + conns: map[string]*grpc.ClientConn{}, + } +} + +func (c *grpcS3BlobCluster) SelfNodeID() string { + if c == nil { + return "" + } + return c.selfNodeID +} + +func (c *grpcS3BlobCluster) ReplicasForChunk(ctx context.Context, chunkKey []byte) ([]S3BlobReplica, error) { + if c == nil || c.members == nil { + return nil, errors.New("s3 blob membership provider is not configured") + } + members, err := c.members.RaftMembersForKey(ctx, chunkKey) + if err != nil { + return nil, errors.WithStack(err) + } + seen := make(map[string]struct{}, len(members)) + replicas := make([]S3BlobReplica, 0, len(members)) + for _, member := range members { + nodeID := strings.TrimSpace(member.NodeID) + address := strings.TrimSpace(member.Address) + if nodeID == "" || address == "" { + return nil, errors.New("s3 blob membership contains an empty node id or address") + } + if _, duplicate := seen[nodeID]; duplicate { + return nil, errors.WithStack(errors.Newf("s3 blob membership contains duplicate node id %q", nodeID)) + } + seen[nodeID] = struct{}{} + replicas = append(replicas, S3BlobReplica{NodeID: nodeID, Address: address, Suffrage: member.Suffrage}) + } + if len(replicas) == 0 { + return nil, errors.New("s3 blob membership is empty") + } + sort.Slice(replicas, func(i, j int) bool { return replicas[i].NodeID < replicas[j].NodeID }) + return replicas, nil +} + +func (c *grpcS3BlobCluster) AllPeersSupportS3BlobOffload(ctx context.Context) bool { + if c == nil || c.adminToken == "" || c.peerToken == "" || c.members == nil || !S3BlobOffloadLocalCapability() { + return false + } + members, err := c.members.RaftMembers(ctx) + if err != nil { + return false + } + replicas, err := s3BlobCapabilityReplicas(members) + if err != nil { + return false + } + fingerprint := s3BlobReplicaFingerprint(replicas) + now := c.now() + if supported, ok := c.cachedCapability(fingerprint, now); ok { + return supported + } + allSupport := c.pollCapabilities(ctx, replicas) + c.storeCapability(fingerprint, now.Add(s3BlobCapabilityRefreshInterval), allSupport) + return allSupport +} + +func s3BlobCapabilityReplicas(members []kv.RaftMember) ([]S3BlobReplica, error) { + seen := make(map[string]struct{}, len(members)) + replicas := make([]S3BlobReplica, 0, len(members)) + for _, member := range members { + nodeID := strings.TrimSpace(member.NodeID) + address := strings.TrimSpace(member.Address) + if nodeID == "" || address == "" { + return nil, errors.New("s3 blob cluster membership contains an empty node id or address") + } + endpoint := nodeID + "\x00" + address + if _, duplicate := seen[endpoint]; duplicate { + continue + } + seen[endpoint] = struct{}{} + replicas = append(replicas, S3BlobReplica{NodeID: nodeID, Address: address, Suffrage: member.Suffrage}) + } + if len(replicas) == 0 { + return nil, errors.New("s3 blob cluster membership is empty") + } + sort.Slice(replicas, func(i, j int) bool { + if replicas[i].NodeID == replicas[j].NodeID { + return replicas[i].Address < replicas[j].Address + } + return replicas[i].NodeID < replicas[j].NodeID + }) + return replicas, nil +} + +func s3BlobReplicaFingerprint(replicas []S3BlobReplica) string { + var b strings.Builder + for _, replica := range replicas { + b.WriteString(replica.NodeID) + b.WriteByte(0) + b.WriteString(replica.Address) + b.WriteByte(0) + b.WriteString(replica.Suffrage) + b.WriteByte('\n') + } + return b.String() +} + +func (c *grpcS3BlobCluster) cachedCapability(fingerprint string, now time.Time) (bool, bool) { + c.mu.Lock() + defer c.mu.Unlock() + if c.capability.fingerprint != fingerprint || !now.Before(c.capability.expiresAt) { + return false, false + } + return c.capability.allSupport, true +} + +func (c *grpcS3BlobCluster) storeCapability(fingerprint string, expiresAt time.Time, supported bool) { + c.mu.Lock() + defer c.mu.Unlock() + c.capability = s3BlobCapabilityCache{fingerprint: fingerprint, expiresAt: expiresAt, allSupport: supported} +} + +func (c *grpcS3BlobCluster) pollCapabilities(ctx context.Context, replicas []S3BlobReplica) bool { + remote := make([]S3BlobReplica, 0, len(replicas)) + for _, replica := range replicas { + if replica.NodeID != c.selfNodeID { + remote = append(remote, replica) + } + } + results := make(chan bool, len(remote)) + for _, replica := range remote { + go func() { + results <- c.peerSupportsS3BlobOffload(ctx, replica) + }() + } + for range remote { + if !<-results { + return false + } + } + return true +} + +func (c *grpcS3BlobCluster) peerSupportsS3BlobOffload(ctx context.Context, replica S3BlobReplica) bool { + conn, err := c.connFor(replica.Address) + if err != nil { + return false + } + timeout := c.capTimeout + if timeout <= 0 { + timeout = s3BlobCapabilityPollTimeout + } + callCtx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + resp, err := pb.NewAdminClient(conn).GetClusterOverview(c.authorizedAdminContext(callCtx), &pb.GetClusterOverviewRequest{}) + return err == nil && resp.GetSelf().GetNodeId() == replica.NodeID && resp.GetCapabilities()[S3BlobOffloadCapabilityName] +} + +func (c *grpcS3BlobCluster) PushChunkBlob(ctx context.Context, replica S3BlobReplica, digest [s3ChunkBlobSHA256Bytes]byte, payload []byte, commitTS uint64) error { + if commitTS == 0 { + return errors.New("s3 chunkblob push requires a commit timestamp") + } + conn, err := c.connFor(replica.Address) + if err != nil { + return err + } + callCtx, cancel := c.peerRPCContext(ctx) + defer cancel() + stream, err := pb.NewS3BlobFetchClient(conn).PushChunkBlob(c.authorizedPeerContext(callCtx)) + if err != nil { + return errors.WithStack(err) + } + if err := sendS3ChunkBlobPushFrames(stream, digest, payload, commitTS); err != nil { + return err + } + resp, err := stream.CloseAndRecv() + if err != nil { + return errors.WithStack(err) + } + if !resp.GetDurable() { + return errors.WithStack(status.Error(codes.DataLoss, "s3 chunkblob peer did not acknowledge durable storage")) + } + return nil +} + +func sendS3ChunkBlobPushFrames( + stream pb.S3BlobFetch_PushChunkBlobClient, + digest [s3ChunkBlobSHA256Bytes]byte, + payload []byte, + commitTS uint64, +) error { + for offset, first := 0, true; first || offset < len(payload); first = false { + end := offset + s3BlobFetchFrameBytes + if end > len(payload) { + end = len(payload) + } + req := &pb.PushChunkBlobRequest{Payload: payload[offset:end], Eof: end == len(payload)} + if first { + req.ContentSha256 = digest[:] + req.CommitTs = commitTS + } + if err := stream.Send(req); err != nil { + return errors.WithStack(err) + } + offset = end + } + return nil +} + +func (c *grpcS3BlobCluster) FetchChunkBlob(ctx context.Context, replica S3BlobReplica, digest [s3ChunkBlobSHA256Bytes]byte) ([]byte, error) { + conn, err := c.connFor(replica.Address) + if err != nil { + return nil, err + } + callCtx, cancel := c.peerRPCContext(ctx) + defer cancel() + stream, err := pb.NewS3BlobFetchClient(conn).FetchChunkBlob(c.authorizedPeerContext(callCtx), &pb.FetchChunkBlobRequest{ContentSha256: digest[:]}) + if err != nil { + return nil, errors.WithStack(err) + } + payload := make([]byte, 0, s3ChunkSize) + seenEOF := false + for { + frame, recvErr := stream.Recv() + if errors.Is(recvErr, io.EOF) { + break + } + if recvErr != nil { + return nil, errors.WithStack(recvErr) + } + if seenEOF { + return nil, errors.WithStack(status.Error(codes.InvalidArgument, "s3 chunkblob fetch returned data after eof")) + } + if len(payload)+len(frame.GetPayload()) > s3ChunkSize { + return nil, errors.WithStack(status.Error(codes.ResourceExhausted, "s3 chunkblob fetch exceeds chunk size")) + } + payload = append(payload, frame.GetPayload()...) + seenEOF = frame.GetEof() + } + if !seenEOF { + return nil, errors.WithStack(status.Error(codes.InvalidArgument, "s3 chunkblob fetch omitted eof")) + } + actual := sha256.Sum256(payload) + if !bytes.Equal(actual[:], digest[:]) { + return nil, errors.WithStack(status.Error(codes.InvalidArgument, "s3 chunkblob fetch sha256 mismatch")) + } + return payload, nil +} + +func (c *grpcS3BlobCluster) peerRPCContext(ctx context.Context) (context.Context, context.CancelFunc) { + timeout := c.rpcTimeout + if timeout <= 0 { + timeout = s3BlobPeerRPCTimeout + } + return context.WithTimeout(ctx, timeout) +} + +func (c *grpcS3BlobCluster) authorizedAdminContext(ctx context.Context) context.Context { + if c == nil || c.adminToken == "" { + return ctx + } + return metadata.AppendToOutgoingContext(ctx, "authorization", "Bearer "+c.adminToken) +} + +func (c *grpcS3BlobCluster) authorizedPeerContext(ctx context.Context) context.Context { + if c == nil || c.peerToken == "" { + return ctx + } + return metadata.AppendToOutgoingContext(ctx, "authorization", "Bearer "+c.peerToken) +} + +func (c *grpcS3BlobCluster) connFor(address string) (*grpc.ClientConn, error) { + address = strings.TrimSpace(address) + if address == "" { + return nil, errors.New("s3 blob peer address is empty") + } + c.mu.Lock() + defer c.mu.Unlock() + if c.closeCalled { + return nil, errors.New("s3 blob cluster is closed") + } + if conn := c.conns[address]; conn != nil { + return conn, nil + } + conn, err := grpc.NewClient(address, internal.GRPCDialOptions()...) + if err != nil { + return nil, errors.WithStack(err) + } + c.conns[address] = conn + return conn, nil +} + +func (c *grpcS3BlobCluster) Close() error { + if c == nil { + return nil + } + c.mu.Lock() + if c.closeCalled { + c.mu.Unlock() + return nil + } + c.closeCalled = true + conns := make([]*grpc.ClientConn, 0, len(c.conns)) + for _, conn := range c.conns { + conns = append(conns, conn) + } + clear(c.conns) + c.mu.Unlock() + var first error + for _, conn := range conns { + if err := conn.Close(); err != nil && first == nil { + first = errors.WithStack(err) + } + } + return first +} diff --git a/adapter/s3_blob_fetch.go b/adapter/s3_blob_fetch.go new file mode 100644 index 000000000..452f18b73 --- /dev/null +++ b/adapter/s3_blob_fetch.go @@ -0,0 +1,473 @@ +package adapter + +import ( + "bytes" + "context" + "crypto/sha256" + "io" + "math" + "time" + + "github.com/bootjp/elastickv/internal/s3keys" + "github.com/bootjp/elastickv/kv" + pb "github.com/bootjp/elastickv/proto" + "github.com/bootjp/elastickv/store" + "github.com/cockroachdb/errors" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +const ( + s3ChunkBlobSHA256Bytes = 32 + s3BlobFetchFrameBytes = 256 * 1024 + + s3BlobFetchRegistrationRetryInitial = 10 * time.Millisecond + s3BlobFetchRegistrationRetryMax = 250 * time.Millisecond + s3BlobFetchRegistrationRetryFactor = 2 +) + +// S3BlobFetchServer serves local content-addressed S3 chunk blobs. It is the +// internal peer-to-peer RPC substrate for the chunkref/chunkblob rollout; it +// deliberately does not enable the public S3 PUT/GET offload path by itself. +type S3BlobFetchServer struct { + store store.MVCCStore + observer S3BlobOffloadObserver + clock *kv.HLC + pushBlocked func() bool + + pb.UnimplementedS3BlobFetchServer +} + +type S3BlobFetchServerOption func(*S3BlobFetchServer) + +type s3ChunkBlobMutationStore interface { + ApplyMutationsPreservingLastCommitTS(ctx context.Context, mutations []*store.KVPairMutation, readKeys [][]byte, startTS, commitTS uint64) error +} + +func WithS3BlobFetchClock(clock *kv.HLC) S3BlobFetchServerOption { + return func(s *S3BlobFetchServer) { + s.clock = clock + } +} + +func WithS3BlobFetchPushBlocked(blocked func() bool) S3BlobFetchServerOption { + return func(s *S3BlobFetchServer) { + s.pushBlocked = blocked + } +} + +func NewS3BlobFetchServer(st store.MVCCStore, observer S3BlobOffloadObserver, opts ...S3BlobFetchServerOption) *S3BlobFetchServer { + server := &S3BlobFetchServer{ + store: st, + observer: observer, + } + for _, opt := range opts { + if opt != nil { + opt(server) + } + } + return server +} + +func (s *S3BlobFetchServer) FetchChunkBlob(req *pb.FetchChunkBlobRequest, stream pb.S3BlobFetch_FetchChunkBlobServer) error { + if s == nil || s.store == nil { + return s3BlobFetchStatus(codes.FailedPrecondition, "s3 blob fetch store is not configured") + } + digest, err := s3ChunkBlobDigest(req.GetContentSha256()) + if err != nil { + return err + } + payload, err := s.fetchChunkBlobPayload(stream.Context(), digest) + if err != nil { + return err + } + if err := s.verifyChunkBlobDigest(digest, payload, codes.InvalidArgument); err != nil { + return err + } + return sendChunkBlobPayload(stream, payload) +} + +func (s *S3BlobFetchServer) fetchChunkBlobPayload(ctx context.Context, digest [s3ChunkBlobSHA256Bytes]byte) ([]byte, error) { + key := s3keys.ChunkBlobKey(digest) + payload, exists, err := s.currentChunkBlobPayload(ctx, key) + if err != nil { + return nil, err + } + if !exists { + return nil, s3BlobFetchStatus(codes.NotFound, "s3 chunkblob not found") + } + return payload, nil +} + +func sendChunkBlobPayload(stream pb.S3BlobFetch_FetchChunkBlobServer, payload []byte) error { + if len(payload) == 0 { + return errors.WithStack(stream.Send(&pb.FetchChunkBlobResponse{Eof: true})) + } + for offset := 0; offset < len(payload); { + end := offset + s3BlobFetchFrameBytes + if end > len(payload) { + end = len(payload) + } + if err := stream.Send(&pb.FetchChunkBlobResponse{ + Payload: payload[offset:end], + Eof: end == len(payload), + }); err != nil { + return errors.WithStack(err) + } + offset = end + } + return nil +} + +func (s *S3BlobFetchServer) PushChunkBlob(stream pb.S3BlobFetch_PushChunkBlobServer) error { + if s == nil || s.store == nil { + return s3BlobFetchStatus(codes.FailedPrecondition, "s3 blob fetch server is not configured") + } + if err := s.ensurePushAllowed(); err != nil { + return err + } + digest, payload, commitTS, err := s.recvChunkBlob(stream) + if err != nil { + return err + } + if err := s.verifyChunkBlobDigest(digest, payload, codes.InvalidArgument); err != nil { + return err + } + if err := s.ensurePushAllowed(); err != nil { + return err + } + if err := s.storeChunkBlob(stream.Context(), digest, payload, commitTS); err != nil { + return err + } + return sendChunkBlobPushAck(stream) +} + +func (s *S3BlobFetchServer) storeChunkBlob( + ctx context.Context, + digest [s3ChunkBlobSHA256Bytes]byte, + payload []byte, + commitTS uint64, +) error { + key := s3keys.ChunkBlobKey(digest) + for { + if err := s.ensurePushAllowed(); err != nil { + return err + } + startTS, err := s.chunkBlobWriteStartTS(ctx, key, digest, payload, commitTS) + if err != nil { + return err + } + if startTS == commitTS { + s.observeCommitTS(commitTS) + return nil + } + if err := s.applyChunkBlobUntilRegistered(ctx, key, payload, startTS, commitTS); err != nil { + done, retry, retryErr := s.retryAfterChunkBlobWriteConflict(ctx, err, key, digest, payload, commitTS) + if retryErr != nil { + return retryErr + } + if done { + s.observeCommitTS(commitTS) + return nil + } + if retry { + continue + } + if code := status.Code(err); code != codes.Unknown { + return err + } + return s3BlobFetchStatusf(codes.Internal, "write s3 chunkblob: %v", err) + } + s.observeCommitTS(commitTS) + return nil + } +} + +func (s *S3BlobFetchServer) retryAfterChunkBlobWriteConflict( + ctx context.Context, + err error, + key []byte, + digest [s3ChunkBlobSHA256Bytes]byte, + payload []byte, + commitTS uint64, +) (bool, bool, error) { + if !errors.Is(err, store.ErrWriteConflict) { + return false, false, nil + } + latestTS, exists, latestErr := s.latestChunkBlobTS(ctx, key) + if latestErr != nil { + return false, false, latestErr + } + stored, storedErr := s.chunkBlobAlreadyStored(ctx, key, digest, payload) + if storedErr != nil { + return false, false, storedErr + } + if stored { + return exists && latestTS >= commitTS, exists && latestTS < commitTS, nil + } + if exists && latestTS < commitTS { + return false, true, nil + } + return false, false, s3BlobFetchStatusf( + codes.FailedPrecondition, + "s3 chunkblob commit timestamp %d is not after latest version %d", + commitTS, + latestTS, + ) +} + +func sendChunkBlobPushAck(stream pb.S3BlobFetch_PushChunkBlobServer) error { + return errors.WithStack(stream.SendAndClose(&pb.PushChunkBlobResponse{Durable: true})) +} + +func (s *S3BlobFetchServer) currentChunkBlobPayload(ctx context.Context, key []byte) ([]byte, bool, error) { + payload, err := s.store.GetAt(ctx, key, math.MaxUint64) + if errors.Is(err, store.ErrKeyNotFound) { + return nil, false, nil + } + if err != nil { + return nil, false, s3BlobFetchStatusf(codes.Internal, "read s3 chunkblob: %v", err) + } + return payload, true, nil +} + +func (s *S3BlobFetchServer) chunkBlobAlreadyStored( + ctx context.Context, + key []byte, + digest [s3ChunkBlobSHA256Bytes]byte, + payload []byte, +) (bool, error) { + existing, exists, err := s.currentChunkBlobPayload(ctx, key) + if err != nil || !exists { + return false, err + } + if bytes.Equal(existing, payload) { + return true, nil + } + if err := s.verifyChunkBlobDigest(digest, existing, codes.InvalidArgument); err != nil { + return false, err + } + return false, s3BlobFetchStatus(codes.InvalidArgument, "s3 chunkblob already exists with different payload") +} + +func (s *S3BlobFetchServer) chunkBlobWriteStartTS( + ctx context.Context, + key []byte, + digest [s3ChunkBlobSHA256Bytes]byte, + payload []byte, + commitTS uint64, +) (uint64, error) { + if commitTS == 0 { + return 0, s3BlobFetchStatus(codes.InvalidArgument, "missing s3 chunkblob commit timestamp") + } + latestTS, exists, err := s.latestChunkBlobTS(ctx, key) + if err != nil { + return 0, err + } + if exists && commitTS <= latestTS { + stored, storedErr := s.chunkBlobAlreadyStored(ctx, key, digest, payload) + if storedErr != nil { + return 0, storedErr + } + if stored { + return commitTS, nil + } + return 0, s3BlobFetchStatusf( + codes.FailedPrecondition, + "s3 chunkblob commit timestamp %d is not after latest version %d", + commitTS, + latestTS, + ) + } + return latestTS, nil +} + +func (s *S3BlobFetchServer) latestChunkBlobTS(ctx context.Context, key []byte) (uint64, bool, error) { + latestTS, exists, err := s.store.LatestCommitTS(ctx, key) + if err != nil { + return 0, false, s3BlobFetchStatusf(codes.Internal, "read s3 chunkblob latest timestamp: %v", err) + } + return latestTS, exists, nil +} + +func (s *S3BlobFetchServer) applyChunkBlob(ctx context.Context, key, payload []byte, startTS, commitTS uint64) error { + chunkStore, ok := s.store.(s3ChunkBlobMutationStore) + if !ok { + return s3BlobFetchStatus(codes.FailedPrecondition, "s3 chunkblob store cannot preserve mvcc watermark") + } + if err := chunkStore.ApplyMutationsPreservingLastCommitTS(ctx, []*store.KVPairMutation{{ + Op: store.OpTypePut, + Key: key, + Value: payload, + }}, nil, startTS, commitTS); err != nil { + return errors.WithStack(err) + } + return nil +} + +func (s *S3BlobFetchServer) applyChunkBlobUntilRegistered(ctx context.Context, key, payload []byte, startTS, commitTS uint64) error { + backoff := s3BlobFetchRegistrationRetryInitial + for { + if err := s.ensurePushAllowed(); err != nil { + return err + } + err := s.applyChunkBlob(ctx, key, payload, startTS, commitTS) + if err == nil || !errors.Is(err, store.ErrWriterNotRegistered) { + return err + } + select { + case <-ctx.Done(): + return s3BlobFetchStatusf(codes.Unavailable, "s3 chunkblob writer registration: %v", ctx.Err()) + case <-time.After(backoff): + backoff *= s3BlobFetchRegistrationRetryFactor + if backoff > s3BlobFetchRegistrationRetryMax { + backoff = s3BlobFetchRegistrationRetryMax + } + } + } +} + +func (s *S3BlobFetchServer) recvChunkBlob(stream pb.S3BlobFetch_PushChunkBlobServer) ([s3ChunkBlobSHA256Bytes]byte, []byte, uint64, error) { + var state s3ChunkBlobReceiveState + for { + if err := s.ensurePushAllowed(); err != nil { + return state.digest, nil, 0, err + } + req, err := stream.Recv() + if errors.Is(err, io.EOF) { + return state.finish() + } + if err != nil { + return state.digest, nil, 0, errors.WithStack(err) + } + if err := state.apply(req); err != nil { + return state.digest, nil, 0, err + } + } +} + +func (s *S3BlobFetchServer) ensurePushAllowed() error { + if s != nil && s.pushBlocked != nil && s.pushBlocked() { + return s3BlobFetchStatus(codes.Unavailable, "startup rotation has not completed") + } + return nil +} + +func (s *S3BlobFetchServer) observeCommitTS(commitTS uint64) { + if s != nil && s.clock != nil { + s.clock.Observe(commitTS) + } +} + +type s3ChunkBlobReceiveState struct { + digest [s3ChunkBlobSHA256Bytes]byte + commitTS uint64 + haveDigest bool + haveCommitTS bool + seenEOF bool + payload bytes.Buffer +} + +func (s *s3ChunkBlobReceiveState) finish() ([s3ChunkBlobSHA256Bytes]byte, []byte, uint64, error) { + if !s.haveDigest { + return s.digest, nil, 0, s3BlobFetchStatus(codes.InvalidArgument, "missing s3 chunkblob sha256") + } + if !s.haveCommitTS { + return s.digest, nil, 0, s3BlobFetchStatus(codes.InvalidArgument, "missing s3 chunkblob commit timestamp") + } + if !s.seenEOF { + return s.digest, nil, 0, s3BlobFetchStatus(codes.InvalidArgument, "missing s3 chunkblob eof") + } + return s.digest, s.payload.Bytes(), s.commitTS, nil +} + +func (s *s3ChunkBlobReceiveState) apply(req *pb.PushChunkBlobRequest) error { + if req == nil { + return s3BlobFetchStatus(codes.InvalidArgument, "nil s3 chunkblob request") + } + if s.seenEOF { + return s3BlobFetchStatus(codes.InvalidArgument, "s3 chunkblob frame after eof") + } + if err := s.applyDigest(req.GetContentSha256()); err != nil { + return err + } + if !s.haveDigest { + return s3BlobFetchStatus(codes.InvalidArgument, "first s3 chunkblob frame must include sha256") + } + if err := s.applyCommitTS(req.GetCommitTs()); err != nil { + return err + } + if !s.haveCommitTS { + return s3BlobFetchStatus(codes.InvalidArgument, "first s3 chunkblob frame must include commit timestamp") + } + if s.payload.Len()+len(req.GetPayload()) > s3ChunkSize { + return s3BlobFetchStatus(codes.ResourceExhausted, "s3 chunkblob payload exceeds chunk size") + } + if _, err := s.payload.Write(req.GetPayload()); err != nil { + return s3BlobFetchStatusf(codes.Internal, "buffer s3 chunkblob: %v", err) + } + if req.GetEof() { + s.seenEOF = true + } + return nil +} + +func (s *s3ChunkBlobReceiveState) applyCommitTS(commitTS uint64) error { + if commitTS == 0 { + return nil + } + if s.haveCommitTS && commitTS != s.commitTS { + return s3BlobFetchStatus(codes.InvalidArgument, "s3 chunkblob commit timestamp changed mid-stream") + } + s.commitTS = commitTS + s.haveCommitTS = true + return nil +} + +func (s *s3ChunkBlobReceiveState) applyDigest(raw []byte) error { + if len(raw) == 0 { + return nil + } + digest, err := s3ChunkBlobDigest(raw) + if err != nil { + return err + } + if s.haveDigest && digest != s.digest { + return s3BlobFetchStatus(codes.InvalidArgument, "s3 chunkblob sha256 changed mid-stream") + } + s.digest = digest + s.haveDigest = true + return nil +} + +func (s *S3BlobFetchServer) verifyChunkBlobDigest(expected [s3ChunkBlobSHA256Bytes]byte, payload []byte, code codes.Code) error { + actual := sha256.Sum256(payload) + if actual == expected { + return nil + } + s.observeSHAMismatch() + return s3BlobFetchStatus(code, "s3 chunkblob sha256 mismatch") +} + +func s3ChunkBlobDigest(raw []byte) ([s3ChunkBlobSHA256Bytes]byte, error) { + var digest [s3ChunkBlobSHA256Bytes]byte + if len(raw) != s3ChunkBlobSHA256Bytes { + return digest, s3BlobFetchStatusf(codes.InvalidArgument, "s3 chunkblob sha256 must be %d bytes", s3ChunkBlobSHA256Bytes) + } + copy(digest[:], raw) + return digest, nil +} + +func (s *S3BlobFetchServer) observeSHAMismatch() { + if s != nil && s.observer != nil { + s.observer.ObserveS3ChunkBlobSHAMismatch() + } +} + +func s3BlobFetchStatus(code codes.Code, msg string) error { + return errors.WithStack(status.Error(code, msg)) +} + +func s3BlobFetchStatusf(code codes.Code, format string, args ...any) error { + return errors.WithStack(status.Errorf(code, format, args...)) +} diff --git a/adapter/s3_blob_fetch_test.go b/adapter/s3_blob_fetch_test.go new file mode 100644 index 000000000..71f50ff58 --- /dev/null +++ b/adapter/s3_blob_fetch_test.go @@ -0,0 +1,758 @@ +package adapter + +import ( + "bytes" + "context" + "crypto/sha256" + "errors" + "io" + "testing" + + "github.com/bootjp/elastickv/internal/s3keys" + "github.com/bootjp/elastickv/kv" + pb "github.com/bootjp/elastickv/proto" + "github.com/bootjp/elastickv/store" + "github.com/stretchr/testify/require" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" +) + +type s3BlobFetchPreservingStore interface { + ApplyMutationsPreservingLastCommitTS(context.Context, []*store.KVPairMutation, [][]byte, uint64, uint64) error +} + +type s3BlobFetchLastCommitTSStore interface { + LastCommitTS() uint64 +} + +func requireS3BlobFetchLastCommitTS(t *testing.T, st store.MVCCStore, want uint64) { + t.Helper() + reader, ok := st.(s3BlobFetchLastCommitTSStore) + require.True(t, ok) + require.Equal(t, want, reader.LastCommitTS()) +} + +func applyS3BlobFetchMutationsPreservingLastCommitTS( + st store.MVCCStore, + ctx context.Context, + mutations []*store.KVPairMutation, + readKeys [][]byte, + startTS uint64, + commitTS uint64, +) error { + preserving, ok := st.(s3BlobFetchPreservingStore) + if !ok { + return errors.New("s3 blob fetch test store cannot preserve last commit timestamp") + } + return preserving.ApplyMutationsPreservingLastCommitTS(ctx, mutations, readKeys, startTS, commitTS) +} + +func TestS3BlobFetchPushStoresAndFetchStreamsPayload(t *testing.T) { + t.Parallel() + + base := store.NewMVCCStore() + st := &recordingS3BlobFetchStore{MVCCStore: base} + observer := &recordingS3BlobOffloadObserver{} + server := NewS3BlobFetchServer(st, observer) + payload := bytes.Repeat([]byte("payload-"), (s3BlobFetchFrameBytes/len("payload-"))+2) + digest := sha256.Sum256(payload) + + push := &recordingS3BlobPushStream{ + ctx: context.Background(), + reqs: []*pb.PushChunkBlobRequest{ + {ContentSha256: digest[:], CommitTs: 100, Payload: payload[:17]}, + {Payload: payload[17:], Eof: true}, + }, + } + require.NoError(t, server.PushChunkBlob(push)) + require.True(t, push.response.GetDurable()) + require.Equal(t, 1, st.applyCalls) + require.Zero(t, st.putCalls) + require.Equal(t, []uint64{100}, st.applyCommitTS) + requireS3BlobFetchLastCommitTS(t, base, 0) + + key := s3keys.ChunkBlobKey(digest) + readTS, exists, err := st.LatestCommitTS(context.Background(), key) + require.NoError(t, err) + require.True(t, exists) + stored, err := st.GetAt(context.Background(), key, readTS) + require.NoError(t, err) + require.Equal(t, payload, stored) + + fetch := &recordingS3BlobFetchStream{ctx: context.Background()} + require.NoError(t, server.FetchChunkBlob(&pb.FetchChunkBlobRequest{ContentSha256: digest[:]}, fetch)) + require.Greater(t, len(fetch.responses), 1) + require.True(t, fetch.responses[len(fetch.responses)-1].GetEof()) + for _, resp := range fetch.responses[:len(fetch.responses)-1] { + require.False(t, resp.GetEof()) + } + var fetched []byte + for _, resp := range fetch.responses { + fetched = append(fetched, resp.GetPayload()...) + } + require.Equal(t, payload, fetched) + require.Zero(t, observer.shaMismatch) +} + +func TestS3BlobFetchRejectsMissingBlob(t *testing.T) { + t.Parallel() + + server := NewS3BlobFetchServer(store.NewMVCCStore(), nil) + digest := sha256.Sum256([]byte("missing")) + fetch := &recordingS3BlobFetchStream{ctx: context.Background()} + + err := server.FetchChunkBlob(&pb.FetchChunkBlobRequest{ContentSha256: digest[:]}, fetch) + require.Equal(t, codes.NotFound, status.Code(err)) + require.Empty(t, fetch.responses) +} + +func TestS3BlobFetchRejectsPushDigestMismatch(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + observer := &recordingS3BlobOffloadObserver{} + server := NewS3BlobFetchServer(st, observer) + wrongDigest := sha256.Sum256([]byte("expected")) + payload := []byte("actual") + push := &recordingS3BlobPushStream{ + ctx: context.Background(), + reqs: []*pb.PushChunkBlobRequest{ + {ContentSha256: wrongDigest[:], CommitTs: 1, Payload: payload, Eof: true}, + }, + } + + err := server.PushChunkBlob(push) + require.Equal(t, codes.InvalidArgument, status.Code(err)) + require.Nil(t, push.response) + _, exists, latestErr := st.LatestCommitTS(context.Background(), s3keys.ChunkBlobKey(wrongDigest)) + require.NoError(t, latestErr) + require.False(t, exists) + require.Equal(t, 1, observer.shaMismatch) +} + +func TestS3BlobFetchRejectsOversizedPush(t *testing.T) { + t.Parallel() + + server := NewS3BlobFetchServer(store.NewMVCCStore(), nil) + payload := bytes.Repeat([]byte{0xab}, s3ChunkSize+1) + digest := sha256.Sum256(payload) + push := &recordingS3BlobPushStream{ + ctx: context.Background(), + reqs: []*pb.PushChunkBlobRequest{ + {ContentSha256: digest[:], CommitTs: 1, Payload: payload, Eof: true}, + }, + } + + err := server.PushChunkBlob(push) + require.Equal(t, codes.ResourceExhausted, status.Code(err)) + require.Nil(t, push.response) +} + +func TestS3BlobFetchRejectsInvalidDigestLength(t *testing.T) { + t.Parallel() + + server := NewS3BlobFetchServer(store.NewMVCCStore(), nil) + fetch := &recordingS3BlobFetchStream{ctx: context.Background()} + err := server.FetchChunkBlob(&pb.FetchChunkBlobRequest{ContentSha256: []byte("short")}, fetch) + require.Equal(t, codes.InvalidArgument, status.Code(err)) + + push := &recordingS3BlobPushStream{ + ctx: context.Background(), + reqs: []*pb.PushChunkBlobRequest{ + {ContentSha256: []byte("short"), CommitTs: 1, Eof: true}, + }, + } + err = server.PushChunkBlob(push) + require.Equal(t, codes.InvalidArgument, status.Code(err)) + require.Nil(t, push.response) +} + +func TestS3BlobFetchRejectsFrameAfterEOF(t *testing.T) { + t.Parallel() + + server := NewS3BlobFetchServer(store.NewMVCCStore(), nil) + digest := sha256.Sum256([]byte("payloadextra")) + push := &recordingS3BlobPushStream{ + ctx: context.Background(), + reqs: []*pb.PushChunkBlobRequest{ + {ContentSha256: digest[:], CommitTs: 1, Payload: []byte("payload"), Eof: true}, + {Payload: []byte("extra")}, + }, + } + + err := server.PushChunkBlob(push) + require.Equal(t, codes.InvalidArgument, status.Code(err)) + require.Nil(t, push.response) +} + +func TestS3BlobFetchPushIsIdempotent(t *testing.T) { + t.Parallel() + + st := &recordingS3BlobFetchStore{MVCCStore: store.NewMVCCStore()} + server := NewS3BlobFetchServer(st, nil) + payload := []byte("idempotent payload") + digest := sha256.Sum256(payload) + + require.NoError(t, server.PushChunkBlob(newS3BlobPushStreamAt(digest, payload, 42))) + require.NoError(t, server.PushChunkBlob(newS3BlobPushStreamAt(digest, payload, 43))) + + require.Equal(t, 2, st.applyCalls) + require.Zero(t, st.putCalls) + require.Equal(t, []uint64{42, 43}, st.applyCommitTS) + + key := s3keys.ChunkBlobKey(digest) + latestTS, exists, err := st.LatestCommitTS(context.Background(), key) + require.NoError(t, err) + require.True(t, exists) + require.Equal(t, uint64(43), latestTS) +} + +func TestS3BlobFetchPushRestampsExistingPayloadAtLeaderCommitTS(t *testing.T) { + t.Parallel() + + base := store.NewMVCCStore() + st := &recordingS3BlobFetchStore{MVCCStore: base} + server := NewS3BlobFetchServer(st, nil) + payload := []byte("already present payload") + digest := sha256.Sum256(payload) + key := s3keys.ChunkBlobKey(digest) + require.NoError(t, base.PutAt(context.Background(), key, payload, 10, 0)) + + require.NoError(t, server.PushChunkBlob(newS3BlobPushStreamAt(digest, payload, 30))) + + require.Equal(t, []uint64{30}, st.applyCommitTS) + requireS3BlobFetchLastCommitTS(t, base, 10) + latestTS, exists, err := st.LatestCommitTS(context.Background(), key) + require.NoError(t, err) + require.True(t, exists) + require.Equal(t, uint64(30), latestTS) +} + +func TestS3BlobFetchPushDoesNotAdvanceMVCCWatermark(t *testing.T) { + t.Parallel() + + base := store.NewMVCCStore() + st := &recordingS3BlobFetchStore{MVCCStore: base} + server := NewS3BlobFetchServer(st, nil) + payload := []byte("remote leader timestamp must not become local watermark") + digest := sha256.Sum256(payload) + + require.NoError(t, server.PushChunkBlob(newS3BlobPushStreamAt(digest, payload, 500))) + + requireS3BlobFetchLastCommitTS(t, base, 0) + key := s3keys.ChunkBlobKey(digest) + latestTS, exists, err := st.LatestCommitTS(context.Background(), key) + require.NoError(t, err) + require.True(t, exists) + require.Equal(t, uint64(500), latestTS) + got, err := st.GetAt(context.Background(), key, 500) + require.NoError(t, err) + require.Equal(t, payload, got) +} + +func TestS3BlobFetchPushDoesNotAdvancePebbleMVCCWatermark(t *testing.T) { + t.Parallel() + + base, err := store.NewPebbleStore(t.TempDir()) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, base.Close()) }) + st := &recordingS3BlobFetchStore{MVCCStore: base} + server := NewS3BlobFetchServer(st, nil) + payload := []byte("remote leader timestamp must not become pebble watermark") + digest := sha256.Sum256(payload) + + require.NoError(t, server.PushChunkBlob(newS3BlobPushStreamAt(digest, payload, 700))) + + require.Equal(t, uint64(0), base.LastCommitTS()) + key := s3keys.ChunkBlobKey(digest) + latestTS, exists, latestErr := st.LatestCommitTS(context.Background(), key) + require.NoError(t, latestErr) + require.True(t, exists) + require.Equal(t, uint64(700), latestTS) + got, getErr := st.GetAt(context.Background(), key, 700) + require.NoError(t, getErr) + require.Equal(t, payload, got) +} + +func TestS3BlobFetchPushDoesNotAdvancePebbleMVCCWatermarkAfterReopen(t *testing.T) { + t.Parallel() + + dir := t.TempDir() + base, err := store.NewPebbleStore(dir) + require.NoError(t, err) + st := &recordingS3BlobFetchStore{MVCCStore: base} + server := NewS3BlobFetchServer(st, nil) + payload := []byte("remote timestamp remains auxiliary after reopen") + digest := sha256.Sum256(payload) + + require.NoError(t, server.PushChunkBlob(newS3BlobPushStreamAt(digest, payload, 900))) + require.NoError(t, base.Close()) + + reopened, err := store.NewPebbleStore(dir) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, reopened.Close()) }) + require.Equal(t, uint64(0), reopened.LastCommitTS()) + got, err := reopened.GetAt(context.Background(), s3keys.ChunkBlobKey(digest), 900) + require.NoError(t, err) + require.Equal(t, payload, got) +} + +func TestS3BlobFetchPushAcknowledgesConcurrentMatchingWrite(t *testing.T) { + t.Parallel() + + st := &conflictOnceS3BlobFetchStore{MVCCStore: store.NewMVCCStore()} + server := NewS3BlobFetchServer(st, nil) + payload := []byte("concurrent payload") + digest := sha256.Sum256(payload) + + require.NoError(t, server.PushChunkBlob(newS3BlobPushStreamAt(digest, payload, 51))) + require.Equal(t, 1, st.applyCalls) +} + +func TestS3BlobFetchPushRetriesAfterConcurrentTombstone(t *testing.T) { + t.Parallel() + + st := &tombstoneConflictS3BlobFetchStore{ + recordingS3BlobFetchStore: recordingS3BlobFetchStore{MVCCStore: store.NewMVCCStore()}, + tombstoneTS: 20, + } + server := NewS3BlobFetchServer(st, nil) + payload := []byte("retry after tombstone") + digest := sha256.Sum256(payload) + + require.NoError(t, server.PushChunkBlob(newS3BlobPushStreamAt(digest, payload, 30))) + require.Equal(t, 2, st.applyCalls) + require.Equal(t, []uint64{0, 20}, st.applyStartTS) + require.Equal(t, []uint64{30, 30}, st.applyCommitTS) +} + +func TestS3BlobFetchPushObservesLeaderCommitTS(t *testing.T) { + t.Parallel() + + clock := kv.NewHLC() + server := NewS3BlobFetchServer( + store.NewMVCCStore(), + nil, + WithS3BlobFetchClock(clock), + ) + payload := []byte("observe commit timestamp") + digest := sha256.Sum256(payload) + + require.NoError(t, server.PushChunkBlob(newS3BlobPushStreamAt(digest, payload, 100))) + require.Equal(t, uint64(100), clock.Current()) +} + +func TestS3BlobFetchPushRechecksGateBetweenFrames(t *testing.T) { + t.Parallel() + + st := &recordingS3BlobFetchStore{MVCCStore: store.NewMVCCStore()} + blocked := false + server := NewS3BlobFetchServer( + st, + nil, + WithS3BlobFetchPushBlocked(func() bool { return blocked }), + ) + payload := []byte("gate closes while streaming") + digest := sha256.Sum256(payload) + push := &recordingS3BlobPushStream{ + ctx: context.Background(), + reqs: []*pb.PushChunkBlobRequest{ + {ContentSha256: digest[:], CommitTs: 10, Payload: payload[:4]}, + {Payload: payload[4:], Eof: true}, + }, + afterRecv: func(n int) { + if n == 1 { + blocked = true + } + }, + } + + err := server.PushChunkBlob(push) + require.Equal(t, codes.Unavailable, status.Code(err)) + require.Zero(t, st.applyCalls) +} + +func TestS3BlobFetchRejectsPushMissingCommitTS(t *testing.T) { + t.Parallel() + + st := &recordingS3BlobFetchStore{MVCCStore: store.NewMVCCStore()} + server := NewS3BlobFetchServer(st, nil) + payload := []byte("missing commit timestamp") + digest := sha256.Sum256(payload) + + err := server.PushChunkBlob(newS3BlobPushStreamAt(digest, payload, 0)) + require.Equal(t, codes.InvalidArgument, status.Code(err)) + require.Zero(t, st.applyCalls) + require.Zero(t, st.putCalls) + _, exists, latestErr := st.LatestCommitTS(context.Background(), s3keys.ChunkBlobKey(digest)) + require.NoError(t, latestErr) + require.False(t, exists) +} + +func TestS3BlobFetchPushRecreatesAfterTombstone(t *testing.T) { + t.Parallel() + + base := store.NewMVCCStore() + st := &recordingS3BlobFetchStore{MVCCStore: base} + server := NewS3BlobFetchServer(st, nil) + payload := []byte("recreated payload") + digest := sha256.Sum256(payload) + key := s3keys.ChunkBlobKey(digest) + require.NoError(t, st.DeleteAt(context.Background(), key, 100)) + + require.NoError(t, server.PushChunkBlob(newS3BlobPushStreamAt(digest, payload, 101))) + require.Equal(t, []uint64{101}, st.applyCommitTS) + requireS3BlobFetchLastCommitTS(t, base, 100) + + stored, err := st.GetAt(context.Background(), key, 101) + require.NoError(t, err) + require.Equal(t, payload, stored) +} + +func TestS3BlobFetchPushRejectsStaleCommitTSAfterTombstone(t *testing.T) { + t.Parallel() + + st := &recordingS3BlobFetchStore{MVCCStore: store.NewMVCCStore()} + server := NewS3BlobFetchServer(st, nil) + payload := []byte("stale payload") + digest := sha256.Sum256(payload) + key := s3keys.ChunkBlobKey(digest) + require.NoError(t, st.DeleteAt(context.Background(), key, 100)) + + err := server.PushChunkBlob(newS3BlobPushStreamAt(digest, payload, 100)) + require.Equal(t, codes.FailedPrecondition, status.Code(err)) + require.Zero(t, st.applyCalls) +} + +func TestS3BlobFetchPushRetriesUntilWriterRegistered(t *testing.T) { + t.Parallel() + + st := ®istrationOnceS3BlobFetchStore{recordingS3BlobFetchStore: recordingS3BlobFetchStore{MVCCStore: store.NewMVCCStore()}} + server := NewS3BlobFetchServer(st, nil) + payload := []byte("registration retry payload") + digest := sha256.Sum256(payload) + + require.NoError(t, server.PushChunkBlob(newS3BlobPushStreamAt(digest, payload, 12))) + require.Equal(t, 2, st.applyCalls) +} + +func TestS3BlobFetchPushRechecksGateDuringRegistrationRetry(t *testing.T) { + t.Parallel() + + blocked := false + st := ®istrationOnceS3BlobFetchStore{ + recordingS3BlobFetchStore: recordingS3BlobFetchStore{MVCCStore: store.NewMVCCStore()}, + afterFirstErr: func() { blocked = true }, + } + server := NewS3BlobFetchServer( + st, + nil, + WithS3BlobFetchPushBlocked(func() bool { return blocked }), + ) + payload := []byte("registration retry gate closes") + digest := sha256.Sum256(payload) + + err := server.PushChunkBlob(newS3BlobPushStreamAt(digest, payload, 12)) + require.Equal(t, codes.Unavailable, status.Code(err)) + require.Equal(t, 1, st.applyCalls) +} + +func TestS3BlobFetchFetchReadsInMemoryAfterCompaction(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + requireFetchAfterCompaction(t, st) +} + +func TestS3BlobFetchFetchReadsPebbleAfterCompaction(t *testing.T) { + t.Parallel() + + st, err := store.NewPebbleStore(t.TempDir()) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, st.Close()) }) + requireFetchAfterCompaction(t, st) +} + +func requireFetchAfterCompaction(t *testing.T, st store.MVCCStore) { + t.Helper() + + payload := []byte("compacted payload") + digest := sha256.Sum256(payload) + key := s3keys.ChunkBlobKey(digest) + require.NoError(t, st.PutAt(context.Background(), key, payload, 10, 0)) + require.NoError(t, st.Compact(context.Background(), 20)) + _, err := st.GetAt(context.Background(), key, 10) + require.ErrorIs(t, err, store.ErrReadTSCompacted) + + server := NewS3BlobFetchServer(st, nil) + fetch := &recordingS3BlobFetchStream{ctx: context.Background()} + require.NoError(t, server.FetchChunkBlob(&pb.FetchChunkBlobRequest{ContentSha256: digest[:]}, fetch)) + require.Len(t, fetch.responses, 1) + require.Equal(t, payload, fetch.responses[0].GetPayload()) + require.True(t, fetch.responses[0].GetEof()) +} + +func TestS3BlobFetchFetchDigestMismatchReturnsInvalidArgument(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + observer := &recordingS3BlobOffloadObserver{} + expected := sha256.Sum256([]byte("expected payload")) + require.NoError(t, st.PutAt(context.Background(), s3keys.ChunkBlobKey(expected), []byte("corrupt payload"), 10, 0)) + + server := NewS3BlobFetchServer(st, observer) + fetch := &recordingS3BlobFetchStream{ctx: context.Background()} + err := server.FetchChunkBlob(&pb.FetchChunkBlobRequest{ContentSha256: expected[:]}, fetch) + require.Equal(t, codes.InvalidArgument, status.Code(err)) + require.Empty(t, fetch.responses) + require.Equal(t, 1, observer.shaMismatch) +} + +func newS3BlobPushStreamAt(digest [s3ChunkBlobSHA256Bytes]byte, payload []byte, commitTS uint64) *recordingS3BlobPushStream { + return &recordingS3BlobPushStream{ + ctx: context.Background(), + reqs: []*pb.PushChunkBlobRequest{ + {ContentSha256: digest[:], CommitTs: commitTS, Payload: payload, Eof: true}, + }, + } +} + +type recordingS3BlobFetchStream struct { + ctx context.Context + responses []*pb.FetchChunkBlobResponse +} + +func (s *recordingS3BlobFetchStream) Send(resp *pb.FetchChunkBlobResponse) error { + s.responses = append(s.responses, &pb.FetchChunkBlobResponse{ + Payload: bytes.Clone(resp.GetPayload()), + Eof: resp.GetEof(), + }) + return nil +} + +func (*recordingS3BlobFetchStream) SetHeader(metadata.MD) error { return nil } + +func (*recordingS3BlobFetchStream) SendHeader(metadata.MD) error { return nil } + +func (*recordingS3BlobFetchStream) SetTrailer(metadata.MD) {} + +func (s *recordingS3BlobFetchStream) Context() context.Context { + if s.ctx != nil { + return s.ctx + } + return context.Background() +} + +func (*recordingS3BlobFetchStream) SendMsg(any) error { return nil } + +func (*recordingS3BlobFetchStream) RecvMsg(any) error { return nil } + +type recordingS3BlobPushStream struct { + ctx context.Context + reqs []*pb.PushChunkBlobRequest + next int + afterRecv func(int) + response *pb.PushChunkBlobResponse +} + +func (s *recordingS3BlobPushStream) Recv() (*pb.PushChunkBlobRequest, error) { + if s.next >= len(s.reqs) { + return nil, io.EOF + } + req := s.reqs[s.next] + s.next++ + if s.afterRecv != nil { + s.afterRecv(s.next) + } + return req, nil +} + +func (s *recordingS3BlobPushStream) SendAndClose(resp *pb.PushChunkBlobResponse) error { + s.response = resp + return nil +} + +func (*recordingS3BlobPushStream) SetHeader(metadata.MD) error { return nil } + +func (*recordingS3BlobPushStream) SendHeader(metadata.MD) error { return nil } + +func (*recordingS3BlobPushStream) SetTrailer(metadata.MD) {} + +func (s *recordingS3BlobPushStream) Context() context.Context { + if s.ctx != nil { + return s.ctx + } + return context.Background() +} + +func (*recordingS3BlobPushStream) SendMsg(any) error { return nil } + +func (*recordingS3BlobPushStream) RecvMsg(any) error { return nil } + +type recordingS3BlobFetchStore struct { + store.MVCCStore + applyCalls int + putCalls int + applyStartTS []uint64 + applyCommitTS []uint64 +} + +func (s *recordingS3BlobFetchStore) PutAt(context.Context, []byte, []byte, uint64, uint64) error { + s.putCalls++ + return errors.New("unexpected PutAt call for S3 chunkblob durability") +} + +func (s *recordingS3BlobFetchStore) ApplyMutations( + ctx context.Context, + mutations []*store.KVPairMutation, + readKeys [][]byte, + startTS uint64, + commitTS uint64, +) error { + s.applyCalls++ + s.applyStartTS = append(s.applyStartTS, startTS) + s.applyCommitTS = append(s.applyCommitTS, commitTS) + return s.MVCCStore.ApplyMutations(ctx, mutations, readKeys, startTS, commitTS) +} + +func (s *recordingS3BlobFetchStore) ApplyMutationsPreservingLastCommitTS( + ctx context.Context, + mutations []*store.KVPairMutation, + readKeys [][]byte, + startTS uint64, + commitTS uint64, +) error { + s.applyCalls++ + s.applyStartTS = append(s.applyStartTS, startTS) + s.applyCommitTS = append(s.applyCommitTS, commitTS) + return applyS3BlobFetchMutationsPreservingLastCommitTS(s.MVCCStore, ctx, mutations, readKeys, startTS, commitTS) +} + +type conflictOnceS3BlobFetchStore struct { + store.MVCCStore + applyCalls int +} + +func (s *conflictOnceS3BlobFetchStore) ApplyMutations( + ctx context.Context, + mutations []*store.KVPairMutation, + readKeys [][]byte, + startTS uint64, + commitTS uint64, +) error { + s.applyCalls++ + if s.applyCalls == 1 { + if err := s.MVCCStore.ApplyMutations(ctx, mutations, readKeys, startTS, commitTS); err != nil { + return err + } + return store.ErrWriteConflict + } + return s.MVCCStore.ApplyMutations(ctx, mutations, readKeys, startTS, commitTS) +} + +func (s *conflictOnceS3BlobFetchStore) ApplyMutationsPreservingLastCommitTS( + ctx context.Context, + mutations []*store.KVPairMutation, + readKeys [][]byte, + startTS uint64, + commitTS uint64, +) error { + s.applyCalls++ + if s.applyCalls == 1 { + if err := applyS3BlobFetchMutationsPreservingLastCommitTS(s.MVCCStore, ctx, mutations, readKeys, startTS, commitTS); err != nil { + return err + } + return store.ErrWriteConflict + } + return applyS3BlobFetchMutationsPreservingLastCommitTS(s.MVCCStore, ctx, mutations, readKeys, startTS, commitTS) +} + +type registrationOnceS3BlobFetchStore struct { + recordingS3BlobFetchStore + afterFirstErr func() +} + +func (s *registrationOnceS3BlobFetchStore) ApplyMutations( + ctx context.Context, + mutations []*store.KVPairMutation, + readKeys [][]byte, + startTS uint64, + commitTS uint64, +) error { + s.applyCalls++ + s.applyStartTS = append(s.applyStartTS, startTS) + if s.applyCalls == 1 { + return store.ErrWriterNotRegistered + } + s.applyCommitTS = append(s.applyCommitTS, commitTS) + return s.MVCCStore.ApplyMutations(ctx, mutations, readKeys, startTS, commitTS) +} + +func (s *registrationOnceS3BlobFetchStore) ApplyMutationsPreservingLastCommitTS( + ctx context.Context, + mutations []*store.KVPairMutation, + readKeys [][]byte, + startTS uint64, + commitTS uint64, +) error { + s.applyCalls++ + s.applyStartTS = append(s.applyStartTS, startTS) + if s.applyCalls == 1 { + if s.afterFirstErr != nil { + s.afterFirstErr() + } + return store.ErrWriterNotRegistered + } + s.applyCommitTS = append(s.applyCommitTS, commitTS) + return applyS3BlobFetchMutationsPreservingLastCommitTS(s.MVCCStore, ctx, mutations, readKeys, startTS, commitTS) +} + +type tombstoneConflictS3BlobFetchStore struct { + recordingS3BlobFetchStore + tombstoneTS uint64 +} + +func (s *tombstoneConflictS3BlobFetchStore) ApplyMutations( + ctx context.Context, + mutations []*store.KVPairMutation, + readKeys [][]byte, + startTS uint64, + commitTS uint64, +) error { + s.applyCalls++ + s.applyStartTS = append(s.applyStartTS, startTS) + s.applyCommitTS = append(s.applyCommitTS, commitTS) + if s.applyCalls == 1 { + if len(mutations) == 0 { + return store.ErrWriteConflict + } + if err := s.DeleteAt(ctx, mutations[0].Key, s.tombstoneTS); err != nil { + return err + } + return store.ErrWriteConflict + } + return s.MVCCStore.ApplyMutations(ctx, mutations, readKeys, startTS, commitTS) +} + +func (s *tombstoneConflictS3BlobFetchStore) ApplyMutationsPreservingLastCommitTS( + ctx context.Context, + mutations []*store.KVPairMutation, + readKeys [][]byte, + startTS uint64, + commitTS uint64, +) error { + s.applyCalls++ + s.applyStartTS = append(s.applyStartTS, startTS) + s.applyCommitTS = append(s.applyCommitTS, commitTS) + if s.applyCalls == 1 { + if len(mutations) == 0 { + return store.ErrWriteConflict + } + if err := s.DeleteAt(ctx, mutations[0].Key, s.tombstoneTS); err != nil { + return err + } + return store.ErrWriteConflict + } + return applyS3BlobFetchMutationsPreservingLastCommitTS(s.MVCCStore, ctx, mutations, readKeys, startTS, commitTS) +} diff --git a/adapter/s3_blob_m1_test.go b/adapter/s3_blob_m1_test.go new file mode 100644 index 000000000..ccc97c86f --- /dev/null +++ b/adapter/s3_blob_m1_test.go @@ -0,0 +1,869 @@ +package adapter + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/xml" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/bootjp/elastickv/internal/s3keys" + "github.com/bootjp/elastickv/kv" + pb "github.com/bootjp/elastickv/proto" + "github.com/bootjp/elastickv/store" + "github.com/cockroachdb/errors" + "github.com/stretchr/testify/require" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" +) + +func TestS3BlobOffloadPutAndProxyOnMissGet(t *testing.T) { + t.Parallel() + + metadataStore := store.NewMVCCStore() + localResolver := &mutableS3BlobLocalStore{store: metadataStore} + cluster := newFakeS3BlobCluster() + server := NewS3Server( + nil, "", metadataStore, newLocalAdapterCoordinator(metadataStore), nil, + WithS3BlobOffloadEnabled(true), + WithS3BlobCluster(cluster), + WithS3BlobLocalStoreResolver(localResolver), + withS3BlobOffloadGCReadyForTest(), + ) + + rec := httptest.NewRecorder() + server.handle(rec, newS3TestRequest(http.MethodPut, "/bucket-a", nil)) + require.Equal(t, http.StatusOK, rec.Code) + + payload := bytes.Repeat([]byte("blob-offload"), s3ChunkSize/len("blob-offload")+100) + rec = httptest.NewRecorder() + server.handle(rec, newS3TestRequest(http.MethodPut, "/bucket-a/object", bytes.NewReader(payload))) + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + + readTS := server.readTS() + legacy, err := metadataStore.ScanAt(context.Background(), []byte(s3keys.BlobPrefix), prefixScanEnd([]byte(s3keys.BlobPrefix)), 10, readTS) + require.NoError(t, err) + require.Empty(t, legacy) + refs, err := metadataStore.ScanAt(context.Background(), []byte(s3keys.ChunkRefPrefix), prefixScanEnd([]byte(s3keys.ChunkRefPrefix)), 10, readTS) + require.NoError(t, err) + require.Len(t, refs, 2) + for _, pair := range refs { + ref, ok := s3keys.DecodeChunkRefValue(pair.Value) + require.True(t, ok) + require.Len(t, ref.ReplicaPeers, len(cluster.replicas)) + for _, peer := range ref.ReplicaPeers { + require.NotEmpty(t, peer.NodeID) + require.NotEmpty(t, peer.Address) + } + } + blobs, err := metadataStore.ScanAt(context.Background(), []byte(s3keys.ChunkBlobPrefix), prefixScanEnd([]byte(s3keys.ChunkBlobPrefix)), 10, readTS) + require.NoError(t, err) + require.Len(t, blobs, 2) + require.GreaterOrEqual(t, cluster.pushCount(), 4) + + // Simulate a follower that applied metadata but has no local chunkblobs. + // GET must fetch from peers, verify the SHA, persist locally, and only then + // send the successful HTTP response. + followerLocal := store.NewMVCCStore() + localResolver.set(followerLocal) + rec = httptest.NewRecorder() + server.handle(rec, newS3TestRequest(http.MethodGet, "/bucket-a/object", nil)) + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + require.Equal(t, payload, rec.Body.Bytes()) + require.GreaterOrEqual(t, cluster.fetchCount(), 2) + fetched, err := followerLocal.ScanAt(context.Background(), []byte(s3keys.ChunkBlobPrefix), prefixScanEnd([]byte(s3keys.ChunkBlobPrefix)), 10, ^uint64(0)) + require.NoError(t, err) + require.Len(t, fetched, 2) +} + +func TestS3BlobOffloadFailsPutWithoutTwoDurableCopies(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + cluster := newFakeS3BlobCluster() + cluster.pushErr = status.Error(codes.Unavailable, "peer unavailable") + server := NewS3Server( + nil, "", st, newLocalAdapterCoordinator(st), nil, + WithS3BlobOffloadEnabled(true), + WithS3BlobCluster(cluster), + WithS3BlobLocalStoreResolver(&mutableS3BlobLocalStore{store: st}), + withS3BlobOffloadGCReadyForTest(), + ) + + rec := httptest.NewRecorder() + server.handle(rec, newS3TestRequest(http.MethodPut, "/bucket-a", nil)) + require.Equal(t, http.StatusOK, rec.Code) + rec = httptest.NewRecorder() + server.handle(rec, newS3TestRequest(http.MethodPut, "/bucket-a/object", bytes.NewReader([]byte("payload")))) + require.Equal(t, http.StatusServiceUnavailable, rec.Code) + + refs, err := st.ScanAt(context.Background(), []byte(s3keys.ChunkRefPrefix), prefixScanEnd([]byte(s3keys.ChunkRefPrefix)), 10, ^uint64(0)) + require.NoError(t, err) + require.Empty(t, refs) +} + +func TestS3BlobOffloadStartsLocalAndRemoteWritesConcurrently(t *testing.T) { + t.Parallel() + + base := store.NewMVCCStore() + localStarted := make(chan struct{}) + release := make(chan struct{}) + blocking := &blockingS3BlobStore{MVCCStore: base, started: localStarted, release: release} + cluster := newFakeS3BlobCluster() + cluster.pushStarted = make(chan struct{}, 2) + cluster.pushRelease = release + server := NewS3Server( + nil, "", base, newLocalAdapterCoordinator(base), nil, + WithS3BlobCluster(cluster), + WithS3BlobLocalStoreResolver(&mutableS3BlobLocalStore{store: blocking}), + ) + + done := make(chan error, 1) + go func() { + payload := []byte("concurrent durability") + digest := sha256.Sum256(payload) + _, err := server.persistS3ChunkBlob(context.Background(), []byte("object-route"), digest, payload, 10) + done <- err + }() + + <-localStarted + <-cluster.pushStarted + close(release) + require.NoError(t, <-done) +} + +func TestS3BlobOffloadRoutesChunkBlobWithChunkRef(t *testing.T) { + t.Parallel() + + server, cluster, localResolver := newS3BlobM1TestServer(t, nil) + putS3BlobM1Object(t, server, "/bucket-route/object", []byte("co-located snapshot payload")) + + require.True(t, bytes.HasPrefix(localResolver.lastRouteKey(), []byte(s3keys.ChunkRefPrefix))) + require.True(t, bytes.HasPrefix(cluster.lastRouteKey(), []byte(s3keys.ChunkRefPrefix))) +} + +func TestS3BlobOffloadLocalReplicationHonorsPushGate(t *testing.T) { + t.Parallel() + + metadataStore := store.NewMVCCStore() + localStore := &recordingS3BlobFetchStore{MVCCStore: store.NewMVCCStore()} + server := NewS3Server( + nil, "", metadataStore, newLocalAdapterCoordinator(metadataStore), nil, + WithS3BlobCluster(newFakeS3BlobCluster()), + WithS3BlobLocalStoreResolver(&mutableS3BlobLocalStore{store: localStore}), + WithS3BlobPushBlocked(func() bool { return true }), + ) + payload := []byte("blocked local replication") + digest := sha256.Sum256(payload) + + _, err := server.persistS3ChunkBlob(context.Background(), []byte("object-route"), digest, payload, 10) + require.Error(t, err) + require.Zero(t, localStore.applyCalls) +} + +func TestS3BlobOffloadProxyRepairHonorsPushGate(t *testing.T) { + t.Parallel() + + metadataStore := store.NewMVCCStore() + localStore := &recordingS3BlobFetchStore{MVCCStore: store.NewMVCCStore()} + cluster := newFakeS3BlobCluster() + payload := []byte("blocked proxy repair") + digest := sha256.Sum256(payload) + cluster.blobs["n2"] = map[[sha256.Size]byte][]byte{digest: payload} + server := NewS3Server( + nil, "", metadataStore, newLocalAdapterCoordinator(metadataStore), nil, + WithS3BlobCluster(cluster), + WithS3BlobLocalStoreResolver(&mutableS3BlobLocalStore{store: localStore}), + WithS3BlobPushBlocked(func() bool { return true }), + ) + ref := s3keys.ChunkRefValue{ + ContentSHA256: digest, + Size: uint64(len(payload)), + SourcePeer: "n2", + } + + _, err := server.fetchAndStoreS3ChunkBlob(context.Background(), []byte("object-route"), ref, 10) + require.Equal(t, codes.Unavailable, status.Code(err)) + require.Zero(t, localStore.applyCalls) +} + +func TestS3BlobMinReplicasFromEnvRejectsLeaderOnly(t *testing.T) { + t.Setenv(s3BlobMinReplicasEnvVar, "1") + _, err := S3BlobMinReplicasFromEnv() + require.Error(t, err) +} + +func TestS3BlobOffloadRangeFetchesOnlyRequestedChunk(t *testing.T) { + t.Parallel() + + server, cluster, localResolver := newS3BlobM1TestServer(t, nil) + payload := bytes.Repeat([]byte("range-offload"), s3ChunkSize/len("range-offload")+100) + putS3BlobM1Object(t, server, "/bucket-range/object", payload) + + localResolver.set(store.NewMVCCStore()) + rec := httptest.NewRecorder() + req := newS3TestRequest(http.MethodGet, "/bucket-range/object", nil) + req.Header.Set("Range", "bytes=0-9") + server.handle(rec, req) + require.Equal(t, http.StatusPartialContent, rec.Code, rec.Body.String()) + require.Equal(t, payload[:10], rec.Body.Bytes()) + require.Equal(t, 1, cluster.fetchCount()) +} + +func TestS3BlobOffloadRepairsCorruptLocalChunk(t *testing.T) { + t.Parallel() + + observer := &recordingS3BlobOffloadObserver{} + server, _, localResolver := newS3BlobM1TestServer(t, observer) + payload := []byte("repair-corrupt-local-copy") + putS3BlobM1Object(t, server, "/bucket-repair/object", payload) + + follower := store.NewMVCCStore() + digest := sha256.Sum256(payload) + require.NoError(t, follower.PutAt(context.Background(), s3keys.ChunkBlobKey(digest), []byte("corrupt"), 1, 0)) + localResolver.set(follower) + rec := httptest.NewRecorder() + server.handle(rec, newS3TestRequest(http.MethodGet, "/bucket-repair/object", nil)) + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + require.Equal(t, payload, rec.Body.Bytes()) + require.GreaterOrEqual(t, observer.shaMismatch, 1) + repaired, err := follower.GetAt(context.Background(), s3keys.ChunkBlobKey(digest), ^uint64(0)) + require.NoError(t, err) + require.Equal(t, payload, repaired) +} + +func TestS3BlobOffloadAllPeersMissingReturnsInternalError(t *testing.T) { + t.Parallel() + + observer := &recordingS3BlobOffloadObserver{} + server, cluster, localResolver := newS3BlobM1TestServer(t, observer) + payload := []byte("unrecoverable-copy") + putS3BlobM1Object(t, server, "/bucket-missing/object", payload) + + require.Eventually(t, func() bool { return cluster.pushCount() == 2 }, time.Second, time.Millisecond) + cluster.clearBlobs() + localResolver.set(store.NewMVCCStore()) + rec := httptest.NewRecorder() + server.handle(rec, newS3TestRequest(http.MethodGet, "/bucket-missing/object", nil)) + require.Equal(t, http.StatusInternalServerError, rec.Code) + require.Equal(t, 1, observer.unrecoverable) +} + +func TestS3BlobFetchUsesWriteTimeReplicaAfterMembershipChange(t *testing.T) { + t.Parallel() + + payload := []byte("write-time replica survives membership replacement") + digest := sha256.Sum256(payload) + cluster := newFakeS3BlobCluster() + cluster.self = "n4" + cluster.replicas = []S3BlobReplica{ + {NodeID: "n4", Address: "n4:50051", Suffrage: "voter"}, + {NodeID: "n5", Address: "n5:50051", Suffrage: "voter"}, + {NodeID: "n6", Address: "n6:50051", Suffrage: "voter"}, + } + cluster.blobs["n2"] = map[[sha256.Size]byte][]byte{digest: bytes.Clone(payload)} + server := &S3Server{blobCluster: cluster} + + got, err := server.fetchS3ChunkBlob(context.Background(), []byte("object-route"), s3keys.ChunkRefValue{ + ContentSHA256: digest, + Size: uint64(len(payload)), + SourcePeer: "n2", + ReplicaPeers: []s3keys.ChunkRefPeer{ + {NodeID: "n1", Address: "n1:50051"}, + {NodeID: "n2", Address: "n2:50051"}, + }, + }) + require.NoError(t, err) + require.Equal(t, payload, got) +} + +func TestOrderedS3BlobPeersTriesCurrentSourceAddressBeforeWriteTimeAddress(t *testing.T) { + t.Parallel() + + writeReplicas := []S3BlobReplica{ + {NodeID: "n1", Address: "old-n1:50051"}, + {NodeID: "n2", Address: "n2:50051"}, + } + currentReplicas := []S3BlobReplica{ + {NodeID: "n1", Address: "new-n1:50051"}, + {NodeID: "n3", Address: "n3:50051"}, + } + + got := orderedS3BlobPeers(writeReplicas, currentReplicas, "n4", "n1") + require.Len(t, got, 4) + require.Equal(t, S3BlobReplica{NodeID: "n1", Address: "new-n1:50051"}, got[0]) + require.Equal(t, S3BlobReplica{NodeID: "n1", Address: "old-n1:50051"}, got[1]) + require.ElementsMatch(t, []S3BlobReplica{ + {NodeID: "n2", Address: "n2:50051"}, + {NodeID: "n3", Address: "n3:50051"}, + }, got[2:]) +} + +func TestS3BlobReadUsesManifestChunkRefVersion(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + oldDigest := sha256.Sum256([]byte("old part")) + newDigest := sha256.Sum256([]byte("new part")) + oldValue, err := s3keys.EncodeChunkRefValue(s3keys.ChunkRefValue{ + ContentSHA256: oldDigest, + Size: uint64(len("old part")), + SourcePeer: "n1", + }) + require.NoError(t, err) + newValue, err := s3keys.EncodeChunkRefValue(s3keys.ChunkRefValue{ + ContentSHA256: newDigest, + Size: uint64(len("new part")), + SourcePeer: "n1", + }) + require.NoError(t, err) + oldKey := s3keys.VersionedChunkRefKey("bucket", 1, "object", "upload", 1, 0, 100) + newKey := s3keys.VersionedChunkRefKey("bucket", 1, "object", "upload", 1, 0, 200) + require.NoError(t, st.PutAt(context.Background(), oldKey, oldValue, 101, 0)) + require.NoError(t, st.PutAt(context.Background(), newKey, newValue, 201, 0)) + + server := &S3Server{store: st} + ref, _, err := server.loadS3ChunkRef( + context.Background(), "bucket", 1, "object", "upload", + 1, 0, 100, uint64(len("old part")), ^uint64(0), + ) + require.NoError(t, err) + require.Equal(t, oldDigest, ref.ContentSHA256) +} + +func TestS3BlobOffloadUploadPartAllocatesUniqueChunkRefVersion(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + require.NoError(t, st.PutAt(context.Background(), []byte("seed"), []byte("value"), 100, 0)) + server := NewS3Server(nil, "", st, newLocalAdapterCoordinator(st), nil) + versions := make(chan s3UploadPartVersion, 2) + errs := make(chan error, 2) + start := make(chan struct{}) + var wg sync.WaitGroup + for range 2 { + wg.Add(1) + go func() { + defer wg.Done() + <-start + version, err := server.allocateS3UploadPartVersionForMode(context.Background(), true) + versions <- version + errs <- err + }() + } + close(start) + wg.Wait() + close(versions) + close(errs) + + for err := range errs { + require.NoError(t, err) + } + got := make([]s3UploadPartVersion, 0, 2) + for version := range versions { + got = append(got, version) + } + require.Len(t, got, 2) + require.Equal(t, got[0].startTS, got[1].startTS) + require.Zero(t, got[0].commitTS) + require.Zero(t, got[1].commitTS) + require.NotZero(t, got[0].chunkRefVersion) + require.NotZero(t, got[1].chunkRefVersion) + require.NotEqual(t, got[0].chunkRefVersion, got[1].chunkRefVersion) +} + +func TestS3BlobOffloadMultipartRoundTrip(t *testing.T) { + t.Parallel() + + server, _, localResolver := newS3BlobM1TestServer(t, nil) + rec := httptest.NewRecorder() + server.handle(rec, newS3TestRequest(http.MethodPut, "/bucket-multipart", nil)) + require.Equal(t, http.StatusOK, rec.Code) + + rec = httptest.NewRecorder() + server.handle(rec, newS3TestRequest(http.MethodPost, "/bucket-multipart/object?uploads=", nil)) + require.Equal(t, http.StatusOK, rec.Code) + var initiated s3InitiateMultipartUploadResult + require.NoError(t, xml.Unmarshal(rec.Body.Bytes(), &initiated)) + require.NotEmpty(t, initiated.UploadId) + + payload := []byte("offloaded multipart payload") + rec = httptest.NewRecorder() + server.handle(rec, newS3TestRequest( + http.MethodPut, + fmt.Sprintf("/bucket-multipart/object?uploadId=%s&partNumber=1", initiated.UploadId), + bytes.NewReader(payload), + )) + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + etag := strings.Trim(rec.Header().Get("ETag"), `"`) + + completeBody := fmt.Sprintf( + `1"%s"`, + etag, + ) + rec = httptest.NewRecorder() + server.handle(rec, newS3TestRequest( + http.MethodPost, + fmt.Sprintf("/bucket-multipart/object?uploadId=%s", initiated.UploadId), + strings.NewReader(completeBody), + )) + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + + localResolver.set(store.NewMVCCStore()) + rec = httptest.NewRecorder() + server.handle(rec, newS3TestRequest(http.MethodGet, "/bucket-multipart/object", nil)) + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + require.Equal(t, payload, rec.Body.Bytes()) +} + +func TestS3BlobDurableTarget(t *testing.T) { + t.Parallel() + + replicas := newFakeS3BlobCluster().replicas + server := &S3Server{} + target, voterTarget, degraded, err := server.s3BlobDurableTarget(replicas, "n1") + require.NoError(t, err) + require.Equal(t, 3, target) + require.Equal(t, 3, voterTarget) + require.False(t, degraded) + + server.blobMinReplicas = 3 + target, voterTarget, degraded, err = server.s3BlobDurableTarget(replicas, "n1") + require.NoError(t, err) + require.Equal(t, 3, target) + require.Equal(t, 3, voterTarget) + require.False(t, degraded) + + target, voterTarget, degraded, err = server.s3BlobDurableTarget(replicas[:2], "n1") + require.NoError(t, err) + require.Equal(t, 2, target) + require.Equal(t, 2, voterTarget) + require.True(t, degraded) + + five := append(append([]S3BlobReplica(nil), replicas...), + S3BlobReplica{NodeID: "n4", Address: "n4:50051", Suffrage: "voter"}, + S3BlobReplica{NodeID: "n5", Address: "n5:50051", Suffrage: "voter"}, + ) + server.blobMinReplicas = 0 + target, voterTarget, degraded, err = server.s3BlobDurableTarget(five, "n1") + require.NoError(t, err) + require.Equal(t, 5, target) + require.Equal(t, 5, voterTarget) + require.False(t, degraded) + + server.blobMinReplicas = 2 + target, voterTarget, degraded, err = server.s3BlobDurableTarget(five, "n1") + require.NoError(t, err) + require.Equal(t, 5, target, "M1 requires every current member") + require.Equal(t, 5, voterTarget) + require.False(t, degraded) + + withLearner := append(append([]S3BlobReplica(nil), replicas...), + S3BlobReplica{NodeID: "n4", Address: "n4:50051", Suffrage: "learner"}, + ) + target, voterTarget, degraded, err = server.s3BlobDurableTarget(withLearner, "n1") + require.NoError(t, err) + require.Equal(t, 4, target) + require.Equal(t, 3, voterTarget) + require.False(t, degraded) +} + +func TestS3BlobReplicationDoesNotCountLearnerTowardVoterQuorum(t *testing.T) { + t.Parallel() + + results := make(chan s3BlobReplicationResult, 3) + results <- s3BlobReplicationResult{ + replica: S3BlobReplica{NodeID: "n1", Address: "n1:50051", Suffrage: "voter"}, + local: true, + } + results <- s3BlobReplicationResult{ + replica: S3BlobReplica{NodeID: "n4", Address: "n4:50051", Suffrage: "learner"}, + } + results <- s3BlobReplicationResult{ + replica: S3BlobReplica{NodeID: "n2", Address: "n2:50051", Suffrage: "voter"}, + } + + summary := (*S3Server)(nil).collectS3BlobReplication(context.Background(), results, 3, 3, 2) + require.Equal(t, 3, summary.durable) + require.Equal(t, 2, summary.voterDurable) + require.True(t, summary.localDurable) + require.Len(t, summary.replicas, 3) +} + +func TestS3BlobOffloadAdminGetUsesProxyOnMiss(t *testing.T) { + t.Parallel() + + server, cluster, localResolver := newS3BlobM1TestServer(t, nil) + payload := bytes.Repeat([]byte("admin-offload"), s3ChunkSize/len("admin-offload")+10) + putS3BlobM1Object(t, server, "/bucket-admin/object", payload) + + localResolver.set(store.NewMVCCStore()) + body, meta, err := server.AdminGetObject( + context.Background(), fullAdminBucketsPrincipal(), "bucket-admin", "object", + ) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, body.Close()) }) + + got, err := io.ReadAll(body) + require.NoError(t, err) + require.Equal(t, payload, got) + require.Equal(t, int64(len(payload)), meta.Size) + require.GreaterOrEqual(t, cluster.fetchCount(), 2) +} + +func TestGRPCS3BlobClusterRequiresAndForwardsBearerToken(t *testing.T) { + t.Parallel() + + withoutToken := requireGRPCS3BlobCluster(t, "") + require.False(t, withoutToken.AllPeersSupportS3BlobOffload(context.Background())) + + cluster := requireGRPCS3BlobCluster(t, "peer-secret") + adminCtx := cluster.authorizedAdminContext(context.Background()) + md, ok := metadata.FromOutgoingContext(adminCtx) + require.True(t, ok) + require.Equal(t, []string{"Bearer peer-secret"}, md.Get("authorization")) + peerCtx := cluster.authorizedPeerContext(context.Background()) + md, ok = metadata.FromOutgoingContext(peerCtx) + require.True(t, ok) + require.Equal(t, []string{"Bearer peer-secret-peer"}, md.Get("authorization")) +} + +func TestGRPCS3BlobClusterAuthenticatedCapabilityPushAndFetch(t *testing.T) { + t.Parallel() + + listener, err := (&net.ListenConfig{}).Listen(context.Background(), "tcp", "127.0.0.1:0") + require.NoError(t, err) + peerStore := store.NewMVCCStore() + admin := NewAdminServer(NodeIdentity{NodeID: "n2", GRPCAddress: listener.Addr().String()}, nil) + admin.SetCapability(S3BlobOffloadCapabilityName, true) + adminUnary, adminStream := AdminTokenAuth("read-only-admin") + peerUnary, peerStream := S3BlobPeerTokenAuth("peer-secret") + server := grpc.NewServer( + grpc.ChainUnaryInterceptor(adminUnary, peerUnary), + grpc.ChainStreamInterceptor(adminStream, peerStream), + ) + pb.RegisterAdminServer(server, admin) + pb.RegisterS3BlobFetchServer(server, NewS3BlobFetchServer(peerStore, nil)) + serveDone := make(chan error, 1) + go func() { serveDone <- server.Serve(listener) }() + t.Cleanup(func() { + server.Stop() + serveErr := <-serveDone + require.True(t, serveErr == nil || errors.Is(serveErr, grpc.ErrServerStopped)) + }) + + members := staticS3BlobMembership{members: []kv.RaftMember{ + {NodeID: "n1", Address: "127.0.0.1:1", Suffrage: "voter"}, + {NodeID: "n2", Address: listener.Addr().String(), Suffrage: "voter"}, + }} + cluster := NewGRPCS3BlobCluster("n1", members, "read-only-admin", "peer-secret") + t.Cleanup(func() { require.NoError(t, cluster.Close()) }) + require.True(t, cluster.AllPeersSupportS3BlobOffload(context.Background())) + + mixedGroups := staticS3BlobMembership{ + members: members.members, + allMembers: append(append([]kv.RaftMember(nil), members.members...), kv.RaftMember{ + NodeID: "n3", Address: "127.0.0.1:1", Suffrage: "voter", + }), + } + mixedCluster := NewGRPCS3BlobCluster("n1", mixedGroups, "read-only-admin", "peer-secret") + t.Cleanup(func() { require.NoError(t, mixedCluster.Close()) }) + require.False(t, mixedCluster.AllPeersSupportS3BlobOffload(context.Background()), + "a peer outside the chunk's sample group must keep rollout fail-closed") + + payload := []byte("authenticated-peer-transfer") + digest := sha256.Sum256(payload) + replica := S3BlobReplica{NodeID: "n2", Address: listener.Addr().String(), Suffrage: "voter"} + require.NoError(t, cluster.PushChunkBlob(context.Background(), replica, digest, payload, 10)) + fetched, err := cluster.FetchChunkBlob(context.Background(), replica, digest) + require.NoError(t, err) + require.Equal(t, payload, fetched) + + unauthenticated := NewGRPCS3BlobCluster("n1", members, "read-only-admin", "") + t.Cleanup(func() { require.NoError(t, unauthenticated.Close()) }) + err = unauthenticated.PushChunkBlob(context.Background(), replica, digest, payload, 11) + require.Equal(t, codes.Unauthenticated, status.Code(err)) +} + +func TestGRPCS3BlobClusterCapabilityPollHasPerPeerTimeout(t *testing.T) { + t.Parallel() + + listener, err := (&net.ListenConfig{}).Listen(context.Background(), "tcp", "127.0.0.1:0") + require.NoError(t, err) + server := grpc.NewServer() + pb.RegisterAdminServer(server, blockingS3BlobAdminServer{}) + serveDone := make(chan error, 1) + go func() { serveDone <- server.Serve(listener) }() + t.Cleanup(func() { + server.Stop() + serveErr := <-serveDone + require.True(t, serveErr == nil || errors.Is(serveErr, grpc.ErrServerStopped)) + }) + + members := staticS3BlobMembership{members: []kv.RaftMember{ + {NodeID: "n1", Address: "127.0.0.1:1", Suffrage: "voter"}, + {NodeID: "n2", Address: listener.Addr().String(), Suffrage: "voter"}, + }} + cluster := requireGRPCS3BlobCluster(t, "peer-secret") + cluster.members = members + cluster.capTimeout = 50 * time.Millisecond + start := time.Now() + require.False(t, cluster.AllPeersSupportS3BlobOffload(context.Background())) + require.Less(t, time.Since(start), time.Second) +} + +func TestGRPCS3BlobClusterFetchHasPeerTimeout(t *testing.T) { + t.Parallel() + + listener, err := (&net.ListenConfig{}).Listen(context.Background(), "tcp", "127.0.0.1:0") + require.NoError(t, err) + server := grpc.NewServer() + pb.RegisterS3BlobFetchServer(server, blockingS3BlobFetchServer{}) + serveDone := make(chan error, 1) + go func() { serveDone <- server.Serve(listener) }() + t.Cleanup(func() { + server.Stop() + serveErr := <-serveDone + require.True(t, serveErr == nil || errors.Is(serveErr, grpc.ErrServerStopped)) + }) + + cluster := requireGRPCS3BlobCluster(t, "peer-secret") + cluster.rpcTimeout = 50 * time.Millisecond + start := time.Now() + _, err = cluster.FetchChunkBlob(context.Background(), S3BlobReplica{ + NodeID: "n2", Address: listener.Addr().String(), Suffrage: "voter", + }, sha256.Sum256([]byte("missing"))) + require.Equal(t, codes.DeadlineExceeded, status.Code(err)) + require.Less(t, time.Since(start), time.Second) +} + +type blockingS3BlobAdminServer struct { + pb.UnimplementedAdminServer +} + +func (blockingS3BlobAdminServer) GetClusterOverview( + ctx context.Context, + _ *pb.GetClusterOverviewRequest, +) (*pb.GetClusterOverviewResponse, error) { + <-ctx.Done() + return nil, errors.WithStack(ctx.Err()) +} + +type blockingS3BlobFetchServer struct { + pb.UnimplementedS3BlobFetchServer +} + +func (blockingS3BlobFetchServer) FetchChunkBlob( + _ *pb.FetchChunkBlobRequest, + stream pb.S3BlobFetch_FetchChunkBlobServer, +) error { + <-stream.Context().Done() + return errors.WithStack(stream.Context().Err()) +} + +type staticS3BlobMembership struct { + members []kv.RaftMember + allMembers []kv.RaftMember +} + +func (m staticS3BlobMembership) RaftMembers(context.Context) ([]kv.RaftMember, error) { + if m.allMembers != nil { + return append([]kv.RaftMember(nil), m.allMembers...), nil + } + return append([]kv.RaftMember(nil), m.members...), nil +} + +func (m staticS3BlobMembership) RaftMembersForKey(context.Context, []byte) ([]kv.RaftMember, error) { + return append([]kv.RaftMember(nil), m.members...), nil +} + +func requireGRPCS3BlobCluster(t *testing.T, token string) *grpcS3BlobCluster { + t.Helper() + cluster, ok := NewGRPCS3BlobCluster("n1", nil, token, token+"-peer").(*grpcS3BlobCluster) + require.True(t, ok) + return cluster +} + +func newS3BlobM1TestServer( + t *testing.T, + observer S3BlobOffloadObserver, +) (*S3Server, *fakeS3BlobCluster, *mutableS3BlobLocalStore) { + t.Helper() + st := store.NewMVCCStore() + cluster := newFakeS3BlobCluster() + localResolver := &mutableS3BlobLocalStore{store: st} + server := NewS3Server( + nil, "", st, newLocalAdapterCoordinator(st), nil, + WithS3BlobOffloadEnabled(true), + WithS3BlobCluster(cluster), + WithS3BlobLocalStoreResolver(localResolver), + WithS3BlobOffloadObserver(observer), + withS3BlobOffloadGCReadyForTest(), + ) + return server, cluster, localResolver +} + +func putS3BlobM1Object(t *testing.T, server *S3Server, path string, payload []byte) { + t.Helper() + parts := strings.Split(strings.TrimPrefix(path, "/"), "/") + require.Len(t, parts, 2) + rec := httptest.NewRecorder() + server.handle(rec, newS3TestRequest(http.MethodPut, "/"+parts[0], nil)) + require.Equal(t, http.StatusOK, rec.Code) + rec = httptest.NewRecorder() + server.handle(rec, newS3TestRequest(http.MethodPut, path, bytes.NewReader(payload))) + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) +} + +type mutableS3BlobLocalStore struct { + mu sync.RWMutex + store store.MVCCStore + routeKey []byte +} + +func (r *mutableS3BlobLocalStore) LocalStoreForKey(key []byte) (store.MVCCStore, bool) { + r.mu.Lock() + defer r.mu.Unlock() + r.routeKey = bytes.Clone(key) + return r.store, r.store != nil +} + +func (r *mutableS3BlobLocalStore) set(st store.MVCCStore) { + r.mu.Lock() + r.store = st + r.mu.Unlock() +} + +func (r *mutableS3BlobLocalStore) lastRouteKey() []byte { + r.mu.RLock() + defer r.mu.RUnlock() + return bytes.Clone(r.routeKey) +} + +type blockingS3BlobStore struct { + store.MVCCStore + started chan struct{} + release chan struct{} + once sync.Once +} + +func (s *blockingS3BlobStore) ApplyMutationsPreservingLastCommitTS(ctx context.Context, mutations []*store.KVPairMutation, readKeys [][]byte, startTS, commitTS uint64) error { + s.once.Do(func() { close(s.started) }) + select { + case <-ctx.Done(): + return errors.WithStack(ctx.Err()) + case <-s.release: + } + preserving, ok := s.MVCCStore.(interface { + ApplyMutationsPreservingLastCommitTS(context.Context, []*store.KVPairMutation, [][]byte, uint64, uint64) error + }) + if !ok { + return errors.New("test store does not preserve last commit timestamp") + } + return preserving.ApplyMutationsPreservingLastCommitTS(ctx, mutations, readKeys, startTS, commitTS) +} + +type fakeS3BlobCluster struct { + mu sync.Mutex + self string + replicas []S3BlobReplica + blobs map[string]map[[sha256.Size]byte][]byte + pushErr error + pushes int + fetches int + routeKey []byte + pushStarted chan struct{} + pushRelease chan struct{} +} + +func newFakeS3BlobCluster() *fakeS3BlobCluster { + return &fakeS3BlobCluster{ + self: "n1", + replicas: []S3BlobReplica{ + {NodeID: "n1", Address: "n1:50051", Suffrage: "voter"}, + {NodeID: "n2", Address: "n2:50051", Suffrage: "voter"}, + {NodeID: "n3", Address: "n3:50051", Suffrage: "voter"}, + }, + blobs: map[string]map[[sha256.Size]byte][]byte{}, + } +} + +func (c *fakeS3BlobCluster) AllPeersSupportS3BlobOffload(context.Context) bool { return true } +func (c *fakeS3BlobCluster) SelfNodeID() string { return c.self } +func (c *fakeS3BlobCluster) Close() error { return nil } + +func (c *fakeS3BlobCluster) ReplicasForChunk(_ context.Context, key []byte) ([]S3BlobReplica, error) { + c.mu.Lock() + defer c.mu.Unlock() + c.routeKey = bytes.Clone(key) + return append([]S3BlobReplica(nil), c.replicas...), nil +} + +func (c *fakeS3BlobCluster) PushChunkBlob(ctx context.Context, replica S3BlobReplica, digest [sha256.Size]byte, payload []byte, _ uint64) error { + if c.pushStarted != nil { + c.pushStarted <- struct{}{} + } + if c.pushRelease != nil { + select { + case <-ctx.Done(): + return errors.WithStack(ctx.Err()) + case <-c.pushRelease: + } + } + c.mu.Lock() + defer c.mu.Unlock() + c.pushes++ + if c.pushErr != nil { + return c.pushErr + } + if c.blobs[replica.NodeID] == nil { + c.blobs[replica.NodeID] = map[[sha256.Size]byte][]byte{} + } + c.blobs[replica.NodeID][digest] = bytes.Clone(payload) + return nil +} + +func (c *fakeS3BlobCluster) FetchChunkBlob(_ context.Context, replica S3BlobReplica, digest [sha256.Size]byte) ([]byte, error) { + c.mu.Lock() + defer c.mu.Unlock() + c.fetches++ + payload := c.blobs[replica.NodeID][digest] + if payload == nil { + return nil, status.Error(codes.NotFound, "missing") + } + return bytes.Clone(payload), nil +} + +func (c *fakeS3BlobCluster) pushCount() int { + c.mu.Lock() + defer c.mu.Unlock() + return c.pushes +} + +func (c *fakeS3BlobCluster) lastRouteKey() []byte { + c.mu.Lock() + defer c.mu.Unlock() + return bytes.Clone(c.routeKey) +} + +func (c *fakeS3BlobCluster) fetchCount() int { + c.mu.Lock() + defer c.mu.Unlock() + return c.fetches +} + +func (c *fakeS3BlobCluster) clearBlobs() { + c.mu.Lock() + c.blobs = map[string]map[[sha256.Size]byte][]byte{} + c.mu.Unlock() +} diff --git a/adapter/s3_blob_offload.go b/adapter/s3_blob_offload.go index fb98dadd5..88ccfd1b9 100644 --- a/adapter/s3_blob_offload.go +++ b/adapter/s3_blob_offload.go @@ -6,6 +6,8 @@ import ( "os" "strconv" "strings" + + "github.com/cockroachdb/errors" ) const ( @@ -13,14 +15,18 @@ const ( // key used by mixed-version PUT admission before writing chunkref metadata. S3BlobOffloadCapabilityName = "feature_s3_blob_offload" - s3BlobOffloadEnvVar = "ELASTICKV_S3_BLOB_OFFLOAD" + s3BlobOffloadEnvVar = "ELASTICKV_S3_BLOB_OFFLOAD" + s3BlobMinReplicasEnvVar = "ELASTICKV_S3_CHUNKBLOB_MIN_REPLICAS" s3BlobOffloadModeLegacy = "legacy" s3BlobOffloadModeOffload = "offload" + s3BlobMinimumDurableCopies = 2 + s3BlobOffloadReasonFlagDisabled = "flag_disabled" s3BlobOffloadReasonCapabilityMissing = "capability_missing" s3BlobOffloadReasonDataPathDisabled = "data_path_disabled" + s3BlobOffloadReasonGCNotReady = "gc_not_ready" s3BlobOffloadReasonEnabled = "enabled" ) @@ -46,11 +52,10 @@ type s3BlobOffloadDecision struct { reason string } -// S3BlobOffloadLocalCapability reports whether this binary can safely serve the -// full offloaded S3 chunkref/chunkblob data path. The scaffolding PR keeps this -// false so future leaders do not mistake these nodes for offload readers. +// S3BlobOffloadLocalCapability reports whether this binary can serve the M1 +// offloaded PUT/GET path, including durable peer push and proxy-on-miss reads. func S3BlobOffloadLocalCapability() bool { - return false + return true } func WithS3BlobOffloadEnabled(enabled bool) S3ServerOption { @@ -80,7 +85,64 @@ func WithS3BlobOffloadObserver(observer S3BlobOffloadObserver) S3ServerOption { } } -func newS3BlobOffloadEnabledFromEnv() bool { +func WithS3BlobCluster(cluster S3BlobCluster) S3ServerOption { + return func(server *S3Server) { + if server == nil { + return + } + server.blobCluster = cluster + server.blobOffloadChecker = cluster + } +} + +func WithS3BlobLocalStoreResolver(resolver S3BlobLocalStoreResolver) S3ServerOption { + return func(server *S3Server) { + if server == nil { + return + } + server.blobLocalStores = resolver + } +} + +func WithS3BlobMinReplicas(minReplicas int) S3ServerOption { + return func(server *S3Server) { + if server == nil { + return + } + server.blobMinReplicas = minReplicas + } +} + +func WithS3BlobPushBlocked(blocked func() bool) S3ServerOption { + return func(server *S3Server) { + if server == nil { + return + } + server.blobPushBlocked = blocked + } +} + +// S3BlobMinReplicasFromEnv returns zero when the dynamic Raft-quorum default +// should be used. Explicit values below two are rejected because a leader-only +// chunkblob is weaker than the legacy Raft data path. +func S3BlobMinReplicasFromEnv() (int, error) { + raw := strings.TrimSpace(os.Getenv(s3BlobMinReplicasEnvVar)) + if raw == "" { + return 0, nil + } + minReplicas, err := strconv.Atoi(raw) + if err != nil { + return 0, errors.Wrap(err, "parse S3 chunkblob minimum replicas") + } + if minReplicas < s3BlobMinimumDurableCopies { + return 0, errors.New("S3 chunkblob minimum replicas must be at least 2") + } + return minReplicas, nil +} + +// S3BlobOffloadEnabledFromEnv returns the local rollout flag used by both the +// S3 write path and the Admin capability advertisement. +func S3BlobOffloadEnabledFromEnv() bool { raw, ok := os.LookupEnv(s3BlobOffloadEnvVar) if !ok { return false @@ -107,12 +169,20 @@ func (s *S3Server) s3BlobOffloadDecision(ctx context.Context) s3BlobOffloadDecis if !S3BlobOffloadLocalCapability() { return s3BlobOffloadDecision{mode: s3BlobOffloadModeLegacy, reason: s3BlobOffloadReasonDataPathDisabled} } + // M1 has a complete PUT/GET path, but content-addressed chunkblobs cannot + // be enabled operationally until M3 installs reference counting, the grace + // queue, and the orphan scanner. M3 sets this readiness only after those + // workers are wired into the server lifecycle. + if !s.blobOffloadGCReady { + return s3BlobOffloadDecision{mode: s3BlobOffloadModeLegacy, reason: s3BlobOffloadReasonGCNotReady} + } return s3BlobOffloadDecision{mode: s3BlobOffloadModeOffload, reason: s3BlobOffloadReasonEnabled} } -func (s *S3Server) observeS3BlobOffloadDecision(ctx context.Context) { +func (s *S3Server) observeS3BlobOffloadDecision(ctx context.Context) s3BlobOffloadDecision { decision := s.s3BlobOffloadDecision(ctx) if s != nil && s.blobOffloadObserver != nil { s.blobOffloadObserver.ObserveS3BlobOffloadDecision(decision.mode, decision.reason) } + return decision } diff --git a/adapter/s3_blob_offload_test.go b/adapter/s3_blob_offload_test.go index 3d3a8fb48..b75f77796 100644 --- a/adapter/s3_blob_offload_test.go +++ b/adapter/s3_blob_offload_test.go @@ -28,6 +28,7 @@ func TestS3BlobOffloadEnvCanEnableGate(t *testing.T) { server := NewS3Server(nil, "", st, newLocalAdapterCoordinator(st), nil) require.True(t, server.blobOffloadEnabled) + require.False(t, server.blobOffloadGCReady) } func TestS3BlobOffloadDecisionFailsClosed(t *testing.T) { @@ -48,7 +49,27 @@ func TestS3BlobOffloadDecisionFailsClosed(t *testing.T) { require.Equal(t, s3BlobOffloadReasonCapabilityMissing, decision.reason) } -func TestS3BlobOffloadDecisionRefusesUntilDataPathIsEnabled(t *testing.T) { +func TestS3BlobOffloadDecisionEnablesImplementedDataPath(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + server := NewS3Server( + nil, + "", + st, + newLocalAdapterCoordinator(st), + nil, + WithS3BlobOffloadEnabled(true), + WithS3BlobOffloadCapabilityChecker(alwaysS3BlobOffloadCapable{}), + withS3BlobOffloadGCReadyForTest(), + ) + + decision := server.s3BlobOffloadDecision(context.Background()) + require.Equal(t, s3BlobOffloadModeOffload, decision.mode) + require.Equal(t, s3BlobOffloadReasonEnabled, decision.reason) +} + +func TestS3BlobOffloadDecisionFailsClosedWithoutGC(t *testing.T) { t.Parallel() st := store.NewMVCCStore() @@ -64,7 +85,7 @@ func TestS3BlobOffloadDecisionRefusesUntilDataPathIsEnabled(t *testing.T) { decision := server.s3BlobOffloadDecision(context.Background()) require.Equal(t, s3BlobOffloadModeLegacy, decision.mode) - require.Equal(t, s3BlobOffloadReasonDataPathDisabled, decision.reason) + require.Equal(t, s3BlobOffloadReasonGCNotReady, decision.reason) } func TestS3Server_PutObjectBlobOffloadFlagFallsBackToLegacyWithoutCapability(t *testing.T) { @@ -111,16 +132,33 @@ func (alwaysS3BlobOffloadCapable) AllPeersSupportS3BlobOffload(context.Context) return true } +func withS3BlobOffloadGCReadyForTest() S3ServerOption { + return func(server *S3Server) { + if server != nil { + server.blobOffloadGCReady = true + } + } +} + type recordingS3BlobOffloadObserver struct { - decisions []s3BlobOffloadDecision + decisions []s3BlobOffloadDecision + replicationDegraded int + shaMismatch int + unrecoverable int } func (o *recordingS3BlobOffloadObserver) ObserveS3BlobOffloadDecision(mode, reason string) { o.decisions = append(o.decisions, s3BlobOffloadDecision{mode: mode, reason: reason}) } -func (o *recordingS3BlobOffloadObserver) ObserveS3ChunkBlobReplicationDegraded() {} +func (o *recordingS3BlobOffloadObserver) ObserveS3ChunkBlobReplicationDegraded() { + o.replicationDegraded++ +} -func (o *recordingS3BlobOffloadObserver) ObserveS3ChunkBlobSHAMismatch() {} +func (o *recordingS3BlobOffloadObserver) ObserveS3ChunkBlobSHAMismatch() { + o.shaMismatch++ +} -func (o *recordingS3BlobOffloadObserver) ObserveS3ChunkBlobUnrecoverable() {} +func (o *recordingS3BlobOffloadObserver) ObserveS3ChunkBlobUnrecoverable() { + o.unrecoverable++ +} diff --git a/adapter/s3_blob_read.go b/adapter/s3_blob_read.go new file mode 100644 index 000000000..0bd21aeaf --- /dev/null +++ b/adapter/s3_blob_read.go @@ -0,0 +1,313 @@ +package adapter + +import ( + "bytes" + "context" + "crypto/sha256" + "math" + rand "math/rand/v2" + + "github.com/bootjp/elastickv/internal/s3keys" + "github.com/bootjp/elastickv/store" + "github.com/cockroachdb/errors" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +func (s *S3Server) ensureS3ObjectRangeLocal( + ctx context.Context, + bucket string, + generation uint64, + objectKey string, + manifest *s3ObjectManifest, + readTS uint64, + offset int64, + length int64, +) error { + remaining := length + position := int64(0) + for _, part := range manifest.Parts { + for chunkIndex, chunkSize := range part.ChunkSizes { + chunkEnd := position + int64(chunkSize) //nolint:gosec // Chunk size is bounded by s3ChunkSize. + if remaining > 0 && chunkEnd > offset { + index, err := uint64FromInt(chunkIndex) + if err != nil { + return err + } + if _, err := s.readS3ObjectChunk(ctx, bucket, generation, objectKey, manifest.UploadID, part, index, chunkSize, readTS); err != nil { + return err + } + consumedStart := position + if consumedStart < offset { + consumedStart = offset + } + consumed := chunkEnd - consumedStart + if consumed > remaining { + consumed = remaining + } + remaining -= consumed + } + position = chunkEnd + if remaining <= 0 { + return nil + } + } + } + return nil +} + +func (s *S3Server) readS3ObjectChunk( + ctx context.Context, + bucket string, + generation uint64, + objectKey string, + uploadID string, + part s3ObjectPart, + chunkIndex uint64, + expectedSize uint64, + readTS uint64, +) ([]byte, error) { + if !part.Offloaded { + return s.readLegacyS3ObjectChunk(ctx, bucket, generation, objectKey, uploadID, part, chunkIndex, readTS) + } + refKey := s3keys.VersionedChunkRefKey( + bucket, generation, objectKey, uploadID, part.PartNo, chunkIndex, part.ChunkRefVersion, + ) + ref, refCommitTS, err := s.loadS3ChunkRef( + ctx, bucket, generation, objectKey, uploadID, + part.PartNo, chunkIndex, part.ChunkRefVersion, expectedSize, readTS, + ) + if err != nil { + return nil, err + } + payload, found, err := s.localS3ChunkBlob(ctx, refKey, ref.ContentSHA256) + if err != nil || found { + return payload, err + } + return s.fetchAndStoreS3ChunkBlob(ctx, refKey, ref, refCommitTS) +} + +func (s *S3Server) readLegacyS3ObjectChunk( + ctx context.Context, + bucket string, + generation uint64, + objectKey string, + uploadID string, + part s3ObjectPart, + chunkIndex uint64, + readTS uint64, +) ([]byte, error) { + key := s3keys.VersionedBlobKey(bucket, generation, objectKey, uploadID, part.PartNo, chunkIndex, part.PartVersion) + payload, err := s.store.GetAt(ctx, key, readTS) + return payload, errors.WithStack(err) +} + +func (s *S3Server) loadS3ChunkRef( + ctx context.Context, + bucket string, + generation uint64, + objectKey string, + uploadID string, + partNo uint64, + chunkIndex uint64, + chunkRefVersion uint64, + expectedSize uint64, + readTS uint64, +) (s3keys.ChunkRefValue, uint64, error) { + refKey := s3keys.VersionedChunkRefKey(bucket, generation, objectKey, uploadID, partNo, chunkIndex, chunkRefVersion) + rawRef, err := s.store.GetAt(ctx, refKey, readTS) + if err != nil { + return s3keys.ChunkRefValue{}, 0, errors.Wrap(err, "read s3 chunkref") + } + ref, ok := s3keys.DecodeChunkRefValue(rawRef) + if !ok { + return s3keys.ChunkRefValue{}, 0, errors.New("decode s3 chunkref: invalid value") + } + if ref.Size != expectedSize { + s.observeS3ChunkBlobMismatch() + return s3keys.ChunkRefValue{}, 0, errors.WithStack(errors.Newf( + "s3 chunkref size %d does not match manifest size %d", + ref.Size, + expectedSize, + )) + } + refCommitTS, exists, err := s.store.LatestCommitTS(ctx, refKey) + if err != nil { + return s3keys.ChunkRefValue{}, 0, errors.Wrap(err, "read s3 chunkref commit timestamp") + } + if !exists || refCommitTS == 0 { + return s3keys.ChunkRefValue{}, 0, errors.New("s3 chunkref commit timestamp is unavailable") + } + return ref, refCommitTS, nil +} + +func (s *S3Server) fetchAndStoreS3ChunkBlob( + ctx context.Context, + refKey []byte, + ref s3keys.ChunkRefValue, + refCommitTS uint64, +) ([]byte, error) { + payload, err := s.fetchS3ChunkBlob(ctx, refKey, ref) + if err != nil { + return nil, err + } + if uint64(len(payload)) != ref.Size { //nolint:gosec // len is bounded by s3ChunkSize. + s.observeS3ChunkBlobMismatch() + return nil, errors.New("fetched s3 chunkblob size does not match chunkref") + } + if err := s.storeFetchedS3ChunkBlob(ctx, refKey, ref.ContentSHA256, payload, refCommitTS); err != nil { + return nil, err + } + return payload, nil +} + +func (s *S3Server) localS3ChunkBlob(ctx context.Context, routeKey []byte, digest [s3ChunkBlobSHA256Bytes]byte) ([]byte, bool, error) { + if s == nil || s.blobLocalStores == nil { + return nil, false, s3BlobUnavailable("s3 chunkblob local store is not configured") + } + key := s3keys.ChunkBlobKey(digest) + localStore, ok := s.blobLocalStores.LocalStoreForKey(routeKey) + if !ok || localStore == nil { + return nil, false, s3BlobUnavailable("s3 chunkblob local store is unavailable") + } + payload, err := localStore.GetAt(ctx, key, math.MaxUint64) + if errors.Is(err, store.ErrKeyNotFound) { + return nil, false, nil + } + if err != nil { + return nil, false, errors.WithStack(err) + } + actual := sha256.Sum256(payload) + if !bytes.Equal(actual[:], digest[:]) { + s.observeS3ChunkBlobMismatch() + return nil, false, nil + } + return payload, true, nil +} + +func (s *S3Server) fetchS3ChunkBlob(ctx context.Context, routeKey []byte, ref s3keys.ChunkRefValue) ([]byte, error) { + if s == nil || s.blobCluster == nil { + return nil, s3BlobUnavailable("s3 chunkblob peer client is not configured") + } + currentReplicas, discoveryErr := s.blobCluster.ReplicasForChunk(ctx, routeKey) + writeReplicas := writeTimeS3BlobReplicas(ref.ReplicaPeers) + if discoveryErr != nil && len(writeReplicas) == 0 { + return nil, errors.Wrap(discoveryErr, "resolve s3 chunkblob fetch replicas") + } + peers := orderedS3BlobPeers(writeReplicas, currentReplicas, s.blobCluster.SelfNodeID(), ref.SourcePeer) + payload, err := s.fetchS3ChunkBlobFromPeers(ctx, peers, ref.ContentSHA256) + if err == nil { + return payload, nil + } + if s.blobOffloadObserver != nil { + s.blobOffloadObserver.ObserveS3ChunkBlobUnrecoverable() + } + return nil, err +} + +func writeTimeS3BlobReplicas(peers []s3keys.ChunkRefPeer) []S3BlobReplica { + replicas := make([]S3BlobReplica, 0, len(peers)) + for _, peer := range peers { + if peer.NodeID == "" || peer.Address == "" { + continue + } + replicas = append(replicas, S3BlobReplica{NodeID: peer.NodeID, Address: peer.Address}) + } + return replicas +} + +func orderedS3BlobPeers(writeReplicas []S3BlobReplica, currentReplicas []S3BlobReplica, selfNodeID, sourcePeer string) []S3BlobReplica { + peers := make([]S3BlobReplica, 0, len(writeReplicas)+len(currentReplicas)) + sources := make([]S3BlobReplica, 0) + seen := make(map[string]struct{}, len(writeReplicas)+len(currentReplicas)) + for _, replica := range currentReplicas { + if replica.NodeID == sourcePeer { + sources = appendUniqueS3BlobPeer(sources, seen, replica, selfNodeID) + } + } + for _, replica := range writeReplicas { + if replica.NodeID == sourcePeer { + sources = appendUniqueS3BlobPeer(sources, seen, replica, selfNodeID) + } + } + for _, replica := range writeReplicas { + if replica.NodeID != sourcePeer { + peers = appendUniqueS3BlobPeer(peers, seen, replica, selfNodeID) + } + } + for _, replica := range currentReplicas { + if replica.NodeID != sourcePeer { + peers = appendUniqueS3BlobPeer(peers, seen, replica, selfNodeID) + } + } + rand.Shuffle(len(peers), func(i, j int) { peers[i], peers[j] = peers[j], peers[i] }) + return append(sources, peers...) +} + +func appendUniqueS3BlobPeer(dst []S3BlobReplica, seen map[string]struct{}, replica S3BlobReplica, selfNodeID string) []S3BlobReplica { + if replica.NodeID == selfNodeID || replica.NodeID == "" || replica.Address == "" { + return dst + } + identity := replica.NodeID + "\x00" + replica.Address + if _, duplicate := seen[identity]; duplicate { + return dst + } + seen[identity] = struct{}{} + return append(dst, replica) +} + +func (s *S3Server) fetchS3ChunkBlobFromPeers( + ctx context.Context, + peers []S3BlobReplica, + digest [s3ChunkBlobSHA256Bytes]byte, +) ([]byte, error) { + for _, peer := range peers { + payload, fetchErr := s.blobCluster.FetchChunkBlob(ctx, peer, digest) + if fetchErr == nil { + return payload, nil + } + if status.Code(fetchErr) == codes.InvalidArgument { + s.observeS3ChunkBlobMismatch() + } + if ctx.Err() != nil { + return nil, errors.WithStack(ctx.Err()) + } + } + return nil, errors.New("s3 chunkblob is unavailable on every replica") +} + +func (s *S3Server) storeFetchedS3ChunkBlob(ctx context.Context, routeKey []byte, digest [s3ChunkBlobSHA256Bytes]byte, payload []byte, commitTS uint64) error { + localStore, ok := s.blobLocalStores.LocalStoreForKey(routeKey) + if !ok || localStore == nil { + return s3BlobUnavailable("s3 chunkblob local store is unavailable") + } + repairTS, err := s.nextTxnCommitTS(ctx, commitTS) + if err != nil { + return errors.Wrap(err, "allocate s3 chunkblob repair timestamp") + } + server := NewS3BlobFetchServer( + localStore, + s.blobOffloadObserver, + WithS3BlobFetchClock(s.clock()), + WithS3BlobFetchPushBlocked(s.blobPushBlocked), + ) + return server.storeChunkBlob(ctx, digest, payload, repairTS) +} + +func (s *S3Server) observeS3ChunkBlobMismatch() { + if s != nil && s.blobOffloadObserver != nil { + s.blobOffloadObserver.ObserveS3ChunkBlobSHAMismatch() + } +} + +func s3ManifestHasOffloadedParts(manifest *s3ObjectManifest) bool { + if manifest == nil { + return false + } + for _, part := range manifest.Parts { + if part.Offloaded { + return true + } + } + return false +} diff --git a/adapter/s3_blob_replicator.go b/adapter/s3_blob_replicator.go new file mode 100644 index 000000000..fb4820d6d --- /dev/null +++ b/adapter/s3_blob_replicator.go @@ -0,0 +1,236 @@ +package adapter + +import ( + "context" + "log/slog" + + "github.com/bootjp/elastickv/internal/s3keys" + "github.com/bootjp/elastickv/store" + "github.com/cockroachdb/errors" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +const s3BlobVoterSuffrage = "voter" + +type s3BlobReplicationResult struct { + replica S3BlobReplica + err error + local bool +} + +type s3BlobReplicationPlan struct { + localStore store.MVCCStore + replicas []S3BlobReplica + selfNodeID string + durableTarget int + voterTarget int + degraded bool +} + +type s3BlobReplicationSummary struct { + durable int + voterDurable int + localDurable bool + replicas []S3BlobReplica +} + +func (s *S3Server) persistS3ChunkBlob( + ctx context.Context, + routeKey []byte, + digest [s3ChunkBlobSHA256Bytes]byte, + payload []byte, + commitTS uint64, +) (s3keys.ChunkRefValue, error) { + plan, err := s.prepareS3BlobReplication(ctx, routeKey) + if err != nil { + return s3keys.ChunkRefValue{}, err + } + replicationCtx, cancel := context.WithCancel(ctx) + defer cancel() + results := s.startS3BlobReplication(replicationCtx, plan, digest, payload, commitTS) + summary := s.collectS3BlobReplication( + ctx, results, len(plan.replicas), plan.durableTarget, plan.voterTarget, + ) + if !summary.localDurable || summary.durable < plan.durableTarget || summary.voterDurable < plan.voterTarget { + return s3keys.ChunkRefValue{}, s3BlobUnavailablef( + "s3 chunkblob reached %d durable replicas (%d voters); require %d including local and %d voters", + summary.durable, + summary.voterDurable, + plan.durableTarget, + plan.voterTarget, + ) + } + if plan.degraded && s.blobOffloadObserver != nil { + s.blobOffloadObserver.ObserveS3ChunkBlobReplicationDegraded() + } + return s3keys.ChunkRefValue{ + ContentSHA256: digest, + Size: uint64(len(payload)), //nolint:gosec // Payload is bounded by s3ChunkSize. + SourcePeer: plan.selfNodeID, + ReplicaPeers: s3ChunkRefPeers(summary.replicas), + }, nil +} + +func s3ChunkRefPeers(replicas []S3BlobReplica) []s3keys.ChunkRefPeer { + peers := make([]s3keys.ChunkRefPeer, 0, len(replicas)) + for _, replica := range replicas { + peers = append(peers, s3keys.ChunkRefPeer{NodeID: replica.NodeID, Address: replica.Address}) + } + return peers +} + +func (s *S3Server) prepareS3BlobReplication( + ctx context.Context, + routeKey []byte, +) (s3BlobReplicationPlan, error) { + if s == nil || s.blobCluster == nil || s.blobLocalStores == nil { + return s3BlobReplicationPlan{}, s3BlobUnavailable("s3 blob offload data path is not configured") + } + localStore, ok := s.blobLocalStores.LocalStoreForKey(routeKey) + if !ok || localStore == nil { + return s3BlobReplicationPlan{}, s3BlobUnavailable("s3 chunkblob local store is unavailable") + } + replicas, err := s.blobCluster.ReplicasForChunk(ctx, routeKey) + if err != nil { + return s3BlobReplicationPlan{}, s3BlobUnavailablef("resolve s3 chunkblob replicas: %v", err) + } + selfNodeID := s.blobCluster.SelfNodeID() + durableTarget, voterTarget, degraded, err := s.s3BlobDurableTarget(replicas, selfNodeID) + if err != nil { + return s3BlobReplicationPlan{}, err + } + return s3BlobReplicationPlan{ + localStore: localStore, + replicas: replicas, + selfNodeID: selfNodeID, + durableTarget: durableTarget, + voterTarget: voterTarget, + degraded: degraded, + }, nil +} + +func (s *S3Server) startS3BlobReplication( + ctx context.Context, + plan s3BlobReplicationPlan, + digest [s3ChunkBlobSHA256Bytes]byte, + payload []byte, + commitTS uint64, +) <-chan s3BlobReplicationResult { + results := make(chan s3BlobReplicationResult, len(plan.replicas)) + for _, replica := range plan.replicas { + if replica.NodeID == plan.selfNodeID { + go func() { + fetchServer := NewS3BlobFetchServer( + plan.localStore, + s.blobOffloadObserver, + WithS3BlobFetchClock(s.clock()), + WithS3BlobFetchPushBlocked(s.blobPushBlocked), + ) + writeErr := fetchServer.storeChunkBlob(ctx, digest, payload, commitTS) + results <- s3BlobReplicationResult{replica: replica, err: writeErr, local: true} + }() + continue + } + go func() { + pushErr := s.blobCluster.PushChunkBlob(ctx, replica, digest, payload, commitTS) + results <- s3BlobReplicationResult{replica: replica, err: pushErr} + }() + } + return results +} + +func (s *S3Server) collectS3BlobReplication( + ctx context.Context, + results <-chan s3BlobReplicationResult, + resultCount int, + durableTarget int, + voterTarget int, +) s3BlobReplicationSummary { + summary := s3BlobReplicationSummary{replicas: make([]S3BlobReplica, 0, durableTarget)} + for range resultCount { + result := <-results + if result.err != nil { + slog.WarnContext(ctx, "s3 chunkblob replica write failed", + "node_id", result.replica.NodeID, + "address", result.replica.Address, + "local", result.local, + "err", result.err, + ) + continue + } + summary.durable++ + if result.replica.Suffrage == s3BlobVoterSuffrage { + summary.voterDurable++ + } + summary.localDurable = summary.localDurable || result.local + summary.replicas = append(summary.replicas, result.replica) + if summary.localDurable && summary.durable >= durableTarget && summary.voterDurable >= voterTarget { + return summary + } + } + return summary +} + +func (s *S3Server) s3BlobDurableTarget(replicas []S3BlobReplica, selfNodeID string) (target int, voterTarget int, degraded bool, err error) { + voters, err := validateS3BlobMembership(replicas, selfNodeID) + if err != nil { + return 0, 0, false, err + } + // M1 has no follower apply/backfill worker yet. Requiring every current + // member keeps each node's offline logical backup self-contained; M2 can + // restore quorum acknowledgement after it guarantees eventual local fetch. + voterTarget = voters + target = len(replicas) + if target < s3BlobMinimumDurableCopies { + return 0, 0, false, errors.New("s3 chunkblob durable target is below two") + } + if s == nil || s.blobMinReplicas <= len(replicas) { + return target, voterTarget, false, nil + } + // Membership has already shrunk below a previously configured target. + // Require every remaining member and emit the degradation signal. A mere + // unreachable peer does not reduce len(replicas), so M1 still fails closed + // during transient outages. + return target, voterTarget, true, nil +} + +func validateS3BlobMembership(replicas []S3BlobReplica, selfNodeID string) (int, error) { + if len(replicas) < s3BlobMinimumDurableCopies { + return 0, s3BlobUnavailable("s3 chunkblob replication requires at least two members") + } + selfFound := false + selfVoter := false + voters := 0 + for _, replica := range replicas { + if replica.NodeID == selfNodeID { + selfFound = true + selfVoter = replica.Suffrage == s3BlobVoterSuffrage + } + switch replica.Suffrage { + case s3BlobVoterSuffrage: + voters++ + case "learner": + default: + return 0, s3BlobUnavailablef("s3 chunkblob member %q has unknown suffrage %q", replica.NodeID, replica.Suffrage) + } + } + if !selfFound { + return 0, s3BlobUnavailable("s3 chunkblob membership does not include this node") + } + if !selfVoter { + return 0, s3BlobUnavailable("s3 chunkblob local member is not a voter") + } + if voters < s3BlobMinimumDurableCopies { + return 0, s3BlobUnavailable("s3 chunkblob replication requires at least two voters") + } + return voters, nil +} + +func s3BlobUnavailable(message string) error { + return errors.WithStack(status.Error(codes.Unavailable, message)) +} + +func s3BlobUnavailablef(format string, args ...any) error { + return errors.WithStack(status.Errorf(codes.Unavailable, format, args...)) +} diff --git a/adapter/s3_chunk_batch_test.go b/adapter/s3_chunk_batch_test.go index 77ee67ab5..af4523be4 100644 --- a/adapter/s3_chunk_batch_test.go +++ b/adapter/s3_chunk_batch_test.go @@ -1,6 +1,8 @@ package adapter import ( + "context" + "crypto/sha256" "strings" "testing" @@ -19,6 +21,16 @@ import ( // silently grows past it. const raftMaxSizePerMsgPostPR593 = 4 << 20 +type recordingS3MetaCoordinator struct { + stubAdapterCoordinator + dispatchSizes []int +} + +func (c *recordingS3MetaCoordinator) Dispatch(_ context.Context, req *kv.OperationGroup[kv.OP]) (*kv.CoordinateResponse, error) { + c.dispatchSizes = append(c.dispatchSizes, len(req.Elems)) + return &kv.CoordinateResponse{}, nil +} + // TestS3ChunkBatchFitsInRaftMaxSize is the byte-budget invariant the // s3ChunkBatchOps comment advertises: a worst-case S3 PutObject / // UploadPart batch must encode strictly under @@ -139,6 +151,62 @@ func TestS3MetaBatchFitsInRaftMaxSize(t *testing.T) { ) } +func TestS3OffloadedChunkRefBatchFitsInRaftMaxSize(t *testing.T) { + t.Parallel() + + bucket := "test-bucket" + objectKey := strings.Repeat("a", 1024) + uploadID := "upload-12345678901234567890" + digest := sha256.Sum256([]byte("chunk")) + refValue, err := s3keys.EncodeChunkRefValue(s3keys.ChunkRefValue{ + ContentSHA256: digest, + Size: s3ChunkSize, + SourcePeer: "node-1", + ReplicaPeers: []s3keys.ChunkRefPeer{ + {NodeID: "node-1", Address: "192.168.0.210:50051"}, + {NodeID: "node-2", Address: "192.168.0.211:50051"}, + {NodeID: "node-3", Address: "192.168.0.212:50051"}, + {NodeID: "node-4", Address: "192.168.0.213:50051"}, + {NodeID: "node-5", Address: "192.168.0.214:50051"}, + }, + }) + require.NoError(t, err) + + muts := make([]*pb.Mutation, 0, s3MetaBatchOps) + for i := uint64(0); i < uint64(s3MetaBatchOps); i++ { + muts = append(muts, &pb.Mutation{ + Op: pb.Op_PUT, + Key: s3keys.VersionedChunkRefKey(bucket, 1, objectKey, uploadID, 1, i, 99), + Value: refValue, + }) + } + encoded, err := proto.Marshal(&pb.Request{Ts: 1234567890, Mutations: muts}) + require.NoError(t, err) + require.Less(t, len(encoded)+1, raftMaxSizePerMsgPostPR593) +} + +func TestS3OffloadedChunkRefsFlushAtMetaBatchLimit(t *testing.T) { + t.Parallel() + + coordinator := &recordingS3MetaCoordinator{} + uploader := &s3ChunkUploader{ + server: &S3Server{coordinator: coordinator}, + ctx: context.Background(), + cfg: s3ChunkUploadConfig{offloaded: true}, + } + for i := 0; i < s3MetaBatchOps-1; i++ { + uploader.pendingBatch = append(uploader.pendingBatch, &kv.Elem[kv.OP]{Op: kv.Put}) + } + require.NoError(t, uploader.flushFullBatch()) + require.Empty(t, coordinator.dispatchSizes) + require.Len(t, uploader.pendingBatch, s3MetaBatchOps-1) + + uploader.pendingBatch = append(uploader.pendingBatch, &kv.Elem[kv.OP]{Op: kv.Put}) + require.NoError(t, uploader.flushFullBatch()) + require.Equal(t, []int{s3MetaBatchOps}, coordinator.dispatchSizes) + require.Empty(t, uploader.pendingBatch) +} + // TestAppendPartBlobKeys_FlushFiresEveryS3MetaBatchOps is the // regression guard for the slice-by-value bug Gemini caught: a // previous version of appendPartBlobKeys took `pending` by value, so @@ -192,3 +260,27 @@ func TestAppendPartBlobKeys_FlushFiresEveryS3MetaBatchOps(t *testing.T) { require.Len(t, pending, 7, "trailing 7 entries should remain for the caller's final flush()") } + +func TestAppendPartBlobKeys_UsesVersionedChunkRefsForOffloadedPart(t *testing.T) { + t.Parallel() + + pending := make([]*kv.Elem[kv.OP], 0, 1) + part := s3ObjectPart{ + PartNo: 7, + ChunkSizes: []uint64{123}, + ChunkRefVersion: 99, + Offloaded: true, + } + + ok := (*S3Server)(nil).appendPartBlobKeys( + &pending, "bucket", 3, "object", "upload", part, func() {}, + ) + require.True(t, ok) + require.Len(t, pending, 1) + require.Equal(t, kv.Del, pending[0].Op) + require.Equal( + t, + s3keys.VersionedChunkRefKey("bucket", 3, "object", "upload", 7, 0, 99), + pending[0].Key, + ) +} diff --git a/adapter/s3_chunk_upload.go b/adapter/s3_chunk_upload.go index 70e025a6f..27ace12ad 100644 --- a/adapter/s3_chunk_upload.go +++ b/adapter/s3_chunk_upload.go @@ -10,6 +10,7 @@ import ( "net/http" "strings" + "github.com/bootjp/elastickv/internal/s3keys" "github.com/bootjp/elastickv/kv" "github.com/cockroachdb/errors" ) @@ -21,6 +22,8 @@ type s3ChunkUploadConfig struct { tooLargeMessage string expectedPayloadSHA string chunkKey func(chunkNo uint64) []byte + chunkRefKey func(chunkNo uint64) []byte + offloaded bool } type s3ChunkUploadResult struct { @@ -29,6 +32,8 @@ type s3ChunkUploadResult struct { ChunkSizes []uint64 PersistedChunks int ChunkCount uint64 + Offloaded bool + ChunkRefElems []*kv.Elem[kv.OP] } type s3ChunkUploader struct { @@ -50,16 +55,21 @@ type s3ChunkUploader struct { // The caller remains responsible for protocol setup, metadata transactions, // and cleanup of any chunks reported in the result. func (s *S3Server) uploadS3Chunks(ctx context.Context, cfg s3ChunkUploadConfig) (s3ChunkUploadResult, *s3PutBodyError, error) { + batchOps := s3ChunkBatchOps + if cfg.offloaded { + batchOps = s3MetaBatchOps + } uploader := &s3ChunkUploader{ server: s, ctx: ctx, cfg: cfg, etagHasher: md5.New(), //nolint:gosec // S3 ETag compatibility requires MD5. buf: make([]byte, s3ChunkSize), - pendingBatch: make([]*kv.Elem[kv.OP], 0, s3ChunkBatchOps), + pendingBatch: make([]*kv.Elem[kv.OP], 0, batchOps), pendingChunkSizes: make([]uint64, 0, s3ChunkBatchOps), pendingAdmission: make([]func(), 0, s3ChunkBatchOps), } + uploader.result.Offloaded = cfg.offloaded if cfg.expectedPayloadSHA != "" { uploader.payloadHasher = sha256.New() } @@ -77,12 +87,17 @@ func (u *s3ChunkUploader) run() (s3ChunkUploadResult, *s3PutBodyError, error) { break } } - if err := u.flushBatch(); err != nil { - return u.result, nil, err + if !u.cfg.offloaded { + if err := u.flushBatch(); err != nil { + return u.result, nil, err + } } if bodyErr := u.verifyChecksums(); bodyErr != nil { return u.result, bodyErr, nil } + if u.cfg.offloaded { + u.result.ChunkRefElems = append([]*kv.Elem[kv.OP](nil), u.pendingBatch...) + } u.result.ETag = hex.EncodeToString(u.etagHasher.Sum(nil)) return u.result, nil, nil } @@ -127,7 +142,11 @@ func (u *s3ChunkUploader) flushBeforeRead() error { } func (u *s3ChunkUploader) flushFullBatch() error { - if len(u.pendingBatch) < s3ChunkBatchOps { + batchOps := s3ChunkBatchOps + if u.cfg.offloaded { + batchOps = s3MetaBatchOps + } + if len(u.pendingBatch) < batchOps { return nil } return u.flushBatch() @@ -141,7 +160,6 @@ func (u *s3ChunkUploader) classifyReadError(err error) (bool, *s3PutBodyError, e } func (u *s3ChunkUploader) appendChunk(data []byte, release func()) error { - u.pendingAdmission = append(u.pendingAdmission, release) chunk := append([]byte(nil), data...) if _, err := u.etagHasher.Write(chunk); err != nil { return errors.WithStack(err) @@ -152,29 +170,75 @@ func (u *s3ChunkUploader) appendChunk(data []byte, release func()) error { } } u.cfg.streamBody.writeDecoded(chunk) - u.pendingBatch = append(u.pendingBatch, &kv.Elem[kv.OP]{ - Op: kv.Put, - Key: u.cfg.chunkKey(u.result.ChunkCount), - Value: chunk, - }) + if u.cfg.offloaded { + if err := u.appendOffloadedChunk(chunk, release); err != nil { + return err + } + } else { + u.appendLegacyChunk(chunk, release) + } chunkSize := uint64(len(data)) u.result.ChunkSizes = append(u.result.ChunkSizes, chunkSize) - u.pendingChunkSizes = append(u.pendingChunkSizes, chunkSize) + if !u.cfg.offloaded { + u.pendingChunkSizes = append(u.pendingChunkSizes, chunkSize) + } u.result.SizeBytes += int64(len(data)) u.result.ChunkCount++ return nil } +func (u *s3ChunkUploader) appendOffloadedChunk(chunk []byte, release func()) error { + defer release() + if u.cfg.chunkRefKey == nil { + return errors.New("s3 chunkref key builder is not configured") + } + digest := sha256.Sum256(chunk) + refKey := u.cfg.chunkRefKey(u.result.ChunkCount) + commitTS, err := u.server.nextTxnCommitTS(u.ctx, 0) + if err != nil { + return errors.WithStack(err) + } + ref, err := u.server.persistS3ChunkBlob(u.ctx, refKey, digest, chunk, commitTS) + if err != nil { + return err + } + refValue, err := s3keys.EncodeChunkRefValue(ref) + if err != nil { + return errors.WithStack(err) + } + u.pendingBatch = append(u.pendingBatch, &kv.Elem[kv.OP]{ + Op: kv.Put, + Key: refKey, + Value: refValue, + }) + u.result.PersistedChunks++ + return nil +} + +func (u *s3ChunkUploader) appendLegacyChunk(chunk []byte, release func()) { + u.pendingAdmission = append(u.pendingAdmission, release) + u.pendingBatch = append(u.pendingBatch, &kv.Elem[kv.OP]{ + Op: kv.Put, + Key: u.cfg.chunkKey(u.result.ChunkCount), + Value: chunk, + }) +} + func (u *s3ChunkUploader) flushBatch() error { if len(u.pendingBatch) == 0 { return nil } + // Full offloaded batches are immutable staged refs. Object readers cannot + // discover them until the final manifest or part descriptor transaction; + // that final metadata commit remains the public linearization point. _, err := u.server.coordinator.Dispatch(u.ctx, &kv.OperationGroup[kv.OP]{Elems: u.pendingBatch}) u.releasePendingAdmission() if err != nil { return errors.WithStack(err) } - u.result.PersistedChunks += len(u.pendingChunkSizes) + if !u.cfg.offloaded { + u.result.PersistedChunks += len(u.pendingChunkSizes) + } u.pendingBatch = u.pendingBatch[:0] u.pendingChunkSizes = u.pendingChunkSizes[:0] return nil diff --git a/adapter/s3_put_object.go b/adapter/s3_put_object.go index bb74e8db6..3b53c2e83 100644 --- a/adapter/s3_put_object.go +++ b/adapter/s3_put_object.go @@ -7,6 +7,8 @@ import ( "github.com/bootjp/elastickv/internal/s3keys" "github.com/bootjp/elastickv/kv" "github.com/cockroachdb/errors" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" ) type s3PutObjectState struct { @@ -53,7 +55,7 @@ func (s *S3Server) prepareS3PutObject(ctx context.Context, request *http.Request return state, nil } -func (s *S3Server) uploadS3ObjectData(ctx context.Context, request *http.Request, streamBody *s3StreamingBody, state *s3PutObjectState, bucket, objectKey, expectedPayloadSHA string) (s3ChunkUploadResult, *s3PutBodyError, error) { +func (s *S3Server) uploadS3ObjectData(ctx context.Context, request *http.Request, streamBody *s3StreamingBody, state *s3PutObjectState, bucket, objectKey, expectedPayloadSHA string, offloaded bool) (s3ChunkUploadResult, *s3PutBodyError, error) { payloadChecksum := expectedPayloadSHA if isS3PayloadMarker(payloadChecksum) { payloadChecksum = "" @@ -67,11 +69,19 @@ func (s *S3Server) uploadS3ObjectData(ctx context.Context, request *http.Request chunkKey: func(chunkNo uint64) []byte { return s3keys.BlobKey(bucket, state.meta.Generation, objectKey, state.uploadID, 1, chunkNo) }, + chunkRefKey: func(chunkNo uint64) []byte { + return s3keys.ChunkRefKey(bucket, state.meta.Generation, objectKey, state.uploadID, 1, chunkNo) + }, + offloaded: offloaded, }) } func (s *S3Server) cleanupS3PutObjectChunks(ctx context.Context, state *s3PutObjectState, upload s3ChunkUploadResult, bucket, objectKey string) { - part := s3ObjectPart{PartNo: 1, ChunkSizes: upload.ChunkSizes[:upload.PersistedChunks]} + chunkSizes := upload.ChunkSizes + if !upload.Offloaded { + chunkSizes = upload.ChunkSizes[:upload.PersistedChunks] + } + part := s3ObjectPart{PartNo: 1, ChunkSizes: chunkSizes, Offloaded: upload.Offloaded} manifest := &s3ObjectManifest{UploadID: state.uploadID} if len(part.ChunkSizes) > 0 { manifest.Parts = []s3ObjectPart{part} @@ -98,7 +108,7 @@ func (s *S3Server) commitS3PutObject(ctx context.Context, request *http.Request, if upload.SizeBytes > 0 { manifest.Parts = []s3ObjectPart{{ PartNo: 1, ETag: upload.ETag, SizeBytes: upload.SizeBytes, - ChunkCount: upload.ChunkCount, ChunkSizes: upload.ChunkSizes, + ChunkCount: upload.ChunkCount, ChunkSizes: upload.ChunkSizes, Offloaded: upload.Offloaded, }} } body, err := encodeS3ObjectManifest(manifest) @@ -109,14 +119,16 @@ func (s *S3Server) commitS3PutObject(ctx context.Context, request *http.Request, if err != nil { return false, errors.WithStack(err) } + elems := append([]*kv.Elem[kv.OP](nil), upload.ChunkRefElems...) + elems = append(elems, + &kv.Elem[kv.OP]{Op: kv.Put, Key: s3keys.BucketMetaKey(bucket), Value: bucketFence}, + &kv.Elem[kv.OP]{Op: kv.Put, Key: state.headKey, Value: body}, + ) _, err = s.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: state.startTS, CommitTS: commitTS, - Elems: []*kv.Elem[kv.OP]{ - {Op: kv.Put, Key: s3keys.BucketMetaKey(bucket), Value: bucketFence}, - {Op: kv.Put, Key: state.headKey, Value: body}, - }, + Elems: elems, }) return true, errors.WithStack(err) } @@ -138,5 +150,9 @@ func (s *S3Server) writeS3ChunkUploadError(w http.ResponseWriter, bodyErr *s3Put writeS3AdmissionError(w, bucket, objectKey, s.s3PutAdmissionRetryAfter()) return } + if status.Code(err) == codes.Unavailable { + writeS3Error(w, http.StatusServiceUnavailable, "ServiceUnavailable", "chunkblob durability is temporarily unavailable", bucket, objectKey) + return + } writeS3ResponseOrInternalError(w, err) } diff --git a/adapter/s3_upload_part.go b/adapter/s3_upload_part.go index 1f491faa4..1a9e154b2 100644 --- a/adapter/s3_upload_part.go +++ b/adapter/s3_upload_part.go @@ -20,6 +20,12 @@ type s3UploadPartState struct { readPin *kv.ActiveTimestampToken } +type s3UploadPartVersion struct { + startTS uint64 + commitTS uint64 + chunkRefVersion uint64 +} + func (s *S3Server) prepareS3UploadPart(ctx context.Context, bucket, objectKey, uploadID, partNumberRaw string) (*s3UploadPartState, error) { partNo, err := parseS3UploadPartNumber(partNumberRaw, bucket, objectKey) if err != nil { @@ -61,8 +67,8 @@ func parseS3UploadPartNumber(raw, bucket, objectKey string) (uint64, error) { return uint64(partNumber), nil } -func (s *S3Server) storeS3UploadPart(ctx context.Context, request *http.Request, streamBody *s3StreamingBody, state *s3UploadPartState, bucket, objectKey, uploadID, admissionProtocol string) (s3ChunkUploadResult, *s3PartDescriptor, *s3PutBodyError, error) { - startTS, commitTS, err := s.allocateS3UploadPartVersion(ctx) +func (s *S3Server) storeS3UploadPart(ctx context.Context, request *http.Request, streamBody *s3StreamingBody, state *s3UploadPartState, bucket, objectKey, uploadID, admissionProtocol string, offloaded bool) (s3ChunkUploadResult, *s3PartDescriptor, *s3PutBodyError, error) { + version, err := s.allocateS3UploadPartVersionForMode(ctx, offloaded) if err != nil { return s3ChunkUploadResult{}, nil, nil, err } @@ -72,19 +78,30 @@ func (s *S3Server) storeS3UploadPart(ctx context.Context, request *http.Request, admissionProtocol: admissionProtocol, tooLargeMessage: "part exceeds maximum allowed size", chunkKey: func(chunkNo uint64) []byte { - return s3keys.VersionedBlobKey(bucket, state.meta.Generation, objectKey, uploadID, state.partNo, chunkNo, commitTS) + return s3keys.VersionedBlobKey(bucket, state.meta.Generation, objectKey, uploadID, state.partNo, chunkNo, version.commitTS) + }, + chunkRefKey: func(chunkNo uint64) []byte { + return s3keys.VersionedChunkRefKey(bucket, state.meta.Generation, objectKey, uploadID, state.partNo, chunkNo, version.chunkRefVersion) }, + offloaded: offloaded, }) committed := false defer func() { if !committed && upload.ChunkCount > 0 { - s.cleanupPartBlobsAsync(bucket, state.meta.Generation, objectKey, uploadID, state.partNo, upload.ChunkCount, commitTS) + s.cleanupPartBlobsAsync( + bucket, state.meta.Generation, objectKey, uploadID, + state.partNo, upload.ChunkCount, version.commitTS, version.chunkRefVersion, upload.Offloaded, + ) } }() if err != nil || bodyErr != nil { return upload, nil, bodyErr, err } - previous, err := s.commitS3UploadPart(ctx, state, upload, bucket, objectKey, uploadID, startTS, commitTS) + version.commitTS, err = s.finalizeS3UploadPartCommitTS(ctx, version.startTS, version.commitTS, offloaded) + if err != nil { + return upload, nil, nil, err + } + previous, err := s.commitS3UploadPart(ctx, state, upload, bucket, objectKey, uploadID, version) if err != nil { return upload, nil, nil, err } @@ -92,6 +109,31 @@ func (s *S3Server) storeS3UploadPart(ctx context.Context, request *http.Request, return upload, previous, nil, nil } +func (s *S3Server) allocateS3UploadPartVersionForMode(ctx context.Context, offloaded bool) (s3UploadPartVersion, error) { + startTS, commitTS, err := s.allocateS3UploadPartVersion(ctx) + if err != nil { + return s3UploadPartVersion{}, err + } + version := s3UploadPartVersion{startTS: startTS, commitTS: commitTS, chunkRefVersion: startTS} + if offloaded { + // Reserve the initial commit timestamp as this attempt's immutable + // chunkref namespace before delaying the descriptor commit timestamp. + version.chunkRefVersion = commitTS + // The descriptor timestamp is allocated after blob durability so the + // chunkrefs cannot commit behind a side-channel write they make reachable. + version.commitTS = 0 + } + return version, nil +} + +func (s *S3Server) finalizeS3UploadPartCommitTS(ctx context.Context, startTS, commitTS uint64, offloaded bool) (uint64, error) { + if !offloaded { + return commitTS, nil + } + ts, err := s.nextTxnCommitTS(ctx, startTS) + return ts, errors.WithStack(err) +} + func (s *S3Server) allocateS3UploadPartVersion(ctx context.Context) (uint64, uint64, error) { readTS := s.readTS() startTS, err := s.txnStartTS(ctx, readTS) @@ -105,10 +147,11 @@ func (s *S3Server) allocateS3UploadPartVersion(ctx context.Context) (uint64, uin return startTS, commitTS, nil } -func (s *S3Server) commitS3UploadPart(ctx context.Context, state *s3UploadPartState, upload s3ChunkUploadResult, bucket, objectKey, uploadID string, startTS, commitTS uint64) (*s3PartDescriptor, error) { +func (s *S3Server) commitS3UploadPart(ctx context.Context, state *s3UploadPartState, upload s3ChunkUploadResult, bucket, objectKey, uploadID string, version s3UploadPartVersion) (*s3PartDescriptor, error) { descriptor := &s3PartDescriptor{ PartNo: state.partNo, ETag: upload.ETag, SizeBytes: upload.SizeBytes, - ChunkCount: upload.ChunkCount, ChunkSizes: upload.ChunkSizes, PartVersion: commitTS, + ChunkCount: upload.ChunkCount, ChunkSizes: upload.ChunkSizes, PartVersion: version.commitTS, + ChunkRefVersion: version.chunkRefVersion, Offloaded: upload.Offloaded, } body, err := json.Marshal(descriptor) if err != nil { @@ -119,9 +162,12 @@ func (s *S3Server) commitS3UploadPart(ctx context.Context, state *s3UploadPartSt if err := s.verifyS3UploadStillExists(ctx, state.uploadMetaKey, bucket, objectKey); err != nil { return nil, err } + elems := make([]*kv.Elem[kv.OP], 0, len(upload.ChunkRefElems)+1) + elems = append(elems, upload.ChunkRefElems...) + elems = append(elems, &kv.Elem[kv.OP]{Op: kv.Put, Key: partKey, Value: body}) _, err = s.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ - IsTxn: true, StartTS: startTS, CommitTS: commitTS, - Elems: []*kv.Elem[kv.OP]{{Op: kv.Put, Key: partKey, Value: body}}, + IsTxn: true, StartTS: version.startTS, CommitTS: version.commitTS, + Elems: elems, }) if err != nil { return nil, errors.WithStack(err) diff --git a/docs/design/2026_04_25_proposed_s3_raft_blob_offload.md b/docs/design/2026_04_25_proposed_s3_raft_blob_offload.md index 12148c0ac..7745d4b8c 100644 --- a/docs/design/2026_04_25_proposed_s3_raft_blob_offload.md +++ b/docs/design/2026_04_25_proposed_s3_raft_blob_offload.md @@ -142,13 +142,26 @@ client ─► HTTP PUT body fsync-acks have returned (= chunkBlobMinReplicas durable copies including the leader) 5. queue ChunkRef into pendingBatch - ─► flushBatch: - coordinator.Dispatch(OperationGroup{ - Elems: [ chunkref Puts ... + manifest Put ], + 6. dispatch each full s3MetaBatchOps chunkref batch + through Raft under immutable attempt-versioned keys + ─► final flushBatch: + coordinator.Dispatch(Transaction{ + Elems: [ final partial chunkref batch + manifest Put ], }) ─► HTTP 200 OK once Dispatch acks ``` +The full chunkref batches in step 6 are staged metadata: readers +cannot discover them until the final manifest transaction commits, +so the manifest remains the linearisation point. Attempt-versioned +multipart keys prevent a retry from overwriting references selected +by an older committed part descriptor. A failed request schedules +best-effort deletion of staged chunkrefs by their unique attempt +prefix; chunkblobs that were already fsynced remain content-addressed +orphans for the M3 orphan scan in §3.5. Batching at `s3MetaBatchOps` +keeps every Raft proposal below `MaxSizePerMsg` even for +multi-terabyte objects. + Step 3 — synchronous chunkblob replication before the chunkref commit — is the difference between "Raft-equivalent durability" and "leader-only durability." Without it, a leader crash between @@ -189,7 +202,17 @@ follower that crashes after acking the push but before the chunkref commits is fine — the chunkref will be retried by the leader on the next attempt because it has not yet entered the Raft log. -**Behaviour during cluster shrink / partial outage.** A controlled +**M1 safety override.** M1 does not yet include M2's follower-apply +and backfill worker. It therefore waits for durable acknowledgements +from **every current Raft member**, including learners, before +committing a chunkref. This makes each current node's offline logical +backup self-contained: a backup taken from any member can resolve +every committed chunkref from its local `chunkblob` keyspace. An +unreachable current member fails the PUT closed. Once M2 guarantees +eventual local fetch on apply/backfill, the implementation may adopt +the quorum policy below without making member-local backups partial. + +**Post-M2 behaviour during cluster shrink / partial outage.** A controlled decommission (5 → 3 nodes) or a transient partition can leave the leader with fewer reachable peers than `(N/2)+1` configured at the last membership commit. Blocking PUTs until the configured minimum @@ -606,6 +629,13 @@ capability-advertising peers exists. Acceptance criteria for M3 (the milestone that flips `ELASTICKV_S3_BLOB_OFFLOAD=true` by default): +Until M3 wires reference counting, the grace queue, and the orphan +scanner into the server lifecycle, the runtime GC-readiness gate +remains false. Setting `ELASTICKV_S3_BLOB_OFFLOAD=true` before that +point still selects the legacy path. This fail-closed gate prevents +M1 chunkblobs from becoming an unbounded production keyspace while +preserving explicit test coverage of the M1 data path. + - WAL growth per GiB of S3 PUT < 1 MiB on a one-week soak test. - Snapshot transfer for a 100 GiB-cluster follower restart completes in < 60 s on a 1 Gbps interconnect. diff --git a/internal/backup/decode.go b/internal/backup/decode.go index eb3db251a..6b06b5fb6 100644 --- a/internal/backup/decode.go +++ b/internal/backup/decode.go @@ -326,6 +326,8 @@ func buildPrefixRoutes() []prefixRoute { {[]byte(S3UploadMetaPrefix), routeS3UploadMeta(S3UploadMetaPrefix)}, {[]byte(S3UploadPartPrefix), routeS3UploadMeta(S3UploadPartPrefix)}, {[]byte(S3BlobPrefix), routeS3Blob}, + {[]byte(S3ChunkRefPrefix), routeS3ChunkRef}, + {[]byte(S3ChunkBlobPrefix), routeS3ChunkBlob}, {[]byte(S3GCUploadPrefix), routeInternalDrop}, {[]byte(S3RoutePrefix), routeInternalDrop}, // SQS @@ -472,6 +474,24 @@ func routeS3Blob(d *dispatcher, k, v []byte) error { return d.s3.HandleBlob(k, v) } +func routeS3ChunkRef(d *dispatcher, k, v []byte) error { + if d.s3 == nil { + d.counters.Internal++ + return nil + } + d.counters.S3++ + return d.s3.HandleChunkRef(k, v) +} + +func routeS3ChunkBlob(d *dispatcher, k, v []byte) error { + if d.s3 == nil { + d.counters.Internal++ + return nil + } + d.counters.S3++ + return d.s3.HandleChunkBlob(k, v) +} + // routeS3UploadMeta returns a handler that forwards a specific // in-flight-multipart prefix into HandleIncompleteUpload. The // indirection captures the prefix label the S3 encoder uses to diff --git a/internal/backup/decode_test.go b/internal/backup/decode_test.go index 2e7885cba..bd4e2de52 100644 --- a/internal/backup/decode_test.go +++ b/internal/backup/decode_test.go @@ -2,10 +2,13 @@ package backup import ( "bytes" + "crypto/sha256" "errors" "os" "path/filepath" "testing" + + "github.com/bootjp/elastickv/internal/s3keys" ) // TestDecodeSnapshot_RejectsEmptyOutRoot pins that DecodeOptions @@ -101,6 +104,51 @@ func TestDecodeSnapshot_RoutesRedisString(t *testing.T) { } } +func TestDecodeSnapshot_RoutesS3OffloadRecords(t *testing.T) { + t.Parallel() + root := t.TempDir() + b := newSnapBuilder(1) + body := []byte("offloaded") + digest := sha256.Sum256(body) + ref, err := s3keys.EncodeChunkRefValue(s3keys.ChunkRefValue{ + ContentSHA256: digest, + Size: uint64(len(body)), + SourcePeer: "n1", + }) + if err != nil { + t.Fatal(err) + } + manifest := encodeS3ManifestValue(t, map[string]any{ + "upload_id": "u", "size_bytes": len(body), "parts": []map[string]any{ + {"part_no": 1, "chunk_count": 1, "chunk_sizes": []uint64{uint64(len(body))}, "offloaded": true}, + }, + }) + b.WriteEntry(s3keys.BucketMetaKey("b"), 1, encodeS3BucketMetaValue(t, map[string]any{ + "bucket_name": "b", "generation": 1, + }), false, 0, snapshotEncStateCleartx) + b.WriteEntry(s3keys.ChunkBlobKey(digest), 1, body, false, 0, snapshotEncStateCleartx) + b.WriteEntry(s3keys.ChunkRefKey("b", 1, "o", "u", 1, 0), 1, ref, false, 0, snapshotEncStateCleartx) + b.WriteEntry(s3keys.ObjectManifestKey("b", 1, "o"), 1, manifest, false, 0, snapshotEncStateCleartx) + + result, err := DecodeSnapshot(bytes.NewReader(b.Bytes()), DecodeOptions{ + OutRoot: root, + Adapters: AdapterSet{S3: true}, + }) + if err != nil { + t.Fatalf("DecodeSnapshot: %v", err) + } + if result.Counters.S3 != 4 { + t.Fatalf("Counters.S3 = %d, want 4", result.Counters.S3) + } + got, err := os.ReadFile(filepath.Join(root, "s3", "b", "o")) //nolint:gosec // test path + if err != nil { + t.Fatal(err) + } + if !bytes.Equal(got, body) { + t.Fatalf("body = %q, want %q", got, body) + } +} + // TestDecodeSnapshot_DropsTombstones pins the up-front tombstone // filter — entries flagged tombstone bypass the prefix router and // only increment Counters.Tombstone. diff --git a/internal/backup/s3.go b/internal/backup/s3.go index 84608f393..c9ebb2c88 100644 --- a/internal/backup/s3.go +++ b/internal/backup/s3.go @@ -1,9 +1,12 @@ package backup import ( + "crypto/sha256" + "encoding/hex" "encoding/json" "fmt" "io" + "math" "os" "path/filepath" "sort" @@ -24,6 +27,8 @@ const ( S3UploadMetaPrefix = s3keys.UploadMetaPrefix S3UploadPartPrefix = s3keys.UploadPartPrefix S3BlobPrefix = s3keys.BlobPrefix + S3ChunkRefPrefix = s3keys.ChunkRefPrefix + S3ChunkBlobPrefix = s3keys.ChunkBlobPrefix S3GCUploadPrefix = s3keys.GCUploadPrefix S3RoutePrefix = s3keys.RoutePrefix ) @@ -55,6 +60,12 @@ var ( // Without this guard a partial / racy snapshot would silently // emit a truncated body. Codex P1 #729. ErrS3IncompleteBlobChunks = errors.New("backup: incomplete blob chunks for manifest-declared part") + // ErrS3InvalidChunkRef is returned when a !s3|chunkref| value cannot + // be decoded or disagrees with its content-addressed chunkblob. + ErrS3InvalidChunkRef = errors.New("backup: invalid S3 chunk reference") + // ErrS3ChunkBlobHashMismatch is returned when a !s3|chunkblob| value + // does not hash to the digest encoded in its key. + ErrS3ChunkBlobHashMismatch = errors.New("backup: S3 chunkblob hash mismatch") ) // verifyChunkCompleteness checks every (partNo, partVersion) entry in @@ -79,7 +90,7 @@ func verifyChunkCompleteness(chunks []s3ChunkKey, declaredParts map[s3PartKey]s3 } got := make(map[s3PartKey]map[uint64]struct{}, len(declaredParts)) for _, k := range chunks { - pk := s3PartKey{partNo: k.partNo, partVersion: k.partVersion} + pk := s3PartKey{partNo: k.partNo, partVersion: k.partVersion, offloaded: k.offloaded} if got[pk] == nil { got[pk] = make(map[uint64]struct{}) } @@ -92,14 +103,14 @@ func verifyChunkCompleteness(chunks []s3ChunkKey, declaredParts map[s3PartKey]s3 seen := got[pk] if uint64(len(seen)) != want.chunkCount { //nolint:gosec // bounded return errors.Wrapf(ErrS3IncompleteBlobChunks, - "partNo=%d partVersion=%d declared chunks=%d, observed unique=%d", - pk.partNo, pk.partVersion, want.chunkCount, len(seen)) + "partNo=%d partVersion=%d offloaded=%t declared chunks=%d, observed unique=%d", + pk.partNo, pk.partVersion, pk.offloaded, want.chunkCount, len(seen)) } for i := uint64(0); i < want.chunkCount; i++ { if _, ok := seen[i]; !ok { return errors.Wrapf(ErrS3IncompleteBlobChunks, - "partNo=%d partVersion=%d declared chunks=%d, missing chunkNo=%d", - pk.partNo, pk.partVersion, want.chunkCount, i) + "partNo=%d partVersion=%d offloaded=%t declared chunks=%d, missing chunkNo=%d", + pk.partNo, pk.partVersion, pk.offloaded, want.chunkCount, i) } } } @@ -117,6 +128,8 @@ func verifyChunkCompleteness(chunks []s3ChunkKey, declaredParts map[s3PartKey]s3 // scratch chunk pool // !s3|bucket|gen|* (bg) -- ignored (operational counter) // !s3|bucket|meta|* (bm) -- buffered until Finalize +// !s3|chunkblob|* (cb) -- validated and spilled once per digest +// !s3|chunkref|* (cr) -- linked to manifest parts at Finalize // !s3|gc|upload|* (g) -- ignored (in-flight cleanup state) // !s3|obj|head|* (o) -- buffered until Finalize // !s3|upload|meta|* (um) -- excluded by default; opt in via @@ -131,8 +144,9 @@ func verifyChunkCompleteness(chunks []s3ChunkKey, declaredParts map[s3PartKey]s3 // assembled body to /s3// with the metadata // sidecar at .elastickv-meta.json. // -// Memory: O(num_objects + num_buckets) buffered metadata. Per-blob -// payloads are streamed to disk as they arrive — never held in memory. +// Memory: O(num_objects + num_buckets + num_chunkrefs + num_chunkblobs) +// buffered metadata and scratch paths. Payloads are streamed to disk as they +// arrive and are not retained in memory by the encoder. type S3Encoder struct { outRoot string scratchRoot string @@ -141,7 +155,11 @@ type S3Encoder struct { renameCollisions bool buckets map[string]*s3BucketState - warn func(event string, fields ...any) + // chunkBlobPaths indexes content-addressed offload payloads spilled + // under scratchRoot. Multiple object chunks can safely share one path. + chunkBlobPaths map[[sha256.Size]byte]string + chunkBlobDirCreated bool + warn func(event string, fields ...any) } type s3BucketState struct { @@ -197,6 +215,10 @@ type s3ObjectState struct { // chunkPaths maps (uploadID, partNo, chunkNo, partVersion) -> // scratch path. chunkPaths map[s3ChunkKey]string + // chunkRefs maps an offloaded object-chunk location to its + // content-addressed payload descriptor. Paths are resolved at finalize + // after all chunkblob records have been scanned. + chunkRefs map[s3ChunkKey]s3keys.ChunkRefValue } type s3ChunkKey struct { @@ -204,6 +226,7 @@ type s3ChunkKey struct { partNo uint64 chunkNo uint64 partVersion uint64 + offloaded bool } // s3PartKey is the manifest-declared part identifier: a (partNo, @@ -213,6 +236,7 @@ type s3ChunkKey struct { type s3PartKey struct { partNo uint64 partVersion uint64 + offloaded bool } // s3DeclaredPart captures what the manifest claims for a part: its @@ -222,6 +246,7 @@ type s3PartKey struct { // rather than a silently-truncated body (Codex P1 #729). type s3DeclaredPart struct { chunkCount uint64 + chunkSizes []uint64 } // s3PublicBucket is the dump-format projection of s3BucketMeta. @@ -293,12 +318,14 @@ type s3LiveManifest struct { } type s3LivePart 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"` - PartVersion uint64 `json:"part_version"` + PartNo uint64 `json:"part_no"` + ETag string `json:"etag"` + SizeBytes int64 `json:"size_bytes"` + ChunkCount uint64 `json:"chunk_count"` + ChunkSizes []uint64 `json:"chunk_sizes"` + PartVersion uint64 `json:"part_version"` + ChunkRefVersion uint64 `json:"chunk_ref_version"` + Offloaded bool `json:"offloaded"` } // NewS3Encoder constructs an encoder rooted at /s3/. Blob @@ -309,9 +336,10 @@ type s3LivePart struct { // Close(). func NewS3Encoder(outRoot, scratchRoot string) *S3Encoder { return &S3Encoder{ - outRoot: outRoot, - scratchRoot: filepath.Join(scratchRoot, "s3"), - buckets: make(map[string]*s3BucketState), + outRoot: outRoot, + scratchRoot: filepath.Join(scratchRoot, "s3"), + buckets: make(map[string]*s3BucketState), + chunkBlobPaths: make(map[[sha256.Size]byte]string), } } @@ -407,8 +435,18 @@ func (s *S3Encoder) HandleObjectManifest(key, value []byte) error { st.uploadID = live.UploadID st.declaredParts = make(map[s3PartKey]s3DeclaredPart, len(live.Parts)) for _, p := range live.Parts { - st.declaredParts[s3PartKey{partNo: p.PartNo, partVersion: p.PartVersion}] = s3DeclaredPart{ + partVersion := p.PartVersion + if p.Offloaded { + if uint64(len(p.ChunkSizes)) != p.ChunkCount { //nolint:gosec // Chunk count is bounded by the object size limit. + return errors.Wrapf(ErrS3InvalidManifest, + "offloaded part %d declares %d chunks but has %d chunk sizes", + p.PartNo, p.ChunkCount, len(p.ChunkSizes)) + } + partVersion = p.ChunkRefVersion + } + st.declaredParts[s3PartKey{partNo: p.PartNo, partVersion: partVersion, offloaded: p.Offloaded}] = s3DeclaredPart{ chunkCount: p.ChunkCount, + chunkSizes: append([]uint64(nil), p.ChunkSizes...), } } st.chunkPaths = ensureChunkPaths(st.chunkPaths) @@ -451,6 +489,55 @@ func (s *S3Encoder) HandleBlob(key, value []byte) error { return nil } +// HandleChunkBlob validates and spills a content-addressed offload payload. +// The payload is stored once per digest and linked to object chunks later by +// HandleChunkRef metadata. +func (s *S3Encoder) HandleChunkBlob(key, value []byte) error { + digest, ok := s3keys.ParseChunkBlobKey(key) + if !ok { + return errors.Wrapf(ErrS3MalformedKey, "chunkblob key: %q", key) + } + if actual := sha256.Sum256(value); actual != digest { + return errors.Wrapf(ErrS3ChunkBlobHashMismatch, "key digest=%x actual=%x", digest, actual) + } + dir := filepath.Join(s.scratchRoot, "chunkblobs") + if !s.chunkBlobDirCreated { + if err := os.MkdirAll(dir, 0o755); err != nil { //nolint:mnd // 0755 == standard dir mode + return errors.WithStack(err) + } + s.chunkBlobDirCreated = true + } + path := filepath.Join(dir, hex.EncodeToString(digest[:])) + if err := writeFileAtomic(path, value); err != nil { + return err + } + s.chunkBlobPaths[digest] = path + return nil +} + +// HandleChunkRef parks an offloaded object-chunk descriptor. The referenced +// chunkblob can appear before or after this record; resolution is deferred to +// Finalize so snapshot scan order does not affect the logical backup. +func (s *S3Encoder) HandleChunkRef(key, value []byte) error { + bucket, gen, object, uploadID, partNo, chunkNo, partVersion, ok := s3keys.ParseVersionedChunkRefKey(key) + if !ok { + return errors.Wrapf(ErrS3MalformedKey, "chunkref key: %q", key) + } + ref, ok := s3keys.DecodeChunkRefValue(value) + if !ok { + return errors.Wrapf(ErrS3InvalidChunkRef, "chunkref key: %q", key) + } + st := s.objectState(bucket, gen, object) + if st.chunkRefs == nil { + st.chunkRefs = make(map[s3ChunkKey]s3keys.ChunkRefValue) + } + st.chunkRefs[s3ChunkKey{ + uploadID: uploadID, partNo: partNo, chunkNo: chunkNo, + partVersion: partVersion, offloaded: true, + }] = ref + return nil +} + // scratchDirForBlob builds the per-(bucket,object) scratch path and // validates it stays under scratchRoot. A bucket or object name of // `.` / `..` would let `filepath.Join` resolve out of scratchRoot @@ -706,6 +793,9 @@ func (s *S3Encoder) flushObjectWithCollision(b *s3BucketState, bucketDir string, if obj.manifest == nil { return s.flushOrphanObject(b, bucketDir, obj) } + if err := s.resolveOffloadedChunkPaths(obj); err != nil { + return err + } objectName, kind, err := s.resolveObjectFilename(b, obj, needsLeafDataRename, objectKeys) if err != nil { return err @@ -732,6 +822,44 @@ func (s *S3Encoder) flushObjectWithCollision(b *s3BucketState, bucketDir string, return nil } +// resolveOffloadedChunkPaths links manifest-selected chunkrefs to their +// content-addressed scratch payloads. Stale refs from overwritten uploads are +// ignored before validation so asynchronous cleanup cannot break a backup of +// the currently committed manifest. +func (s *S3Encoder) resolveOffloadedChunkPaths(obj *s3ObjectState) error { + obj.chunkPaths = ensureChunkPaths(obj.chunkPaths) + for key, ref := range obj.chunkRefs { + if key.uploadID != obj.uploadID { + continue + } + partKey := s3PartKey{partNo: key.partNo, partVersion: key.partVersion, offloaded: true} + declared, ok := obj.declaredParts[partKey] + if !ok { + continue + } + path, ok := s.chunkBlobPaths[ref.ContentSHA256] + if !ok { + continue + } + info, err := os.Stat(path) + if err != nil { + return errors.WithStack(err) + } + if ref.Size > math.MaxInt64 || info.Size() != int64(ref.Size) { //nolint:gosec // ref.Size is bounded above. + return errors.Wrapf(ErrS3InvalidChunkRef, + "partNo=%d chunkNo=%d version=%d declared size=%d actual size=%d", + key.partNo, key.chunkNo, key.partVersion, ref.Size, info.Size()) + } + if key.chunkNo >= uint64(len(declared.chunkSizes)) || declared.chunkSizes[key.chunkNo] != ref.Size { //nolint:gosec // Manifest validation bounds len to chunkCount. + return errors.Wrapf(ErrS3InvalidChunkRef, + "partNo=%d chunkNo=%d version=%d ref size=%d does not match manifest chunk size", + key.partNo, key.chunkNo, key.partVersion, ref.Size) + } + obj.chunkPaths[key] = path + } + return nil +} + // flushOrphanObject handles objects with chunks but no manifest. By // default they emit only a warning. With --include-orphans on, the // chunks are written under /_orphans// as @@ -1012,7 +1140,7 @@ func filterChunksForManifest(m map[s3ChunkKey]string, manifestUploadID string, d continue } if declaredParts != nil { - declared, ok := declaredParts[s3PartKey{partNo: k.partNo, partVersion: k.partVersion}] + declared, ok := declaredParts[s3PartKey{partNo: k.partNo, partVersion: k.partVersion, offloaded: k.offloaded}] if !ok { continue } @@ -1033,20 +1161,23 @@ func filterChunksForManifest(m map[s3ChunkKey]string, manifestUploadID string, d } keys = append(keys, k) } - sort.SliceStable(keys, func(i, j int) bool { - a, b := keys[i], keys[j] - switch { - case a.partNo != b.partNo: - return a.partNo < b.partNo - case a.partVersion != b.partVersion: - return a.partVersion < b.partVersion - default: - return a.chunkNo < b.chunkNo - } - }) + sort.SliceStable(keys, func(i, j int) bool { return lessS3ChunkKey(keys[i], keys[j]) }) return keys } +func lessS3ChunkKey(a, b s3ChunkKey) bool { + switch { + case a.partNo != b.partNo: + return a.partNo < b.partNo + case a.partVersion != b.partVersion: + return a.partVersion < b.partVersion + case a.offloaded != b.offloaded: + return !a.offloaded + default: + return a.chunkNo < b.chunkNo + } +} + func appendFile(dst io.Writer, srcPath string) error { f, err := os.Open(srcPath) //nolint:gosec // srcPath composed from scratch root if err != nil { diff --git a/internal/backup/s3_test.go b/internal/backup/s3_test.go index 86131a4f3..6b1823c1e 100644 --- a/internal/backup/s3_test.go +++ b/internal/backup/s3_test.go @@ -2,6 +2,7 @@ package backup import ( "bytes" + "crypto/sha256" "encoding/json" "os" "path/filepath" @@ -65,6 +66,26 @@ func emitObject(t *testing.T, enc *S3Encoder, bucket string, gen uint64, object } } +func emitOffloadedChunk(t *testing.T, enc *S3Encoder, bucket string, gen uint64, object, uploadID string, partNo, chunkNo, version uint64, body []byte) { + t.Helper() + digest := sha256.Sum256(body) + if err := enc.HandleChunkBlob(s3keys.ChunkBlobKey(digest), body); err != nil { + t.Fatalf("HandleChunkBlob(%d): %v", chunkNo, err) + } + ref, err := s3keys.EncodeChunkRefValue(s3keys.ChunkRefValue{ + ContentSHA256: digest, + Size: uint64(len(body)), + SourcePeer: "n1", + }) + if err != nil { + t.Fatal(err) + } + key := s3keys.VersionedChunkRefKey(bucket, gen, object, uploadID, partNo, chunkNo, version) + if err := enc.HandleChunkRef(key, ref); err != nil { + t.Fatalf("HandleChunkRef(%d): %v", chunkNo, err) + } +} + func readJSONFile[T any](t *testing.T, path string, into *T) { t.Helper() body, err := os.ReadFile(path) //nolint:gosec // test path @@ -148,6 +169,128 @@ func TestS3_MultipartObjectAssemblesInPartChunkOrder(t *testing.T) { } } +func TestS3_OffloadedObjectAssemblesFromChunkRefs(t *testing.T) { + t.Parallel() + enc, root := newS3Encoder(t) + bucket, object, uploadID := "logs", "offloaded.bin", "u-offload" + const gen, chunkRefVersion = uint64(3), uint64(91) + chunks := []struct { + no uint64 + body []byte + }{ + {no: 0, body: []byte("hello ")}, + {no: 1, body: []byte("world")}, + } + + if err := enc.HandleBucketMeta(s3keys.BucketMetaKey(bucket), encodeS3BucketMetaValue(t, map[string]any{ + "bucket_name": bucket, "generation": gen, + })); err != nil { + t.Fatal(err) + } + for _, chunk := range chunks { + emitOffloadedChunk(t, enc, bucket, gen, object, uploadID, 1, chunk.no, chunkRefVersion, chunk.body) + } + // A legacy blob with the same numeric part version must not satisfy an + // offloaded manifest part. + if err := enc.HandleBlob(s3keys.VersionedBlobKey(bucket, gen, object, uploadID, 1, 0, chunkRefVersion), []byte("wrong")); err != nil { + t.Fatal(err) + } + if err := enc.HandleObjectManifest(s3keys.ObjectManifestKey(bucket, gen, object), encodeS3ManifestValue(t, map[string]any{ + "upload_id": uploadID, "size_bytes": 11, "parts": []map[string]any{ + { + "part_no": 1, "size_bytes": 11, "chunk_count": 2, + "chunk_sizes": []uint64{6, 5}, "offloaded": true, "chunk_ref_version": chunkRefVersion, + }, + }, + })); err != nil { + t.Fatal(err) + } + if err := enc.Finalize(); err != nil { + t.Fatalf("Finalize: %v", err) + } + body, err := os.ReadFile(filepath.Join(root, "s3", bucket, object)) //nolint:gosec // test path + if err != nil { + t.Fatal(err) + } + if string(body) != "hello world" { + t.Fatalf("body = %q, want %q", body, "hello world") + } +} + +func TestS3_OffloadedObjectFailsClosedWhenChunkBlobIsMissing(t *testing.T) { + t.Parallel() + enc, _ := newS3Encoder(t) + body := []byte("missing") + digest := sha256.Sum256(body) + ref, err := s3keys.EncodeChunkRefValue(s3keys.ChunkRefValue{ + ContentSHA256: digest, + Size: uint64(len(body)), + SourcePeer: "n1", + }) + if err != nil { + t.Fatal(err) + } + if err := enc.HandleChunkRef(s3keys.ChunkRefKey("b", 1, "o", "u", 1, 0), ref); err != nil { + t.Fatal(err) + } + if err := enc.HandleObjectManifest(s3keys.ObjectManifestKey("b", 1, "o"), encodeS3ManifestValue(t, map[string]any{ + "upload_id": "u", "size_bytes": len(body), "parts": []map[string]any{ + {"part_no": 1, "chunk_count": 1, "chunk_sizes": []uint64{uint64(len(body))}, "offloaded": true}, + }, + })); err != nil { + t.Fatal(err) + } + if err := enc.Finalize(); !errors.Is(err, ErrS3IncompleteBlobChunks) { + t.Fatalf("Finalize error = %v, want ErrS3IncompleteBlobChunks", err) + } +} + +func TestS3_OffloadedObjectRejectsManifestChunkSizeMismatch(t *testing.T) { + t.Parallel() + enc, _ := newS3Encoder(t) + body := []byte("content") + const chunkRefVersion = uint64(7) + const mismatchedChunkSize = uint64(8) + emitOffloadedChunk(t, enc, "b", 1, "o", "u", 1, 0, chunkRefVersion, body) + manifest := encodeS3ManifestValue(t, map[string]any{ + "upload_id": "u", "size_bytes": len(body), "parts": []map[string]any{ + { + "part_no": 1, "chunk_count": 1, "chunk_sizes": []uint64{mismatchedChunkSize}, + "offloaded": true, "chunk_ref_version": chunkRefVersion, + }, + }, + }) + if err := enc.HandleObjectManifest(s3keys.ObjectManifestKey("b", 1, "o"), manifest); err != nil { + t.Fatal(err) + } + if err := enc.Finalize(); !errors.Is(err, ErrS3InvalidChunkRef) { + t.Fatalf("Finalize error = %v, want ErrS3InvalidChunkRef", err) + } +} + +func TestS3_OffloadedObjectRejectsMissingManifestChunkSizes(t *testing.T) { + t.Parallel() + enc, _ := newS3Encoder(t) + manifest := encodeS3ManifestValue(t, map[string]any{ + "upload_id": "u", "size_bytes": 1, "parts": []map[string]any{ + {"part_no": 1, "chunk_count": 1, "offloaded": true}, + }, + }) + if err := enc.HandleObjectManifest(s3keys.ObjectManifestKey("b", 1, "o"), manifest); !errors.Is(err, ErrS3InvalidManifest) { + t.Fatalf("HandleObjectManifest error = %v, want ErrS3InvalidManifest", err) + } +} + +func TestS3_ChunkBlobRejectsContentHashMismatch(t *testing.T) { + t.Parallel() + enc, _ := newS3Encoder(t) + digest := sha256.Sum256([]byte("expected")) + err := enc.HandleChunkBlob(s3keys.ChunkBlobKey(digest), []byte("different")) + if !errors.Is(err, ErrS3ChunkBlobHashMismatch) { + t.Fatalf("HandleChunkBlob error = %v, want ErrS3ChunkBlobHashMismatch", err) + } +} + func TestS3_OrphanChunksWarn(t *testing.T) { t.Parallel() enc, _ := newS3Encoder(t) diff --git a/internal/raftadmin/health.go b/internal/raftadmin/health.go index 1e1567656..b19de6036 100644 --- a/internal/raftadmin/health.go +++ b/internal/raftadmin/health.go @@ -36,6 +36,20 @@ func RegisterOperationalServicesWithInterceptor( engine raftengine.Engine, serviceNames []string, interceptor MembershipChangeInterceptor, +) { + RegisterOperationalServicesWithInterceptorAndStaticServing(ctx, gs, engine, serviceNames, nil, interceptor) +} + +// RegisterOperationalServicesWithInterceptorAndStaticServing registers +// leader-gated operational service health plus peer-local services that should +// stay SERVING on followers. +func RegisterOperationalServicesWithInterceptorAndStaticServing( + ctx context.Context, + gs *grpc.Server, + engine raftengine.Engine, + leaderServiceNames []string, + staticServingServiceNames []string, + interceptor MembershipChangeInterceptor, ) { if gs == nil { return @@ -48,7 +62,18 @@ func RegisterOperationalServicesWithInterceptor( healthSrv := health.NewServer() healthpb.RegisterHealthServer(gs, healthSrv) - go observeLeaderHealth(ctx, engine, healthSrv, serviceNames, healthPollInterval()) + go observeStaticServingHealth(ctx, healthSrv, staticServingServiceNames) + go observeLeaderHealth(ctx, engine, healthSrv, leaderServiceNames, healthPollInterval()) +} + +func observeStaticServingHealth(ctx context.Context, healthSrv *health.Server, serviceNames []string) { + services := dedupeNamedServices(serviceNames) + if len(services) == 0 { + return + } + setHealthStatus(healthSrv, services, healthpb.HealthCheckResponse_SERVING) + <-ctx.Done() + setHealthStatus(healthSrv, services, healthpb.HealthCheckResponse_NOT_SERVING) } func observeLeaderHealth(ctx context.Context, engine raftengine.Engine, healthSrv *health.Server, serviceNames []string, pollInterval time.Duration) { @@ -125,6 +150,22 @@ func dedupeServices(serviceNames []string) []string { return services } +func dedupeNamedServices(serviceNames []string) []string { + seen := map[string]struct{}{} + services := make([]string, 0, len(serviceNames)) + for _, name := range serviceNames { + if name == "" { + continue + } + if _, ok := seen[name]; ok { + continue + } + seen[name] = struct{}{} + services = append(services, name) + } + return services +} + func setHealthStatus(healthSrv *health.Server, services []string, status healthpb.HealthCheckResponse_ServingStatus) { if healthSrv == nil { return diff --git a/internal/raftadmin/server_test.go b/internal/raftadmin/server_test.go index 4ec02aba6..3aa100823 100644 --- a/internal/raftadmin/server_test.go +++ b/internal/raftadmin/server_test.go @@ -360,6 +360,45 @@ func TestRegisterOperationalServicesPublishesLeaderHealth(t *testing.T) { }, 5*time.Second, 50*time.Millisecond) } +func TestRegisterOperationalServicesPublishesStaticServingHealth(t *testing.T) { + t.Parallel() + + engine := &fakeEngine{ + status: raftengine.Status{State: raftengine.StateFollower}, + serving: false, + } + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + + listener := bufconn.Listen(1024 * 1024) + server := grpc.NewServer() + RegisterOperationalServicesWithInterceptorAndStaticServing(ctx, server, engine, []string{"RawKV"}, []string{"S3BlobFetch"}, nil) + go func() { + _ = server.Serve(listener) + }() + t.Cleanup(server.Stop) + + conn, err := grpc.NewClient( + "passthrough:///bufnet", + grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) { + return listener.Dial() + }), + grpc.WithTransportCredentials(insecure.NewCredentials()), + ) + require.NoError(t, err) + t.Cleanup(func() { _ = conn.Close() }) + + client := healthpb.NewHealthClient(conn) + require.Eventually(t, func() bool { + resp, checkErr := client.Check(context.Background(), &healthpb.HealthCheckRequest{Service: "S3BlobFetch"}) + return checkErr == nil && resp.Status == healthpb.HealthCheckResponse_SERVING + }, 5*time.Second, 50*time.Millisecond) + require.Eventually(t, func() bool { + resp, checkErr := client.Check(context.Background(), &healthpb.HealthCheckRequest{Service: "RawKV"}) + return checkErr == nil && resp.Status == healthpb.HealthCheckResponse_NOT_SERVING + }, 5*time.Second, 50*time.Millisecond) +} + type stateOnlyEngine struct { state raftengine.State } diff --git a/internal/s3keys/chunkref.go b/internal/s3keys/chunkref.go index b14d6baf9..ef9e92c86 100644 --- a/internal/s3keys/chunkref.go +++ b/internal/s3keys/chunkref.go @@ -11,33 +11,64 @@ const ( chunkBlobSHA256Bytes = 32 chunkBlobSHA256HexBytes = chunkBlobSHA256Bytes * 2 chunkRefValueVersionV1 = byte(1) + chunkRefValueVersionV2 = byte(2) chunkRefValueFixedBytes = 1 + chunkBlobSHA256Bytes + u64Bytes + 2 chunkRefSourcePeerLenSize = 2 + chunkRefReplicaCountSize = 2 + chunkRefPeerFieldLenSize = 2 maxChunkRefSourcePeerLen = int(^uint16(0)) + maxChunkRefReplicaCount = int(^uint16(0)) ) // ErrChunkRefSourcePeerTooLong is returned when a ChunkRefValue source peer // cannot be encoded into the fixed-width value header. var ErrChunkRefSourcePeerTooLong = errors.New("s3 chunkref source peer is too long") +// ErrChunkRefReplicaSetTooLarge is returned when a write-time replica set +// cannot be represented by the V2 chunkref format. +var ErrChunkRefReplicaSetTooLarge = errors.New("s3 chunkref replica set is too large") + +// ErrChunkRefPeerFieldTooLong is returned when a replica node ID or address +// cannot be represented by the V2 chunkref format. +var ErrChunkRefPeerFieldTooLong = errors.New("s3 chunkref replica field is too long") + +// ChunkRefPeer records a peer that durably acknowledged the chunkblob before +// its reference became reachable through Raft. +type ChunkRefPeer struct { + NodeID string + Address string +} + // ChunkRefValue is the stored value for a !s3|chunkref| entry. It points from // an object part chunk to the content-addressed !s3|chunkblob| payload. type ChunkRefValue struct { ContentSHA256 [chunkBlobSHA256Bytes]byte Size uint64 SourcePeer string + ReplicaPeers []ChunkRefPeer } // ParseChunkRefKey decodes a !s3|chunkref| key into its object-part location. func ParseChunkRefKey(key []byte) (bucket string, generation uint64, object string, uploadID string, partNo uint64, chunkNo uint64, ok bool) { + bucket, generation, object, uploadID, partNo, chunkNo, _, ok = ParseVersionedChunkRefKey(key) + return bucket, generation, object, uploadID, partNo, chunkNo, ok +} + +// ParseVersionedChunkRefKey decodes both the original key shape and the +// attempt-versioned shape. partVersion is zero for original keys. +func ParseVersionedChunkRefKey(key []byte) (bucket string, generation uint64, object string, uploadID string, partNo uint64, chunkNo uint64, partVersion uint64, ok bool) { if !bytes.HasPrefix(key, chunkRefPrefixBytes) { - return "", 0, "", "", 0, 0, false + return "", 0, "", "", 0, 0, 0, false + } + parts, ok := parseObjectChunkKeyHead(key, len(chunkRefPrefixBytes)) + if !ok { + return "", 0, "", "", 0, 0, 0, false } - parts, ok := parseObjectChunkKey(key, len(chunkRefPrefixBytes)) + partVersion, ok = parseOptionalPartVersion(key, parts.next) if !ok { - return "", 0, "", "", 0, 0, false + return "", 0, "", "", 0, 0, 0, false } - return parts.bucket, parts.generation, parts.object, parts.uploadID, parts.partNo, parts.chunkNo, true + return parts.bucket, parts.generation, parts.object, parts.uploadID, parts.partNo, parts.chunkNo, partVersion, true } // ChunkBlobKey builds a content-addressed blob key from a SHA-256 digest. @@ -71,14 +102,48 @@ func EncodeChunkRefValue(ref ChunkRefValue) ([]byte, error) { if err != nil { return nil, err } + version := chunkRefValueVersionV1 + if len(ref.ReplicaPeers) > 0 { + version = chunkRefValueVersionV2 + } out := make([]byte, 0, chunkRefValueFixedBytes+len(ref.SourcePeer)) - out = append(out, chunkRefValueVersionV1) + out = append(out, version) out = append(out, ref.ContentSHA256[:]...) out = appendU64(out, ref.Size) var sourceLen [chunkRefSourcePeerLenSize]byte binary.BigEndian.PutUint16(sourceLen[:], sourcePeerLen) out = append(out, sourceLen[:]...) out = append(out, ref.SourcePeer...) + if version == chunkRefValueVersionV1 { + return out, nil + } + if len(ref.ReplicaPeers) > maxChunkRefReplicaCount { + return nil, ErrChunkRefReplicaSetTooLarge + } + var replicaCount [chunkRefReplicaCountSize]byte + binary.BigEndian.PutUint16(replicaCount[:], uint16(len(ref.ReplicaPeers))) //nolint:gosec // Bounded above. + out = append(out, replicaCount[:]...) + for _, peer := range ref.ReplicaPeers { + out, err = appendChunkRefPeerField(out, peer.NodeID) + if err != nil { + return nil, err + } + out, err = appendChunkRefPeerField(out, peer.Address) + if err != nil { + return nil, err + } + } + return out, nil +} + +func appendChunkRefPeerField(out []byte, value string) ([]byte, error) { + if len(value) > int(^uint16(0)) { + return nil, ErrChunkRefPeerFieldTooLong + } + var encodedLen [chunkRefPeerFieldLenSize]byte + binary.BigEndian.PutUint16(encodedLen[:], uint16(len(value))) //nolint:gosec // Bounded above. + out = append(out, encodedLen[:]...) + out = append(out, value...) return out, nil } @@ -92,26 +157,78 @@ func chunkRefSourcePeerLen(n int) (uint16, error) { // DecodeChunkRefValue decodes a chunk reference value. It returns ok=false for // unknown versions, truncated values, or trailing-length mismatches. func DecodeChunkRefValue(value []byte) (ChunkRefValue, bool) { + ref, version, next, ok := decodeChunkRefHeader(value) + if !ok { + return ChunkRefValue{}, false + } + if version == chunkRefValueVersionV1 { + return ref, next == len(value) + } + return decodeChunkRefV2(value, ref, next) +} + +func decodeChunkRefHeader(value []byte) (ChunkRefValue, byte, int, bool) { var ref ChunkRefValue - if len(value) < chunkRefValueFixedBytes || value[0] != chunkRefValueVersionV1 { - return ref, false + if len(value) < chunkRefValueFixedBytes || (value[0] != chunkRefValueVersionV1 && value[0] != chunkRefValueVersionV2) { + return ref, 0, 0, false } copy(ref.ContentSHA256[:], value[1:1+chunkBlobSHA256Bytes]) sizeOffset := 1 + chunkBlobSHA256Bytes var ok bool ref.Size, sizeOffset, ok = readU64(value, sizeOffset) if !ok || len(value)-sizeOffset < chunkRefSourcePeerLenSize { - return ChunkRefValue{}, false + return ChunkRefValue{}, 0, 0, false } sourceLen := int(binary.BigEndian.Uint16(value[sizeOffset : sizeOffset+chunkRefSourcePeerLenSize])) sourceOffset := sizeOffset + chunkRefSourcePeerLenSize - if len(value)-sourceOffset != sourceLen { + if len(value)-sourceOffset < sourceLen { + return ChunkRefValue{}, 0, 0, false + } + ref.SourcePeer = string(value[sourceOffset : sourceOffset+sourceLen]) + return ref, value[0], sourceOffset + sourceLen, true +} + +func decodeChunkRefV2(value []byte, ref ChunkRefValue, next int) (ChunkRefValue, bool) { + if len(value)-next < chunkRefReplicaCountSize { + return ChunkRefValue{}, false + } + replicaCount := int(binary.BigEndian.Uint16(value[next : next+chunkRefReplicaCountSize])) + next += chunkRefReplicaCountSize + const minimumEncodedPeerBytes = 2 * chunkRefPeerFieldLenSize + if replicaCount > (len(value)-next)/minimumEncodedPeerBytes { + return ChunkRefValue{}, false + } + ref.ReplicaPeers = make([]ChunkRefPeer, 0, replicaCount) + for range replicaCount { + nodeID, afterNode, valid := readChunkRefPeerField(value, next) + if !valid { + return ChunkRefValue{}, false + } + address, afterAddress, valid := readChunkRefPeerField(value, afterNode) + if !valid { + return ChunkRefValue{}, false + } + ref.ReplicaPeers = append(ref.ReplicaPeers, ChunkRefPeer{NodeID: nodeID, Address: address}) + next = afterAddress + } + if next != len(value) { return ChunkRefValue{}, false } - ref.SourcePeer = string(value[sourceOffset:]) return ref, true } +func readChunkRefPeerField(value []byte, offset int) (string, int, bool) { + if len(value)-offset < chunkRefPeerFieldLenSize { + return "", offset, false + } + fieldLen := int(binary.BigEndian.Uint16(value[offset : offset+chunkRefPeerFieldLenSize])) + start := offset + chunkRefPeerFieldLenSize + if len(value)-start < fieldLen { + return "", offset, false + } + return string(value[start : start+fieldLen]), start + fieldLen, true +} + type parsedObjectChunkKey struct { bucket string generation uint64 @@ -119,9 +236,10 @@ type parsedObjectChunkKey struct { uploadID string partNo uint64 chunkNo uint64 + next int } -func parseObjectChunkKey(key []byte, offset int) (parsedObjectChunkKey, bool) { +func parseObjectChunkKeyHead(key []byte, offset int) (parsedObjectChunkKey, bool) { var p parsedObjectChunkKey bucketRaw, next, ok := decodeSegment(key, offset) if !ok { @@ -141,11 +259,12 @@ func parseObjectChunkKey(key []byte, offset int) (parsedObjectChunkKey, bool) { if p.partNo, next, ok = readU64(key, next); !ok { return p, false } - if p.chunkNo, next, ok = readU64(key, next); !ok || next != len(key) { + if p.chunkNo, next, ok = readU64(key, next); !ok { return p, false } p.bucket = string(bucketRaw) p.object = string(objectRaw) p.uploadID = string(uploadIDRaw) + p.next = next return p, true } diff --git a/internal/s3keys/keys.go b/internal/s3keys/keys.go index 651f5834d..1ee2a09e6 100644 --- a/internal/s3keys/keys.go +++ b/internal/s3keys/keys.go @@ -104,6 +104,25 @@ func ChunkRefKey(bucket string, generation uint64, object string, uploadID strin return buildObjectKey(chunkRefPrefixBytes, bucket, generation, object, uploadID, partNo, chunkNo) } +// VersionedChunkRefKey returns the chunk reference key for one immutable part +// attempt. A zero version preserves the original key shape for existing data +// and single PUTs whose upload ID is already unique. +func VersionedChunkRefKey(bucket string, generation uint64, object string, uploadID string, partNo uint64, chunkNo uint64, partVersion uint64) []byte { + if partVersion == 0 { + return ChunkRefKey(bucket, generation, object, uploadID, partNo, chunkNo) + } + out := make([]byte, 0, len(ChunkRefPrefix)+len(bucket)+len(object)+len(uploadID)+buildObjectExtraBytes+4*u64Bytes) + out = append(out, chunkRefPrefixBytes...) + out = append(out, EncodeSegment([]byte(bucket))...) + out = appendU64(out, generation) + out = append(out, EncodeSegment([]byte(object))...) + out = append(out, EncodeSegment([]byte(uploadID))...) + out = appendU64(out, partNo) + out = appendU64(out, chunkNo) + out = appendU64(out, partVersion) + return out +} + // VersionedBlobKey returns the blob key for a specific part attempt identified by // partVersion (typically the part's commit timestamp). When partVersion is 0 the // result is identical to BlobKey, preserving backward compatibility with data diff --git a/internal/s3keys/keys_test.go b/internal/s3keys/keys_test.go index 56664138c..97431d286 100644 --- a/internal/s3keys/keys_test.go +++ b/internal/s3keys/keys_test.go @@ -295,6 +295,25 @@ func TestChunkRefKey_RoundTripAndRouteKey(t *testing.T) { require.Equal(t, RouteKey(bucket, 11, object), ExtractRouteKey(key)) } +func TestVersionedChunkRefKey_RoundTrip(t *testing.T) { + t.Parallel() + + key := VersionedChunkRefKey("bucket", 11, "object", "upload", 7, 3, 99) + bucket, generation, object, uploadID, partNo, chunkNo, partVersion, ok := ParseVersionedChunkRefKey(key) + require.True(t, ok) + require.Equal(t, "bucket", bucket) + require.Equal(t, uint64(11), generation) + require.Equal(t, "object", object) + require.Equal(t, "upload", uploadID) + require.Equal(t, uint64(7), partNo) + require.Equal(t, uint64(3), chunkNo) + require.Equal(t, uint64(99), partVersion) + require.Equal(t, RouteKey("bucket", 11, "object"), ExtractRouteKey(key)) + + unversioned := VersionedChunkRefKey("bucket", 11, "object", "upload", 7, 3, 0) + require.Equal(t, ChunkRefKey("bucket", 11, "object", "upload", 7, 3), unversioned) +} + func TestChunkBlobKey_RoundTrip(t *testing.T) { t.Parallel() @@ -325,6 +344,28 @@ func TestChunkRefValue_RoundTrip(t *testing.T) { require.Equal(t, "node-a", got.SourcePeer) } +func TestChunkRefValueV2_RoundTripWriteTimeReplicas(t *testing.T) { + t.Parallel() + + sum := sha256.Sum256([]byte("chunk data")) + want := ChunkRefValue{ + ContentSHA256: sum, + Size: 1234, + SourcePeer: "node-a", + ReplicaPeers: []ChunkRefPeer{ + {NodeID: "node-a", Address: "10.0.0.1:50051"}, + {NodeID: "node-b", Address: "10.0.0.2:50051"}, + }, + } + value, err := EncodeChunkRefValue(want) + require.NoError(t, err) + require.Equal(t, chunkRefValueVersionV2, value[0]) + + got, ok := DecodeChunkRefValue(value) + require.True(t, ok) + require.Equal(t, want, got) +} + func TestChunkRefValue_RejectsMalformedValues(t *testing.T) { t.Parallel() @@ -342,6 +383,19 @@ func TestChunkRefValue_RejectsMalformedValues(t *testing.T) { _, ok := DecodeChunkRefValue(malformed) require.False(t, ok, "value %x should be rejected", malformed) } + + v2, err := EncodeChunkRefValue(ChunkRefValue{ + ContentSHA256: sum, + Size: 1, + SourcePeer: "n1", + ReplicaPeers: []ChunkRefPeer{{NodeID: "n1", Address: "n1:50051"}}, + }) + require.NoError(t, err) + replicaCountOffset := chunkRefValueFixedBytes + len("n1") + v2[replicaCountOffset] = 0xff + v2[replicaCountOffset+1] = 0xff + _, ok := DecodeChunkRefValue(v2) + require.False(t, ok, "impossible replica count must be rejected before allocation") } // TestPerBucketPrefixes_IsolateByBucketAndGeneration covers the diff --git a/kv/coordinator.go b/kv/coordinator.go index 5ee5c4f46..7fa2ff18f 100644 --- a/kv/coordinator.go +++ b/kv/coordinator.go @@ -323,6 +323,24 @@ type GroupRoutableCoordinator interface { EngineGroupIDForKey(key []byte) uint64 } +// RaftMember describes one member of the Raft group that owns a key. The +// adapter layer uses this read-only view for peer-local side channels whose +// payloads deliberately do not enter the Raft log. +type RaftMember struct { + NodeID string + Address string + Suffrage string +} + +// RaftMembershipCoordinator is the optional capability implemented by +// coordinators that can expose both cluster-wide and per-key Raft membership. +// It does not establish a read fence; callers must use it only for peer +// discovery and keep correctness decisions behind their own quorum checks. +type RaftMembershipCoordinator interface { + RaftMembers(ctx context.Context) ([]RaftMember, error) + RaftMembersForKey(ctx context.Context, key []byte) ([]RaftMember, error) +} + // LeaseReadGroupKey returns a representative key per distinct owning // group for the supplied keys, so callers can issue one lease read per // group rather than one per key. The returned slice preserves the order diff --git a/kv/shard_store.go b/kv/shard_store.go index 14c72332f..4d071dd35 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -2415,6 +2415,18 @@ func (s *ShardStore) groupForKey(key []byte) (*ShardGroup, bool) { return g, ok } +// LocalStoreForKey returns this process's store for the key's owning group +// without a leader fence or network proxy. It is reserved for node-local +// auxiliary state such as content-addressed S3 chunk blobs; replicated state +// must continue through the normal ShardStore or Coordinator paths. +func (s *ShardStore) LocalStoreForKey(key []byte) (store.MVCCStore, bool) { + g, ok := s.groupForKey(key) + if !ok || g == nil || g.Store == nil { + return nil, false + } + return g.Store, true +} + func (s *ShardStore) proxyRawGet(ctx context.Context, g *ShardGroup, key []byte, ts uint64, groupID uint64) ([]byte, error) { engine := engineForGroup(g) if engine == nil { diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index 37b963c90..14fd4c763 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -1865,6 +1865,67 @@ func (c *ShardedCoordinator) Clock() *HLC { return c.clock } +// RaftMembers returns every configured Raft endpoint across all local groups. +// Endpoints, rather than node IDs, are deduplicated because one process can +// expose a distinct listener per group. The stable sort keeps capability-cache +// fingerprints deterministic despite map iteration order. +func (c *ShardedCoordinator) RaftMembers(ctx context.Context) ([]RaftMember, error) { + if c == nil || len(c.groups) == 0 { + return nil, errors.WithStack(ErrLeaderNotFound) + } + membersByEndpoint := make(map[string]RaftMember) + for _, group := range c.groups { + members, err := raftMembersForGroup(ctx, group) + if err != nil { + return nil, err + } + for _, member := range members { + endpoint := member.NodeID + "\x00" + member.Address + membersByEndpoint[endpoint] = member + } + } + endpoints := make([]string, 0, len(membersByEndpoint)) + for endpoint := range membersByEndpoint { + endpoints = append(endpoints, endpoint) + } + slices.Sort(endpoints) + members := make([]RaftMember, 0, len(endpoints)) + for _, endpoint := range endpoints { + members = append(members, membersByEndpoint[endpoint]) + } + return members, nil +} + +// RaftMembersForKey returns a stable value copy of the current configuration +// for the key's owning group. It is intentionally membership-only: callers do +// not gain access to the mutable ShardGroup or its local store. +func (c *ShardedCoordinator) RaftMembersForKey(ctx context.Context, key []byte) ([]RaftMember, error) { + g, ok := c.groupForKey(key) + if !ok { + return nil, errors.WithStack(ErrLeaderNotFound) + } + return raftMembersForGroup(ctx, g) +} + +func raftMembersForGroup(ctx context.Context, group *ShardGroup) ([]RaftMember, error) { + if group == nil || group.Engine == nil { + return nil, errors.WithStack(ErrLeaderNotFound) + } + cfg, err := group.Engine.Configuration(ctx) + if err != nil { + return nil, errors.WithStack(err) + } + members := make([]RaftMember, 0, len(cfg.Servers)) + for _, server := range cfg.Servers { + members = append(members, RaftMember{ + NodeID: server.ID, + Address: server.Address, + Suffrage: server.Suffrage, + }) + } + return members, nil +} + func (c *ShardedCoordinator) groupForKey(key []byte) (*ShardGroup, bool) { gid, ok := c.router.ResolveGroup(key) if !ok { diff --git a/kv/sharded_membership_test.go b/kv/sharded_membership_test.go new file mode 100644 index 000000000..956840edd --- /dev/null +++ b/kv/sharded_membership_test.go @@ -0,0 +1,54 @@ +package kv + +import ( + "context" + "testing" + + "github.com/bootjp/elastickv/distribution" + "github.com/bootjp/elastickv/internal/raftengine" + "github.com/stretchr/testify/require" +) + +type membershipConfigEngine struct { + noopEngine + config raftengine.Configuration +} + +func (e membershipConfigEngine) Configuration(context.Context) (raftengine.Configuration, error) { + return e.config, nil +} + +func TestShardedCoordinatorRaftMembersIncludesEveryGroupEndpoint(t *testing.T) { + t.Parallel() + + routes := distribution.NewEngine() + routes.UpdateRoute(nil, []byte("m"), 1) + routes.UpdateRoute([]byte("m"), nil, 2) + groups := map[uint64]*ShardGroup{ + 1: {Engine: membershipConfigEngine{config: raftengine.Configuration{Servers: []raftengine.Server{ + {ID: "n1", Address: "127.0.0.1:5101", Suffrage: "voter"}, + {ID: "n2", Address: "127.0.0.1:5102", Suffrage: "voter"}, + }}}}, + 2: {Engine: membershipConfigEngine{config: raftengine.Configuration{Servers: []raftengine.Server{ + {ID: "n1", Address: "127.0.0.1:5201", Suffrage: "voter"}, + {ID: "n3", Address: "127.0.0.1:5203", Suffrage: "learner"}, + }}}}, + } + coordinator := NewShardedCoordinator(routes, groups, 1, NewHLC(), nil) + + members, err := coordinator.RaftMembers(context.Background()) + require.NoError(t, err) + require.Equal(t, []RaftMember{ + {NodeID: "n1", Address: "127.0.0.1:5101", Suffrage: "voter"}, + {NodeID: "n1", Address: "127.0.0.1:5201", Suffrage: "voter"}, + {NodeID: "n2", Address: "127.0.0.1:5102", Suffrage: "voter"}, + {NodeID: "n3", Address: "127.0.0.1:5203", Suffrage: "learner"}, + }, members) + + keyMembers, err := coordinator.RaftMembersForKey(context.Background(), []byte("z")) + require.NoError(t, err) + require.Equal(t, []RaftMember{ + {NodeID: "n1", Address: "127.0.0.1:5201", Suffrage: "voter"}, + {NodeID: "n3", Address: "127.0.0.1:5203", Suffrage: "learner"}, + }, keyMembers) +} diff --git a/main.go b/main.go index 479bcc69b..5bc27a830 100644 --- a/main.go +++ b/main.go @@ -162,6 +162,7 @@ var ( // SigV4 access keys). adminTokenFile = flag.String("adminTokenFile", "", "Path to a file containing the read-only bearer token required on the Admin gRPC service (leave blank with --adminInsecureNoAuth off to disable the Admin service)") adminInsecureNoAuth = flag.Bool("adminInsecureNoAuth", false, "Register the Admin gRPC service without bearer-token authentication; development only") + s3BlobPeerTokenFile = flag.String("s3BlobPeerTokenFile", "", "Path to the cluster-shared bearer token for peer S3 chunkblob RPCs") // Admin HTTP listener flags (PR #545's parallel work merged into // main; serves the cookie/SigV4-authenticated admin dashboard). @@ -2076,6 +2077,8 @@ func setupAdminService( srv, icept, err := configureAdminService( *adminTokenFile, *adminInsecureNoAuth, + *s3BlobPeerTokenFile, + adapter.S3BlobOffloadEnabledFromEnv(), adapter.NodeIdentity{NodeID: nodeID, GRPCAddress: selfAddr}, members, ) @@ -2083,7 +2086,7 @@ func setupAdminService( return nil, adminGRPCInterceptors{}, err } if srv == nil { - return nil, adminGRPCInterceptors{}, nil + return nil, icept, nil } for _, rt := range runtimes { srv.RegisterGroup(rt.spec.id, rt.engine) @@ -2155,8 +2158,9 @@ func adminMembersFromBootstrap(selfID string, servers []raftengine.Server) []ada // ChainUnaryInterceptor call, so using grpc.UnaryInterceptor alongside risks // silent overwrites (gRPC-Go: last option of the same type wins). type adminGRPCInterceptors struct { - unary []grpc.UnaryServerInterceptor - stream []grpc.StreamServerInterceptor + unary []grpc.UnaryServerInterceptor + stream []grpc.StreamServerInterceptor + s3BlobFetchEnabled bool } func (a adminGRPCInterceptors) empty() bool { @@ -2172,6 +2176,7 @@ var _ kv.Coordinator = (*startupGatedCoordinator)(nil) var _ kv.LeaseReadableCoordinator = (*startupGatedCoordinator)(nil) var _ kv.AllGroupsLeaseReadableCoordinator = (*startupGatedCoordinator)(nil) var _ kv.GroupRoutableCoordinator = (*startupGatedCoordinator)(nil) +var _ kv.RaftMembershipCoordinator = (*startupGatedCoordinator)(nil) func (c startupGatedCoordinator) Dispatch(ctx context.Context, reqs *kv.OperationGroup[kv.OP]) (*kv.CoordinateResponse, error) { if c.gate != nil && c.gate.blocked() { @@ -2231,6 +2236,22 @@ func (c startupGatedCoordinator) EngineGroupIDForKey(key []byte) uint64 { return 0 } +func (c startupGatedCoordinator) RaftMembers(ctx context.Context) ([]kv.RaftMember, error) { + provider, ok := c.inner.(kv.RaftMembershipCoordinator) + if !ok { + return nil, errors.WithStack(kv.ErrLeaderNotFound) + } + return provider.RaftMembers(ctx) //nolint:wrapcheck // Preserve membership lookup errors. +} + +func (c startupGatedCoordinator) RaftMembersForKey(ctx context.Context, key []byte) ([]kv.RaftMember, error) { + provider, ok := c.inner.(kv.RaftMembershipCoordinator) + if !ok { + return nil, errors.WithStack(kv.ErrLeaderNotFound) + } + return provider.RaftMembersForKey(ctx, key) //nolint:wrapcheck // Preserve membership lookup errors. +} + type startupPublicKVGate struct { ready atomic.Bool blockMutator func() bool @@ -2256,6 +2277,20 @@ func (g *startupPublicKVGate) unaryInterceptor( return handler(ctx, req) } +func (g *startupPublicKVGate) streamInterceptor( + srv interface{}, + stream grpc.ServerStream, + info *grpc.StreamServerInfo, + handler grpc.StreamHandler, +) error { + if g != nil && info != nil && startupRotationGatedMethod(info.FullMethod) && g.blocked() { + // Return a raw gRPC status so clients and retry policy see Unavailable. + //nolint:wrapcheck + return status.Error(codes.Unavailable, "startup rotation has not completed") + } + return handler(srv, stream) +} + func (g *startupPublicKVGate) blocked() bool { if g == nil { return false @@ -2281,7 +2316,8 @@ func startupRotationGatedMethod(fullMethod string) bool { pb.EncryptionAdmin_RegisterEncryptionWriter_FullMethodName, pb.EncryptionAdmin_ResyncSidecar_FullMethodName, pb.EncryptionAdmin_EnableStorageEnvelope_FullMethodName, - pb.EncryptionAdmin_EnableRaftEnvelope_FullMethodName: + pb.EncryptionAdmin_EnableRaftEnvelope_FullMethodName, + pb.S3BlobFetch_PushChunkBlob_FullMethodName: return true default: return strings.HasPrefix(fullMethod, "/RawKV/") || @@ -2297,34 +2333,79 @@ func startupRotationGatedMethod(fullMethod string) bool { func configureAdminService( tokenPath string, insecureNoAuth bool, + s3BlobPeerTokenPath string, + s3BlobOffloadEnabled bool, self adapter.NodeIdentity, members []adapter.NodeIdentity, ) (*adapter.AdminServer, adminGRPCInterceptors, error) { - if tokenPath == "" && !insecureNoAuth { - return nil, adminGRPCInterceptors{}, nil + auth, err := loadNodeGRPCAuth(tokenPath, insecureNoAuth, s3BlobPeerTokenPath) + if err != nil { + return nil, adminGRPCInterceptors{}, err + } + var srv *adapter.AdminServer + if auth.adminEnabled { + srv = adapter.NewAdminServer(self, members) + srv.SetCapability( + adapter.S3BlobOffloadCapabilityName, + s3BlobOffloadAdminCapability(auth.adminToken, auth.peerToken, s3BlobOffloadEnabled), + ) } + var icept adminGRPCInterceptors + adminUnary, adminStream := adapter.AdminTokenAuth(auth.adminToken) + appendGRPCBearerInterceptors(&icept, adminUnary, adminStream) + peerUnary, peerStream := adapter.S3BlobPeerTokenAuth(auth.peerToken) + appendGRPCBearerInterceptors(&icept, peerUnary, peerStream) + icept.s3BlobFetchEnabled = auth.peerToken != "" + return srv, icept, nil +} + +type nodeGRPCAuthConfig struct { + adminToken string + peerToken string + adminEnabled bool +} + +func loadNodeGRPCAuth(tokenPath string, insecureNoAuth bool, peerTokenPath string) (nodeGRPCAuthConfig, error) { if tokenPath != "" && insecureNoAuth { - return nil, adminGRPCInterceptors{}, errors.New("--adminInsecureNoAuth and --adminTokenFile are mutually exclusive") + return nodeGRPCAuthConfig{}, errors.New("--adminInsecureNoAuth and --adminTokenFile are mutually exclusive") } - token := "" + var cfg nodeGRPCAuthConfig + cfg.adminEnabled = tokenPath != "" || insecureNoAuth if tokenPath != "" { loaded, err := loadAdminTokenFile(tokenPath) if err != nil { - return nil, adminGRPCInterceptors{}, err + return nodeGRPCAuthConfig{}, err } - token = loaded + cfg.adminToken = loaded } - srv := adapter.NewAdminServer(self, members) - srv.SetCapability(adapter.S3BlobOffloadCapabilityName, adapter.S3BlobOffloadLocalCapability()) - unary, stream := adapter.AdminTokenAuth(token) - var icept adminGRPCInterceptors + if peerTokenPath != "" { + loaded, err := loadS3BlobPeerTokenFile(peerTokenPath) + if err != nil { + return nodeGRPCAuthConfig{}, err + } + cfg.peerToken = loaded + } + if cfg.adminToken != "" && cfg.peerToken == cfg.adminToken { + return nodeGRPCAuthConfig{}, errors.New("S3 blob peer token must differ from the read-only admin token") + } + return cfg, nil +} + +func appendGRPCBearerInterceptors( + icept *adminGRPCInterceptors, + unary grpc.UnaryServerInterceptor, + stream grpc.StreamServerInterceptor, +) { if unary != nil { icept.unary = append(icept.unary, unary) } if stream != nil { icept.stream = append(icept.stream, stream) } - return srv, icept, nil +} + +func s3BlobOffloadAdminCapability(adminToken, peerToken string, enabled bool) bool { + return adapter.S3BlobOffloadLocalCapability() && enabled && adminToken != "" && peerToken != "" } // loadAdminTokenFile materialises --adminTokenFile with a strict upper bound @@ -2339,6 +2420,14 @@ func loadAdminTokenFile(path string) (string, error) { return tok, nil } +func loadS3BlobPeerTokenFile(path string) (string, error) { + tok, err := internalutil.LoadBearerTokenFile(path, adminTokenMaxBytes, "S3 blob peer token") + if err != nil { + return "", errors.Wrap(err, "load S3 blob peer token") + } + return tok, nil +} + // startMemoryWatchdog optionally starts the memwatch goroutine. The // watcher is off by default; it is enabled only when the operator sets // ELASTICKV_MEMORY_SHUTDOWN_THRESHOLD_MB. On threshold crossing the @@ -2486,6 +2575,26 @@ func fsmCompactionRuntimes(runtimes []*raftGroupRuntime) []kv.FSMCompactRuntime return out } +func registerS3BlobFetchServer( + registrar grpc.ServiceRegistrar, + enabled bool, + st store.MVCCStore, + observer adapter.S3BlobOffloadObserver, + clock *kv.HLC, + pushBlocked func() bool, +) []string { + if !enabled { + return nil + } + pb.RegisterS3BlobFetchServer(registrar, adapter.NewS3BlobFetchServer( + st, + observer, + adapter.WithS3BlobFetchClock(clock), + adapter.WithS3BlobFetchPushBlocked(pushBlocked), + )) + return []string{"S3BlobFetch"} +} + func startRaftServers( ctx context.Context, lc *net.ListenConfig, @@ -2502,6 +2611,8 @@ func startRaftServers( forwardDeps adminForwardServerDeps, confChangeInterceptor internalraftadmin.MembershipChangeInterceptor, encWiring encryptionWriteWiring, + s3BlobObserver adapter.S3BlobOffloadObserver, + s3BlobPushBlocked func() bool, ) error { forwardLogger := slog.Default().With(slog.String("component", "admin")) // extraOptsCap reserves slots for the unary + stream admin interceptor @@ -2530,6 +2641,14 @@ func startRaftServers( grpcSvc := adapter.NewGRPCServer(shardStore, coordinate) pb.RegisterRawKVServer(gs, grpcSvc) pb.RegisterTransactionalKVServer(gs, grpcSvc) + staticServingServices := registerS3BlobFetchServer( + gs, + adminGRPCOpts.s3BlobFetchEnabled, + rt.store, + s3BlobObserver, + coordinate.Clock(), + s3BlobPushBlocked, + ) pb.RegisterInternalServer(gs, adapter.NewInternalWithEngine( trx, rt.engine, @@ -2573,7 +2692,14 @@ func startRaftServers( // Stage 7c §3.1: pass the encryption-aware pre-register hook // (nil when encryption is not wired); raftadmin.Server invokes // it before AddVoter/AddLearner propose the conf-change. - internalraftadmin.RegisterOperationalServicesWithInterceptor(ctx, gs, rt.engine, []string{"RawKV"}, confChangeInterceptor) + internalraftadmin.RegisterOperationalServicesWithInterceptorAndStaticServing( + ctx, + gs, + rt.engine, + []string{"RawKV"}, + staticServingServices, + confChangeInterceptor, + ) reflection.Register(gs) grpcSock, err := lc.Listen(ctx, "tcp", rt.spec.address) @@ -2975,6 +3101,11 @@ func (r *runtimeServerRunner) startRaftTransport() error { adminGRPCOpts := r.adminGRPCOpts if r.publicKVGate != nil { adminGRPCOpts.unary = append(adminGRPCOpts.unary, r.publicKVGate.unaryInterceptor) + adminGRPCOpts.stream = append(adminGRPCOpts.stream, r.publicKVGate.streamInterceptor) + } + var s3BlobPushBlocked func() bool + if r.publicKVGate != nil { + s3BlobPushBlocked = r.publicKVGate.blocked } forwardDeps := adminForwardServerDeps{ tables: newDynamoTablesSource(r.dynamoServer), @@ -2999,6 +3130,8 @@ func (r *runtimeServerRunner) startRaftTransport() error { forwardDeps, r.encryptionConfChangeInterceptor, r.encWiring, + r.metricsRegistry.S3BlobOffloadObserver(), + s3BlobPushBlocked, ); err != nil { return r.startupFailure(err) } @@ -3007,19 +3140,55 @@ func (r *runtimeServerRunner) startRaftTransport() error { func (r *runtimeServerRunner) prepareAdminForwardServers() error { r.dynamoServer = newDynamoDBServer(r.shardStore, r.coordinate, r.leaderDynamo, r.metricsRegistry, r.readTracker) + blobCluster, err := r.newS3BlobCluster() + if err != nil { + return err + } s3Server, err := newS3Server( r.s3Address, r.shardStore, r.coordinate, r.leaderS3, r.s3Region, r.s3CredsFile, r.s3PathStyleOnly, r.readTracker, r.metricsRegistry.S3PutAdmissionObserver(), r.metricsRegistry.S3BlobOffloadObserver(), + blobCluster, + r.publicKVGate.blocked, ) if err != nil { + if blobCluster != nil { + _ = blobCluster.Close() + } return err } r.s3Server = s3Server return nil } +func (r *runtimeServerRunner) newS3BlobCluster() (adapter.S3BlobCluster, error) { + if r == nil || strings.TrimSpace(r.s3Address) == "" { + return nil, nil + } + members, ok := r.coordinate.(kv.RaftMembershipCoordinator) + if !ok { + return nil, errors.New("S3 blob offload requires raft membership discovery") + } + adminToken := "" + if strings.TrimSpace(*adminTokenFile) != "" { + loaded, err := loadAdminTokenFile(*adminTokenFile) + if err != nil { + return nil, err + } + adminToken = loaded + } + peerToken := "" + if strings.TrimSpace(*s3BlobPeerTokenFile) != "" { + loaded, err := loadS3BlobPeerTokenFile(*s3BlobPeerTokenFile) + if err != nil { + return nil, err + } + peerToken = loaded + } + return adapter.NewGRPCS3BlobCluster(*raftId, members, adminToken, peerToken), nil +} + func (r *runtimeServerRunner) preparePublicServices() error { redisServer, redisCompactor, redisListener, err := prepareRedisServer( r.ctx, r.lc, r.redisAddress, r.shardStore, r.coordinate, diff --git a/main_admin_test.go b/main_admin_test.go index f12d7314d..d9191d7a6 100644 --- a/main_admin_test.go +++ b/main_admin_test.go @@ -26,13 +26,14 @@ import ( "github.com/bootjp/elastickv/internal/admin" "github.com/bootjp/elastickv/internal/raftengine" "github.com/bootjp/elastickv/kv" + proto "github.com/bootjp/elastickv/proto" "github.com/stretchr/testify/require" "golang.org/x/sync/errgroup" ) func TestConfigureAdminServiceDisabledByDefault(t *testing.T) { t.Parallel() - srv, icept, err := configureAdminService("", false, adapter.NodeIdentity{NodeID: "n1"}, nil) + srv, icept, err := configureAdminService("", false, "", false, adapter.NodeIdentity{NodeID: "n1"}, nil) if err != nil { t.Fatalf("disabled-by-default should not error: %v", err) } @@ -48,7 +49,7 @@ func TestConfigureAdminServiceRejectsMutualExclusion(t *testing.T) { if err := os.WriteFile(tokPath, []byte("x"), 0o600); err != nil { t.Fatal(err) } - if _, _, err := configureAdminService(tokPath, true, adapter.NodeIdentity{}, nil); err == nil { + if _, _, err := configureAdminService(tokPath, true, "", false, adapter.NodeIdentity{}, nil); err == nil { t.Fatal("expected mutual-exclusion error") } } @@ -57,25 +58,80 @@ func TestConfigureAdminServiceTokenFile(t *testing.T) { t.Parallel() dir := t.TempDir() tokPath := filepath.Join(dir, "t") + peerPath := filepath.Join(dir, "peer") if err := os.WriteFile(tokPath, []byte("hunter2\n"), 0o600); err != nil { t.Fatal(err) } - srv, icept, err := configureAdminService(tokPath, false, adapter.NodeIdentity{NodeID: "n1"}, nil) + if err := os.WriteFile(peerPath, []byte("peer-secret\n"), 0o600); err != nil { + t.Fatal(err) + } + srv, icept, err := configureAdminService(tokPath, false, peerPath, true, adapter.NodeIdentity{NodeID: "n1"}, nil) if err != nil { t.Fatalf("configureAdminService: %v", err) } if srv == nil { t.Fatal("expected an AdminServer instance") } - // Expect one unary + one stream interceptor for the admin-token gate. - if len(icept.unary) != 1 || len(icept.stream) != 1 { - t.Fatalf("expected 1 unary + 1 stream interceptor, got %d + %d", len(icept.unary), len(icept.stream)) + if len(icept.unary) != 2 || len(icept.stream) != 2 { + t.Fatalf("expected separate admin and peer interceptors, got %d + %d", len(icept.unary), len(icept.stream)) + } + if !icept.s3BlobFetchEnabled { + t.Fatal("token-authenticated configuration should enable S3BlobFetch") + } + overview, err := srv.GetClusterOverview(context.Background(), &proto.GetClusterOverviewRequest{}) + if err != nil { + t.Fatalf("GetClusterOverview: %v", err) + } + if !overview.GetCapabilities()[adapter.S3BlobOffloadCapabilityName] { + t.Fatal("enabled authenticated configuration should advertise S3 blob offload") + } +} + +func TestConfigureAdminServiceRejectsSharedAdminAndPeerToken(t *testing.T) { + t.Parallel() + dir := t.TempDir() + adminPath := filepath.Join(dir, "admin") + peerPath := filepath.Join(dir, "peer") + for _, path := range []string{adminPath, peerPath} { + if err := os.WriteFile(path, []byte("same-secret\n"), 0o600); err != nil { + t.Fatal(err) + } + } + if _, _, err := configureAdminService(adminPath, false, peerPath, true, adapter.NodeIdentity{}, nil); err == nil { + t.Fatal("expected shared admin and peer token to be rejected") + } +} + +func TestConfigureAdminServiceDoesNotAdvertiseDisabledS3BlobOffload(t *testing.T) { + t.Parallel() + dir := t.TempDir() + tokPath := filepath.Join(dir, "t") + peerPath := filepath.Join(dir, "peer") + if err := os.WriteFile(tokPath, []byte("hunter2\n"), 0o600); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(peerPath, []byte("peer-secret\n"), 0o600); err != nil { + t.Fatal(err) + } + srv, icept, err := configureAdminService(tokPath, false, peerPath, false, adapter.NodeIdentity{NodeID: "n1"}, nil) + if err != nil { + t.Fatalf("configureAdminService: %v", err) + } + if !icept.s3BlobFetchEnabled { + t.Fatal("rollback mode must retain authenticated blob reads") + } + overview, err := srv.GetClusterOverview(context.Background(), &proto.GetClusterOverviewRequest{}) + if err != nil { + t.Fatalf("GetClusterOverview: %v", err) + } + if overview.GetCapabilities()[adapter.S3BlobOffloadCapabilityName] { + t.Fatal("disabled rollout flag must remove S3 blob offload capability") } } func TestConfigureAdminServiceInsecureNoAuth(t *testing.T) { t.Parallel() - srv, icept, err := configureAdminService("", true, adapter.NodeIdentity{NodeID: "n1"}, nil) + srv, icept, err := configureAdminService("", true, "", true, adapter.NodeIdentity{NodeID: "n1"}, nil) if err != nil { t.Fatalf("insecure mode should succeed: %v", err) } @@ -85,6 +141,16 @@ func TestConfigureAdminServiceInsecureNoAuth(t *testing.T) { if !icept.empty() { t.Fatalf("insecure mode should not attach interceptors, got %+v", icept) } + if icept.s3BlobFetchEnabled { + t.Fatal("insecure admin mode must not expose S3BlobFetch") + } + overview, err := srv.GetClusterOverview(context.Background(), &proto.GetClusterOverviewRequest{}) + if err != nil { + t.Fatalf("GetClusterOverview: %v", err) + } + if overview.GetCapabilities()[adapter.S3BlobOffloadCapabilityName] { + t.Fatal("insecure admin mode must not advertise S3 blob offload capability") + } } func TestAdminMembersFromBootstrapExcludesSelf(t *testing.T) { diff --git a/main_encryption_rotate_on_startup_test.go b/main_encryption_rotate_on_startup_test.go index 5e16348ad..9d8063deb 100644 --- a/main_encryption_rotate_on_startup_test.go +++ b/main_encryption_rotate_on_startup_test.go @@ -432,6 +432,7 @@ func TestStartupPublicKVGate_BlocksMutatorsUntilReady(t *testing.T) { pb.EncryptionAdmin_ResyncSidecar_FullMethodName, pb.EncryptionAdmin_EnableStorageEnvelope_FullMethodName, pb.EncryptionAdmin_EnableRaftEnvelope_FullMethodName, + pb.S3BlobFetch_PushChunkBlob_FullMethodName, } for _, method := range blockedMethods { @@ -503,6 +504,42 @@ func TestStartupPublicKVGate_BlocksMutatorsUntilReady(t *testing.T) { } } +func TestStartupPublicKVGate_BlocksMutatorStreamsUntilReady(t *testing.T) { + t.Parallel() + gate := &startupPublicKVGate{} + + handlerCalled := false + handler := func(interface{}, grpc.ServerStream) error { + handlerCalled = true + return nil + } + err := gate.streamInterceptor( + nil, + nil, + &grpc.StreamServerInfo{FullMethod: pb.S3BlobFetch_PushChunkBlob_FullMethodName}, + handler, + ) + if status.Code(err) != codes.Unavailable { + t.Fatalf("PushChunkBlob before ready err=%v, want Unavailable", err) + } + if handlerCalled { + t.Fatal("stream handler ran before startup gate opened") + } + + gate.markReady() + if err := gate.streamInterceptor( + nil, + nil, + &grpc.StreamServerInfo{FullMethod: pb.S3BlobFetch_PushChunkBlob_FullMethodName}, + handler, + ); err != nil { + t.Fatalf("PushChunkBlob after ready err=%v", err) + } + if !handlerCalled { + t.Fatal("stream handler did not run after startup gate opened") + } +} + func TestStartupPublicKVGate_BlocksAfterReadyWhenStartupRotationPending(t *testing.T) { t.Parallel() blocked := true @@ -587,6 +624,9 @@ func TestRuntimeServerRunner_PrepareAdminForwardServersDoesNotBindPublicListener s3Address: heldS3.Addr().String(), s3PathStyleOnly: true, metricsRegistry: monitoring.NewRegistry("test-node", "127.0.0.1:0"), + coordinate: startupGatedCoordinator{ + inner: &stubStartupCoordinator{clock: kv.NewHLC()}, + }, } if err := runner.prepareAdminForwardServers(); err != nil { t.Fatalf("prepareAdminForwardServers: %v", err) diff --git a/main_s3.go b/main_s3.go index f23b58a3a..e80aedfba 100644 --- a/main_s3.go +++ b/main_s3.go @@ -63,7 +63,7 @@ func prepareS3Server( ) (*adapter.S3Server, net.Listener, error) { s3Server, err := newS3Server( s3Addr, shardStore, coordinate, leaderS3, region, credentialsFile, - pathStyleOnly, readTracker, putAdmissionObserver, blobOffloadObserver, + pathStyleOnly, readTracker, putAdmissionObserver, blobOffloadObserver, nil, nil, ) if err != nil { return nil, nil, err @@ -86,6 +86,8 @@ func newS3Server( readTracker *kv.ActiveTimestampTracker, putAdmissionObserver adapter.S3PutAdmissionObserver, blobOffloadObserver adapter.S3BlobOffloadObserver, + blobCluster adapter.S3BlobCluster, + blobPushBlocked func() bool, ) (*adapter.S3Server, error) { s3Addr = strings.TrimSpace(s3Addr) if s3Addr == "" { @@ -102,17 +104,29 @@ func newS3Server( if err != nil { return nil, err } + minReplicas, err := adapter.S3BlobMinReplicasFromEnv() + if err != nil { + return nil, errors.WithStack(err) + } + options := []adapter.S3ServerOption{ + adapter.WithS3Region(region), + adapter.WithS3StaticCredentials(staticCreds), + adapter.WithS3ActiveTimestampTracker(readTracker), + adapter.WithS3PutAdmissionObserver(putAdmissionObserver), + adapter.WithS3BlobOffloadObserver(blobOffloadObserver), + adapter.WithS3BlobMinReplicas(minReplicas), + adapter.WithS3BlobPushBlocked(blobPushBlocked), + } + if blobCluster != nil { + options = append(options, adapter.WithS3BlobCluster(blobCluster)) + } s3Server := adapter.NewS3Server( nil, s3Addr, shardStore, coordinate, leaderS3, - adapter.WithS3Region(region), - adapter.WithS3StaticCredentials(staticCreds), - adapter.WithS3ActiveTimestampTracker(readTracker), - adapter.WithS3PutAdmissionObserver(putAdmissionObserver), - adapter.WithS3BlobOffloadObserver(blobOffloadObserver), + options..., ) return s3Server, nil } diff --git a/proto/service.pb.go b/proto/service.pb.go index 00000ae2f..04d7cce82 100644 --- a/proto/service.pb.go +++ b/proto/service.pb.go @@ -2223,6 +2223,214 @@ func (x *TransferTarget) GetTargetAddress() string { return "" } +type FetchChunkBlobRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + ContentSha256 []byte `protobuf:"bytes,1,opt,name=content_sha256,json=contentSha256,proto3" json:"content_sha256,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *FetchChunkBlobRequest) Reset() { + *x = FetchChunkBlobRequest{} + mi := &file_service_proto_msgTypes[39] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *FetchChunkBlobRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*FetchChunkBlobRequest) ProtoMessage() {} + +func (x *FetchChunkBlobRequest) ProtoReflect() protoreflect.Message { + mi := &file_service_proto_msgTypes[39] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use FetchChunkBlobRequest.ProtoReflect.Descriptor instead. +func (*FetchChunkBlobRequest) Descriptor() ([]byte, []int) { + return file_service_proto_rawDescGZIP(), []int{39} +} + +func (x *FetchChunkBlobRequest) GetContentSha256() []byte { + if x != nil { + return x.ContentSha256 + } + return nil +} + +type FetchChunkBlobResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Payload []byte `protobuf:"bytes,1,opt,name=payload,proto3" json:"payload,omitempty"` + Eof bool `protobuf:"varint,2,opt,name=eof,proto3" json:"eof,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *FetchChunkBlobResponse) Reset() { + *x = FetchChunkBlobResponse{} + mi := &file_service_proto_msgTypes[40] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *FetchChunkBlobResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*FetchChunkBlobResponse) ProtoMessage() {} + +func (x *FetchChunkBlobResponse) ProtoReflect() protoreflect.Message { + mi := &file_service_proto_msgTypes[40] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use FetchChunkBlobResponse.ProtoReflect.Descriptor instead. +func (*FetchChunkBlobResponse) Descriptor() ([]byte, []int) { + return file_service_proto_rawDescGZIP(), []int{40} +} + +func (x *FetchChunkBlobResponse) GetPayload() []byte { + if x != nil { + return x.Payload + } + return nil +} + +func (x *FetchChunkBlobResponse) GetEof() bool { + if x != nil { + return x.Eof + } + return false +} + +type PushChunkBlobRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + ContentSha256 []byte `protobuf:"bytes,1,opt,name=content_sha256,json=contentSha256,proto3" json:"content_sha256,omitempty"` + Payload []byte `protobuf:"bytes,2,opt,name=payload,proto3" json:"payload,omitempty"` + Eof bool `protobuf:"varint,3,opt,name=eof,proto3" json:"eof,omitempty"` + CommitTs uint64 `protobuf:"varint,4,opt,name=commit_ts,json=commitTs,proto3" json:"commit_ts,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *PushChunkBlobRequest) Reset() { + *x = PushChunkBlobRequest{} + mi := &file_service_proto_msgTypes[41] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *PushChunkBlobRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*PushChunkBlobRequest) ProtoMessage() {} + +func (x *PushChunkBlobRequest) ProtoReflect() protoreflect.Message { + mi := &file_service_proto_msgTypes[41] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use PushChunkBlobRequest.ProtoReflect.Descriptor instead. +func (*PushChunkBlobRequest) Descriptor() ([]byte, []int) { + return file_service_proto_rawDescGZIP(), []int{41} +} + +func (x *PushChunkBlobRequest) GetContentSha256() []byte { + if x != nil { + return x.ContentSha256 + } + return nil +} + +func (x *PushChunkBlobRequest) GetPayload() []byte { + if x != nil { + return x.Payload + } + return nil +} + +func (x *PushChunkBlobRequest) GetEof() bool { + if x != nil { + return x.Eof + } + return false +} + +func (x *PushChunkBlobRequest) GetCommitTs() uint64 { + if x != nil { + return x.CommitTs + } + return 0 +} + +type PushChunkBlobResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Durable bool `protobuf:"varint,1,opt,name=durable,proto3" json:"durable,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *PushChunkBlobResponse) Reset() { + *x = PushChunkBlobResponse{} + mi := &file_service_proto_msgTypes[42] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *PushChunkBlobResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*PushChunkBlobResponse) ProtoMessage() {} + +func (x *PushChunkBlobResponse) ProtoReflect() protoreflect.Message { + mi := &file_service_proto_msgTypes[42] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use PushChunkBlobResponse.ProtoReflect.Descriptor instead. +func (*PushChunkBlobResponse) Descriptor() ([]byte, []int) { + return file_service_proto_rawDescGZIP(), []int{42} +} + +func (x *PushChunkBlobResponse) GetDurable() bool { + if x != nil { + return x.Durable + } + return false +} + type RaftAdminTransferLeadershipResponse struct { state protoimpl.MessageState `protogen:"open.v1"` unknownFields protoimpl.UnknownFields @@ -2231,7 +2439,7 @@ type RaftAdminTransferLeadershipResponse struct { func (x *RaftAdminTransferLeadershipResponse) Reset() { *x = RaftAdminTransferLeadershipResponse{} - mi := &file_service_proto_msgTypes[39] + mi := &file_service_proto_msgTypes[43] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -2243,7 +2451,7 @@ func (x *RaftAdminTransferLeadershipResponse) String() string { func (*RaftAdminTransferLeadershipResponse) ProtoMessage() {} func (x *RaftAdminTransferLeadershipResponse) ProtoReflect() protoreflect.Message { - mi := &file_service_proto_msgTypes[39] + mi := &file_service_proto_msgTypes[43] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -2256,7 +2464,7 @@ func (x *RaftAdminTransferLeadershipResponse) ProtoReflect() protoreflect.Messag // Deprecated: Use RaftAdminTransferLeadershipResponse.ProtoReflect.Descriptor instead. func (*RaftAdminTransferLeadershipResponse) Descriptor() ([]byte, []int) { - return file_service_proto_rawDescGZIP(), []int{39} + return file_service_proto_rawDescGZIP(), []int{43} } var File_service_proto protoreflect.FileDescriptor @@ -2398,7 +2606,19 @@ const file_service_proto_rawDesc = "" + "\x11target_candidates\x18\x05 \x03(\v2\x0f.TransferTargetR\x10targetCandidates\"T\n" + "\x0eTransferTarget\x12\x1b\n" + "\ttarget_id\x18\x01 \x01(\tR\btargetId\x12%\n" + - "\x0etarget_address\x18\x02 \x01(\tR\rtargetAddress\"%\n" + + "\x0etarget_address\x18\x02 \x01(\tR\rtargetAddress\">\n" + + "\x15FetchChunkBlobRequest\x12%\n" + + "\x0econtent_sha256\x18\x01 \x01(\fR\rcontentSha256\"D\n" + + "\x16FetchChunkBlobResponse\x12\x18\n" + + "\apayload\x18\x01 \x01(\fR\apayload\x12\x10\n" + + "\x03eof\x18\x02 \x01(\bR\x03eof\"\x86\x01\n" + + "\x14PushChunkBlobRequest\x12%\n" + + "\x0econtent_sha256\x18\x01 \x01(\fR\rcontentSha256\x12\x18\n" + + "\apayload\x18\x02 \x01(\fR\apayload\x12\x10\n" + + "\x03eof\x18\x03 \x01(\bR\x03eof\x12\x1b\n" + + "\tcommit_ts\x18\x04 \x01(\x04R\bcommitTs\"1\n" + + "\x15PushChunkBlobResponse\x12\x18\n" + + "\adurable\x18\x01 \x01(\bR\adurable\"%\n" + "#RaftAdminTransferLeadershipResponse*\xa9\x01\n" + "\x0eRaftAdminState\x12\x1c\n" + "\x18RAFT_ADMIN_STATE_UNKNOWN\x10\x00\x12\x1d\n" + @@ -2428,7 +2648,10 @@ const file_service_proto_rawDesc = "" + "AddLearner\x12\x1b.RaftAdminAddLearnerRequest\x1a%.RaftAdminConfigurationChangeResponse\"\x00\x12Z\n" + "\x0ePromoteLearner\x12\x1f.RaftAdminPromoteLearnerRequest\x1a%.RaftAdminConfigurationChangeResponse\"\x00\x12V\n" + "\fRemoveServer\x12\x1d.RaftAdminRemoveServerRequest\x1a%.RaftAdminConfigurationChangeResponse\"\x00\x12a\n" + - "\x12TransferLeadership\x12#.RaftAdminTransferLeadershipRequest\x1a$.RaftAdminTransferLeadershipResponse\"\x00B#Z!github.com/bootjp/elastickv/protob\x06proto3" + "\x12TransferLeadership\x12#.RaftAdminTransferLeadershipRequest\x1a$.RaftAdminTransferLeadershipResponse\"\x002\x98\x01\n" + + "\vS3BlobFetch\x12E\n" + + "\x0eFetchChunkBlob\x12\x16.FetchChunkBlobRequest\x1a\x17.FetchChunkBlobResponse\"\x000\x01\x12B\n" + + "\rPushChunkBlob\x12\x15.PushChunkBlobRequest\x1a\x16.PushChunkBlobResponse\"\x00(\x01B#Z!github.com/bootjp/elastickv/protob\x06proto3" var ( file_service_proto_rawDescOnce sync.Once @@ -2443,7 +2666,7 @@ func file_service_proto_rawDescGZIP() []byte { } var file_service_proto_enumTypes = make([]protoimpl.EnumInfo, 1) -var file_service_proto_msgTypes = make([]protoimpl.MessageInfo, 40) +var file_service_proto_msgTypes = make([]protoimpl.MessageInfo, 44) var file_service_proto_goTypes = []any{ (RaftAdminState)(0), // 0: RaftAdminState (*RawPutRequest)(nil), // 1: RawPutRequest @@ -2485,7 +2708,11 @@ var file_service_proto_goTypes = []any{ (*RaftAdminConfigurationChangeResponse)(nil), // 37: RaftAdminConfigurationChangeResponse (*RaftAdminTransferLeadershipRequest)(nil), // 38: RaftAdminTransferLeadershipRequest (*TransferTarget)(nil), // 39: TransferTarget - (*RaftAdminTransferLeadershipResponse)(nil), // 40: RaftAdminTransferLeadershipResponse + (*FetchChunkBlobRequest)(nil), // 40: FetchChunkBlobRequest + (*FetchChunkBlobResponse)(nil), // 41: FetchChunkBlobResponse + (*PushChunkBlobRequest)(nil), // 42: PushChunkBlobRequest + (*PushChunkBlobResponse)(nil), // 43: PushChunkBlobResponse + (*RaftAdminTransferLeadershipResponse)(nil), // 44: RaftAdminTransferLeadershipResponse } var file_service_proto_depIdxs = []int32{ 10, // 0: RawScanAtResponse.kv:type_name -> RawKVPair @@ -2516,27 +2743,31 @@ var file_service_proto_depIdxs = []int32{ 35, // 25: RaftAdmin.PromoteLearner:input_type -> RaftAdminPromoteLearnerRequest 36, // 26: RaftAdmin.RemoveServer:input_type -> RaftAdminRemoveServerRequest 38, // 27: RaftAdmin.TransferLeadership:input_type -> RaftAdminTransferLeadershipRequest - 2, // 28: RawKV.RawPut:output_type -> RawPutResponse - 4, // 29: RawKV.RawGet:output_type -> RawGetResponse - 6, // 30: RawKV.RawDelete:output_type -> RawDeleteResponse - 8, // 31: RawKV.RawLatestCommitTS:output_type -> RawLatestCommitTSResponse - 11, // 32: RawKV.RawScanAt:output_type -> RawScanAtResponse - 13, // 33: TransactionalKV.Put:output_type -> PutResponse - 17, // 34: TransactionalKV.Get:output_type -> GetResponse - 15, // 35: TransactionalKV.Delete:output_type -> DeleteResponse - 21, // 36: TransactionalKV.Scan:output_type -> ScanResponse - 23, // 37: TransactionalKV.PreWrite:output_type -> PreCommitResponse - 25, // 38: TransactionalKV.Commit:output_type -> CommitResponse - 27, // 39: TransactionalKV.Rollback:output_type -> RollbackResponse - 29, // 40: RaftAdmin.Status:output_type -> RaftAdminStatusResponse - 32, // 41: RaftAdmin.Configuration:output_type -> RaftAdminConfigurationResponse - 37, // 42: RaftAdmin.AddVoter:output_type -> RaftAdminConfigurationChangeResponse - 37, // 43: RaftAdmin.AddLearner:output_type -> RaftAdminConfigurationChangeResponse - 37, // 44: RaftAdmin.PromoteLearner:output_type -> RaftAdminConfigurationChangeResponse - 37, // 45: RaftAdmin.RemoveServer:output_type -> RaftAdminConfigurationChangeResponse - 40, // 46: RaftAdmin.TransferLeadership:output_type -> RaftAdminTransferLeadershipResponse - 28, // [28:47] is the sub-list for method output_type - 9, // [9:28] is the sub-list for method input_type + 40, // 28: S3BlobFetch.FetchChunkBlob:input_type -> FetchChunkBlobRequest + 42, // 29: S3BlobFetch.PushChunkBlob:input_type -> PushChunkBlobRequest + 2, // 30: RawKV.RawPut:output_type -> RawPutResponse + 4, // 31: RawKV.RawGet:output_type -> RawGetResponse + 6, // 32: RawKV.RawDelete:output_type -> RawDeleteResponse + 8, // 33: RawKV.RawLatestCommitTS:output_type -> RawLatestCommitTSResponse + 11, // 34: RawKV.RawScanAt:output_type -> RawScanAtResponse + 13, // 35: TransactionalKV.Put:output_type -> PutResponse + 17, // 36: TransactionalKV.Get:output_type -> GetResponse + 15, // 37: TransactionalKV.Delete:output_type -> DeleteResponse + 21, // 38: TransactionalKV.Scan:output_type -> ScanResponse + 23, // 39: TransactionalKV.PreWrite:output_type -> PreCommitResponse + 25, // 40: TransactionalKV.Commit:output_type -> CommitResponse + 27, // 41: TransactionalKV.Rollback:output_type -> RollbackResponse + 29, // 42: RaftAdmin.Status:output_type -> RaftAdminStatusResponse + 32, // 43: RaftAdmin.Configuration:output_type -> RaftAdminConfigurationResponse + 37, // 44: RaftAdmin.AddVoter:output_type -> RaftAdminConfigurationChangeResponse + 37, // 45: RaftAdmin.AddLearner:output_type -> RaftAdminConfigurationChangeResponse + 37, // 46: RaftAdmin.PromoteLearner:output_type -> RaftAdminConfigurationChangeResponse + 37, // 47: RaftAdmin.RemoveServer:output_type -> RaftAdminConfigurationChangeResponse + 44, // 48: RaftAdmin.TransferLeadership:output_type -> RaftAdminTransferLeadershipResponse + 41, // 49: S3BlobFetch.FetchChunkBlob:output_type -> FetchChunkBlobResponse + 43, // 50: S3BlobFetch.PushChunkBlob:output_type -> PushChunkBlobResponse + 30, // [30:51] is the sub-list for method output_type + 9, // [9:30] is the sub-list for method input_type 9, // [9:9] is the sub-list for extension type_name 9, // [9:9] is the sub-list for extension extendee 0, // [0:9] is the sub-list for field type_name @@ -2553,9 +2784,9 @@ func file_service_proto_init() { GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_service_proto_rawDesc), len(file_service_proto_rawDesc)), NumEnums: 1, - NumMessages: 40, + NumMessages: 44, NumExtensions: 0, - NumServices: 3, + NumServices: 4, }, GoTypes: file_service_proto_goTypes, DependencyIndexes: file_service_proto_depIdxs, diff --git a/proto/service.proto b/proto/service.proto index 68d12d5f7..a2212b101 100644 --- a/proto/service.proto +++ b/proto/service.proto @@ -31,6 +31,11 @@ service RaftAdmin { rpc TransferLeadership(RaftAdminTransferLeadershipRequest) returns (RaftAdminTransferLeadershipResponse) {} } +service S3BlobFetch { + rpc FetchChunkBlob(FetchChunkBlobRequest) returns (stream FetchChunkBlobResponse) {} + rpc PushChunkBlob(stream PushChunkBlobRequest) returns (PushChunkBlobResponse) {} +} + message RawPutRequest { bytes key = 1; bytes value = 2; @@ -261,4 +266,24 @@ message TransferTarget { string target_address = 2; } +message FetchChunkBlobRequest { + bytes content_sha256 = 1; +} + +message FetchChunkBlobResponse { + bytes payload = 1; + bool eof = 2; +} + +message PushChunkBlobRequest { + bytes content_sha256 = 1; + bytes payload = 2; + bool eof = 3; + uint64 commit_ts = 4; +} + +message PushChunkBlobResponse { + bool durable = 1; +} + message RaftAdminTransferLeadershipResponse {} diff --git a/proto/service_grpc.pb.go b/proto/service_grpc.pb.go index 484d04c64..1254bcbf9 100644 --- a/proto/service_grpc.pb.go +++ b/proto/service_grpc.pb.go @@ -931,3 +931,139 @@ var RaftAdmin_ServiceDesc = grpc.ServiceDesc{ Streams: []grpc.StreamDesc{}, Metadata: "service.proto", } + +const ( + S3BlobFetch_FetchChunkBlob_FullMethodName = "/S3BlobFetch/FetchChunkBlob" + S3BlobFetch_PushChunkBlob_FullMethodName = "/S3BlobFetch/PushChunkBlob" +) + +// S3BlobFetchClient is the client API for S3BlobFetch service. +// +// For semantics around ctx use and closing/ending streaming RPCs, please refer to https://pkg.go.dev/google.golang.org/grpc/?tab=doc#ClientConn.NewStream. +type S3BlobFetchClient interface { + FetchChunkBlob(ctx context.Context, in *FetchChunkBlobRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[FetchChunkBlobResponse], error) + PushChunkBlob(ctx context.Context, opts ...grpc.CallOption) (grpc.ClientStreamingClient[PushChunkBlobRequest, PushChunkBlobResponse], error) +} + +type s3BlobFetchClient struct { + cc grpc.ClientConnInterface +} + +func NewS3BlobFetchClient(cc grpc.ClientConnInterface) S3BlobFetchClient { + return &s3BlobFetchClient{cc} +} + +func (c *s3BlobFetchClient) FetchChunkBlob(ctx context.Context, in *FetchChunkBlobRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[FetchChunkBlobResponse], error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + stream, err := c.cc.NewStream(ctx, &S3BlobFetch_ServiceDesc.Streams[0], S3BlobFetch_FetchChunkBlob_FullMethodName, cOpts...) + if err != nil { + return nil, err + } + x := &grpc.GenericClientStream[FetchChunkBlobRequest, FetchChunkBlobResponse]{ClientStream: stream} + if err := x.ClientStream.SendMsg(in); err != nil { + return nil, err + } + if err := x.ClientStream.CloseSend(); err != nil { + return nil, err + } + return x, nil +} + +// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name. +type S3BlobFetch_FetchChunkBlobClient = grpc.ServerStreamingClient[FetchChunkBlobResponse] + +func (c *s3BlobFetchClient) PushChunkBlob(ctx context.Context, opts ...grpc.CallOption) (grpc.ClientStreamingClient[PushChunkBlobRequest, PushChunkBlobResponse], error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + stream, err := c.cc.NewStream(ctx, &S3BlobFetch_ServiceDesc.Streams[1], S3BlobFetch_PushChunkBlob_FullMethodName, cOpts...) + if err != nil { + return nil, err + } + x := &grpc.GenericClientStream[PushChunkBlobRequest, PushChunkBlobResponse]{ClientStream: stream} + return x, nil +} + +// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name. +type S3BlobFetch_PushChunkBlobClient = grpc.ClientStreamingClient[PushChunkBlobRequest, PushChunkBlobResponse] + +// S3BlobFetchServer is the server API for S3BlobFetch service. +// All implementations must embed UnimplementedS3BlobFetchServer +// for forward compatibility. +type S3BlobFetchServer interface { + FetchChunkBlob(*FetchChunkBlobRequest, grpc.ServerStreamingServer[FetchChunkBlobResponse]) error + PushChunkBlob(grpc.ClientStreamingServer[PushChunkBlobRequest, PushChunkBlobResponse]) error + mustEmbedUnimplementedS3BlobFetchServer() +} + +// UnimplementedS3BlobFetchServer must be embedded to have +// forward compatible implementations. +// +// NOTE: this should be embedded by value instead of pointer to avoid a nil +// pointer dereference when methods are called. +type UnimplementedS3BlobFetchServer struct{} + +func (UnimplementedS3BlobFetchServer) FetchChunkBlob(*FetchChunkBlobRequest, grpc.ServerStreamingServer[FetchChunkBlobResponse]) error { + return status.Error(codes.Unimplemented, "method FetchChunkBlob not implemented") +} +func (UnimplementedS3BlobFetchServer) PushChunkBlob(grpc.ClientStreamingServer[PushChunkBlobRequest, PushChunkBlobResponse]) error { + return status.Error(codes.Unimplemented, "method PushChunkBlob not implemented") +} +func (UnimplementedS3BlobFetchServer) mustEmbedUnimplementedS3BlobFetchServer() {} +func (UnimplementedS3BlobFetchServer) testEmbeddedByValue() {} + +// UnsafeS3BlobFetchServer may be embedded to opt out of forward compatibility for this service. +// Use of this interface is not recommended, as added methods to S3BlobFetchServer will +// result in compilation errors. +type UnsafeS3BlobFetchServer interface { + mustEmbedUnimplementedS3BlobFetchServer() +} + +func RegisterS3BlobFetchServer(s grpc.ServiceRegistrar, srv S3BlobFetchServer) { + // If the following call panics, it indicates UnimplementedS3BlobFetchServer was + // embedded by pointer and is nil. This will cause panics if an + // unimplemented method is ever invoked, so we test this at initialization + // time to prevent it from happening at runtime later due to I/O. + if t, ok := srv.(interface{ testEmbeddedByValue() }); ok { + t.testEmbeddedByValue() + } + s.RegisterService(&S3BlobFetch_ServiceDesc, srv) +} + +func _S3BlobFetch_FetchChunkBlob_Handler(srv interface{}, stream grpc.ServerStream) error { + m := new(FetchChunkBlobRequest) + if err := stream.RecvMsg(m); err != nil { + return err + } + return srv.(S3BlobFetchServer).FetchChunkBlob(m, &grpc.GenericServerStream[FetchChunkBlobRequest, FetchChunkBlobResponse]{ServerStream: stream}) +} + +// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name. +type S3BlobFetch_FetchChunkBlobServer = grpc.ServerStreamingServer[FetchChunkBlobResponse] + +func _S3BlobFetch_PushChunkBlob_Handler(srv interface{}, stream grpc.ServerStream) error { + return srv.(S3BlobFetchServer).PushChunkBlob(&grpc.GenericServerStream[PushChunkBlobRequest, PushChunkBlobResponse]{ServerStream: stream}) +} + +// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name. +type S3BlobFetch_PushChunkBlobServer = grpc.ClientStreamingServer[PushChunkBlobRequest, PushChunkBlobResponse] + +// S3BlobFetch_ServiceDesc is the grpc.ServiceDesc for S3BlobFetch service. +// It's only intended for direct use with grpc.RegisterService, +// and not to be introspected or modified (even as a copy) +var S3BlobFetch_ServiceDesc = grpc.ServiceDesc{ + ServiceName: "S3BlobFetch", + HandlerType: (*S3BlobFetchServer)(nil), + Methods: []grpc.MethodDesc{}, + Streams: []grpc.StreamDesc{ + { + StreamName: "FetchChunkBlob", + Handler: _S3BlobFetch_FetchChunkBlob_Handler, + ServerStreams: true, + }, + { + StreamName: "PushChunkBlob", + Handler: _S3BlobFetch_PushChunkBlob_Handler, + ClientStreams: true, + }, + }, + Metadata: "service.proto", +} diff --git a/store/lsm_store.go b/store/lsm_store.go index 665b6396b..d104c0eec 100644 --- a/store/lsm_store.go +++ b/store/lsm_store.go @@ -176,6 +176,10 @@ func resolvePebbleCacheBytes(envVal string) int64 { } var metaLastCommitTSBytes = []byte(metaLastCommitTS) + +var localAuxiliaryMVCCPrefixes = [][]byte{ + []byte("!s3|chunkblob|"), +} var metaMinRetainedTSBytes = []byte(metaMinRetainedTS) var metaPendingMinRetainedTSBytes = []byte(metaPendingMinRetainedTS) var metaAppliedIndexBytes = []byte(metaAppliedIndex) @@ -659,6 +663,9 @@ func (s *pebbleStore) findMaxDataCommitTS() (uint64, error) { if userKey == nil { continue } + if isLocalAuxiliaryMVCCKey(userKey) { + continue + } if ts > maxTS { maxTS = ts } @@ -669,6 +676,15 @@ func (s *pebbleStore) findMaxDataCommitTS() (uint64, error) { return maxTS, nil } +func isLocalAuxiliaryMVCCKey(userKey []byte) bool { + for _, prefix := range localAuxiliaryMVCCPrefixes { + if bytes.HasPrefix(userKey, prefix) { + return true + } + } + return false +} + func (s *pebbleStore) findMinRetainedTS() (uint64, error) { return readPebbleUint64(s.db, metaMinRetainedTSBytes) } @@ -2126,7 +2142,14 @@ func (s *pebbleStore) ApplyMutations(ctx context.Context, mutations []*KVPairMut // registration commits. // appliedIndex=0: direct path has no raft index; the leaf treats 0 as // "do not write metaAppliedIndex" so the meta key stays unchanged. - return s.applyMutationsWithOpts(ctx, mutations, readKeys, startTS, commitTS, s.directApplyWriteOpts(), true, 0) + return s.applyMutationsWithOpts(ctx, mutations, readKeys, startTS, commitTS, s.directApplyWriteOpts(), true, 0, true) +} + +// ApplyMutationsPreservingLastCommitTS is a direct, durable write path for +// local auxiliary MVCC keys that must not advance the store-wide safe snapshot +// watermark. It still uses pebble.Sync and the direct writer-registration gate. +func (s *pebbleStore) ApplyMutationsPreservingLastCommitTS(ctx context.Context, mutations []*KVPairMutation, readKeys [][]byte, startTS, commitTS uint64) error { + return s.applyMutationsWithOpts(ctx, mutations, readKeys, startTS, commitTS, s.directApplyWriteOpts(), true, 0, false) } // ApplyMutationsRaft is the raft-apply commit path. Durability is governed @@ -2155,7 +2178,7 @@ func (s *pebbleStore) ApplyMutationsRaft(ctx context.Context, mutations []*KVPai // land here; their LastAppliedIndex() will stay behind the snapshot // pointer and the skip optimisation will fall back to full restore // for them. Preferred path is ApplyMutationsRaftAt. - return s.applyMutationsWithOpts(ctx, mutations, readKeys, startTS, commitTS, s.raftApplyWriteOpts(), false, 0) + return s.applyMutationsWithOpts(ctx, mutations, readKeys, startTS, commitTS, s.raftApplyWriteOpts(), false, 0, true) } // ApplyMutationsRaftAt is ApplyMutationsRaft with the raft entry @@ -2167,7 +2190,7 @@ func (s *pebbleStore) ApplyMutationsRaft(ctx context.Context, mutations []*KVPai // Production callers (kvFSM.applyXxx with f.pendingApplyIdx) SHOULD // pass the entry.Index value the engine delivered via SetApplyIndex. func (s *pebbleStore) ApplyMutationsRaftAt(ctx context.Context, mutations []*KVPairMutation, readKeys [][]byte, startTS, commitTS, appliedIndex uint64) error { - return s.applyMutationsWithOpts(ctx, mutations, readKeys, startTS, commitTS, s.raftApplyWriteOpts(), false, appliedIndex) + return s.applyMutationsWithOpts(ctx, mutations, readKeys, startTS, commitTS, s.raftApplyWriteOpts(), false, appliedIndex, true) } func (s *pebbleStore) raftApplyAlreadyLandedLocked(mutations []*KVPairMutation, commitTS uint64) (bool, error) { @@ -2260,8 +2283,8 @@ func (s *pebbleStore) hasCommitsAfter(startTS uint64) bool { return s.lastCommitTS > startTS } -func (s *pebbleStore) checkApplyConflicts(ctx context.Context, mutations []*KVPairMutation, readKeys [][]byte, startTS uint64) error { - if !s.hasCommitsAfter(startTS) { +func (s *pebbleStore) checkApplyConflicts(ctx context.Context, mutations []*KVPairMutation, readKeys [][]byte, startTS uint64, force bool) error { + if !force && !s.hasCommitsAfter(startTS) { return nil } if err := s.checkConflicts(ctx, mutations, startTS); err != nil { @@ -2270,7 +2293,7 @@ func (s *pebbleStore) checkApplyConflicts(ctx context.Context, mutations []*KVPa return s.checkReadConflicts(ctx, readKeys, startTS) } -func (s *pebbleStore) applyMutationsWithOpts(ctx context.Context, mutations []*KVPairMutation, readKeys [][]byte, startTS, commitTS uint64, writeOpts *pebble.WriteOptions, gateRegistration bool, appliedIndex uint64) error { +func (s *pebbleStore) applyMutationsWithOpts(ctx context.Context, mutations []*KVPairMutation, readKeys [][]byte, startTS, commitTS uint64, writeOpts *pebble.WriteOptions, gateRegistration bool, appliedIndex uint64, advanceLastCommitTS bool) error { s.dbMu.RLock() defer s.dbMu.RUnlock() @@ -2293,7 +2316,7 @@ func (s *pebbleStore) applyMutationsWithOpts(ctx context.Context, mutations []*K // Keep OCC at the store boundary: callers pass the read snapshot via // startTS/readKeys, and leader-issued commit timestamps alone do not prove // the read set is still current. - if err := s.checkApplyConflicts(ctx, mutations, readKeys, startTS); err != nil { + if err := s.checkApplyConflicts(ctx, mutations, readKeys, startTS, !advanceLastCommitTS); err != nil { return err } @@ -2301,39 +2324,51 @@ func (s *pebbleStore) applyMutationsWithOpts(ctx context.Context, mutations []*K return err } - // Hold mtx across read → batch-set → commit → in-memory update so that a - // concurrent alignCommitTS (PutAt/DeleteAt/ExpireAt) cannot advance+persist - // metaLastCommitTS between our read and batch commit, which would let this - // batch overwrite the meta key with a smaller value. - s.mtx.Lock() - newLastTS := s.lastCommitTS - if commitTS > newLastTS { - newLastTS = commitTS - } - if err := setPebbleUint64InBatch(b, metaLastCommitTSBytes, newLastTS); err != nil { - s.mtx.Unlock() + newLastTS, unlockLastCommitTS, err := s.stageLastCommitTSInBatch(b, commitTS, advanceLastCommitTS) + if err != nil { return err } + defer unlockLastCommitTS() // Bundle metaAppliedIndex in the same batch as the data + commitTS // meta key so a crash either commits all three atomically or none. // appliedIndex==0 is the legacy / non-raft callers (ApplyMutations // or ApplyMutationsRaft); they leave the key unchanged. if appliedIndex > 0 { if err := setPebbleUint64InBatch(b, metaAppliedIndexBytes, appliedIndex); err != nil { - s.mtx.Unlock() return err } } if err := b.Commit(writeOpts); err != nil { - s.mtx.Unlock() return errors.WithStack(err) } - s.updateLastCommitTS(newLastTS) - s.mtx.Unlock() + if advanceLastCommitTS { + s.updateLastCommitTS(newLastTS) + } return nil } +func (s *pebbleStore) stageLastCommitTSInBatch(b *pebble.Batch, commitTS uint64, advance bool) (uint64, func(), error) { + if !advance { + return 0, func() {}, nil + } + + // Hold mtx across read → batch-set → commit → in-memory update so that a + // concurrent alignCommitTS (PutAt/DeleteAt/ExpireAt) cannot advance+persist + // metaLastCommitTS between our read and batch commit, which would let this + // batch overwrite the meta key with a smaller value. + s.mtx.Lock() + newLastTS := s.lastCommitTS + if commitTS > newLastTS { + newLastTS = commitTS + } + if err := setPebbleUint64InBatch(b, metaLastCommitTSBytes, newLastTS); err != nil { + s.mtx.Unlock() + return 0, nil, err + } + return newLastTS, s.mtx.Unlock, nil +} + // DeletePrefixAt atomically deletes all visible keys matching prefix by writing // tombstone versions at commitTS. An empty prefix deletes all keys. Keys // matching excludePrefix are preserved. Uses an iterator-based approach that diff --git a/store/mvcc_store.go b/store/mvcc_store.go index 99737db0e..b1302d447 100644 --- a/store/mvcc_store.go +++ b/store/mvcc_store.go @@ -683,6 +683,41 @@ func (s *mvccStore) ApplyMutations(ctx context.Context, mutations []*KVPairMutat return nil } +// ApplyMutationsPreservingLastCommitTS validates and writes versions without +// advancing the store-wide LastCommitTS watermark. This is for local auxiliary +// data whose timestamps mirror a remote Raft leader's commit order but must not +// advertise that this replica has applied ordinary Raft entries up to that +// timestamp. +func (s *mvccStore) ApplyMutationsPreservingLastCommitTS(ctx context.Context, mutations []*KVPairMutation, readKeys [][]byte, startTS, commitTS uint64) error { + s.mtx.Lock() + defer s.mtx.Unlock() + + if err := s.checkConflictsLocked(mutations, readKeys, startTS); err != nil { + return err + } + + for _, mut := range mutations { + switch mut.Op { + case OpTypePut: + if err := validateValueSize(mut.Value); err != nil { + return err + } + s.putVersionLocked(mut.Key, mut.Value, commitTS, mut.ExpireAt) + case OpTypeDelete: + s.deleteVersionLocked(mut.Key, commitTS) + default: + return errors.WithStack(ErrUnknownOp) + } + s.log.InfoContext(ctx, "apply mutation preserving last commit timestamp", + slog.String("key", string(mut.Key)), + slog.Uint64("commit_ts", commitTS), + slog.Bool("delete", mut.Op == OpTypeDelete), + ) + } + + return nil +} + func (s *mvccStore) checkConflictsLocked(mutations []*KVPairMutation, readKeys [][]byte, startTS uint64) error { for _, mut := range mutations { if latestVer, ok := s.latestVersionLocked(mut.Key); ok && latestVer.TS > startTS {