From 8eadac85b1eb9ecf0e2bcc1823a33ffb5f4e5c82 Mon Sep 17 00:00:00 2001
From: Marek Siarkowicz <siarkowicz@google.com>
Date: Fri, 29 Nov 2024 12:15:14 +0100
Subject: [PATCH] Separate memory and disk snapshot indexes

Signed-off-by: Marek Siarkowicz <siarkowicz@google.com>
---
 server/etcdserver/server.go      | 20 ++++++++++++--------
 server/etcdserver/server_test.go |  3 +++
 2 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/server/etcdserver/server.go b/server/etcdserver/server.go
index 45ba1a676799..f141731bf18f 100644
--- a/server/etcdserver/server.go
+++ b/server/etcdserver/server.go
@@ -741,10 +741,11 @@ func (s *EtcdServer) ReportSnapshot(id uint64, status raft.SnapshotStatus) {
 }
 
 type etcdProgress struct {
-	confState         raftpb.ConfState
-	diskSnapshotIndex uint64
-	appliedt          uint64
-	appliedi          uint64
+	confState           raftpb.ConfState
+	diskSnapshotIndex   uint64
+	memorySnapshotIndex uint64
+	appliedt            uint64
+	appliedi            uint64
 }
 
 // raftReadyHandler contains a set of EtcdServer operations to be called by raftNode,
@@ -809,10 +810,11 @@ func (s *EtcdServer) run() {
 	s.r.start(rh)
 
 	ep := etcdProgress{
-		confState:         sn.Metadata.ConfState,
-		diskSnapshotIndex: sn.Metadata.Index,
-		appliedt:          sn.Metadata.Term,
-		appliedi:          sn.Metadata.Index,
+		confState:           sn.Metadata.ConfState,
+		diskSnapshotIndex:   sn.Metadata.Index,
+		memorySnapshotIndex: sn.Metadata.Index,
+		appliedt:            sn.Metadata.Term,
+		appliedi:            sn.Metadata.Index,
 	}
 
 	defer func() {
@@ -1133,6 +1135,7 @@ func (s *EtcdServer) applySnapshot(ep *etcdProgress, toApply *toApply) {
 	ep.appliedt = toApply.snapshot.Metadata.Term
 	ep.appliedi = toApply.snapshot.Metadata.Index
 	ep.diskSnapshotIndex = ep.appliedi
+	ep.memorySnapshotIndex = ep.appliedi
 	ep.confState = toApply.snapshot.Metadata.ConfState
 
 	// As backends and implementations like alarmsStore changed, we need
@@ -2160,6 +2163,7 @@ func (s *EtcdServer) snapshot(ep *etcdProgress) {
 		lg.Panic("failed to save snapshot", zap.Error(err))
 	}
 	ep.diskSnapshotIndex = ep.appliedi
+	ep.memorySnapshotIndex = ep.appliedi
 	if err = s.r.storage.Release(snap); err != nil {
 		lg.Panic("failed to release wal", zap.Error(err))
 	}
diff --git a/server/etcdserver/server_test.go b/server/etcdserver/server_test.go
index dcd73a7996b7..8724cdbe2649 100644
--- a/server/etcdserver/server_test.go
+++ b/server/etcdserver/server_test.go
@@ -688,6 +688,9 @@ func TestSnapshot(t *testing.T) {
 	if ep.diskSnapshotIndex != 1 {
 		t.Errorf("ep.diskSnapshotIndex = %d, want 1", ep.diskSnapshotIndex)
 	}
+	if ep.memorySnapshotIndex != 1 {
+		t.Errorf("ep.memorySnapshotIndex = %d, want 1", ep.memorySnapshotIndex)
+	}
 }
 
 // TestSnapshotOrdering ensures raft persists snapshot onto disk before