Allow unbounded distribution routes

Author: bootjp

adapter/distribution_server.go

@@ -0,0 +1,43 @@
+package adapter
+
+import (
+	"context"
+
+	"github.com/bootjp/elastickv/distribution"
+	pb "github.com/bootjp/elastickv/proto"
+)
+
+// DistributionServer serves distribution related gRPC APIs.
+type DistributionServer struct {
+	engine *distribution.Engine
+	pb.UnimplementedDistributionServer
+}
+
+// NewDistributionServer creates a new server.
+func NewDistributionServer(e *distribution.Engine) *DistributionServer {
+	return &DistributionServer{engine: e}
+}
+
+// UpdateRoute allows updating route information.
+func (s *DistributionServer) UpdateRoute(start, end []byte, group uint64) {
+	s.engine.UpdateRoute(start, end, group)
+}
+
+// GetRoute returns route for a key.
+func (s *DistributionServer) GetRoute(ctx context.Context, req *pb.GetRouteRequest) (*pb.GetRouteResponse, error) {
+	r, ok := s.engine.GetRoute(req.Key)
+	if !ok {
+		return &pb.GetRouteResponse{}, nil
+	}
+	return &pb.GetRouteResponse{
+		Start:       r.Start,
+		End:         r.End,
+		RaftGroupId: r.GroupID,
+	}, nil
+}
+
+// GetTimestamp returns monotonically increasing timestamp.
+func (s *DistributionServer) GetTimestamp(ctx context.Context, req *pb.GetTimestampRequest) (*pb.GetTimestampResponse, error) {
+	ts := s.engine.NextTimestamp()
+	return &pb.GetTimestampResponse{Timestamp: ts}, nil
+}

cmd/server/demo.go

@@ -11,6 +11,7 @@ import (
 	transport "github.com/Jille/raft-grpc-transport"
 	"github.com/Jille/raftadmin"
 	"github.com/bootjp/elastickv/adapter"
+	"github.com/bootjp/elastickv/distribution"
 	"github.com/bootjp/elastickv/kv"
 	pb "github.com/bootjp/elastickv/proto"
 	"github.com/bootjp/elastickv/store"
@@ -89,10 +90,15 @@ func run(eg *errgroup.Group) error {
 		trx := kv.NewTransaction(r)
 		coordinator := kv.NewCoordinator(trx, r)
 		gs := adapter.NewGRPCServer(st, coordinator)
+		distEngine := distribution.NewEngine()
+		distSrv := adapter.NewDistributionServer(distEngine)
+		// example route for demo purposes
+		distSrv.UpdateRoute([]byte("a"), []byte("z"), uint64(i))
 		tm.Register(s)
 		pb.RegisterRawKVServer(s, gs)
 		pb.RegisterTransactionalKVServer(s, gs)
 		pb.RegisterInternalServer(s, adapter.NewInternal(trx, r))
+		pb.RegisterDistributionServer(s, distSrv)
 		leaderhealth.Setup(r, s, []string{"RawKV"})
 		raftadmin.Register(s, r)
 

distribution/engine.go

@@ -0,0 +1,60 @@
+package distribution
+
+import (
+	"bytes"
+	"sync"
+	"sync/atomic"
+)
+
+// Route represents mapping from key range to raft group.
+// The key range is half-open: [Start, End)
+// meaning Start is inclusive and End is exclusive. If End is nil or empty
+// the range is unbounded and covers all keys >= Start.
+type Route struct {
+	// Start marks the inclusive beginning of the range.
+	Start []byte
+        // End marks the exclusive end of the range. When empty, the range has
+        // no upper bound.
+        End     []byte
+	GroupID uint64
+}
+
+// Engine holds in-memory metadata of routes and provides timestamp generation.
+type Engine struct {
+	mu     sync.RWMutex
+	routes []Route
+	ts     uint64
+}
+
+// NewEngine creates an Engine.
+func NewEngine() *Engine {
+	return &Engine{routes: make([]Route, 0)}
+}
+
+// UpdateRoute registers or updates a route. The range is [start, end). If end
+// is nil or empty the range is treated as unbounded above.
+func (e *Engine) UpdateRoute(start, end []byte, group uint64) {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+	e.routes = append(e.routes, Route{Start: start, End: end, GroupID: group})
+}
+
+// GetRoute finds a route for the given key. The search uses half-open ranges
+// [Start, End). If End is empty the range is treated as unbounded above.
+func (e *Engine) GetRoute(key []byte) (Route, bool) {
+	e.mu.RLock()
+	defer e.mu.RUnlock()
+        for _, r := range e.routes {
+                if bytes.Compare(key, r.Start) >= 0 {
+                        if len(r.End) == 0 || bytes.Compare(key, r.End) < 0 {
+                                return r, true
+                        }
+                }
+        }
+	return Route{}, false
+}
+
+// NextTimestamp returns a monotonic increasing timestamp.
+func (e *Engine) NextTimestamp() uint64 {
+	return atomic.AddUint64(&e.ts, 1)
+}

distribution/engine_test.go

@@ -0,0 +1,44 @@
+package distribution
+
+import "testing"
+
+func TestEngineRouteLookup(t *testing.T) {
+        e := NewEngine()
+        e.UpdateRoute([]byte("a"), []byte("m"), 1)
+        // second range has no upper bound
+        e.UpdateRoute([]byte("m"), nil, 2)
+
+	cases := []struct {
+		key    []byte
+		group  uint64
+		expect bool
+	}{
+		{[]byte("a"), 1, true}, // start is inclusive
+		{[]byte("b"), 1, true},
+                {[]byte("m"), 2, true}, // end is exclusive, so m belongs to second range
+                {[]byte("x"), 2, true},
+                {[]byte("za"), 2, true}, // unbounded range covers keys beyond z
+        }
+
+	for _, c := range cases {
+		r, ok := e.GetRoute(c.key)
+		if ok != c.expect {
+			t.Fatalf("key %q expected ok=%v, got %v", c.key, c.expect, ok)
+		}
+		if ok && r.GroupID != c.group {
+			t.Fatalf("key %q expected group %d, got %d", c.key, c.group, r.GroupID)
+		}
+	}
+}
+
+func TestEngineTimestampMonotonic(t *testing.T) {
+	e := NewEngine()
+	last := e.NextTimestamp()
+	for i := 0; i < 100; i++ {
+		ts := e.NextTimestamp()
+		if ts <= last {
+			t.Fatalf("timestamp not monotonic: %d <= %d", ts, last)
+		}
+		last = ts
+	}
+}

proto/Makefile

@@ -1,10 +1,12 @@
+all: gen
+
 gen:
 	protoc --go_out=. --go_opt=paths=source_relative \
-		--go-grpc_out=. --go-grpc_opt=paths=source_relative \
-		service.proto
+	    --go-grpc_out=. --go-grpc_opt=paths=source_relative \
+	    service.proto
 	protoc --go_out=. --go_opt=paths=source_relative \
-		--go-grpc_out=. --go-grpc_opt=paths=source_relative \
-		internal.proto
-
-
-
+	    --go-grpc_out=. --go-grpc_opt=paths=source_relative \
+	    internal.proto
+	protoc --go_out=. --go_opt=paths=source_relative \
+	    --go-grpc_out=. --go-grpc_opt=paths=source_relative \
+	    distribution.proto