From dee248b91387cec537603d3aaf97b2bd1ae1db5e Mon Sep 17 00:00:00 2001
From: Milos Zivkovic <milos.zivkovic@mvpworkshop.co>
Date: Tue, 11 Oct 2022 23:30:12 +0400
Subject: [PATCH] Replace the binary search method for finding the closest
 snapshot

---
 validators/store/snapshot/types.go      | 44 ++++++++++-
 validators/store/snapshot/types_test.go | 99 +++++++++++++++++--------
 2 files changed, 109 insertions(+), 34 deletions(-)

diff --git a/validators/store/snapshot/types.go b/validators/store/snapshot/types.go
index a7f7279801..d7a0f19e9d 100644
--- a/validators/store/snapshot/types.go
+++ b/validators/store/snapshot/types.go
@@ -308,11 +308,47 @@ func (s *snapshotStore) deleteLower(num uint64) {
 	s.Lock()
 	defer s.Unlock()
 
-	i := sort.Search(len(s.list), func(i int) bool {
-		return s.list[i].Number >= num
-	})
+	pruneIndex := s.findClosestSnapshotIndex(num)
+	s.list = s.list[pruneIndex:]
+}
+
+// findClosestSnapshotIndex finds the closest snapshot index for the specified
+// block number
+func (s *snapshotStore) findClosestSnapshotIndex(blockNum uint64) int {
+	// Check if the block number is lower than the highest saved snapshot
+	if blockNum < s.list[0].Number {
+		return 0
+	}
+
+	// Check if the block number if higher than the highest saved snapshot
+	if blockNum > s.list[len(s.list)-1].Number {
+		return len(s.list) - 1
+	}
+
+	var (
+		low  = 0
+		high = len(s.list) - 1
+	)
+
+	// Find the closest value using binary search
+	for low <= high {
+		mid := (high + low) / 2
+
+		if blockNum < s.list[mid].Number {
+			high = mid - 1
+		} else if blockNum > s.list[mid].Number {
+			low = mid + 1
+		} else {
+			return mid
+		}
+	}
+
+	// Check which of the two positions is closest (and has a higher block num)
+	if s.list[low].Number-blockNum < blockNum-s.list[high].Number {
+		return high
+	}
 
-	s.list = s.list[i:]
+	return low
 }
 
 // find returns the index of the first closest snapshot to the number specified
diff --git a/validators/store/snapshot/types_test.go b/validators/store/snapshot/types_test.go
index 82d48a31ec..9119eca745 100644
--- a/validators/store/snapshot/types_test.go
+++ b/validators/store/snapshot/types_test.go
@@ -1096,41 +1096,80 @@ func Test_snapshotStore_updateLastBlock(t *testing.T) {
 func Test_snapshotStore_deleteLower(t *testing.T) {
 	t.Parallel()
 
-	var (
-		metadata = &SnapshotMetadata{
-			LastBlock: 10,
-		}
+	metadata := &SnapshotMetadata{
+		LastBlock: 10,
+	}
 
-		snapshots = []*Snapshot{
-			{Number: 10},
-			{Number: 19},
-			{Number: 20},
-			{Number: 21},
-			{Number: 30},
-		}
+	testTable := []struct {
+		name              string
+		snapshots         []*Snapshot
+		boundary          uint64
+		expectedSnapshots []*Snapshot
+	}{
+		{
+			"Drop lower-number snapshots",
+			[]*Snapshot{
+				{Number: 10},
+				{Number: 19},
+				{Number: 25},
+				{Number: 30},
+			},
+			uint64(20),
+			[]*Snapshot{
+				{Number: 25},
+				{Number: 30},
+			},
+		},
+		{
+			"Higher block value",
+			[]*Snapshot{
+				{Number: 10},
+				{Number: 11},
+				{Number: 12},
+				{Number: 13},
+				{Number: 14},
+			},
+			uint64(15),
+			[]*Snapshot{
+				{Number: 14},
+			},
+		},
+		{
+			// Single snapshots shouldn't be dropped
+			"Single snapshot",
+			[]*Snapshot{
+				{Number: 10},
+			},
+			uint64(15),
+			[]*Snapshot{
+				{Number: 10},
+			},
+		},
+	}
 
-		boundary = uint64(20)
-	)
+	for _, testCase := range testTable {
+		testCase := testCase
 
-	store := newSnapshotStore(
-		metadata,
-		snapshots,
-	)
+		t.Run(testCase.name, func(t *testing.T) {
+			t.Parallel()
 
-	store.deleteLower(boundary)
+			store := newSnapshotStore(
+				metadata,
+				testCase.snapshots,
+			)
 
-	assert.Equal(
-		t,
-		&snapshotStore{
-			lastNumber: metadata.LastBlock,
-			list: []*Snapshot{
-				{Number: 20},
-				{Number: 21},
-				{Number: 30},
-			},
-		},
-		store,
-	)
+			store.deleteLower(testCase.boundary)
+
+			assert.Equal(
+				t,
+				&snapshotStore{
+					lastNumber: metadata.LastBlock,
+					list:       testCase.expectedSnapshots,
+				},
+				store,
+			)
+		})
+	}
 }
 
 func Test_snapshotStore_find(t *testing.T) {