route gets through shard router

bootjp · bootjp · commit 8d82b6d4f80b · 2025-08-31T00:24:38.000+09:00
diff --git a/kv/shard_router.go b/kv/shard_router.go
@@ -1,10 +1,12 @@
 package kv
 
 import (
+	"context"
 	"sync"
 
 	"github.com/bootjp/elastickv/distribution"
 	pb "github.com/bootjp/elastickv/proto"
+	"github.com/bootjp/elastickv/store"
 	"github.com/cockroachdb/errors"
 )
 
@@ -14,50 +16,55 @@ import (
 type ShardRouter struct {
 	engine *distribution.Engine
 	mu     sync.RWMutex
-	groups map[uint64]Transactional
+	groups map[uint64]*routerGroup
+}
+
+type routerGroup struct {
+	tm    Transactional
+	store store.Store
 }
 
 // NewShardRouter creates a new router.
 func NewShardRouter(e *distribution.Engine) *ShardRouter {
 	return &ShardRouter{
 		engine: e,
-		groups: make(map[uint64]Transactional),
+		groups: make(map[uint64]*routerGroup),
 	}
 }
 
-// Register associates a raft group ID with a Transactional.
-func (s *ShardRouter) Register(group uint64, tm Transactional) {
+// Register associates a raft group ID with its transactional manager and store.
+func (s *ShardRouter) Register(group uint64, tm Transactional, st store.Store) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
-	s.groups[group] = tm
+	s.groups[group] = &routerGroup{tm: tm, store: st}
 }
 
 func (s *ShardRouter) Commit(reqs []*pb.Request) (*TransactionResponse, error) {
-	return s.process(reqs, func(tm Transactional, rs []*pb.Request) (*TransactionResponse, error) {
-		return tm.Commit(rs)
+	return s.process(reqs, func(g *routerGroup, rs []*pb.Request) (*TransactionResponse, error) {
+		return g.tm.Commit(rs)
 	})
 }
 
 // Abort dispatches aborts to the correct raft group.
 func (s *ShardRouter) Abort(reqs []*pb.Request) (*TransactionResponse, error) {
-	return s.process(reqs, func(tm Transactional, rs []*pb.Request) (*TransactionResponse, error) {
-		return tm.Abort(rs)
+	return s.process(reqs, func(g *routerGroup, rs []*pb.Request) (*TransactionResponse, error) {
+		return g.tm.Abort(rs)
 	})
 }
 
-func (s *ShardRouter) process(reqs []*pb.Request, fn func(Transactional, []*pb.Request) (*TransactionResponse, error)) (*TransactionResponse, error) {
+func (s *ShardRouter) process(reqs []*pb.Request, fn func(*routerGroup, []*pb.Request) (*TransactionResponse, error)) (*TransactionResponse, error) {
 	grouped, err := s.groupRequests(reqs)
 	if err != nil {
 		return nil, errors.WithStack(err)
 	}
 
 	var max uint64
 	for gid, rs := range grouped {
-		tm, ok := s.getGroup(gid)
+		g, ok := s.getGroup(gid)
 		if !ok {
 			return nil, errors.Wrapf(ErrInvalidRequest, "unknown group %d", gid)
 		}
-		r, err := fn(tm, rs)
+		r, err := fn(g, rs)
 		if err != nil {
 			return nil, errors.WithStack(err)
 		}
@@ -68,11 +75,11 @@ func (s *ShardRouter) process(reqs []*pb.Request, fn func(Transactional, []*pb.R
 	return &TransactionResponse{CommitIndex: max}, nil
 }
 
-func (s *ShardRouter) getGroup(id uint64) (Transactional, bool) {
+func (s *ShardRouter) getGroup(id uint64) (*routerGroup, bool) {
 	s.mu.RLock()
 	defer s.mu.RUnlock()
-	tm, ok := s.groups[id]
-	return tm, ok
+	g, ok := s.groups[id]
+	return g, ok
 }
 
 func (s *ShardRouter) groupRequests(reqs []*pb.Request) (map[uint64][]*pb.Request, error) {
@@ -91,4 +98,21 @@ func (s *ShardRouter) groupRequests(reqs []*pb.Request) (map[uint64][]*pb.Reques
 	return batches, nil
 }
 
+// Get retrieves a key routed to the correct shard.
+func (s *ShardRouter) Get(ctx context.Context, key []byte) ([]byte, error) {
+	route, ok := s.engine.GetRoute(key)
+	if !ok {
+		return nil, errors.Wrapf(ErrInvalidRequest, "no route for key %q", key)
+	}
+	g, ok := s.getGroup(route.GroupID)
+	if !ok {
+		return nil, errors.Wrapf(ErrInvalidRequest, "unknown group %d", route.GroupID)
+	}
+	v, err := g.store.Get(ctx, key)
+	if err != nil {
+		return nil, errors.WithStack(err)
+	}
+	return v, nil
+}
+
 var _ Transactional = (*ShardRouter)(nil)
diff --git a/kv/shard_router_test.go b/kv/shard_router_test.go
@@ -93,6 +93,8 @@ func newTestRaft(t *testing.T, id string, fsm raft.FSM) (*raft.Raft, func()) {
 }
 
 func TestShardRouterCommit(t *testing.T) {
+	ctx := context.Background()
+
 	e := distribution.NewEngine()
 	e.UpdateRoute([]byte("a"), []byte("m"), 1)
 	e.UpdateRoute([]byte("m"), nil, 2)
@@ -104,32 +106,31 @@ func TestShardRouterCommit(t *testing.T) {
 	l1 := store.NewRbMemoryStoreWithExpire(time.Minute)
 	r1, stop1 := newTestRaft(t, "1", NewKvFSM(s1, l1))
 	defer stop1()
-	router.Register(1, NewTransaction(r1))
+	router.Register(1, NewTransaction(r1), s1)
 
 	// group 2
 	s2 := store.NewRbMemoryStore()
 	l2 := store.NewRbMemoryStoreWithExpire(time.Minute)
 	r2, stop2 := newTestRaft(t, "2", NewKvFSM(s2, l2))
 	defer stop2()
-	router.Register(2, NewTransaction(r2))
+	router.Register(2, NewTransaction(r2), s2)
 
 	reqs := []*pb.Request{
 		{IsTxn: false, Phase: pb.Phase_NONE, Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: []byte("b"), Value: []byte("v1")}}},
 		{IsTxn: false, Phase: pb.Phase_NONE, Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: []byte("x"), Value: []byte("v2")}}},
 	}
 
-	_, err := router.Commit(reqs)
-	if err != nil {
+	if _, err := router.Commit(reqs); err != nil {
 		t.Fatalf("commit: %v", err)
 	}
 
-	v, err := s1.Get(context.Background(), []byte("b"))
+	v, err := router.Get(ctx, []byte("b"))
 	if err != nil || string(v) != "v1" {
-		t.Fatalf("group1 value: %v %v", v, err)
+		t.Fatalf("group1 get: %v %v", v, err)
 	}
-	v, err = s2.Get(context.Background(), []byte("x"))
+	v, err = router.Get(ctx, []byte("x"))
 	if err != nil || string(v) != "v2" {
-		t.Fatalf("group2 value: %v %v", v, err)
+		t.Fatalf("group2 get: %v %v", v, err)
 	}
 }
 
@@ -147,14 +148,14 @@ func TestShardRouterSplitAndMerge(t *testing.T) {
 	l1 := store.NewRbMemoryStoreWithExpire(time.Minute)
 	r1, stop1 := newTestRaft(t, "1", NewKvFSM(s1, l1))
 	defer stop1()
-	router.Register(1, NewTransaction(r1))
+	router.Register(1, NewTransaction(r1), s1)
 
 	// group 2 (will be used after split)
 	s2 := store.NewRbMemoryStore()
 	l2 := store.NewRbMemoryStoreWithExpire(time.Minute)
 	r2, stop2 := newTestRaft(t, "2", NewKvFSM(s2, l2))
 	defer stop2()
-	router.Register(2, NewTransaction(r2))
+	router.Register(2, NewTransaction(r2), s2)
 
 	// initial write routed to group 1
 	req := []*pb.Request{
@@ -163,7 +164,7 @@ func TestShardRouterSplitAndMerge(t *testing.T) {
 	if _, err := router.Commit(req); err != nil {
 		t.Fatalf("commit group1: %v", err)
 	}
-	v, err := s1.Get(ctx, []byte("b"))
+	v, err := router.Get(ctx, []byte("b"))
 	if err != nil || string(v) != "v1" {
 		t.Fatalf("group1 value before split: %v %v", v, err)
 	}
@@ -181,7 +182,7 @@ func TestShardRouterSplitAndMerge(t *testing.T) {
 	if _, err := router.Commit(req); err != nil {
 		t.Fatalf("commit group2: %v", err)
 	}
-	v, err = s2.Get(ctx, []byte("x"))
+	v, err = router.Get(ctx, []byte("x"))
 	if err != nil || string(v) != "v2" {
 		t.Fatalf("group2 value after split: %v %v", v, err)
 	}
@@ -198,7 +199,7 @@ func TestShardRouterSplitAndMerge(t *testing.T) {
 	if _, err := router.Commit(req); err != nil {
 		t.Fatalf("commit after merge: %v", err)
 	}
-	v, err = s1.Get(ctx, []byte("z"))
+	v, err = router.Get(ctx, []byte("z"))
 	if err != nil || string(v) != "v3" {
 		t.Fatalf("group1 value after merge: %v %v", v, err)
 	}
@@ -232,8 +233,8 @@ func TestShardRouterCommitFailure(t *testing.T) {
 
 	ok := &fakeTM{}
 	fail := &fakeTM{commitErr: true}
-	router.Register(1, ok)
-	router.Register(2, fail)
+	router.Register(1, ok, nil)
+	router.Register(2, fail, nil)
 
 	reqs := []*pb.Request{
 		{IsTxn: false, Phase: pb.Phase_NONE, Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: []byte("b"), Value: []byte("v1")}}},
