From 0d53a94ca8aa601f84f1d19d628e8d8d53967fbd Mon Sep 17 00:00:00 2001
From: garethgeorge <garethgeorge97@gmail.com>
Date: Fri, 3 May 2024 01:15:23 -0700
Subject: [PATCH] feat: add seek support to join iterator for better
 performance

---
 internal/oplog/indexutil/indexutil.go | 50 ++++++++++++++++++++++++---
 1 file changed, 45 insertions(+), 5 deletions(-)

diff --git a/internal/oplog/indexutil/indexutil.go b/internal/oplog/indexutil/indexutil.go
index c00d52cf..8b236632 100644
--- a/internal/oplog/indexutil/indexutil.go
+++ b/internal/oplog/indexutil/indexutil.go
@@ -30,12 +30,19 @@ type IndexIterator interface {
 	Next() (int64, bool)
 }
 
+type SeekableIndexIterator interface {
+	IndexIterator
+	Seek(int64) (int64, bool) // seek to the first recordId >= id and return it or return false.
+}
+
 type IndexSearchIterator struct {
 	c      *bolt.Cursor
 	k      []byte
 	prefix []byte
 }
 
+var _ SeekableIndexIterator = &IndexSearchIterator{}
+
 func newSearchIterator(b *bolt.Bucket, prefix []byte) IndexIterator {
 	c := b.Cursor()
 	k, _ := c.Seek(prefix)
@@ -59,13 +66,38 @@ func (i *IndexSearchIterator) Next() (int64, bool) {
 	return id, true
 }
 
+func (i *IndexSearchIterator) Seek(id int64) (int64, bool) {
+	seekTo := []byte{}
+	seekTo = append(seekTo, i.prefix...)
+	seekTo = append(seekTo, serializationutil.Itob(id)...)
+	k, _ := i.c.Seek(seekTo)
+	if k == nil || !bytes.HasPrefix(k, i.prefix) {
+		return 0, false
+	}
+	id, err := serializationutil.Btoi(k[len(i.prefix):])
+	if err != nil {
+		return 0, false
+	}
+	return id, true
+}
+
 type JoinIterator struct {
-	iters []IndexIterator
+	iters     []IndexIterator
+	seekables []SeekableIndexIterator
 }
 
 func NewJoinIterator(iters ...IndexIterator) *JoinIterator {
+	seekables := make([]SeekableIndexIterator, 0, len(iters))
+	for _, iter := range iters {
+		if seekable, ok := iter.(SeekableIndexIterator); ok {
+			seekables = append(seekables, seekable)
+		} else {
+			seekables = append(seekables, nil)
+		}
+	}
 	return &JoinIterator{
-		iters: iters,
+		iters:     iters,
+		seekables: seekables,
 	}
 }
 
@@ -104,9 +136,17 @@ func (j *JoinIterator) Next() (int64, bool) {
 			if id == nexts[maxIdx] {
 				continue
 			}
-			nexts[idx], ok = j.iters[idx].Next()
-			if !ok {
-				return 0, false
+
+			if j.seekables[idx] != nil {
+				nexts[idx], ok = j.seekables[idx].Seek(nexts[maxIdx])
+				if !ok {
+					return 0, false
+				}
+			} else {
+				nexts[idx], ok = j.iters[idx].Next()
+				if !ok {
+					return 0, false
+				}
 			}
 		}
 	}