membuffer: protect snapshot getter with mutex

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(),
+	}
 }

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(),
+	}
 }

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)
+}