diff --git a/adapter/s3_blob_fetch.go b/adapter/s3_blob_fetch.go new file mode 100644 index 000000000..d0923fd2c --- /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, digest, payload, commitTS); err != nil { + return err + } + return sendChunkBlobPushAck(stream) +} + +func (s *S3BlobFetchServer) storeChunkBlob( + stream pb.S3BlobFetch_PushChunkBlobServer, + 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(stream.Context(), key, digest, payload, commitTS) + if err != nil { + return err + } + if startTS == commitTS { + s.observeCommitTS(commitTS) + return nil + } + if err := s.applyChunkBlobUntilRegistered(stream.Context(), key, payload, startTS, commitTS); err != nil { + done, retry, retryErr := s.retryAfterChunkBlobWriteConflict(stream.Context(), 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..cc422bd3d --- /dev/null +++ b/adapter/s3_blob_fetch_test.go @@ -0,0 +1,735 @@ +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 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_offload_test.go b/adapter/s3_blob_offload_test.go index 3d3a8fb48..9b8d00fbb 100644 --- a/adapter/s3_blob_offload_test.go +++ b/adapter/s3_blob_offload_test.go @@ -112,15 +112,24 @@ func (alwaysS3BlobOffloadCapable) AllPeersSupportS3BlobOffload(context.Context) } 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/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/main.go b/main.go index c7a1bae3a..7575972ab 100644 --- a/main.go +++ b/main.go @@ -1815,6 +1815,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 @@ -1840,7 +1854,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/") || @@ -2061,6 +2076,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 @@ -2089,6 +2106,12 @@ func startRaftServers( grpcSvc := adapter.NewGRPCServer(shardStore, coordinate) pb.RegisterRawKVServer(gs, grpcSvc) pb.RegisterTransactionalKVServer(gs, grpcSvc) + pb.RegisterS3BlobFetchServer(gs, adapter.NewS3BlobFetchServer( + rt.store, + s3BlobObserver, + adapter.WithS3BlobFetchClock(coordinate.Clock()), + adapter.WithS3BlobFetchPushBlocked(s3BlobPushBlocked), + )) pb.RegisterInternalServer(gs, adapter.NewInternalWithEngine( trx, rt.engine, @@ -2132,7 +2155,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"}, + []string{"S3BlobFetch"}, + confChangeInterceptor, + ) reflection.Register(gs) grpcSock, err := lc.Listen(ctx, "tcp", rt.spec.address) @@ -2534,6 +2564,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), @@ -2558,6 +2593,8 @@ func (r *runtimeServerRunner) startRaftTransport() error { forwardDeps, r.encryptionConfChangeInterceptor, r.encWiring, + r.metricsRegistry.S3BlobOffloadObserver(), + s3BlobPushBlocked, ); err != nil { return r.startupFailure(err) } diff --git a/main_encryption_rotate_on_startup_test.go b/main_encryption_rotate_on_startup_test.go index 5e16348ad..7fc5a066e 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 diff --git a/proto/service.pb.go b/proto/service.pb.go index a869e0c1d..c5f1fc016 100644 --- a/proto/service.pb.go +++ b/proto/service.pb.go @@ -2215,6 +2215,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 @@ -2223,7 +2431,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) } @@ -2235,7 +2443,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 { @@ -2248,7 +2456,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 @@ -2389,7 +2597,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" + @@ -2419,7 +2639,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 @@ -2434,7 +2657,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 @@ -2476,7 +2699,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 @@ -2507,27 +2734,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 @@ -2544,9 +2775,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 b1ac9105f..051a491fa 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; @@ -260,4 +265,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 1bf9f4574..743f024e1 100644 --- a/store/lsm_store.go +++ b/store/lsm_store.go @@ -1402,7 +1402,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 @@ -1431,7 +1438,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 @@ -1443,10 +1450,10 @@ 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) 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() @@ -1469,6 +1476,34 @@ func (s *pebbleStore) applyMutationsWithOpts(ctx context.Context, mutations []*K return err } + 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 err := stageAppliedIndexInBatch(b, appliedIndex); err != nil { + return err + } + if err := b.Commit(writeOpts); err != nil { + return errors.WithStack(err) + } + 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 @@ -1480,26 +1515,16 @@ func (s *pebbleStore) applyMutationsWithOpts(ctx context.Context, mutations []*K } if err := setPebbleUint64InBatch(b, metaLastCommitTSBytes, newLastTS); err != nil { s.mtx.Unlock() - return err - } - // 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) + return 0, nil, err } - s.updateLastCommitTS(newLastTS) - s.mtx.Unlock() + return newLastTS, s.mtx.Unlock, nil +} - return nil +func stageAppliedIndexInBatch(b *pebble.Batch, appliedIndex uint64) error { + if appliedIndex == 0 { + return nil + } + return setPebbleUint64InBatch(b, metaAppliedIndexBytes, appliedIndex) } // DeletePrefixAt atomically deletes all visible keys matching prefix by writing diff --git a/store/mvcc_store.go b/store/mvcc_store.go index 1886fbf3b..9389be6f3 100644 --- a/store/mvcc_store.go +++ b/store/mvcc_store.go @@ -538,6 +538,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 {