fix: transfer leader on close

kv/kv.go

@@ -204,14 +204,20 @@ func (kv *KV) Join(id, addr string) error {
 
 // Leave removes a node from the Raft cluster.
 func (kv *KV) Leave(id string) error {
-	removeFuture := kv.raft.RemoveServer(raft.ServerID(id), 0, 0)
-	return removeFuture.Error()
+	future := kv.raft.RemoveServer(raft.ServerID(id), 0, 0)
+	return future.Error()
 }
 
 // Close gracefully shuts down the Raft node.
 func (kv *KV) Close() error {
-	f := kv.raft.Shutdown()
-	return f.Error()
+	if kv.isLeader() {
+		future := kv.raft.LeadershipTransfer()
+		if err := future.Error(); err != nil {
+			return err
+		}
+	}
+	future := kv.raft.Shutdown()
+	return future.Error()
 }
 
 // WaitForLeader blocks until the Raft node detects a leader in the cluster or times out.
@@ -230,6 +236,9 @@ func (kv *KV) WaitForLeader(timeout time.Duration) error {
 		}
 	}
 }
+func (kv *KV) isLeader() bool {
+	return kv.raft.State() == raft.Leader
+}
 
 // apply encodes the request and submits it to Raft.
 func (kv *KV) apply(reqType RequestType, req proto.Message) (any, error) {

mokv/mokv.go

@@ -47,6 +47,7 @@ type Storer interface {
 	Set(key string, value []byte) error
 	Delete(key string) error
 	List() <-chan []byte
+	Close() error
 }
 
 // GetEnv defines a function signature for retrieving environment variables.
@@ -227,6 +228,10 @@ func (r *MOKV) close(ctx context.Context) error {
 		errs = append(errs, fmt.Errorf("membership leave error: %w", err))
 	}
 
+	if err := r.kv.Close(); err != nil {
+		errs = append(errs, fmt.Errorf("KV close error: %w", err))
+	}
+
 	r.grpcServer.GracefulStop()
 
 	if err := r.meterProvider.Shutdown(ctx); err != nil {

server/server_test.go

@@ -29,6 +29,10 @@ func (t *testKV) GetServers() ([]*api.Server, error) {
 	return nil, nil
 }
 
+func (t *testKV) Close() error {
+	return nil
+}
+
 // setupTestServer creates and starts a test server, returning cleanup function
 func setupTestServer(t *testing.T) (api.KVClient, func()) {
 	t.Helper()