From 64f537335660dbdcfc520d4aeb8d468ef7e10dda Mon Sep 17 00:00:00 2001
From: you06 <you1474600@gmail.com>
Date: Tue, 17 Dec 2024 15:20:28 +0900
Subject: [PATCH] snapshot getter with mutex

Signed-off-by: you06 <you1474600@gmail.com>
---
 internal/unionstore/memdb_art.go   |  6 ++++--
 internal/unionstore/memdb_rbt.go   |  5 ++++-
 internal/unionstore/union_store.go | 12 ++++++++++++
 3 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/internal/unionstore/memdb_art.go b/internal/unionstore/memdb_art.go
index c7c1b21d98..dcc3e85995 100644
--- a/internal/unionstore/memdb_art.go
+++ b/internal/unionstore/memdb_art.go
@@ -161,7 +161,9 @@ func (db *artDBWithContext) SnapshotIterReverse(upper, lower []byte) Iterator {
 	return db.ART.SnapshotIterReverse(upper, lower)
 }
 
-// SnapshotGetter returns a Getter for a snapshot of MemBuffer.
 func (db *artDBWithContext) SnapshotGetter() Getter {
-	return db.ART.SnapshotGetter()
+	return &SnapshotGetter{
+		mu:     &db.RWMutex,
+		getter: db.ART.SnapshotGetter(),
+	}
 }
diff --git a/internal/unionstore/memdb_rbt.go b/internal/unionstore/memdb_rbt.go
index c805f49935..687d632a3a 100644
--- a/internal/unionstore/memdb_rbt.go
+++ b/internal/unionstore/memdb_rbt.go
@@ -173,5 +173,8 @@ func (db *rbtDBWithContext) SnapshotIterReverse(upper, lower []byte) Iterator {
 
 // SnapshotGetter returns a Getter for a snapshot of MemBuffer.
 func (db *rbtDBWithContext) SnapshotGetter() Getter {
-	return db.RBT.SnapshotGetter()
+	return &SnapshotGetter{
+		mu:     &db.RWMutex,
+		getter: db.RBT.SnapshotGetter(),
+	}
 }
diff --git a/internal/unionstore/union_store.go b/internal/unionstore/union_store.go
index 1a5f1a36b9..91060c0856 100644
--- a/internal/unionstore/union_store.go
+++ b/internal/unionstore/union_store.go
@@ -36,6 +36,7 @@ package unionstore
 
 import (
 	"context"
+	"sync"
 	"time"
 
 	tikverr "github.com/tikv/client-go/v2/error"
@@ -254,3 +255,14 @@ var (
 	_ MemBuffer = &rbtDBWithContext{}
 	_ MemBuffer = &artDBWithContext{}
 )
+
+type SnapshotGetter struct {
+	mu     *sync.RWMutex
+	getter Getter
+}
+
+func (getter *SnapshotGetter) Get(ctx context.Context, key []byte) ([]byte, error) {
+	getter.mu.RLock()
+	defer getter.mu.RUnlock()
+	return getter.getter.Get(ctx, key)
+}