Refactor coordinator logic and add sharded support

.gitignore

@@ -29,3 +29,11 @@ jepsen/.lein-*
 jepsen/.nrepl-port
 .m2/
 jepsen/store/
+
+# Jepsen local SSH keys (generated locally; never commit)
+jepsen/docker/id_rsa
+jepsen/.ssh/
+
+# Build and lint cache directories
+.cache/
+.golangci-cache/

Dockerfile

@@ -3,7 +3,7 @@ FROM golang:latest AS build
 WORKDIR $GOPATH/src/app
 COPY . .
 
-RUN CGO_ENABLED=0 go build -o /app main.go
+RUN CGO_ENABLED=0 go build -o /app .
 
 FROM gcr.io/distroless/static:latest
 COPY --from=build /app /app

adapter/dynamodb.go

@@ -32,7 +32,7 @@ type DynamoDBServer struct {
 	httpServer       *http.Server
 }
 
-func NewDynamoDBServer(listen net.Listener, st store.MVCCStore, coordinate *kv.Coordinate) *DynamoDBServer {
+func NewDynamoDBServer(listen net.Listener, st store.MVCCStore, coordinate kv.Coordinator) *DynamoDBServer {
 	d := &DynamoDBServer{
 		listen:           listen,
 		store:            st,
@@ -85,7 +85,7 @@ func (d *DynamoDBServer) putItem(w http.ResponseWriter, r *http.Request) {
 		http.Error(w, err.Error(), http.StatusBadRequest)
 		return
 	}
-	if _, err = d.coordinator.Dispatch(reqs); err != nil {
+	if _, err = d.coordinator.Dispatch(r.Context(), reqs); err != nil {
 		http.Error(w, err.Error(), http.StatusInternalServerError)
 		return
 	}
@@ -195,7 +195,7 @@ func (d *DynamoDBServer) updateItem(w http.ResponseWriter, r *http.Request) {
 		IsTxn: false,
 		Elems: []*kv.Elem[kv.OP]{elem},
 	}
-	if _, err = d.coordinator.Dispatch(req); err != nil {
+	if _, err = d.coordinator.Dispatch(r.Context(), req); err != nil {
 		http.Error(w, err.Error(), http.StatusInternalServerError)
 		return
 	}
@@ -214,7 +214,7 @@ func (d *DynamoDBServer) transactWriteItems(w http.ResponseWriter, r *http.Reque
 		http.Error(w, err.Error(), http.StatusBadRequest)
 		return
 	}
-	if _, err = d.coordinator.Dispatch(reqs); err != nil {
+	if _, err = d.coordinator.Dispatch(r.Context(), reqs); err != nil {
 		http.Error(w, err.Error(), http.StatusInternalServerError)
 		return
 	}

adapter/grpc.go

@@ -4,15 +4,14 @@ import (
 	"context"
 	"log/slog"
 	"os"
+	"sync"
 
 	"github.com/bootjp/elastickv/internal"
 	"github.com/bootjp/elastickv/kv"
 	pb "github.com/bootjp/elastickv/proto"
 	"github.com/bootjp/elastickv/store"
 	"github.com/cockroachdb/errors"
 	"github.com/spaolacci/murmur3"
-	"google.golang.org/grpc"
-	"google.golang.org/grpc/credentials/insecure"
 )
 
 var _ pb.RawKVServer = (*GRPCServer)(nil)
@@ -24,96 +23,151 @@ type GRPCServer struct {
 	coordinator    kv.Coordinator
 	store          store.MVCCStore
 
+	closeStore bool
+	closeOnce  sync.Once
+	closeErr   error
+
 	pb.UnimplementedRawKVServer
 	pb.UnimplementedTransactionalKVServer
 }
 
-func NewGRPCServer(store store.MVCCStore, coordinate *kv.Coordinate) *GRPCServer {
-	return &GRPCServer{
+type GRPCServerOption func(*GRPCServer)
+
+func WithCloseStore() GRPCServerOption {
+	return func(s *GRPCServer) {
+		s.closeStore = true
+	}
+}
+
+func NewGRPCServer(store store.MVCCStore, coordinate kv.Coordinator, opts ...GRPCServerOption) *GRPCServer {
+	s := &GRPCServer{
 		log: slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{
 			Level: slog.LevelWarn,
 		})),
 		grpcTranscoder: newGrpcGrpcTranscoder(),
 		coordinator:    coordinate,
 		store:          store,
 	}
+	for _, opt := range opts {
+		if opt == nil {
+			continue
+		}
+		opt(s)
+	}
+	return s
 }
 
-func (r GRPCServer) RawGet(ctx context.Context, req *pb.RawGetRequest) (*pb.RawGetResponse, error) {
-	readTS := req.GetTs()
-	if readTS == 0 {
-		readTS = snapshotTS(r.coordinator.Clock(), r.store)
+func (r *GRPCServer) Close() error {
+	if r == nil {
+		return nil
 	}
-
-	if r.coordinator.IsLeader() {
-		v, err := r.store.GetAt(ctx, req.Key, readTS)
-		if err != nil {
-			switch {
-			case errors.Is(err, store.ErrKeyNotFound):
-				return &pb.RawGetResponse{
-					Value: nil,
-				}, nil
-			default:
-				return nil, errors.WithStack(err)
-			}
+	r.closeOnce.Do(func() {
+		if !r.closeStore || r.store == nil {
+			return
+		}
+		if err := r.store.Close(); err != nil {
+			r.closeErr = errors.WithStack(err)
 		}
-		r.log.InfoContext(ctx, "Get",
-			slog.String("key", string(req.Key)),
-			slog.String("value", string(v)))
+	})
+	return r.closeErr
+}
+
+func (r *GRPCServer) clock() *kv.HLC {
+	if r == nil || r.coordinator == nil {
+		return nil
+	}
+	return r.coordinator.Clock()
+}
 
-		return &pb.RawGetResponse{
-			Value: v,
-		}, nil
+func (r *GRPCServer) RawGet(ctx context.Context, req *pb.RawGetRequest) (*pb.RawGetResponse, error) {
+	readTS := req.GetTs()
+	if readTS == 0 {
+		readTS = snapshotTS(r.clock(), r.store)
 	}
 
-	v, err := r.tryLeaderGet(req.Key)
+	v, err := r.store.GetAt(ctx, req.Key, readTS)
+	if errors.Is(err, store.ErrKeyNotFound) {
+		return &pb.RawGetResponse{Value: nil}, nil
+	}
 	if err != nil {
-		return &pb.RawGetResponse{
-			Value: nil,
-		}, err
+		return nil, errors.WithStack(err)
 	}
 
 	r.log.InfoContext(ctx, "Get",
 		slog.String("key", string(req.Key)),
 		slog.String("value", string(v)))
 
-	return &pb.RawGetResponse{
-		Value: v,
-	}, nil
+	return &pb.RawGetResponse{Value: v}, nil
 }
 
-func (r GRPCServer) tryLeaderGet(key []byte) ([]byte, error) {
-	addr := r.coordinator.RaftLeader()
-	if addr == "" {
-		return nil, ErrLeaderNotFound
+func (r *GRPCServer) RawLatestCommitTS(ctx context.Context, req *pb.RawLatestCommitTSRequest) (*pb.RawLatestCommitTSResponse, error) {
+	key := req.GetKey()
+	if len(key) == 0 {
+		return nil, errors.WithStack(kv.ErrInvalidRequest)
 	}
 
-	conn, err := grpc.NewClient(string(addr),
-		grpc.WithTransportCredentials(insecure.NewCredentials()),
-		grpc.WithDefaultCallOptions(grpc.WaitForReady(true)),
-	)
+	ts, exists, err := r.store.LatestCommitTS(ctx, key)
 	if err != nil {
 		return nil, errors.WithStack(err)
 	}
-	defer conn.Close()
+	return &pb.RawLatestCommitTSResponse{
+		Ts:     ts,
+		Exists: exists,
+	}, nil
+}
 
-	cli := pb.NewRawKVClient(conn)
-	ts := snapshotTS(r.coordinator.Clock(), r.store)
-	resp, err := cli.RawGet(context.Background(), &pb.RawGetRequest{Key: key, Ts: ts})
+func (r *GRPCServer) RawScanAt(ctx context.Context, req *pb.RawScanAtRequest) (*pb.RawScanAtResponse, error) {
+	limit64 := req.GetLimit()
+	limit, err := rawScanLimit(limit64)
 	if err != nil {
-		return nil, errors.WithStack(err)
+		return &pb.RawScanAtResponse{Kv: nil}, err
 	}
 
-	return resp.Value, nil
+	readTS := req.GetTs()
+	if readTS == 0 {
+		readTS = snapshotTS(r.clock(), r.store)
+	}
+
+	res, err := r.store.ScanAt(ctx, req.StartKey, req.EndKey, limit, readTS)
+	if err != nil {
+		return &pb.RawScanAtResponse{Kv: nil}, errors.WithStack(err)
+	}
+
+	return &pb.RawScanAtResponse{Kv: rawKvPairs(res)}, nil
+}
+
+func rawScanLimit(limit64 int64) (int, error) {
+	if limit64 < 0 {
+		return 0, errors.WithStack(kv.ErrInvalidRequest)
+	}
+	maxInt64 := int64(^uint(0) >> 1)
+	if limit64 > maxInt64 {
+		return 0, errors.WithStack(internal.ErrIntOverflow)
+	}
+	return int(limit64), nil
+}
+
+func rawKvPairs(res []*store.KVPair) []*pb.RawKVPair {
+	out := make([]*pb.RawKVPair, 0, len(res))
+	for _, kvp := range res {
+		if kvp == nil {
+			continue
+		}
+		out = append(out, &pb.RawKVPair{
+			Key:   kvp.Key,
+			Value: kvp.Value,
+		})
+	}
+	return out
 }
 
-func (r GRPCServer) RawPut(_ context.Context, req *pb.RawPutRequest) (*pb.RawPutResponse, error) {
+func (r *GRPCServer) RawPut(ctx context.Context, req *pb.RawPutRequest) (*pb.RawPutResponse, error) {
 	m, err := r.grpcTranscoder.RawPutToRequest(req)
 	if err != nil {
 		return nil, errors.WithStack(err)
 	}
 
-	res, err := r.coordinator.Dispatch(m)
+	res, err := r.coordinator.Dispatch(ctx, m)
 	if err != nil {
 		return &pb.RawPutResponse{
 			CommitIndex: uint64(0),
@@ -127,13 +181,13 @@ func (r GRPCServer) RawPut(_ context.Context, req *pb.RawPutRequest) (*pb.RawPut
 	}, nil
 }
 
-func (r GRPCServer) RawDelete(ctx context.Context, req *pb.RawDeleteRequest) (*pb.RawDeleteResponse, error) {
+func (r *GRPCServer) RawDelete(ctx context.Context, req *pb.RawDeleteRequest) (*pb.RawDeleteResponse, error) {
 	m, err := r.grpcTranscoder.RawDeleteToRequest(req)
 	if err != nil {
 		return nil, errors.WithStack(err)
 	}
 
-	res, err := r.coordinator.Dispatch(m)
+	res, err := r.coordinator.Dispatch(ctx, m)
 	if err != nil {
 		return &pb.RawDeleteResponse{
 			CommitIndex: uint64(0),
@@ -147,27 +201,27 @@ func (r GRPCServer) RawDelete(ctx context.Context, req *pb.RawDeleteRequest) (*p
 	}, nil
 }
 
-func (r GRPCServer) PreWrite(ctx context.Context, req *pb.PreWriteRequest) (*pb.PreCommitResponse, error) {
+func (r *GRPCServer) PreWrite(ctx context.Context, req *pb.PreWriteRequest) (*pb.PreCommitResponse, error) {
 	return nil, kv.ErrNotImplemented
 }
 
-func (r GRPCServer) Commit(ctx context.Context, req *pb.CommitRequest) (*pb.CommitResponse, error) {
+func (r *GRPCServer) Commit(ctx context.Context, req *pb.CommitRequest) (*pb.CommitResponse, error) {
 	return nil, kv.ErrNotImplemented
 }
 
-func (r GRPCServer) Rollback(ctx context.Context, req *pb.RollbackRequest) (*pb.RollbackResponse, error) {
+func (r *GRPCServer) Rollback(ctx context.Context, req *pb.RollbackRequest) (*pb.RollbackResponse, error) {
 	return nil, kv.ErrNotImplemented
 }
 
-func (r GRPCServer) Put(ctx context.Context, req *pb.PutRequest) (*pb.PutResponse, error) {
+func (r *GRPCServer) Put(ctx context.Context, req *pb.PutRequest) (*pb.PutResponse, error) {
 	reqs, err := r.grpcTranscoder.TransactionalPutToRequests(req)
 	if err != nil {
 		return nil, errors.WithStack(err)
 	}
 
 	r.log.InfoContext(ctx, "Put", slog.Any("reqs", reqs))
 
-	res, err := r.coordinator.Dispatch(reqs)
+	res, err := r.coordinator.Dispatch(ctx, reqs)
 	if err != nil {
 		return &pb.PutResponse{
 			CommitIndex: uint64(0),
@@ -180,13 +234,13 @@ func (r GRPCServer) Put(ctx context.Context, req *pb.PutRequest) (*pb.PutRespons
 	}, nil
 }
 
-func (r GRPCServer) Get(ctx context.Context, req *pb.GetRequest) (*pb.GetResponse, error) {
+func (r *GRPCServer) Get(ctx context.Context, req *pb.GetRequest) (*pb.GetResponse, error) {
 	h := murmur3.New64()
 	if _, err := h.Write(req.Key); err != nil {
 		return nil, errors.WithStack(err)
 	}
 
-	readTS := snapshotTS(r.coordinator.Clock(), r.store)
+	readTS := snapshotTS(r.clock(), r.store)
 	v, err := r.store.GetAt(ctx, req.Key, readTS)
 	if err != nil {
 		switch {
@@ -206,15 +260,15 @@ func (r GRPCServer) Get(ctx context.Context, req *pb.GetRequest) (*pb.GetRespons
 	}, nil
 }
 
-func (r GRPCServer) Delete(ctx context.Context, req *pb.DeleteRequest) (*pb.DeleteResponse, error) {
+func (r *GRPCServer) Delete(ctx context.Context, req *pb.DeleteRequest) (*pb.DeleteResponse, error) {
 	reqs, err := r.grpcTranscoder.TransactionalDeleteToRequests(req)
 	if err != nil {
 		return nil, errors.WithStack(err)
 	}
 
 	r.log.InfoContext(ctx, "Delete", slog.Any("reqs", reqs))
 
-	res, err := r.coordinator.Dispatch(reqs)
+	res, err := r.coordinator.Dispatch(ctx, reqs)
 	if err != nil {
 		return &pb.DeleteResponse{
 			CommitIndex: uint64(0),
@@ -227,14 +281,14 @@ func (r GRPCServer) Delete(ctx context.Context, req *pb.DeleteRequest) (*pb.Dele
 	}, nil
 }
 
-func (r GRPCServer) Scan(ctx context.Context, req *pb.ScanRequest) (*pb.ScanResponse, error) {
+func (r *GRPCServer) Scan(ctx context.Context, req *pb.ScanRequest) (*pb.ScanResponse, error) {
 	limit, err := internal.Uint64ToInt(req.Limit)
 	if err != nil {
 		return &pb.ScanResponse{
 			Kv: nil,
 		}, errors.WithStack(err)
 	}
-	readTS := snapshotTS(r.coordinator.Clock(), r.store)
+	readTS := snapshotTS(r.clock(), r.store)
 	res, err := r.store.ScanAt(ctx, req.StartKey, req.EndKey, limit, readTS)
 	if err != nil {
 		return &pb.ScanResponse{