Add MultiRaft sharding manager

Author: bootjp

kv/sharded_transaction.go

@@ -0,0 +1,104 @@
+package kv
+
+import (
+	"sync"
+
+	"github.com/bootjp/elastickv/distribution"
+	pb "github.com/bootjp/elastickv/proto"
+	"github.com/cockroachdb/errors"
+)
+
+// ShardedTransactionManager routes requests to multiple raft groups based on key ranges.
+type ShardedTransactionManager struct {
+	engine *distribution.Engine
+	mu     sync.RWMutex
+	groups map[uint64]Transactional
+}
+
+// NewShardedTransactionManager creates a new manager.
+func NewShardedTransactionManager(e *distribution.Engine) *ShardedTransactionManager {
+	return &ShardedTransactionManager{
+		engine: e,
+		groups: make(map[uint64]Transactional),
+	}
+}
+
+// Register associates a raft group ID with a Transactional.
+func (s *ShardedTransactionManager) Register(group uint64, tm Transactional) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.groups[group] = tm
+}
+
+// Commit dispatches requests to the correct raft group.
+func (s *ShardedTransactionManager) Commit(reqs []*pb.Request) (*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)
+		if !ok {
+			return nil, errors.Wrapf(ErrInvalidRequest, "unknown group %d", gid)
+		}
+		r, err := tm.Commit(rs)
+		if err != nil {
+			return nil, errors.WithStack(err)
+		}
+		if r.CommitIndex > max {
+			max = r.CommitIndex
+		}
+	}
+	return &TransactionResponse{CommitIndex: max}, nil
+}
+
+// Abort dispatches aborts to the correct raft group.
+func (s *ShardedTransactionManager) Abort(reqs []*pb.Request) (*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)
+		if !ok {
+			return nil, errors.Wrapf(ErrInvalidRequest, "unknown group %d", gid)
+		}
+		r, err := tm.Abort(rs)
+		if err != nil {
+			return nil, errors.WithStack(err)
+		}
+		if r.CommitIndex > max {
+			max = r.CommitIndex
+		}
+	}
+	return &TransactionResponse{CommitIndex: max}, nil
+}
+
+func (s *ShardedTransactionManager) getGroup(id uint64) (Transactional, bool) {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	tm, ok := s.groups[id]
+	return tm, ok
+}
+
+func (s *ShardedTransactionManager) groupRequests(reqs []*pb.Request) (map[uint64][]*pb.Request, error) {
+	batches := make(map[uint64][]*pb.Request)
+	for _, r := range reqs {
+		if len(r.Mutations) == 0 {
+			return nil, ErrInvalidRequest
+		}
+		key := r.Mutations[0].Key
+		route, ok := s.engine.GetRoute(key)
+		if !ok {
+			return nil, errors.Wrapf(ErrInvalidRequest, "no route for key %q", key)
+		}
+		batches[route.GroupID] = append(batches[route.GroupID], r)
+	}
+	return batches, nil
+}
+
+var _ Transactional = (*ShardedTransactionManager)(nil)

kv/sharded_transaction_test.go

@@ -0,0 +1,84 @@
+package kv
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/bootjp/elastickv/distribution"
+	pb "github.com/bootjp/elastickv/proto"
+	"github.com/bootjp/elastickv/store"
+	"github.com/hashicorp/raft"
+)
+
+// helper to create single-node raft
+func newTestRaft(t *testing.T, id string, fsm raft.FSM) *raft.Raft {
+	t.Helper()
+	c := raft.DefaultConfig()
+	c.LocalID = raft.ServerID(id)
+	ldb := raft.NewInmemStore()
+	sdb := raft.NewInmemStore()
+	fss := raft.NewInmemSnapshotStore()
+	addr, trans := raft.NewInmemTransport(raft.ServerAddress(id))
+	r, err := raft.NewRaft(c, fsm, ldb, sdb, fss, trans)
+	if err != nil {
+		t.Fatalf("new raft: %v", err)
+	}
+	cfg := raft.Configuration{Servers: []raft.Server{{ID: raft.ServerID(id), Address: addr}}}
+	if err := r.BootstrapCluster(cfg).Error(); err != nil {
+		t.Fatalf("bootstrap: %v", err)
+	}
+
+	// single node should eventually become leader
+	for i := 0; i < 100; i++ {
+		if r.State() == raft.Leader {
+			break
+		}
+		time.Sleep(50 * time.Millisecond)
+	}
+	if r.State() != raft.Leader {
+		t.Fatalf("node %s is not leader", id)
+	}
+	return r
+}
+
+func TestShardedTransactionManagerCommit(t *testing.T) {
+	e := distribution.NewEngine()
+	e.UpdateRoute([]byte("a"), []byte("m"), 1)
+	e.UpdateRoute([]byte("m"), nil, 2)
+
+	stm := NewShardedTransactionManager(e)
+
+	// group 1
+	s1 := store.NewRbMemoryStore()
+	l1 := store.NewRbMemoryStoreWithExpire(time.Minute)
+	r1 := newTestRaft(t, "1", NewKvFSM(s1, l1))
+	defer r1.Shutdown()
+	stm.Register(1, NewTransaction(r1))
+
+	// group 2
+	s2 := store.NewRbMemoryStore()
+	l2 := store.NewRbMemoryStoreWithExpire(time.Minute)
+	r2 := newTestRaft(t, "2", NewKvFSM(s2, l2))
+	defer r2.Shutdown()
+	stm.Register(2, NewTransaction(r2))
+
+	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 := stm.Commit(reqs)
+	if err != nil {
+		t.Fatalf("commit: %v", err)
+	}
+
+	v, err := s1.Get(context.Background(), []byte("b"))
+	if err != nil || string(v) != "v1" {
+		t.Fatalf("group1 value: %v %v", v, err)
+	}
+	v, err = s2.Get(context.Background(), []byte("x"))
+	if err != nil || string(v) != "v2" {
+		t.Fatalf("group2 value: %v %v", v, err)
+	}
+}