s3: add blob fetch rpc

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++
+}

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

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
 }

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)
 	}

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