Complete standalone hotspot split automation

adapter/admin_grpc.go

@@ -33,6 +33,13 @@ type KeyVizSampler interface {
 	Snapshot(from, to time.Time) []keyviz.MatrixColumn
 }
 
+// AutoSplitRuntime is the atomic runtime control exposed by the authenticated
+// Admin service.
+type AutoSplitRuntime interface {
+	Enabled() bool
+	SetEnabled(enabled bool)
+}
+
 // AdminGroup exposes per-Raft-group state to the Admin service. It is a narrow
 // subset of raftengine.Engine so tests can supply an in-memory fake without
 // standing up a real Raft cluster. Configuration is polled on each
@@ -77,7 +84,8 @@ type AdminServer struct {
 	// Nil means keyviz is disabled — the RPC returns Unavailable.
 	// Guarded by groupsMu (same lock as groups/now) so RegisterSampler
 	// pairs atomically with concurrent RPC reads.
-	sampler KeyVizSampler
+	sampler          KeyVizSampler
+	autoSplitRuntime AutoSplitRuntime
 
 	pb.UnimplementedAdminServer
 }
@@ -145,6 +153,32 @@ func (s *AdminServer) SetCapability(name string, enabled bool) {
 	s.groupsMu.Unlock()
 }
 
+// RegisterAutoSplitRuntime wires the process-local automatic split switch.
+func (s *AdminServer) RegisterAutoSplitRuntime(runtime AutoSplitRuntime) {
+	s.groupsMu.Lock()
+	s.autoSplitRuntime = runtime
+	s.groupsMu.Unlock()
+}
+
+// SetAutoSplitEnabled atomically changes whether the scheduler may issue new
+// SplitRange calls. Detector observation continues while disabled.
+func (s *AdminServer) SetAutoSplitEnabled(
+	_ context.Context,
+	req *pb.SetAutoSplitEnabledRequest,
+) (*pb.SetAutoSplitEnabledResponse, error) {
+	if req == nil {
+		return nil, grpcStatusError(codes.InvalidArgument, "request is required")
+	}
+	s.groupsMu.RLock()
+	runtime := s.autoSplitRuntime
+	s.groupsMu.RUnlock()
+	if runtime == nil {
+		return nil, grpcStatusError(codes.Unavailable, "automatic splitting is not configured")
+	}
+	runtime.SetEnabled(req.GetEnabled())
+	return &pb.SetAutoSplitEnabledResponse{Enabled: runtime.Enabled()}, nil
+}
+
 // GetClusterOverview returns the local node identity, the current member
 // list, and per-group leader identity collected from the engines registered
 // via RegisterGroup. The member list is the union of (a) the bootstrap seed

adapter/admin_grpc_autosplit_test.go

@@ -0,0 +1,60 @@
+package adapter
+
+import (
+	"context"
+	"sync"
+	"testing"
+
+	"github.com/bootjp/elastickv/distribution/autosplit"
+	pb "github.com/bootjp/elastickv/proto"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+func TestAdminSetAutoSplitEnabledControlsRuntimeSwitch(t *testing.T) {
+	t.Parallel()
+	server := NewAdminServer(NodeIdentity{NodeID: "node-a"}, nil)
+
+	_, err := server.SetAutoSplitEnabled(context.Background(), &pb.SetAutoSplitEnabledRequest{Enabled: false})
+	require.Equal(t, codes.Unavailable, status.Code(err))
+
+	runtime := autosplit.NewRuntimeSwitch(true)
+	server.RegisterAutoSplitRuntime(runtime)
+	disabled, err := server.SetAutoSplitEnabled(context.Background(), &pb.SetAutoSplitEnabledRequest{Enabled: false})
+	require.NoError(t, err)
+	require.False(t, disabled.GetEnabled())
+	require.True(t, runtime.KillSwitch(context.Background()))
+
+	enabled, err := server.SetAutoSplitEnabled(context.Background(), &pb.SetAutoSplitEnabledRequest{Enabled: true})
+	require.NoError(t, err)
+	require.True(t, enabled.GetEnabled())
+	require.False(t, runtime.KillSwitch(context.Background()))
+}
+
+func TestAdminSetAutoSplitEnabledIsConcurrentSafe(t *testing.T) {
+	t.Parallel()
+	server := NewAdminServer(NodeIdentity{NodeID: "node-a"}, nil)
+	runtime := autosplit.NewRuntimeSwitch(true)
+	server.RegisterAutoSplitRuntime(runtime)
+
+	var wg sync.WaitGroup
+	errs := make(chan error, 100)
+	for i := range 100 {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			_, err := server.SetAutoSplitEnabled(context.Background(), &pb.SetAutoSplitEnabledRequest{Enabled: i%2 == 0})
+			errs <- err
+		}()
+	}
+	wg.Wait()
+	close(errs)
+	for err := range errs {
+		require.NoError(t, err)
+	}
+
+	_, err := server.SetAutoSplitEnabled(context.Background(), &pb.SetAutoSplitEnabledRequest{Enabled: true})
+	require.NoError(t, err)
+	require.True(t, runtime.Enabled())
+}

adapter/distribution_server.go

@@ -649,6 +649,9 @@ func (s *DistributionServer) applyEngineSnapshot(snapshot distribution.CatalogSn
 		return grpcStatusError(codes.FailedPrecondition, errDistributionEngineNotConfigured.Error())
 	}
 	if err := s.engine.ApplySnapshot(snapshot); err != nil {
+		if errors.Is(err, distribution.ErrEngineSnapshotVersionStale) {
+			return nil
+		}
 		return grpcStatusErrorf(codes.Internal, "apply engine snapshot: %v", err)
 	}
 	return nil

adapter/distribution_server_test.go

@@ -653,6 +653,29 @@ func TestDistributionServerSplitRange_ReturnsExactCommittedSplitVersion(t *testi
 	require.Equal(t, uint64(3), latest.Version)
 }
 
+func TestDistributionServerApplyEngineSnapshotAcceptsNewerEngine(t *testing.T) {
+	t.Parallel()
+
+	engine := distribution.NewEngine()
+	recent := distribution.CatalogSnapshot{
+		Version: 3,
+		Routes: []distribution.RouteDescriptor{{
+			RouteID: 1,
+			Start:   []byte(""),
+			End:     nil,
+			GroupID: 1,
+			State:   distribution.RouteStateActive,
+		}},
+	}
+	require.NoError(t, engine.ApplySnapshot(recent))
+
+	s := NewDistributionServer(engine, nil)
+	stale := recent
+	stale.Version = 2
+	require.NoError(t, s.applyEngineSnapshot(stale))
+	require.Equal(t, uint64(3), engine.Version())
+}
+
 func TestDistributionServerSplitRange_RetriesCatalogReloadUntilVisible(t *testing.T) {
 	t.Parallel()
 

adapter/internal.go

@@ -18,6 +18,17 @@ func WithInternalTimestampAllocator(alloc kv.TimestampAllocator) InternalOption
 	}
 }
 
+// ForwardWriteObserver observes successfully committed leader-side forwarded
+// writes. It lets the autosplit sampler account for writes that entered through
+// followers without coupling adapter.Internal to keyviz.
+type ForwardWriteObserver func([]*pb.Request)
+
+func WithInternalForwardWriteObserver(observer ForwardWriteObserver) InternalOption {
+	return func(i *Internal) {
+		i.forwardWriteObserver = observer
+	}
+}
+
 func NewInternalWithEngine(txm kv.Transactional, leader raftengine.LeaderView, clock *kv.HLC, relay *RedisPubSubRelay, opts ...InternalOption) *Internal {
 	i := &Internal{
 		leader:             leader,
@@ -32,11 +43,12 @@ func NewInternalWithEngine(txm kv.Transactional, leader raftengine.LeaderView, c
 }
 
 type Internal struct {
-	leader             raftengine.LeaderView
-	transactionManager kv.Transactional
-	clock              *kv.HLC
-	tsAllocator        kv.TimestampAllocator
-	relay              *RedisPubSubRelay
+	leader               raftengine.LeaderView
+	transactionManager   kv.Transactional
+	clock                *kv.HLC
+	tsAllocator          kv.TimestampAllocator
+	relay                *RedisPubSubRelay
+	forwardWriteObserver ForwardWriteObserver
 
 	pb.UnimplementedInternalServer
 }
@@ -70,6 +82,9 @@ func (i *Internal) Forward(ctx context.Context, req *pb.ForwardRequest) (*pb.For
 			CommitIndex: 0,
 		}, errors.WithStack(err)
 	}
+	if i.forwardWriteObserver != nil {
+		i.forwardWriteObserver(req.GetRequests())
+	}
 
 	return &pb.ForwardResponse{
 		Success:     true,

adapter/internal_test.go

@@ -5,11 +5,40 @@ import (
 	"encoding/binary"
 	"testing"
 
+	"github.com/bootjp/elastickv/internal/raftengine"
 	"github.com/bootjp/elastickv/kv"
 	pb "github.com/bootjp/elastickv/proto"
 	"github.com/stretchr/testify/require"
 )
 
+func TestInternalForwardObservesCommittedWrites(t *testing.T) {
+	t.Parallel()
+
+	reqs := []*pb.Request{{
+		Mutations: []*pb.Mutation{{
+			Op:    pb.Op_PUT,
+			Key:   []byte("hot"),
+			Value: []byte("value"),
+		}},
+	}}
+	txn := &forwardObserverTxn{}
+	var observed []*pb.Request
+	internal := NewInternalWithEngine(txn, forwardObserverLeader{}, nil, nil, WithInternalForwardWriteObserver(func(reqs []*pb.Request) {
+		observed = reqs
+	}))
+
+	resp, err := internal.Forward(context.Background(), &pb.ForwardRequest{Requests: reqs})
+
+	require.NoError(t, err)
+	require.True(t, resp.Success)
+	require.Equal(t, uint64(9), resp.CommitIndex)
+	require.Len(t, observed, 1)
+	require.Same(t, reqs[0], observed[0])
+	require.Len(t, txn.reqs, 1)
+	require.Same(t, reqs[0], txn.reqs[0])
+	require.Equal(t, uint64(1), reqs[0].Ts)
+}
+
 func TestStampTxnTimestamps_RejectsMaxStartTS(t *testing.T) {
 	t.Parallel()
 
@@ -241,3 +270,25 @@ func TestStampTxnTimestamps_UsesSingleTxnStartTS(t *testing.T) {
 	require.Greater(t, meta.CommitTS, uint64(9))
 	require.Equal(t, meta.CommitTS, commitTS)
 }
+
+type forwardObserverLeader struct{}
+
+func (forwardObserverLeader) State() raftengine.State { return raftengine.StateLeader }
+func (forwardObserverLeader) Leader() raftengine.LeaderInfo {
+	return raftengine.LeaderInfo{ID: "self", Address: "127.0.0.1:0"}
+}
+func (forwardObserverLeader) VerifyLeader(context.Context) error               { return nil }
+func (forwardObserverLeader) LinearizableRead(context.Context) (uint64, error) { return 0, nil }
+
+type forwardObserverTxn struct {
+	reqs []*pb.Request
+}
+
+func (t *forwardObserverTxn) Commit(_ context.Context, reqs []*pb.Request) (*kv.TransactionResponse, error) {
+	t.reqs = reqs
+	return &kv.TransactionResponse{CommitIndex: 9}, nil
+}
+
+func (t *forwardObserverTxn) Abort(context.Context, []*pb.Request) (*kv.TransactionResponse, error) {
+	return &kv.TransactionResponse{}, nil
+}