Add Raft leader proxying for GET requests

Author: bootjp

adapter/grpc.go

@@ -140,7 +140,12 @@ func (r GRPCServer) Get(ctx context.Context, req *pb.GetRequest) (*pb.GetRespons
 
 	v, err := r.store.Get(ctx, req.Key)
 	if err != nil {
-		return nil, errors.WithStack(err)
+		switch {
+		case errors.Is(err, store.ErrKeyNotFound):
+			return &pb.GetResponse{Value: nil}, nil
+		default:
+			return nil, errors.WithStack(err)
+		}
 	}
 
 	r.log.InfoContext(ctx, "Get",

adapter/redis.go

@@ -8,9 +8,12 @@ import (
 	"strings"
 
 	"github.com/bootjp/elastickv/kv"
+	pb "github.com/bootjp/elastickv/proto"
 	"github.com/bootjp/elastickv/store"
 	"github.com/cockroachdb/errors"
 	"github.com/tidwall/redcon"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
 )
 
 //nolint:mnd
@@ -113,6 +116,15 @@ func (r *RedisServer) set(conn redcon.Conn, cmd redcon.Command) {
 }
 
 func (r *RedisServer) get(conn redcon.Conn, cmd redcon.Command) {
+	if v, ok := r.tryLeaderGet(cmd.Args[1]); ok {
+		if v == nil {
+			conn.WriteNull()
+		} else {
+			conn.WriteBulk(v)
+		}
+		return
+	}
+
 	v, err := r.store.Get(context.Background(), cmd.Args[1])
 	if err != nil {
 		conn.WriteNull()
@@ -190,3 +202,32 @@ func (r *RedisServer) keys(conn redcon.Conn, cmd redcon.Command) {
 		conn.WriteBulk(kvPair.Key)
 	}
 }
+
+// tryLeaderGet proxies a GET to the current Raft leader, returning the value and
+// whether the proxy succeeded.
+func (r *RedisServer) tryLeaderGet(key []byte) ([]byte, bool) {
+	c, ok := r.coordinator.(*kv.Coordinate)
+	if !ok || r.coordinator.IsLeader() {
+		return nil, false
+	}
+	addr := c.RaftLeader()
+	if addr == "" {
+		return nil, false
+	}
+
+	conn, err := grpc.NewClient(string(addr),
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithDefaultCallOptions(grpc.WaitForReady(true)),
+	)
+	if err != nil {
+		return nil, false
+	}
+	defer conn.Close()
+
+	cli := pb.NewRawKVClient(conn)
+	resp, err := cli.RawGet(context.Background(), &pb.RawGetRequest{Key: key})
+	if err != nil {
+		return nil, false
+	}
+	return resp.Value, true
+}

kv/coordinator.go

@@ -2,6 +2,7 @@ package kv
 
 import (
 	"context"
+	"time"
 
 	pb "github.com/bootjp/elastickv/proto"
 	"github.com/cockroachdb/errors"
@@ -49,6 +50,12 @@ func (c *Coordinate) IsLeader() bool {
 	return c.raft.State() == raft.Leader
 }
 
+// RaftLeader returns the current leader's address as known by this node.
+func (c *Coordinate) RaftLeader() raft.ServerAddress {
+	addr, _ := c.raft.LeaderWithID()
+	return addr
+}
+
 func (c *Coordinate) dispatchTxn(reqs []*Elem[OP]) (*CoordinateResponse, error) {
 	var logs []*pb.Request
 	for _, req := range reqs {
@@ -202,11 +209,41 @@ func (c *Coordinate) redirect(reqs *OperationGroup[OP]) (*CoordinateResponse, er
 		return nil, ErrInvalidRequest
 	}
 
+	if err := c.waitApplied(r.CommitIndex); err != nil {
+		return nil, errors.WithStack(err)
+	}
+
 	return &CoordinateResponse{
 		CommitIndex: r.CommitIndex,
 	}, nil
 }
 
+// waitApplied blocks until the local raft FSM has applied at least the given
+// log index. This is used on followers after redirecting a write to the leader
+// so that immediately subsequent reads served locally observe the new value.
+func (c *Coordinate) waitApplied(idx uint64) error {
+	// Fast path: already applied.
+	if c.raft.AppliedIndex() >= idx {
+		return nil
+	}
+
+	const (
+		pollInterval = 5 * time.Millisecond
+		// Keep generous but finite to avoid hanging client calls.
+		timeout = 2 * time.Second
+	)
+
+	deadline := time.Now().Add(timeout)
+	for time.Now().Before(deadline) {
+		if c.raft.AppliedIndex() >= idx {
+			return nil
+		}
+		time.Sleep(pollInterval)
+	}
+
+	return errors.Wrapf(ErrInvalidRequest, "timeout waiting for apply: want=%d applied=%d", idx, c.raft.AppliedIndex())
+}
+
 func (c *Coordinate) toForwardRequest(reqs []*pb.Request) *pb.ForwardRequest {
 	if len(reqs) == 0 {
 		return nil

kv/transaction.go

@@ -28,6 +28,30 @@ type TransactionResponse struct {
 	CommitIndex uint64
 }
 
+// applyAndBarrier submits a log entry, waits for it to be applied, and
+// surfaces both Raft transport errors and errors returned from FSM.Apply.
+// HashiCorp Raft delivers FSM errors via ApplyFuture.Response(), not Error(),
+// so we must inspect the response to avoid silently treating failed writes as
+// successes.
+func applyAndBarrier(r *raft.Raft, b []byte) (uint64, error) {
+	af := r.Apply(b, time.Second)
+	if err := af.Error(); err != nil {
+		return 0, errors.WithStack(err)
+	}
+
+	if resp := af.Response(); resp != nil {
+		if err, ok := resp.(error); ok && err != nil {
+			return 0, errors.WithStack(err)
+		}
+	}
+
+	if f := r.Barrier(time.Second); f.Error() != nil {
+		return 0, errors.WithStack(f.Error())
+	}
+
+	return af.Index(), nil
+}
+
 func (t *TransactionManager) Commit(reqs []*pb.Request) (*TransactionResponse, error) {
 	commitIndex, err := func() (uint64, error) {
 		commitIndex := uint64(0)
@@ -37,15 +61,11 @@ func (t *TransactionManager) Commit(reqs []*pb.Request) (*TransactionResponse, e
 				return 0, errors.WithStack(err)
 			}
 
-			af := t.raft.Apply(b, time.Second)
-			if af.Error() != nil {
-				return 0, errors.WithStack(af.Error())
-			}
-			f := t.raft.Barrier(time.Second)
-			if f.Error() != nil {
-				return 0, errors.WithStack(f.Error())
+			idx, err := applyAndBarrier(t.raft, b)
+			if err != nil {
+				return 0, err
 			}
-			commitIndex = af.Index()
+			commitIndex = idx
 		}
 
 		return commitIndex, nil
@@ -83,15 +103,11 @@ func (t *TransactionManager) Abort(reqs []*pb.Request) (*TransactionResponse, er
 			return nil, errors.WithStack(err)
 		}
 
-		af := t.raft.Apply(b, time.Second)
-		if af.Error() != nil {
-			return nil, errors.WithStack(af.Error())
-		}
-		f := t.raft.Barrier(time.Second)
-		if f.Error() != nil {
-			return nil, errors.WithStack(f.Error())
+		idx, err := applyAndBarrier(t.raft, b)
+		if err != nil {
+			return nil, err
 		}
-		commitIndex = af.Index()
+		commitIndex = idx
 	}
 
 	return &TransactionResponse{

store/bolt_store.go

@@ -61,7 +61,14 @@ func (s *boltStore) Get(ctx context.Context, key []byte) ([]byte, error) {
 		return nil
 	})
 
-	return v, errors.WithStack(err)
+	if err != nil {
+		return nil, errors.WithStack(err)
+	}
+	if v == nil {
+		return nil, ErrKeyNotFound
+	}
+
+	return v, nil
 }
 
 func (s *boltStore) Scan(ctx context.Context, start []byte, end []byte, limit int) ([]*KVPair, error) {