diff --git a/cmd/milvus/mck.go b/cmd/milvus/mck.go
index 6581356c252b6..b22b28114c856 100644
--- a/cmd/milvus/mck.go
+++ b/cmd/milvus/mck.go
@@ -107,7 +107,7 @@ func (c *mck) execute(args []string, flags *flag.FlagSet) {
 func (c *mck) run() {
 	c.connectMinio()
 
-	_, values, err := c.metaKV.LoadWithPrefix(segmentPrefix)
+	_, values, err := c.metaKV.LoadWithPrefix(context.TODO(), segmentPrefix)
 	if err != nil {
 		log.Fatal("failed to list the segment info", zap.String("key", segmentPrefix), zap.Error(err))
 	}
@@ -121,7 +121,7 @@ func (c *mck) run() {
 		c.segmentIDMap[info.ID] = info
 	}
 
-	_, values, err = c.metaKV.LoadWithPrefix(collectionPrefix)
+	_, values, err = c.metaKV.LoadWithPrefix(context.TODO(), collectionPrefix)
 	if err != nil {
 		log.Fatal("failed to list the collection info", zap.String("key", collectionPrefix), zap.Error(err))
 	}
@@ -150,13 +150,13 @@ func (c *mck) run() {
 	}
 	log.Info("partition ids", zap.Int64s("ids", ids))
 
-	keys, values, err := c.metaKV.LoadWithPrefix(triggerTaskPrefix)
+	keys, values, err := c.metaKV.LoadWithPrefix(context.TODO(), triggerTaskPrefix)
 	if err != nil {
 		log.Fatal("failed to list the trigger task info", zap.Error(err))
 	}
 	c.extractTask(triggerTaskPrefix, keys, values)
 
-	keys, values, err = c.metaKV.LoadWithPrefix(activeTaskPrefix)
+	keys, values, err = c.metaKV.LoadWithPrefix(context.TODO(), activeTaskPrefix)
 	if err != nil {
 		log.Fatal("failed to list the active task info", zap.Error(err))
 	}
@@ -259,7 +259,7 @@ func getConfigValue(a string, b string, name string) string {
 }
 
 func (c *mck) cleanTrash() {
-	keys, _, err := c.metaKV.LoadWithPrefix(MckTrash)
+	keys, _, err := c.metaKV.LoadWithPrefix(context.TODO(), MckTrash)
 	if err != nil {
 		log.Error("failed to load backup info", zap.Error(err))
 		return
@@ -273,7 +273,7 @@ func (c *mck) cleanTrash() {
 	deleteAll := ""
 	fmt.Scanln(&deleteAll)
 	if deleteAll == "Y" {
-		err = c.metaKV.RemoveWithPrefix(MckTrash)
+		err = c.metaKV.RemoveWithPrefix(context.TODO(), MckTrash)
 		if err != nil {
 			log.Error("failed to remove backup infos", zap.String("key", MckTrash), zap.Error(err))
 			return
@@ -395,31 +395,31 @@ func (c *mck) extractTask(prefix string, keys []string, values []string) {
 func (c *mck) removeTask(invalidTask int64) bool {
 	taskType := c.taskNameMap[invalidTask]
 	key := c.taskKeyMap[invalidTask]
-	err := c.metaKV.Save(getTrashKey(taskType, key), c.allTaskInfo[key])
+	err := c.metaKV.Save(context.TODO(), getTrashKey(taskType, key), c.allTaskInfo[key])
 	if err != nil {
 		log.Warn("failed to backup task", zap.String("key", getTrashKey(taskType, key)), zap.Int64("task_id", invalidTask), zap.Error(err))
 		return false
 	}
 	fmt.Printf("Back up task successfully, back path: %s\n", getTrashKey(taskType, key))
-	err = c.metaKV.Remove(key)
+	err = c.metaKV.Remove(context.TODO(), key)
 	if err != nil {
 		log.Warn("failed to remove task", zap.Int64("task_id", invalidTask), zap.Error(err))
 		return false
 	}
 
 	key = fmt.Sprintf("%s/%d", taskInfoPrefix, invalidTask)
-	taskInfo, err := c.metaKV.Load(key)
+	taskInfo, err := c.metaKV.Load(context.TODO(), key)
 	if err != nil {
 		log.Warn("failed to load task info", zap.Int64("task_id", invalidTask), zap.Error(err))
 		return false
 	}
-	err = c.metaKV.Save(getTrashKey(taskType, key), taskInfo)
+	err = c.metaKV.Save(context.TODO(), getTrashKey(taskType, key), taskInfo)
 	if err != nil {
 		log.Warn("failed to backup task info", zap.Int64("task_id", invalidTask), zap.Error(err))
 		return false
 	}
 	fmt.Printf("Back up task info successfully, back path: %s\n", getTrashKey(taskType, key))
-	err = c.metaKV.Remove(key)
+	err = c.metaKV.Remove(context.TODO(), key)
 	if err != nil {
 		log.Warn("failed to remove task info", zap.Int64("task_id", invalidTask), zap.Error(err))
 	}
diff --git a/internal/kv/etcd/embed_etcd_config_test.go b/internal/kv/etcd/embed_etcd_config_test.go
index 1a121f6a09bf9..72c9ee156df3a 100644
--- a/internal/kv/etcd/embed_etcd_config_test.go
+++ b/internal/kv/etcd/embed_etcd_config_test.go
@@ -17,6 +17,7 @@
 package etcdkv_test
 
 import (
+	"context"
 	"os"
 	"testing"
 
@@ -49,7 +50,7 @@ func TestEtcdConfigLoad(te *testing.T) {
 		require.NoError(t, err)
 
 		defer metaKv.Close()
-		defer metaKv.RemoveWithPrefix("")
+		defer metaKv.RemoveWithPrefix(context.TODO(), "")
 
 		kv := metaKv.(*embed_etcd_kv.EmbedEtcdKV)
 		assert.Equal(t, kv.GetConfig().SnapshotCount, uint64(1000))
diff --git a/internal/kv/etcd/embed_etcd_kv.go b/internal/kv/etcd/embed_etcd_kv.go
index 754fb81e052ae..072304fb8e3bb 100644
--- a/internal/kv/etcd/embed_etcd_kv.go
+++ b/internal/kv/etcd/embed_etcd_kv.go
@@ -123,7 +123,7 @@ func (kv *EmbedEtcdKV) GetPath(key string) string {
 	return path.Join(kv.rootPath, key)
 }
 
-func (kv *EmbedEtcdKV) WalkWithPrefix(prefix string, paginationSize int, fn func([]byte, []byte) error) error {
+func (kv *EmbedEtcdKV) WalkWithPrefix(ctx context.Context, prefix string, paginationSize int, fn func([]byte, []byte) error) error {
 	prefix = path.Join(kv.rootPath, prefix)
 	ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout)
 	defer cancel()
@@ -158,7 +158,7 @@ func (kv *EmbedEtcdKV) WalkWithPrefix(prefix string, paginationSize int, fn func
 }
 
 // LoadWithPrefix returns all the keys and values with the given key prefix
-func (kv *EmbedEtcdKV) LoadWithPrefix(key string) ([]string, []string, error) {
+func (kv *EmbedEtcdKV) LoadWithPrefix(ctx context.Context, key string) ([]string, []string, error) {
 	key = path.Join(kv.rootPath, key)
 	log.Debug("LoadWithPrefix ", zap.String("prefix", key))
 	ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout)
@@ -178,7 +178,7 @@ func (kv *EmbedEtcdKV) LoadWithPrefix(key string) ([]string, []string, error) {
 	return keys, values, nil
 }
 
-func (kv *EmbedEtcdKV) Has(key string) (bool, error) {
+func (kv *EmbedEtcdKV) Has(ctx context.Context, key string) (bool, error) {
 	key = path.Join(kv.rootPath, key)
 	log.Debug("Has", zap.String("key", key))
 	ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout)
@@ -190,7 +190,7 @@ func (kv *EmbedEtcdKV) Has(key string) (bool, error) {
 	return resp.Count != 0, nil
 }
 
-func (kv *EmbedEtcdKV) HasPrefix(prefix string) (bool, error) {
+func (kv *EmbedEtcdKV) HasPrefix(ctx context.Context, prefix string) (bool, error) {
 	prefix = path.Join(kv.rootPath, prefix)
 	log.Debug("HasPrefix", zap.String("prefix", prefix))
 
@@ -206,7 +206,7 @@ func (kv *EmbedEtcdKV) HasPrefix(prefix string) (bool, error) {
 }
 
 // LoadBytesWithPrefix returns all the keys and values with the given key prefix
-func (kv *EmbedEtcdKV) LoadBytesWithPrefix(key string) ([]string, [][]byte, error) {
+func (kv *EmbedEtcdKV) LoadBytesWithPrefix(ctx context.Context, key string) ([]string, [][]byte, error) {
 	key = path.Join(kv.rootPath, key)
 	log.Debug("LoadBytesWithPrefix ", zap.String("prefix", key))
 	ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout)
@@ -226,7 +226,7 @@ func (kv *EmbedEtcdKV) LoadBytesWithPrefix(key string) ([]string, [][]byte, erro
 }
 
 // LoadBytesWithPrefix2 returns all the keys and values with versions by the given key prefix
-func (kv *EmbedEtcdKV) LoadBytesWithPrefix2(key string) ([]string, [][]byte, []int64, error) {
+func (kv *EmbedEtcdKV) LoadBytesWithPrefix2(ctx context.Context, key string) ([]string, [][]byte, []int64, error) {
 	key = path.Join(kv.rootPath, key)
 	log.Debug("LoadBytesWithPrefix2 ", zap.String("prefix", key))
 	ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout)
@@ -248,7 +248,7 @@ func (kv *EmbedEtcdKV) LoadBytesWithPrefix2(key string) ([]string, [][]byte, []i
 }
 
 // Load returns value of the given key
-func (kv *EmbedEtcdKV) Load(key string) (string, error) {
+func (kv *EmbedEtcdKV) Load(ctx context.Context, key string) (string, error) {
 	key = path.Join(kv.rootPath, key)
 	ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout)
 	defer cancel()
@@ -264,7 +264,7 @@ func (kv *EmbedEtcdKV) Load(key string) (string, error) {
 }
 
 // LoadBytes returns value of the given key
-func (kv *EmbedEtcdKV) LoadBytes(key string) ([]byte, error) {
+func (kv *EmbedEtcdKV) LoadBytes(ctx context.Context, key string) ([]byte, error) {
 	key = path.Join(kv.rootPath, key)
 	ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout)
 	defer cancel()
@@ -280,7 +280,7 @@ func (kv *EmbedEtcdKV) LoadBytes(key string) ([]byte, error) {
 }
 
 // MultiLoad returns values of a set of keys
-func (kv *EmbedEtcdKV) MultiLoad(keys []string) ([]string, error) {
+func (kv *EmbedEtcdKV) MultiLoad(ctx context.Context, keys []string) ([]string, error) {
 	ops := make([]clientv3.Op, 0, len(keys))
 	for _, keyLoad := range keys {
 		ops = append(ops, clientv3.OpGet(path.Join(kv.rootPath, keyLoad)))
@@ -316,7 +316,7 @@ func (kv *EmbedEtcdKV) MultiLoad(keys []string) ([]string, error) {
 }
 
 // MultiLoadBytes returns values of a set of keys
-func (kv *EmbedEtcdKV) MultiLoadBytes(keys []string) ([][]byte, error) {
+func (kv *EmbedEtcdKV) MultiLoadBytes(ctx context.Context, keys []string) ([][]byte, error) {
 	ops := make([]clientv3.Op, 0, len(keys))
 	for _, keyLoad := range keys {
 		ops = append(ops, clientv3.OpGet(path.Join(kv.rootPath, keyLoad)))
@@ -352,7 +352,7 @@ func (kv *EmbedEtcdKV) MultiLoadBytes(keys []string) ([][]byte, error) {
 }
 
 // LoadBytesWithRevision returns keys, values and revision with given key prefix.
-func (kv *EmbedEtcdKV) LoadBytesWithRevision(key string) ([]string, [][]byte, int64, error) {
+func (kv *EmbedEtcdKV) LoadBytesWithRevision(ctx context.Context, key string) ([]string, [][]byte, int64, error) {
 	key = path.Join(kv.rootPath, key)
 	log.Debug("LoadBytesWithRevision ", zap.String("prefix", key))
 	ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout)
@@ -372,7 +372,7 @@ func (kv *EmbedEtcdKV) LoadBytesWithRevision(key string) ([]string, [][]byte, in
 }
 
 // Save saves the key-value pair.
-func (kv *EmbedEtcdKV) Save(key, value string) error {
+func (kv *EmbedEtcdKV) Save(ctx context.Context, key, value string) error {
 	key = path.Join(kv.rootPath, key)
 	ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout)
 	defer cancel()
@@ -381,7 +381,7 @@ func (kv *EmbedEtcdKV) Save(key, value string) error {
 }
 
 // SaveBytes saves the key-value pair.
-func (kv *EmbedEtcdKV) SaveBytes(key string, value []byte) error {
+func (kv *EmbedEtcdKV) SaveBytes(ctx context.Context, key string, value []byte) error {
 	key = path.Join(kv.rootPath, key)
 	ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout)
 	defer cancel()
@@ -390,7 +390,7 @@ func (kv *EmbedEtcdKV) SaveBytes(key string, value []byte) error {
 }
 
 // SaveBytesWithLease is a function to put value in etcd with etcd lease options.
-func (kv *EmbedEtcdKV) SaveBytesWithLease(key string, value []byte, id clientv3.LeaseID) error {
+func (kv *EmbedEtcdKV) SaveBytesWithLease(ctx context.Context, key string, value []byte, id clientv3.LeaseID) error {
 	key = path.Join(kv.rootPath, key)
 	ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout)
 	defer cancel()
@@ -399,7 +399,7 @@ func (kv *EmbedEtcdKV) SaveBytesWithLease(key string, value []byte, id clientv3.
 }
 
 // MultiSave saves the key-value pairs in a transaction.
-func (kv *EmbedEtcdKV) MultiSave(kvs map[string]string) error {
+func (kv *EmbedEtcdKV) MultiSave(ctx context.Context, kvs map[string]string) error {
 	ops := make([]clientv3.Op, 0, len(kvs))
 	for key, value := range kvs {
 		ops = append(ops, clientv3.OpPut(path.Join(kv.rootPath, key), value))
@@ -413,7 +413,7 @@ func (kv *EmbedEtcdKV) MultiSave(kvs map[string]string) error {
 }
 
 // MultiSaveBytes saves the key-value pairs in a transaction.
-func (kv *EmbedEtcdKV) MultiSaveBytes(kvs map[string][]byte) error {
+func (kv *EmbedEtcdKV) MultiSaveBytes(ctx context.Context, kvs map[string][]byte) error {
 	ops := make([]clientv3.Op, 0, len(kvs))
 	for key, value := range kvs {
 		ops = append(ops, clientv3.OpPut(path.Join(kv.rootPath, key), string(value)))
@@ -427,7 +427,7 @@ func (kv *EmbedEtcdKV) MultiSaveBytes(kvs map[string][]byte) error {
 }
 
 // RemoveWithPrefix removes the keys with given prefix.
-func (kv *EmbedEtcdKV) RemoveWithPrefix(prefix string) error {
+func (kv *EmbedEtcdKV) RemoveWithPrefix(ctx context.Context, prefix string) error {
 	key := path.Join(kv.rootPath, prefix)
 	ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout)
 	defer cancel()
@@ -437,7 +437,7 @@ func (kv *EmbedEtcdKV) RemoveWithPrefix(prefix string) error {
 }
 
 // Remove removes the key.
-func (kv *EmbedEtcdKV) Remove(key string) error {
+func (kv *EmbedEtcdKV) Remove(ctx context.Context, key string) error {
 	key = path.Join(kv.rootPath, key)
 	ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout)
 	defer cancel()
@@ -447,7 +447,7 @@ func (kv *EmbedEtcdKV) Remove(key string) error {
 }
 
 // MultiRemove removes the keys in a transaction.
-func (kv *EmbedEtcdKV) MultiRemove(keys []string) error {
+func (kv *EmbedEtcdKV) MultiRemove(ctx context.Context, keys []string) error {
 	ops := make([]clientv3.Op, 0, len(keys))
 	for _, key := range keys {
 		ops = append(ops, clientv3.OpDelete(path.Join(kv.rootPath, key)))
@@ -461,7 +461,7 @@ func (kv *EmbedEtcdKV) MultiRemove(keys []string) error {
 }
 
 // MultiSaveAndRemove saves the key-value pairs and removes the keys in a transaction.
-func (kv *EmbedEtcdKV) MultiSaveAndRemove(saves map[string]string, removals []string, preds ...predicates.Predicate) error {
+func (kv *EmbedEtcdKV) MultiSaveAndRemove(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error {
 	cmps, err := parsePredicates(kv.rootPath, preds...)
 	if err != nil {
 		return err
@@ -491,7 +491,7 @@ func (kv *EmbedEtcdKV) MultiSaveAndRemove(saves map[string]string, removals []st
 }
 
 // MultiSaveBytesAndRemove saves the key-value pairs and removes the keys in a transaction.
-func (kv *EmbedEtcdKV) MultiSaveBytesAndRemove(saves map[string][]byte, removals []string) error {
+func (kv *EmbedEtcdKV) MultiSaveBytesAndRemove(ctx context.Context, saves map[string][]byte, removals []string) error {
 	ops := make([]clientv3.Op, 0, len(saves)+len(removals))
 	for key, value := range saves {
 		ops = append(ops, clientv3.OpPut(path.Join(kv.rootPath, key), string(value)))
@@ -508,26 +508,26 @@ func (kv *EmbedEtcdKV) MultiSaveBytesAndRemove(saves map[string][]byte, removals
 	return err
 }
 
-func (kv *EmbedEtcdKV) Watch(key string) clientv3.WatchChan {
+func (kv *EmbedEtcdKV) Watch(ctx context.Context, key string) clientv3.WatchChan {
 	key = path.Join(kv.rootPath, key)
 	rch := kv.client.Watch(context.Background(), key, clientv3.WithCreatedNotify())
 	return rch
 }
 
-func (kv *EmbedEtcdKV) WatchWithPrefix(key string) clientv3.WatchChan {
+func (kv *EmbedEtcdKV) WatchWithPrefix(ctx context.Context, key string) clientv3.WatchChan {
 	key = path.Join(kv.rootPath, key)
 	rch := kv.client.Watch(context.Background(), key, clientv3.WithPrefix(), clientv3.WithCreatedNotify())
 	return rch
 }
 
-func (kv *EmbedEtcdKV) WatchWithRevision(key string, revision int64) clientv3.WatchChan {
+func (kv *EmbedEtcdKV) WatchWithRevision(ctx context.Context, key string, revision int64) clientv3.WatchChan {
 	key = path.Join(kv.rootPath, key)
 	rch := kv.client.Watch(context.Background(), key, clientv3.WithPrefix(), clientv3.WithPrevKV(), clientv3.WithRev(revision))
 	return rch
 }
 
 // MultiSaveAndRemoveWithPrefix saves kv in @saves and removes the keys with given prefix in @removals.
-func (kv *EmbedEtcdKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string, preds ...predicates.Predicate) error {
+func (kv *EmbedEtcdKV) MultiSaveAndRemoveWithPrefix(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error {
 	cmps, err := parsePredicates(kv.rootPath, preds...)
 	if err != nil {
 		return err
@@ -557,7 +557,7 @@ func (kv *EmbedEtcdKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, rem
 }
 
 // MultiSaveBytesAndRemoveWithPrefix saves kv in @saves and removes the keys with given prefix in @removals.
-func (kv *EmbedEtcdKV) MultiSaveBytesAndRemoveWithPrefix(saves map[string][]byte, removals []string) error {
+func (kv *EmbedEtcdKV) MultiSaveBytesAndRemoveWithPrefix(ctx context.Context, saves map[string][]byte, removals []string) error {
 	ops := make([]clientv3.Op, 0, len(saves)+len(removals))
 	for key, value := range saves {
 		ops = append(ops, clientv3.OpPut(path.Join(kv.rootPath, key), string(value)))
@@ -576,7 +576,7 @@ func (kv *EmbedEtcdKV) MultiSaveBytesAndRemoveWithPrefix(saves map[string][]byte
 
 // CompareVersionAndSwap compares the existing key-value's version with version, and if
 // they are equal, the target is stored in etcd.
-func (kv *EmbedEtcdKV) CompareVersionAndSwap(key string, version int64, target string) (bool, error) {
+func (kv *EmbedEtcdKV) CompareVersionAndSwap(ctx context.Context, key string, version int64, target string) (bool, error) {
 	ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout)
 	defer cancel()
 	resp, err := kv.client.Txn(ctx).If(
@@ -593,7 +593,7 @@ func (kv *EmbedEtcdKV) CompareVersionAndSwap(key string, version int64, target s
 
 // CompareVersionAndSwapBytes compares the existing key-value's version with version, and if
 // they are equal, the target is stored in etcd.
-func (kv *EmbedEtcdKV) CompareVersionAndSwapBytes(key string, version int64, target []byte, opts ...clientv3.OpOption) (bool, error) {
+func (kv *EmbedEtcdKV) CompareVersionAndSwapBytes(ctx context.Context, key string, version int64, target []byte, opts ...clientv3.OpOption) (bool, error) {
 	ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout)
 	defer cancel()
 	resp, err := kv.client.Txn(ctx).If(
diff --git a/internal/kv/etcd/embed_etcd_kv_test.go b/internal/kv/etcd/embed_etcd_kv_test.go
index 7e01b5b5516b9..f18351417e4bc 100644
--- a/internal/kv/etcd/embed_etcd_kv_test.go
+++ b/internal/kv/etcd/embed_etcd_kv_test.go
@@ -17,6 +17,7 @@
 package etcdkv_test
 
 import (
+	"context"
 	"fmt"
 	"path"
 	"sort"
@@ -55,7 +56,7 @@ func TestEmbedEtcd(te *testing.T) {
 		require.NoError(t, err)
 
 		defer metaKv.Close()
-		defer metaKv.RemoveWithPrefix("")
+		defer metaKv.RemoveWithPrefix(context.TODO(), "")
 
 		saveAndLoadTests := []struct {
 			key   string
@@ -69,11 +70,11 @@ func TestEmbedEtcd(te *testing.T) {
 
 		for i, test := range saveAndLoadTests {
 			if i < 4 {
-				err = metaKv.Save(test.key, test.value)
+				err = metaKv.Save(context.TODO(), test.key, test.value)
 				assert.NoError(t, err)
 			}
 
-			val, err := metaKv.Load(test.key)
+			val, err := metaKv.Load(context.TODO(), test.key)
 			assert.NoError(t, err)
 			assert.Equal(t, test.value, val)
 		}
@@ -87,7 +88,7 @@ func TestEmbedEtcd(te *testing.T) {
 		}
 
 		for _, test := range invalidLoadTests {
-			val, err := metaKv.Load(test.invalidKey)
+			val, err := metaKv.Load(context.TODO(), test.invalidKey)
 			assert.Error(t, err)
 			assert.Zero(t, val)
 		}
@@ -124,7 +125,7 @@ func TestEmbedEtcd(te *testing.T) {
 		}
 
 		for _, test := range loadPrefixTests {
-			actualKeys, actualValues, err := metaKv.LoadWithPrefix(test.prefix)
+			actualKeys, actualValues, err := metaKv.LoadWithPrefix(context.TODO(), test.prefix)
 			assert.ElementsMatch(t, test.expectedKeys, actualKeys)
 			assert.ElementsMatch(t, test.expectedValues, actualValues)
 			assert.Equal(t, test.expectedError, err)
@@ -141,15 +142,15 @@ func TestEmbedEtcd(te *testing.T) {
 		}
 
 		for _, test := range removeTests {
-			err = metaKv.Remove(test.validKey)
+			err = metaKv.Remove(context.TODO(), test.validKey)
 			assert.NoError(t, err)
 
-			_, err = metaKv.Load(test.validKey)
+			_, err = metaKv.Load(context.TODO(), test.validKey)
 			assert.Error(t, err)
 
-			err = metaKv.Remove(test.validKey)
+			err = metaKv.Remove(context.TODO(), test.validKey)
 			assert.NoError(t, err)
-			err = metaKv.Remove(test.invalidKey)
+			err = metaKv.Remove(context.TODO(), test.invalidKey)
 			assert.NoError(t, err)
 		}
 	})
@@ -163,7 +164,7 @@ func TestEmbedEtcd(te *testing.T) {
 		require.NoError(t, err)
 
 		defer metaKv.Close()
-		defer metaKv.RemoveWithPrefix("")
+		defer metaKv.RemoveWithPrefix(context.TODO(), "")
 
 		saveAndLoadTests := []struct {
 			key   string
@@ -177,11 +178,11 @@ func TestEmbedEtcd(te *testing.T) {
 
 		for i, test := range saveAndLoadTests {
 			if i < 4 {
-				err = metaKv.SaveBytes(test.key, test.value)
+				err = metaKv.SaveBytes(context.TODO(), test.key, test.value)
 				assert.NoError(t, err)
 			}
 
-			val, err := metaKv.LoadBytes(test.key)
+			val, err := metaKv.LoadBytes(context.TODO(), test.key)
 			assert.NoError(t, err)
 			assert.Equal(t, test.value, val)
 		}
@@ -195,7 +196,7 @@ func TestEmbedEtcd(te *testing.T) {
 		}
 
 		for _, test := range invalidLoadTests {
-			val, err := metaKv.LoadBytes(test.invalidKey)
+			val, err := metaKv.LoadBytes(context.TODO(), test.invalidKey)
 			assert.Error(t, err)
 			assert.Zero(t, val)
 		}
@@ -232,12 +233,12 @@ func TestEmbedEtcd(te *testing.T) {
 		}
 
 		for _, test := range loadPrefixTests {
-			actualKeys, actualValues, err := metaKv.LoadBytesWithPrefix(test.prefix)
+			actualKeys, actualValues, err := metaKv.LoadBytesWithPrefix(context.TODO(), test.prefix)
 			assert.ElementsMatch(t, test.expectedKeys, actualKeys)
 			assert.ElementsMatch(t, test.expectedValues, actualValues)
 			assert.Equal(t, test.expectedError, err)
 
-			actualKeys, actualValues, versions, err := metaKv.LoadBytesWithPrefix2(test.prefix)
+			actualKeys, actualValues, versions, err := metaKv.LoadBytesWithPrefix2(context.TODO(), test.prefix)
 			assert.ElementsMatch(t, test.expectedKeys, actualKeys)
 			assert.ElementsMatch(t, test.expectedValues, actualValues)
 			assert.NotZero(t, versions)
@@ -255,15 +256,15 @@ func TestEmbedEtcd(te *testing.T) {
 		}
 
 		for _, test := range removeTests {
-			err = metaKv.Remove(test.validKey)
+			err = metaKv.Remove(context.TODO(), test.validKey)
 			assert.NoError(t, err)
 
-			_, err = metaKv.Load(test.validKey)
+			_, err = metaKv.Load(context.TODO(), test.validKey)
 			assert.Error(t, err)
 
-			err = metaKv.Remove(test.validKey)
+			err = metaKv.Remove(context.TODO(), test.validKey)
 			assert.NoError(t, err)
-			err = metaKv.Remove(test.invalidKey)
+			err = metaKv.Remove(context.TODO(), test.invalidKey)
 			assert.NoError(t, err)
 		}
 	})
@@ -275,7 +276,7 @@ func TestEmbedEtcd(te *testing.T) {
 		assert.NoError(t, err)
 
 		defer metaKv.Close()
-		defer metaKv.RemoveWithPrefix("")
+		defer metaKv.RemoveWithPrefix(context.TODO(), "")
 
 		prepareKV := []struct {
 			inKey   string
@@ -289,7 +290,7 @@ func TestEmbedEtcd(te *testing.T) {
 		}
 
 		for _, test := range prepareKV {
-			err = metaKv.SaveBytes(test.inKey, test.inValue)
+			err = metaKv.SaveBytes(context.TODO(), test.inKey, test.inValue)
 			require.NoError(t, err)
 		}
 
@@ -305,7 +306,7 @@ func TestEmbedEtcd(te *testing.T) {
 		}
 
 		for _, test := range loadWithRevisionTests {
-			keys, values, revision, err := metaKv.LoadBytesWithRevision(test.inKey)
+			keys, values, revision, err := metaKv.LoadBytesWithRevision(context.TODO(), test.inKey)
 			assert.NoError(t, err)
 			assert.Equal(t, test.expectedKeyNo, len(keys))
 			assert.ElementsMatch(t, test.expectedValues, values)
@@ -319,7 +320,7 @@ func TestEmbedEtcd(te *testing.T) {
 		assert.NoError(t, err)
 
 		defer metaKv.Close()
-		defer metaKv.RemoveWithPrefix("")
+		defer metaKv.RemoveWithPrefix(context.TODO(), "")
 
 		multiSaveTests := map[string]string{
 			"key_1":      "value_1",
@@ -330,10 +331,10 @@ func TestEmbedEtcd(te *testing.T) {
 			"_":          "other",
 		}
 
-		err = metaKv.MultiSave(multiSaveTests)
+		err = metaKv.MultiSave(context.TODO(), multiSaveTests)
 		assert.NoError(t, err)
 		for k, v := range multiSaveTests {
-			actualV, err := metaKv.Load(k)
+			actualV, err := metaKv.Load(context.TODO(), k)
 			assert.NoError(t, err)
 			assert.Equal(t, v, actualV)
 		}
@@ -349,7 +350,7 @@ func TestEmbedEtcd(te *testing.T) {
 		}
 
 		for _, test := range multiLoadTests {
-			vs, err := metaKv.MultiLoad(test.inputKeys)
+			vs, err := metaKv.MultiLoad(context.TODO(), test.inputKeys)
 			assert.NoError(t, err)
 			assert.Equal(t, test.expectedValues, vs)
 		}
@@ -365,7 +366,7 @@ func TestEmbedEtcd(te *testing.T) {
 		}
 
 		for _, test := range invalidMultiLoad {
-			vs, err := metaKv.MultiLoad(test.invalidKeys)
+			vs, err := metaKv.MultiLoad(context.TODO(), test.invalidKeys)
 			assert.Error(t, err)
 			assert.Equal(t, test.expectedValues, vs)
 		}
@@ -376,10 +377,10 @@ func TestEmbedEtcd(te *testing.T) {
 		}
 
 		for _, k := range removeWithPrefixTests {
-			err = metaKv.RemoveWithPrefix(k)
+			err = metaKv.RemoveWithPrefix(context.TODO(), k)
 			assert.NoError(t, err)
 
-			ks, vs, err := metaKv.LoadWithPrefix(k)
+			ks, vs, err := metaKv.LoadWithPrefix(context.TODO(), k)
 			assert.Empty(t, ks)
 			assert.Empty(t, vs)
 			assert.NoError(t, err)
@@ -392,10 +393,10 @@ func TestEmbedEtcd(te *testing.T) {
 			"_",
 		}
 
-		err = metaKv.MultiRemove(multiRemoveTests)
+		err = metaKv.MultiRemove(context.TODO(), multiRemoveTests)
 		assert.NoError(t, err)
 
-		ks, vs, err := metaKv.LoadWithPrefix("")
+		ks, vs, err := metaKv.LoadWithPrefix(context.TODO(), "")
 		assert.NoError(t, err)
 		assert.Empty(t, ks)
 		assert.Empty(t, vs)
@@ -412,11 +413,11 @@ func TestEmbedEtcd(te *testing.T) {
 			{make(map[string]string), []string{"multikey_2"}},
 		}
 		for _, test := range multiSaveAndRemoveTests {
-			err = metaKv.MultiSaveAndRemove(test.multiSaves, test.multiRemoves)
+			err = metaKv.MultiSaveAndRemove(context.TODO(), test.multiSaves, test.multiRemoves)
 			assert.NoError(t, err)
 		}
 
-		ks, vs, err = metaKv.LoadWithPrefix("")
+		ks, vs, err = metaKv.LoadWithPrefix(context.TODO(), "")
 		assert.NoError(t, err)
 		assert.Empty(t, ks)
 		assert.Empty(t, vs)
@@ -429,7 +430,7 @@ func TestEmbedEtcd(te *testing.T) {
 		assert.NoError(t, err)
 
 		defer metaKv.Close()
-		defer metaKv.RemoveWithPrefix("")
+		defer metaKv.RemoveWithPrefix(context.TODO(), "")
 
 		multiSaveTests := map[string][]byte{
 			"key_1":      []byte("value_1"),
@@ -440,10 +441,10 @@ func TestEmbedEtcd(te *testing.T) {
 			"_":          []byte("other"),
 		}
 
-		err = metaKv.MultiSaveBytes(multiSaveTests)
+		err = metaKv.MultiSaveBytes(context.TODO(), multiSaveTests)
 		assert.NoError(t, err)
 		for k, v := range multiSaveTests {
-			actualV, err := metaKv.LoadBytes(k)
+			actualV, err := metaKv.LoadBytes(context.TODO(), k)
 			assert.NoError(t, err)
 			assert.Equal(t, v, actualV)
 		}
@@ -459,7 +460,7 @@ func TestEmbedEtcd(te *testing.T) {
 		}
 
 		for _, test := range multiLoadTests {
-			vs, err := metaKv.MultiLoadBytes(test.inputKeys)
+			vs, err := metaKv.MultiLoadBytes(context.TODO(), test.inputKeys)
 			assert.NoError(t, err)
 			assert.Equal(t, test.expectedValues, vs)
 		}
@@ -475,7 +476,7 @@ func TestEmbedEtcd(te *testing.T) {
 		}
 
 		for _, test := range invalidMultiLoad {
-			vs, err := metaKv.MultiLoadBytes(test.invalidKeys)
+			vs, err := metaKv.MultiLoadBytes(context.TODO(), test.invalidKeys)
 			assert.Error(t, err)
 			assert.Equal(t, test.expectedValues, vs)
 		}
@@ -486,10 +487,10 @@ func TestEmbedEtcd(te *testing.T) {
 		}
 
 		for _, k := range removeWithPrefixTests {
-			err = metaKv.RemoveWithPrefix(k)
+			err = metaKv.RemoveWithPrefix(context.TODO(), k)
 			assert.NoError(t, err)
 
-			ks, vs, err := metaKv.LoadBytesWithPrefix(k)
+			ks, vs, err := metaKv.LoadBytesWithPrefix(context.TODO(), k)
 			assert.Empty(t, ks)
 			assert.Empty(t, vs)
 			assert.NoError(t, err)
@@ -502,10 +503,10 @@ func TestEmbedEtcd(te *testing.T) {
 			"_",
 		}
 
-		err = metaKv.MultiRemove(multiRemoveTests)
+		err = metaKv.MultiRemove(context.TODO(), multiRemoveTests)
 		assert.NoError(t, err)
 
-		ks, vs, err := metaKv.LoadBytesWithPrefix("")
+		ks, vs, err := metaKv.LoadBytesWithPrefix(context.TODO(), "")
 		assert.NoError(t, err)
 		assert.Empty(t, ks)
 		assert.Empty(t, vs)
@@ -522,11 +523,11 @@ func TestEmbedEtcd(te *testing.T) {
 			{map[string][]byte{}, []string{"multikey_2"}},
 		}
 		for _, test := range multiSaveAndRemoveTests {
-			err = metaKv.MultiSaveBytesAndRemove(test.multiSaves, test.multiRemoves)
+			err = metaKv.MultiSaveBytesAndRemove(context.TODO(), test.multiSaves, test.multiRemoves)
 			assert.NoError(t, err)
 		}
 
-		ks, vs, err = metaKv.LoadBytesWithPrefix("")
+		ks, vs, err = metaKv.LoadBytesWithPrefix(context.TODO(), "")
 		assert.NoError(t, err)
 		assert.Empty(t, ks)
 		assert.Empty(t, vs)
@@ -538,7 +539,7 @@ func TestEmbedEtcd(te *testing.T) {
 		require.NoError(t, err)
 
 		defer metaKv.Close()
-		defer metaKv.RemoveWithPrefix("")
+		defer metaKv.RemoveWithPrefix(context.TODO(), "")
 
 		prepareTests := map[string]string{
 			"x/abc/1": "1",
@@ -550,7 +551,7 @@ func TestEmbedEtcd(te *testing.T) {
 		}
 
 		// MultiSaveAndRemoveWithPrefix
-		err = metaKv.MultiSave(prepareTests)
+		err = metaKv.MultiSave(context.TODO(), prepareTests)
 		require.NoError(t, err)
 		multiSaveAndRemoveWithPrefixTests := []struct {
 			multiSave map[string]string
@@ -569,14 +570,14 @@ func TestEmbedEtcd(te *testing.T) {
 		}
 
 		for _, test := range multiSaveAndRemoveWithPrefixTests {
-			k, _, err := metaKv.LoadWithPrefix(test.loadPrefix)
+			k, _, err := metaKv.LoadWithPrefix(context.TODO(), test.loadPrefix)
 			assert.NoError(t, err)
 			assert.Equal(t, test.lengthBeforeRemove, len(k))
 
-			err = metaKv.MultiSaveAndRemoveWithPrefix(test.multiSave, test.prefix)
+			err = metaKv.MultiSaveAndRemoveWithPrefix(context.TODO(), test.multiSave, test.prefix)
 			assert.NoError(t, err)
 
-			k, _, err = metaKv.LoadWithPrefix(test.loadPrefix)
+			k, _, err = metaKv.LoadWithPrefix(context.TODO(), test.loadPrefix)
 			assert.NoError(t, err)
 			assert.Equal(t, test.lengthAfterRemove, len(k))
 		}
@@ -589,7 +590,7 @@ func TestEmbedEtcd(te *testing.T) {
 		require.NoError(t, err)
 
 		defer metaKv.Close()
-		defer metaKv.RemoveWithPrefix("")
+		defer metaKv.RemoveWithPrefix(context.TODO(), "")
 
 		prepareTests := map[string][]byte{
 			"x/abc/1": []byte("1"),
@@ -600,13 +601,13 @@ func TestEmbedEtcd(te *testing.T) {
 			"x/den/2": []byte("200"),
 		}
 
-		k, v, err := metaKv.LoadBytesWithPrefix("/")
+		k, v, err := metaKv.LoadBytesWithPrefix(context.TODO(), "/")
 		assert.NoError(t, err)
 		assert.Zero(t, len(k))
 		assert.Zero(t, len(v))
 
 		// MultiSaveAndRemoveWithPrefix
-		err = metaKv.MultiSaveBytes(prepareTests)
+		err = metaKv.MultiSaveBytes(context.TODO(), prepareTests)
 		require.NoError(t, err)
 		multiSaveAndRemoveWithPrefixTests := []struct {
 			multiSave map[string][]byte
@@ -624,14 +625,14 @@ func TestEmbedEtcd(te *testing.T) {
 		}
 
 		for _, test := range multiSaveAndRemoveWithPrefixTests {
-			k, _, err = metaKv.LoadBytesWithPrefix(test.loadPrefix)
+			k, _, err = metaKv.LoadBytesWithPrefix(context.TODO(), test.loadPrefix)
 			assert.NoError(t, err)
 			assert.Equal(t, test.lengthBeforeRemove, len(k))
 
-			err = metaKv.MultiSaveBytesAndRemoveWithPrefix(test.multiSave, test.prefix)
+			err = metaKv.MultiSaveBytesAndRemoveWithPrefix(context.TODO(), test.multiSave, test.prefix)
 			assert.NoError(t, err)
 
-			k, _, err = metaKv.LoadBytesWithPrefix(test.loadPrefix)
+			k, _, err = metaKv.LoadBytesWithPrefix(context.TODO(), test.loadPrefix)
 			assert.NoError(t, err)
 			assert.Equal(t, test.lengthAfterRemove, len(k))
 		}
@@ -643,13 +644,13 @@ func TestEmbedEtcd(te *testing.T) {
 		assert.NoError(t, err)
 
 		defer watchKv.Close()
-		defer watchKv.RemoveWithPrefix("")
+		defer watchKv.RemoveWithPrefix(context.TODO(), "")
 
-		ch := watchKv.Watch("x")
+		ch := watchKv.Watch(context.TODO(), "x")
 		resp := <-ch
 		assert.True(t, resp.Created)
 
-		ch = watchKv.WatchWithPrefix("x")
+		ch = watchKv.WatchWithPrefix(context.TODO(), "x")
 		resp = <-ch
 		assert.True(t, resp.Created)
 	})
@@ -661,7 +662,7 @@ func TestEmbedEtcd(te *testing.T) {
 		assert.NoError(t, err)
 
 		defer metaKv.Close()
-		defer metaKv.RemoveWithPrefix("")
+		defer metaKv.RemoveWithPrefix(context.TODO(), "")
 
 		revisionTests := []struct {
 			inKey       string
@@ -674,13 +675,13 @@ func TestEmbedEtcd(te *testing.T) {
 		}
 
 		for _, test := range revisionTests {
-			err = metaKv.SaveBytes(test.inKey, test.fistValue)
+			err = metaKv.SaveBytes(context.TODO(), test.inKey, test.fistValue)
 			require.NoError(t, err)
 
-			_, _, revision, _ := metaKv.LoadBytesWithRevision(test.inKey)
-			ch := metaKv.WatchWithRevision(test.inKey, revision+1)
+			_, _, revision, _ := metaKv.LoadBytesWithRevision(context.TODO(), test.inKey)
+			ch := metaKv.WatchWithRevision(context.TODO(), test.inKey, revision+1)
 
-			err = metaKv.SaveBytes(test.inKey, test.secondValue)
+			err = metaKv.SaveBytes(context.TODO(), test.inKey, test.secondValue)
 			require.NoError(t, err)
 
 			resp := <-ch
@@ -689,15 +690,15 @@ func TestEmbedEtcd(te *testing.T) {
 			assert.Equal(t, revision+1, resp.Header.Revision)
 		}
 
-		success, err := metaKv.CompareVersionAndSwapBytes("a/b/c", 0, []byte("1"))
+		success, err := metaKv.CompareVersionAndSwapBytes(context.TODO(), "a/b/c", 0, []byte("1"))
 		assert.NoError(t, err)
 		assert.True(t, success)
 
-		value, err := metaKv.LoadBytes("a/b/c")
+		value, err := metaKv.LoadBytes(context.TODO(), "a/b/c")
 		assert.NoError(t, err)
 		assert.Equal(t, value, []byte("1"))
 
-		success, err = metaKv.CompareVersionAndSwapBytes("a/b/c", 0, []byte("1"))
+		success, err = metaKv.CompareVersionAndSwapBytes(context.TODO(), "a/b/c", 0, []byte("1"))
 		assert.NoError(t, err)
 		assert.False(t, success)
 	})
@@ -708,7 +709,7 @@ func TestEmbedEtcd(te *testing.T) {
 		assert.NoError(t, err)
 
 		defer metaKv.Close()
-		defer metaKv.RemoveWithPrefix("")
+		defer metaKv.RemoveWithPrefix(context.TODO(), "")
 
 		kvs := map[string]string{
 			"A/100":    "v1",
@@ -718,23 +719,23 @@ func TestEmbedEtcd(te *testing.T) {
 			"B/100":    "v5",
 		}
 
-		err = metaKv.MultiSave(kvs)
+		err = metaKv.MultiSave(context.TODO(), kvs)
 		assert.NoError(t, err)
 		for k, v := range kvs {
-			actualV, err := metaKv.Load(k)
+			actualV, err := metaKv.Load(context.TODO(), k)
 			assert.NoError(t, err)
 			assert.Equal(t, v, actualV)
 		}
 
 		t.Run("apply function error ", func(t *testing.T) {
-			err = metaKv.WalkWithPrefix("A", 5, func(key []byte, value []byte) error {
+			err = metaKv.WalkWithPrefix(context.TODO(), "A", 5, func(key []byte, value []byte) error {
 				return errors.New("error")
 			})
 			assert.Error(t, err)
 		})
 
 		t.Run("get with non-exist prefix ", func(t *testing.T) {
-			err = metaKv.WalkWithPrefix("non-exist-prefix", 5, func(key []byte, value []byte) error {
+			err = metaKv.WalkWithPrefix(context.TODO(), "non-exist-prefix", 5, func(key []byte, value []byte) error {
 				return nil
 			})
 			assert.NoError(t, err)
@@ -755,7 +756,7 @@ func TestEmbedEtcd(te *testing.T) {
 				ret := make(map[string]string)
 				actualSortedKey := make([]string, 0)
 
-				err = metaKv.WalkWithPrefix("A", pagination, func(key []byte, value []byte) error {
+				err = metaKv.WalkWithPrefix(context.TODO(), "A", pagination, func(key []byte, value []byte) error {
 					k := string(key)
 					k = k[len(rootPath)+1:]
 					ret[k] = string(value)
@@ -783,23 +784,23 @@ func TestEmbedEtcd(te *testing.T) {
 		assert.NoError(t, err)
 
 		defer kv.Close()
-		defer kv.RemoveWithPrefix("")
+		defer kv.RemoveWithPrefix(context.TODO(), "")
 
-		has, err := kv.Has("key1")
+		has, err := kv.Has(context.TODO(), "key1")
 		assert.NoError(t, err)
 		assert.False(t, has)
 
-		err = kv.Save("key1", "value1")
+		err = kv.Save(context.TODO(), "key1", "value1")
 		assert.NoError(t, err)
 
-		has, err = kv.Has("key1")
+		has, err = kv.Has(context.TODO(), "key1")
 		assert.NoError(t, err)
 		assert.True(t, has)
 
-		err = kv.Remove("key1")
+		err = kv.Remove(context.TODO(), "key1")
 		assert.NoError(t, err)
 
-		has, err = kv.Has("key1")
+		has, err = kv.Has(context.TODO(), "key1")
 		assert.NoError(t, err)
 		assert.False(t, has)
 	})
@@ -810,23 +811,23 @@ func TestEmbedEtcd(te *testing.T) {
 		assert.NoError(t, err)
 
 		defer kv.Close()
-		defer kv.RemoveWithPrefix("")
+		defer kv.RemoveWithPrefix(context.TODO(), "")
 
-		has, err := kv.HasPrefix("key")
+		has, err := kv.HasPrefix(context.TODO(), "key")
 		assert.NoError(t, err)
 		assert.False(t, has)
 
-		err = kv.Save("key1", "value1")
+		err = kv.Save(context.TODO(), "key1", "value1")
 		assert.NoError(t, err)
 
-		has, err = kv.HasPrefix("key")
+		has, err = kv.HasPrefix(context.TODO(), "key")
 		assert.NoError(t, err)
 		assert.True(t, has)
 
-		err = kv.Remove("key1")
+		err = kv.Remove(context.TODO(), "key1")
 		assert.NoError(t, err)
 
-		has, err = kv.HasPrefix("key")
+		has, err = kv.HasPrefix(context.TODO(), "key")
 		assert.NoError(t, err)
 		assert.False(t, has)
 	})
@@ -865,7 +866,7 @@ func (s *EmbedEtcdKVSuite) SetupTest() {
 
 func (s *EmbedEtcdKVSuite) TearDownTest() {
 	if s.kv != nil {
-		s.kv.RemoveWithPrefix("")
+		s.kv.RemoveWithPrefix(context.TODO(), "")
 		s.kv.Close()
 		s.kv = nil
 	}
@@ -879,7 +880,7 @@ func (s *EmbedEtcdKVSuite) TestTxnWithPredicates() {
 		"lease2": "2",
 	}
 
-	err := etcdKV.MultiSave(prepareKV)
+	err := etcdKV.MultiSave(context.TODO(), prepareKV)
 	s.Require().NoError(err)
 
 	badPredicate := predicates.NewMockPredicate(s.T())
@@ -899,13 +900,13 @@ func (s *EmbedEtcdKVSuite) TestTxnWithPredicates() {
 
 	for _, test := range multiSaveAndRemovePredTests {
 		s.Run(test.tag, func() {
-			err := etcdKV.MultiSaveAndRemove(test.multiSave, nil, test.preds...)
+			err := etcdKV.MultiSaveAndRemove(context.TODO(), test.multiSave, nil, test.preds...)
 			if test.expectSuccess {
 				s.NoError(err)
 			} else {
 				s.Error(err)
 			}
-			err = etcdKV.MultiSaveAndRemoveWithPrefix(test.multiSave, nil, test.preds...)
+			err = etcdKV.MultiSaveAndRemoveWithPrefix(context.TODO(), test.multiSave, nil, test.preds...)
 			if test.expectSuccess {
 				s.NoError(err)
 			} else {
diff --git a/internal/kv/etcd/embed_etcd_restart_test.go b/internal/kv/etcd/embed_etcd_restart_test.go
index 97fa8600a33d1..a93db26e7efe1 100644
--- a/internal/kv/etcd/embed_etcd_restart_test.go
+++ b/internal/kv/etcd/embed_etcd_restart_test.go
@@ -17,6 +17,7 @@
 package etcdkv_test
 
 import (
+	"context"
 	"os"
 	"testing"
 
@@ -49,7 +50,7 @@ func TestEtcdRestartLoad(te *testing.T) {
 		require.NoError(t, err)
 
 		defer metaKv.Close()
-		defer metaKv.RemoveWithPrefix("")
+		defer metaKv.RemoveWithPrefix(context.TODO(), "")
 
 		saveAndLoadTests := []struct {
 			key   string
@@ -64,14 +65,14 @@ func TestEtcdRestartLoad(te *testing.T) {
 		// save some data
 		for i, test := range saveAndLoadTests {
 			if i < 4 {
-				err = metaKv.Save(test.key, test.value)
+				err = metaKv.Save(context.TODO(), test.key, test.value)
 				assert.NoError(t, err)
 			}
 		}
 
 		// check test result
 		for _, test := range saveAndLoadTests {
-			val, err := metaKv.Load(test.key)
+			val, err := metaKv.Load(context.TODO(), test.key)
 			assert.NoError(t, err)
 			assert.Equal(t, test.value, val)
 		}
@@ -83,7 +84,7 @@ func TestEtcdRestartLoad(te *testing.T) {
 		metaKv, _ = embed_etcd_kv.NewMetaKvFactory(rootPath, &param.EtcdCfg)
 
 		for _, test := range saveAndLoadTests {
-			val, err := metaKv.Load(test.key)
+			val, err := metaKv.Load(context.TODO(), test.key)
 			assert.NoError(t, err)
 			assert.Equal(t, test.value, val)
 		}
diff --git a/internal/kv/etcd/etcd_kv.go b/internal/kv/etcd/etcd_kv.go
index 5003d0fae4b1f..16f1b66af4420 100644
--- a/internal/kv/etcd/etcd_kv.go
+++ b/internal/kv/etcd/etcd_kv.go
@@ -20,6 +20,7 @@ import (
 	"context"
 	"encoding/binary"
 	"fmt"
+	"github.com/milvus-io/milvus/pkg/kv"
 	"path"
 	"time"
 
@@ -38,6 +39,9 @@ const (
 	defaultRequestTimeout = 10 * time.Second
 )
 
+// implementation assertion
+var _ kv.WatchKV = (*etcdKV)(nil)
+
 // etcdKV implements TxnKV interface, it supports to process multiple kvs in a transaction.
 type etcdKV struct {
 	client   *clientv3.Client
@@ -71,7 +75,7 @@ func (kv *etcdKV) GetPath(key string) string {
 	return path.Join(kv.rootPath, key)
 }
 
-func (kv *etcdKV) WalkWithPrefix(prefix string, paginationSize int, fn func([]byte, []byte) error) error {
+func (kv *etcdKV) WalkWithPrefix(ctx context.Context, prefix string, paginationSize int, fn func([]byte, []byte) error) error {
 	start := time.Now()
 	prefix = path.Join(kv.rootPath, prefix)
 
@@ -109,7 +113,7 @@ func (kv *etcdKV) WalkWithPrefix(prefix string, paginationSize int, fn func([]by
 }
 
 // LoadWithPrefix returns all the keys and values with the given key prefix.
-func (kv *etcdKV) LoadWithPrefix(key string) ([]string, []string, error) {
+func (kv *etcdKV) LoadWithPrefix(ctx context.Context, key string) ([]string, []string, error) {
 	start := time.Now()
 	key = path.Join(kv.rootPath, key)
 	ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout)
@@ -129,7 +133,7 @@ func (kv *etcdKV) LoadWithPrefix(key string) ([]string, []string, error) {
 	return keys, values, nil
 }
 
-func (kv *etcdKV) Has(key string) (bool, error) {
+func (kv *etcdKV) Has(ctx context.Context, key string) (bool, error) {
 	start := time.Now()
 	key = path.Join(kv.rootPath, key)
 	ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout)
@@ -144,7 +148,7 @@ func (kv *etcdKV) Has(key string) (bool, error) {
 	return resp.Count != 0, nil
 }
 
-func (kv *etcdKV) HasPrefix(prefix string) (bool, error) {
+func (kv *etcdKV) HasPrefix(ctx context.Context, prefix string) (bool, error) {
 	start := time.Now()
 	prefix = path.Join(kv.rootPath, prefix)
 	ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout)
@@ -160,7 +164,7 @@ func (kv *etcdKV) HasPrefix(prefix string) (bool, error) {
 }
 
 // LoadBytesWithPrefix returns all the keys and values with the given key prefix.
-func (kv *etcdKV) LoadBytesWithPrefix(key string) ([]string, [][]byte, error) {
+func (kv *etcdKV) LoadBytesWithPrefix(ctx context.Context, key string) ([]string, [][]byte, error) {
 	start := time.Now()
 	key = path.Join(kv.rootPath, key)
 	ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout)
@@ -181,7 +185,7 @@ func (kv *etcdKV) LoadBytesWithPrefix(key string) ([]string, [][]byte, error) {
 }
 
 // LoadBytesWithPrefix2 returns all the keys,values and key versions with the given key prefix.
-func (kv *etcdKV) LoadBytesWithPrefix2(key string) ([]string, [][]byte, []int64, error) {
+func (kv *etcdKV) LoadBytesWithPrefix2(ctx context.Context, key string) ([]string, [][]byte, []int64, error) {
 	start := time.Now()
 	key = path.Join(kv.rootPath, key)
 	ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout)
@@ -204,7 +208,7 @@ func (kv *etcdKV) LoadBytesWithPrefix2(key string) ([]string, [][]byte, []int64,
 }
 
 // Load returns value of the key.
-func (kv *etcdKV) Load(key string) (string, error) {
+func (kv *etcdKV) Load(ctx context.Context, key string) (string, error) {
 	start := time.Now()
 	key = path.Join(kv.rootPath, key)
 	ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout)
@@ -221,7 +225,7 @@ func (kv *etcdKV) Load(key string) (string, error) {
 }
 
 // LoadBytes returns value of the key.
-func (kv *etcdKV) LoadBytes(key string) ([]byte, error) {
+func (kv *etcdKV) LoadBytes(ctx context.Context, key string) ([]byte, error) {
 	start := time.Now()
 	key = path.Join(kv.rootPath, key)
 	ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout)
@@ -238,7 +242,7 @@ func (kv *etcdKV) LoadBytes(key string) ([]byte, error) {
 }
 
 // MultiLoad gets the values of the keys in a transaction.
-func (kv *etcdKV) MultiLoad(keys []string) ([]string, error) {
+func (kv *etcdKV) MultiLoad(ctx context.Context, keys []string) ([]string, error) {
 	start := time.Now()
 	ops := make([]clientv3.Op, 0, len(keys))
 	for _, keyLoad := range keys {
@@ -273,7 +277,7 @@ func (kv *etcdKV) MultiLoad(keys []string) ([]string, error) {
 }
 
 // MultiLoadBytes gets the values of the keys in a transaction.
-func (kv *etcdKV) MultiLoadBytes(keys []string) ([][]byte, error) {
+func (kv *etcdKV) MultiLoadBytes(ctx context.Context, keys []string) ([][]byte, error) {
 	start := time.Now()
 	ops := make([]clientv3.Op, 0, len(keys))
 	for _, keyLoad := range keys {
@@ -308,7 +312,7 @@ func (kv *etcdKV) MultiLoadBytes(keys []string) ([][]byte, error) {
 }
 
 // LoadBytesWithRevision returns keys, values and revision with given key prefix.
-func (kv *etcdKV) LoadBytesWithRevision(key string) ([]string, [][]byte, int64, error) {
+func (kv *etcdKV) LoadBytesWithRevision(ctx context.Context, key string) ([]string, [][]byte, int64, error) {
 	start := time.Now()
 	key = path.Join(kv.rootPath, key)
 	ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout)
@@ -329,7 +333,7 @@ func (kv *etcdKV) LoadBytesWithRevision(key string) ([]string, [][]byte, int64,
 }
 
 // Save saves the key-value pair.
-func (kv *etcdKV) Save(key, value string) error {
+func (kv *etcdKV) Save(ctx context.Context, key, value string) error {
 	start := time.Now()
 	key = path.Join(kv.rootPath, key)
 	ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout)
@@ -341,7 +345,7 @@ func (kv *etcdKV) Save(key, value string) error {
 }
 
 // SaveBytes saves the key-value pair.
-func (kv *etcdKV) SaveBytes(key string, value []byte) error {
+func (kv *etcdKV) SaveBytes(ctx context.Context, key string, value []byte) error {
 	start := time.Now()
 	key = path.Join(kv.rootPath, key)
 	ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout)
@@ -353,7 +357,7 @@ func (kv *etcdKV) SaveBytes(key string, value []byte) error {
 }
 
 // SaveBytesWithLease is a function to put value in etcd with etcd lease options.
-func (kv *etcdKV) SaveBytesWithLease(key string, value []byte, id clientv3.LeaseID) error {
+func (kv *etcdKV) SaveBytesWithLease(ctx context.Context, key string, value []byte, id clientv3.LeaseID) error {
 	start := time.Now()
 	key = path.Join(kv.rootPath, key)
 	ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout)
@@ -365,7 +369,7 @@ func (kv *etcdKV) SaveBytesWithLease(key string, value []byte, id clientv3.Lease
 }
 
 // MultiSave saves the key-value pairs in a transaction.
-func (kv *etcdKV) MultiSave(kvs map[string]string) error {
+func (kv *etcdKV) MultiSave(ctx context.Context, kvs map[string]string) error {
 	start := time.Now()
 	ops := make([]clientv3.Op, 0, len(kvs))
 	var keys []string
@@ -387,7 +391,7 @@ func (kv *etcdKV) MultiSave(kvs map[string]string) error {
 }
 
 // MultiSaveBytes saves the key-value pairs in a transaction.
-func (kv *etcdKV) MultiSaveBytes(kvs map[string][]byte) error {
+func (kv *etcdKV) MultiSaveBytes(ctx context.Context, kvs map[string][]byte) error {
 	start := time.Now()
 	ops := make([]clientv3.Op, 0, len(kvs))
 	var keys []string
@@ -409,7 +413,7 @@ func (kv *etcdKV) MultiSaveBytes(kvs map[string][]byte) error {
 }
 
 // RemoveWithPrefix removes the keys with given prefix.
-func (kv *etcdKV) RemoveWithPrefix(prefix string) error {
+func (kv *etcdKV) RemoveWithPrefix(ctx context.Context, prefix string) error {
 	start := time.Now()
 	key := path.Join(kv.rootPath, prefix)
 	ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout)
@@ -421,7 +425,7 @@ func (kv *etcdKV) RemoveWithPrefix(prefix string) error {
 }
 
 // Remove removes the key.
-func (kv *etcdKV) Remove(key string) error {
+func (kv *etcdKV) Remove(ctx context.Context, key string) error {
 	start := time.Now()
 	key = path.Join(kv.rootPath, key)
 	ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout)
@@ -433,7 +437,7 @@ func (kv *etcdKV) Remove(key string) error {
 }
 
 // MultiRemove removes the keys in a transaction.
-func (kv *etcdKV) MultiRemove(keys []string) error {
+func (kv *etcdKV) MultiRemove(ctx context.Context, keys []string) error {
 	start := time.Now()
 	ops := make([]clientv3.Op, 0, len(keys))
 	for _, key := range keys {
@@ -452,7 +456,7 @@ func (kv *etcdKV) MultiRemove(keys []string) error {
 }
 
 // MultiSaveAndRemove saves the key-value pairs and removes the keys in a transaction.
-func (kv *etcdKV) MultiSaveAndRemove(saves map[string]string, removals []string, preds ...predicates.Predicate) error {
+func (kv *etcdKV) MultiSaveAndRemove(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error {
 	cmps, err := parsePredicates(kv.rootPath, preds...)
 	if err != nil {
 		return err
@@ -492,7 +496,7 @@ func (kv *etcdKV) MultiSaveAndRemove(saves map[string]string, removals []string,
 }
 
 // MultiSaveBytesAndRemove saves the key-value pairs and removes the keys in a transaction.
-func (kv *etcdKV) MultiSaveBytesAndRemove(saves map[string][]byte, removals []string) error {
+func (kv *etcdKV) MultiSaveBytesAndRemove(ctx context.Context, saves map[string][]byte, removals []string) error {
 	start := time.Now()
 	ops := make([]clientv3.Op, 0, len(saves)+len(removals))
 	var keys []string
@@ -522,7 +526,7 @@ func (kv *etcdKV) MultiSaveBytesAndRemove(saves map[string][]byte, removals []st
 }
 
 // Watch starts watching a key, returns a watch channel.
-func (kv *etcdKV) Watch(key string) clientv3.WatchChan {
+func (kv *etcdKV) Watch(ctx context.Context, key string) clientv3.WatchChan {
 	start := time.Now()
 	key = path.Join(kv.rootPath, key)
 	rch := kv.client.Watch(context.Background(), key, clientv3.WithCreatedNotify())
@@ -531,7 +535,7 @@ func (kv *etcdKV) Watch(key string) clientv3.WatchChan {
 }
 
 // WatchWithPrefix starts watching a key with prefix, returns a watch channel.
-func (kv *etcdKV) WatchWithPrefix(key string) clientv3.WatchChan {
+func (kv *etcdKV) WatchWithPrefix(ctx context.Context, key string) clientv3.WatchChan {
 	start := time.Now()
 	key = path.Join(kv.rootPath, key)
 	rch := kv.client.Watch(context.Background(), key, clientv3.WithPrefix(), clientv3.WithCreatedNotify())
@@ -540,7 +544,7 @@ func (kv *etcdKV) WatchWithPrefix(key string) clientv3.WatchChan {
 }
 
 // WatchWithRevision starts watching a key with revision, returns a watch channel.
-func (kv *etcdKV) WatchWithRevision(key string, revision int64) clientv3.WatchChan {
+func (kv *etcdKV) WatchWithRevision(ctx context.Context, key string, revision int64) clientv3.WatchChan {
 	start := time.Now()
 	key = path.Join(kv.rootPath, key)
 	rch := kv.client.Watch(context.Background(), key, clientv3.WithPrefix(), clientv3.WithPrevKV(), clientv3.WithRev(revision))
@@ -549,7 +553,7 @@ func (kv *etcdKV) WatchWithRevision(key string, revision int64) clientv3.WatchCh
 }
 
 // MultiSaveAndRemoveWithPrefix saves kv in @saves and removes the keys with given prefix in @removals.
-func (kv *etcdKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string, preds ...predicates.Predicate) error {
+func (kv *etcdKV) MultiSaveAndRemoveWithPrefix(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error {
 	cmps, err := parsePredicates(kv.rootPath, preds...)
 	if err != nil {
 		return err
@@ -588,7 +592,7 @@ func (kv *etcdKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals
 }
 
 // MultiSaveBytesAndRemoveWithPrefix saves kv in @saves and removes the keys with given prefix in @removals.
-func (kv *etcdKV) MultiSaveBytesAndRemoveWithPrefix(saves map[string][]byte, removals []string) error {
+func (kv *etcdKV) MultiSaveBytesAndRemoveWithPrefix(ctx context.Context, saves map[string][]byte, removals []string) error {
 	start := time.Now()
 	ops := make([]clientv3.Op, 0, len(saves))
 	var keys []string
@@ -619,7 +623,7 @@ func (kv *etcdKV) MultiSaveBytesAndRemoveWithPrefix(saves map[string][]byte, rem
 
 // CompareVersionAndSwap compares the existing key-value's version with version, and if
 // they are equal, the target is stored in etcd.
-func (kv *etcdKV) CompareVersionAndSwap(key string, source int64, target string) (bool, error) {
+func (kv *etcdKV) CompareVersionAndSwap(ctx context.Context, key string, source int64, target string) (bool, error) {
 	start := time.Now()
 	ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout)
 	defer cancel()
@@ -635,7 +639,7 @@ func (kv *etcdKV) CompareVersionAndSwap(key string, source int64, target string)
 
 // CompareVersionAndSwapBytes compares the existing key-value's version with version, and if
 // they are equal, the target is stored in etcd.
-func (kv *etcdKV) CompareVersionAndSwapBytes(key string, source int64, target []byte, opts ...clientv3.OpOption) (bool, error) {
+func (kv *etcdKV) CompareVersionAndSwapBytes(ctx context.Context, key string, source int64, target []byte, opts ...clientv3.OpOption) (bool, error) {
 	start := time.Now()
 	ctx, cancel := context.WithTimeout(context.TODO(), kv.requestTimeout)
 	defer cancel()
diff --git a/internal/kv/etcd/etcd_kv_test.go b/internal/kv/etcd/etcd_kv_test.go
index 93e42d95f347a..123a0a209ce13 100644
--- a/internal/kv/etcd/etcd_kv_test.go
+++ b/internal/kv/etcd/etcd_kv_test.go
@@ -17,6 +17,7 @@
 package etcdkv
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"path"
@@ -79,7 +80,7 @@ func (s *EtcdKVSuite) SetupTest() {
 }
 
 func (s *EtcdKVSuite) TearDownTest() {
-	s.etcdKV.RemoveWithPrefix("")
+	s.etcdKV.RemoveWithPrefix(context.TODO(), "")
 	s.etcdKV.Close()
 }
 
@@ -97,11 +98,11 @@ func (s *EtcdKVSuite) TestSaveLoad() {
 
 	for i, test := range saveAndLoadTests {
 		if i < 4 {
-			err := etcdKV.Save(test.key, test.value)
+			err := etcdKV.Save(context.TODO(), test.key, test.value)
 			s.Require().NoError(err)
 		}
 
-		val, err := etcdKV.Load(test.key)
+		val, err := etcdKV.Load(context.TODO(), test.key)
 		s.Require().NoError(err)
 		s.Equal(test.value, val)
 	}
@@ -115,7 +116,7 @@ func (s *EtcdKVSuite) TestSaveLoad() {
 	}
 
 	for _, test := range invalidLoadTests {
-		val, err := etcdKV.Load(test.invalidKey)
+		val, err := etcdKV.Load(context.TODO(), test.invalidKey)
 		s.Error(err)
 		s.Zero(val)
 	}
@@ -152,7 +153,7 @@ func (s *EtcdKVSuite) TestSaveLoad() {
 	}
 
 	for _, test := range loadPrefixTests {
-		actualKeys, actualValues, err := etcdKV.LoadWithPrefix(test.prefix)
+		actualKeys, actualValues, err := etcdKV.LoadWithPrefix(context.TODO(), test.prefix)
 		s.ElementsMatch(test.expectedKeys, actualKeys)
 		s.ElementsMatch(test.expectedValues, actualValues)
 		s.Equal(test.expectedError, err)
@@ -169,15 +170,15 @@ func (s *EtcdKVSuite) TestSaveLoad() {
 	}
 
 	for _, test := range removeTests {
-		err := etcdKV.Remove(test.validKey)
+		err := etcdKV.Remove(context.TODO(), test.validKey)
 		s.NoError(err)
 
-		_, err = etcdKV.Load(test.validKey)
+		_, err = etcdKV.Load(context.TODO(), test.validKey)
 		s.Error(err)
 
-		err = etcdKV.Remove(test.validKey)
+		err = etcdKV.Remove(context.TODO(), test.validKey)
 		s.NoError(err)
-		err = etcdKV.Remove(test.invalidKey)
+		err = etcdKV.Remove(context.TODO(), test.invalidKey)
 		s.NoError(err)
 	}
 }
@@ -197,11 +198,11 @@ func (s *EtcdKVSuite) TestSaveAndLoadBytes() {
 
 	for i, test := range saveAndLoadTests {
 		if i < 4 {
-			err := etcdKV.SaveBytes(test.key, []byte(test.value))
+			err := etcdKV.SaveBytes(context.TODO(), test.key, []byte(test.value))
 			s.Require().NoError(err)
 		}
 
-		val, err := etcdKV.LoadBytes(test.key)
+		val, err := etcdKV.LoadBytes(context.TODO(), test.key)
 		s.NoError(err)
 		s.Equal(test.value, string(val))
 	}
@@ -215,7 +216,7 @@ func (s *EtcdKVSuite) TestSaveAndLoadBytes() {
 	}
 
 	for _, test := range invalidLoadTests {
-		val, err := etcdKV.LoadBytes(test.invalidKey)
+		val, err := etcdKV.LoadBytes(context.TODO(), test.invalidKey)
 		s.Error(err)
 		s.Zero(string(val))
 	}
@@ -252,7 +253,7 @@ func (s *EtcdKVSuite) TestSaveAndLoadBytes() {
 	}
 
 	for _, test := range loadPrefixTests {
-		actualKeys, actualValues, err := etcdKV.LoadBytesWithPrefix(test.prefix)
+		actualKeys, actualValues, err := etcdKV.LoadBytesWithPrefix(context.TODO(), test.prefix)
 		actualStringValues := make([]string, len(actualValues))
 		for i := range actualValues {
 			actualStringValues[i] = string(actualValues[i])
@@ -261,7 +262,7 @@ func (s *EtcdKVSuite) TestSaveAndLoadBytes() {
 		s.ElementsMatch(test.expectedValues, actualStringValues)
 		s.Equal(test.expectedError, err)
 
-		actualKeys, actualValues, versions, err := etcdKV.LoadBytesWithPrefix2(test.prefix)
+		actualKeys, actualValues, versions, err := etcdKV.LoadBytesWithPrefix2(context.TODO(), test.prefix)
 		actualStringValues = make([]string, len(actualValues))
 		for i := range actualValues {
 			actualStringValues[i] = string(actualValues[i])
@@ -283,15 +284,15 @@ func (s *EtcdKVSuite) TestSaveAndLoadBytes() {
 	}
 
 	for _, test := range removeTests {
-		err := etcdKV.Remove(test.validKey)
+		err := etcdKV.Remove(context.TODO(), test.validKey)
 		s.NoError(err)
 
-		_, err = etcdKV.Load(test.validKey)
+		_, err = etcdKV.Load(context.TODO(), test.validKey)
 		s.Error(err)
 
-		err = etcdKV.Remove(test.validKey)
+		err = etcdKV.Remove(context.TODO(), test.validKey)
 		s.NoError(err)
-		err = etcdKV.Remove(test.invalidKey)
+		err = etcdKV.Remove(context.TODO(), test.invalidKey)
 		s.NoError(err)
 	}
 }
@@ -311,7 +312,7 @@ func (s *EtcdKVSuite) TestLoadBytesWithRevision() {
 	}
 
 	for _, test := range prepareKV {
-		err := etcdKV.SaveBytes(test.inKey, []byte(test.inValue))
+		err := etcdKV.SaveBytes(context.TODO(), test.inKey, []byte(test.inValue))
 		s.NoError(err)
 	}
 
@@ -327,7 +328,7 @@ func (s *EtcdKVSuite) TestLoadBytesWithRevision() {
 	}
 
 	for _, test := range loadWithRevisionTests {
-		keys, values, revision, err := etcdKV.LoadBytesWithRevision(test.inKey)
+		keys, values, revision, err := etcdKV.LoadBytesWithRevision(context.TODO(), test.inKey)
 		s.NoError(err)
 		s.Equal(test.expectedKeyNo, len(keys))
 		stringValues := make([]string, len(values))
@@ -350,10 +351,10 @@ func (s *EtcdKVSuite) TestMultiSaveAndMultiLoad() {
 		"_":          "other",
 	}
 
-	err := etcdKV.MultiSave(multiSaveTests)
+	err := etcdKV.MultiSave(context.TODO(), multiSaveTests)
 	s.Require().NoError(err)
 	for k, v := range multiSaveTests {
-		actualV, err := etcdKV.Load(k)
+		actualV, err := etcdKV.Load(context.TODO(), k)
 		s.NoError(err)
 		s.Equal(v, actualV)
 	}
@@ -369,7 +370,7 @@ func (s *EtcdKVSuite) TestMultiSaveAndMultiLoad() {
 	}
 
 	for _, test := range multiLoadTests {
-		vs, err := etcdKV.MultiLoad(test.inputKeys)
+		vs, err := etcdKV.MultiLoad(context.TODO(), test.inputKeys)
 		s.NoError(err)
 		s.Equal(test.expectedValues, vs)
 	}
@@ -385,7 +386,7 @@ func (s *EtcdKVSuite) TestMultiSaveAndMultiLoad() {
 	}
 
 	for _, test := range invalidMultiLoad {
-		vs, err := etcdKV.MultiLoad(test.invalidKeys)
+		vs, err := etcdKV.MultiLoad(context.TODO(), test.invalidKeys)
 		s.Error(err)
 		s.Equal(test.expectedValues, vs)
 	}
@@ -396,10 +397,10 @@ func (s *EtcdKVSuite) TestMultiSaveAndMultiLoad() {
 	}
 
 	for _, k := range removeWithPrefixTests {
-		err = etcdKV.RemoveWithPrefix(k)
+		err = etcdKV.RemoveWithPrefix(context.TODO(), k)
 		s.NoError(err)
 
-		ks, vs, err := etcdKV.LoadWithPrefix(k)
+		ks, vs, err := etcdKV.LoadWithPrefix(context.TODO(), k)
 		s.Empty(ks)
 		s.Empty(vs)
 		s.NoError(err)
@@ -412,10 +413,10 @@ func (s *EtcdKVSuite) TestMultiSaveAndMultiLoad() {
 		"_",
 	}
 
-	err = etcdKV.MultiRemove(multiRemoveTests)
+	err = etcdKV.MultiRemove(context.TODO(), multiRemoveTests)
 	s.NoError(err)
 
-	ks, vs, err := etcdKV.LoadWithPrefix("")
+	ks, vs, err := etcdKV.LoadWithPrefix(context.TODO(), "")
 	s.NoError(err)
 	s.Empty(ks)
 	s.Empty(vs)
@@ -432,11 +433,11 @@ func (s *EtcdKVSuite) TestMultiSaveAndMultiLoad() {
 		{make(map[string]string), []string{"multikey_2"}},
 	}
 	for _, test := range multiSaveAndRemoveTests {
-		err = etcdKV.MultiSaveAndRemove(test.multiSaves, test.multiRemoves)
+		err = etcdKV.MultiSaveAndRemove(context.TODO(), test.multiSaves, test.multiRemoves)
 		s.NoError(err)
 	}
 
-	ks, vs, err = etcdKV.LoadWithPrefix("")
+	ks, vs, err = etcdKV.LoadWithPrefix(context.TODO(), "")
 	s.NoError(err)
 	s.Empty(ks)
 	s.Empty(vs)
@@ -458,10 +459,10 @@ func (s *EtcdKVSuite) TestMultiSaveBytesAndMultiLoadBytes() {
 		multiSaveBytesTests[k] = []byte(v)
 	}
 
-	err := etcdKV.MultiSaveBytes(multiSaveBytesTests)
+	err := etcdKV.MultiSaveBytes(context.TODO(), multiSaveBytesTests)
 	s.Require().NoError(err)
 	for k, v := range multiSaveTests {
-		actualV, err := etcdKV.LoadBytes(k)
+		actualV, err := etcdKV.LoadBytes(context.TODO(), k)
 		s.NoError(err)
 		s.Equal(v, string(actualV))
 	}
@@ -477,7 +478,7 @@ func (s *EtcdKVSuite) TestMultiSaveBytesAndMultiLoadBytes() {
 	}
 
 	for _, test := range multiLoadTests {
-		vs, err := etcdKV.MultiLoadBytes(test.inputKeys)
+		vs, err := etcdKV.MultiLoadBytes(context.TODO(), test.inputKeys)
 		stringVs := make([]string, len(vs))
 		for i := range vs {
 			stringVs[i] = string(vs[i])
@@ -497,7 +498,7 @@ func (s *EtcdKVSuite) TestMultiSaveBytesAndMultiLoadBytes() {
 	}
 
 	for _, test := range invalidMultiLoad {
-		vs, err := etcdKV.MultiLoadBytes(test.invalidKeys)
+		vs, err := etcdKV.MultiLoadBytes(context.TODO(), test.invalidKeys)
 		stringVs := make([]string, len(vs))
 		for i := range vs {
 			stringVs[i] = string(vs[i])
@@ -512,10 +513,10 @@ func (s *EtcdKVSuite) TestMultiSaveBytesAndMultiLoadBytes() {
 	}
 
 	for _, k := range removeWithPrefixTests {
-		err = etcdKV.RemoveWithPrefix(k)
+		err = etcdKV.RemoveWithPrefix(context.TODO(), k)
 		s.NoError(err)
 
-		ks, vs, err := etcdKV.LoadBytesWithPrefix(k)
+		ks, vs, err := etcdKV.LoadBytesWithPrefix(context.TODO(), k)
 		s.Empty(ks)
 		s.Empty(vs)
 		s.NoError(err)
@@ -528,10 +529,10 @@ func (s *EtcdKVSuite) TestMultiSaveBytesAndMultiLoadBytes() {
 		"_",
 	}
 
-	err = etcdKV.MultiRemove(multiRemoveTests)
+	err = etcdKV.MultiRemove(context.TODO(), multiRemoveTests)
 	s.NoError(err)
 
-	ks, vs, err := etcdKV.LoadBytesWithPrefix("")
+	ks, vs, err := etcdKV.LoadBytesWithPrefix(context.TODO(), "")
 	s.NoError(err)
 	s.Empty(ks)
 	s.Empty(vs)
@@ -549,11 +550,11 @@ func (s *EtcdKVSuite) TestMultiSaveBytesAndMultiLoadBytes() {
 	}
 
 	for _, test := range multiSaveAndRemoveTests {
-		err = etcdKV.MultiSaveBytesAndRemove(test.multiSaves, test.multiRemoves)
+		err = etcdKV.MultiSaveBytesAndRemove(context.TODO(), test.multiSaves, test.multiRemoves)
 		s.NoError(err)
 	}
 
-	ks, vs, err = etcdKV.LoadBytesWithPrefix("")
+	ks, vs, err = etcdKV.LoadBytesWithPrefix(context.TODO(), "")
 	s.NoError(err)
 	s.Empty(ks)
 	s.Empty(vs)
@@ -567,7 +568,7 @@ func (s *EtcdKVSuite) TestTxnWithPredicates() {
 		"lease2": "2",
 	}
 
-	err := etcdKV.MultiSave(prepareKV)
+	err := etcdKV.MultiSave(context.TODO(), prepareKV)
 	s.Require().NoError(err)
 
 	badPredicate := predicates.NewMockPredicate(s.T())
@@ -587,13 +588,13 @@ func (s *EtcdKVSuite) TestTxnWithPredicates() {
 
 	for _, test := range multiSaveAndRemovePredTests {
 		s.Run(test.tag, func() {
-			err := etcdKV.MultiSaveAndRemove(test.multiSave, nil, test.preds...)
+			err := etcdKV.MultiSaveAndRemove(context.TODO(), test.multiSave, nil, test.preds...)
 			if test.expectSuccess {
 				s.NoError(err)
 			} else {
 				s.Error(err)
 			}
-			err = etcdKV.MultiSaveAndRemoveWithPrefix(test.multiSave, nil, test.preds...)
+			err = etcdKV.MultiSaveAndRemoveWithPrefix(context.TODO(), test.multiSave, nil, test.preds...)
 			if test.expectSuccess {
 				s.NoError(err)
 			} else {
@@ -616,7 +617,7 @@ func (s *EtcdKVSuite) TestMultiSaveAndRemoveWithPrefix() {
 	}
 
 	// MultiSaveAndRemoveWithPrefix
-	err := etcdKV.MultiSave(prepareTests)
+	err := etcdKV.MultiSave(context.TODO(), prepareTests)
 	s.Require().NoError(err)
 	multiSaveAndRemoveWithPrefixTests := []struct {
 		multiSave map[string]string
@@ -634,14 +635,14 @@ func (s *EtcdKVSuite) TestMultiSaveAndRemoveWithPrefix() {
 	}
 
 	for _, test := range multiSaveAndRemoveWithPrefixTests {
-		k, _, err := etcdKV.LoadWithPrefix(test.loadPrefix)
+		k, _, err := etcdKV.LoadWithPrefix(context.TODO(), test.loadPrefix)
 		s.NoError(err)
 		s.Equal(test.lengthBeforeRemove, len(k))
 
-		err = etcdKV.MultiSaveAndRemoveWithPrefix(test.multiSave, test.prefix)
+		err = etcdKV.MultiSaveAndRemoveWithPrefix(context.TODO(), test.multiSave, test.prefix)
 		s.NoError(err)
 
-		k, _, err = etcdKV.LoadWithPrefix(test.loadPrefix)
+		k, _, err = etcdKV.LoadWithPrefix(context.TODO(), test.loadPrefix)
 		s.NoError(err)
 		s.Equal(test.lengthAfterRemove, len(k))
 	}
@@ -650,11 +651,11 @@ func (s *EtcdKVSuite) TestMultiSaveAndRemoveWithPrefix() {
 func (s *EtcdKVSuite) TestWatch() {
 	etcdKV := s.etcdKV
 
-	ch := etcdKV.Watch("x")
+	ch := etcdKV.Watch(context.TODO(), "x")
 	resp := <-ch
 	s.True(resp.Created)
 
-	ch = etcdKV.WatchWithPrefix("x")
+	ch = etcdKV.WatchWithPrefix(context.TODO(), "x")
 	resp = <-ch
 	s.True(resp.Created)
 }
@@ -673,13 +674,13 @@ func (s *EtcdKVSuite) TestRevisionBytes() {
 	}
 
 	for _, test := range revisionTests {
-		err := etcdKV.SaveBytes(test.inKey, test.fistValue)
+		err := etcdKV.SaveBytes(context.TODO(), test.inKey, test.fistValue)
 		s.Require().NoError(err)
 
-		_, _, revision, _ := etcdKV.LoadBytesWithRevision(test.inKey)
-		ch := etcdKV.WatchWithRevision(test.inKey, revision+1)
+		_, _, revision, _ := etcdKV.LoadBytesWithRevision(context.TODO(), test.inKey)
+		ch := etcdKV.WatchWithRevision(context.TODO(), test.inKey, revision+1)
 
-		err = etcdKV.SaveBytes(test.inKey, test.secondValue)
+		err = etcdKV.SaveBytes(context.TODO(), test.inKey, test.secondValue)
 		s.Require().NoError(err)
 
 		resp := <-ch
@@ -688,15 +689,15 @@ func (s *EtcdKVSuite) TestRevisionBytes() {
 		s.Equal(revision+1, resp.Header.Revision)
 	}
 
-	success, err := etcdKV.CompareVersionAndSwapBytes("a/b/c", 0, []byte("1"))
+	success, err := etcdKV.CompareVersionAndSwapBytes(context.TODO(), "a/b/c", 0, []byte("1"))
 	s.NoError(err)
 	s.True(success)
 
-	value, err := etcdKV.LoadBytes("a/b/c")
+	value, err := etcdKV.LoadBytes(context.TODO(), "a/b/c")
 	s.NoError(err)
 	s.Equal(string(value), "1")
 
-	success, err = etcdKV.CompareVersionAndSwapBytes("a/b/c", 0, []byte("1"))
+	success, err = etcdKV.CompareVersionAndSwapBytes(context.TODO(), "a/b/c", 0, []byte("1"))
 	s.NoError(err)
 	s.False(success)
 }
@@ -721,7 +722,7 @@ func Test_WalkWithPagination(t *testing.T) {
 	etcdKV := NewEtcdKV(etcdCli, rootPath)
 
 	defer etcdKV.Close()
-	defer etcdKV.RemoveWithPrefix("")
+	defer etcdKV.RemoveWithPrefix(context.TODO(), "")
 
 	kvs := map[string]string{
 		"A/100":    "v1",
@@ -731,23 +732,23 @@ func Test_WalkWithPagination(t *testing.T) {
 		"B/100":    "v5",
 	}
 
-	err = etcdKV.MultiSave(kvs)
+	err = etcdKV.MultiSave(context.TODO(), kvs)
 	assert.NoError(t, err)
 	for k, v := range kvs {
-		actualV, err := etcdKV.Load(k)
+		actualV, err := etcdKV.Load(context.TODO(), k)
 		assert.NoError(t, err)
 		assert.Equal(t, v, actualV)
 	}
 
 	t.Run("apply function error ", func(t *testing.T) {
-		err = etcdKV.WalkWithPrefix("A", 5, func(key []byte, value []byte) error {
+		err = etcdKV.WalkWithPrefix(context.TODO(), "A", 5, func(key []byte, value []byte) error {
 			return errors.New("error")
 		})
 		assert.Error(t, err)
 	})
 
 	t.Run("get with non-exist prefix ", func(t *testing.T) {
-		err = etcdKV.WalkWithPrefix("non-exist-prefix", 5, func(key []byte, value []byte) error {
+		err = etcdKV.WalkWithPrefix(context.TODO(), "non-exist-prefix", 5, func(key []byte, value []byte) error {
 			return nil
 		})
 		assert.NoError(t, err)
@@ -768,7 +769,7 @@ func Test_WalkWithPagination(t *testing.T) {
 			ret := make(map[string]string)
 			actualSortedKey := make([]string, 0)
 
-			err = etcdKV.WalkWithPrefix("A", pagination, func(key []byte, value []byte) error {
+			err = etcdKV.WalkWithPrefix(context.TODO(), "A", pagination, func(key []byte, value []byte) error {
 				k := string(key)
 				k = k[len(rootPath)+1:]
 				ret[k] = string(value)
@@ -844,27 +845,27 @@ func TestHas(t *testing.T) {
 	assert.NoError(t, err)
 	rootPath := "/etcd/test/root/has"
 	kv := NewEtcdKV(etcdCli, rootPath)
-	err = kv.RemoveWithPrefix("")
+	err = kv.RemoveWithPrefix(context.TODO(), "")
 	require.NoError(t, err)
 
 	defer kv.Close()
-	defer kv.RemoveWithPrefix("")
+	defer kv.RemoveWithPrefix(context.TODO(), "")
 
-	has, err := kv.Has("key1")
+	has, err := kv.Has(context.TODO(), "key1")
 	assert.NoError(t, err)
 	assert.False(t, has)
 
-	err = kv.Save("key1", "value1")
+	err = kv.Save(context.TODO(), "key1", "value1")
 	assert.NoError(t, err)
 
-	has, err = kv.Has("key1")
+	has, err = kv.Has(context.TODO(), "key1")
 	assert.NoError(t, err)
 	assert.True(t, has)
 
-	err = kv.Remove("key1")
+	err = kv.Remove(context.TODO(), "key1")
 	assert.NoError(t, err)
 
-	has, err = kv.Has("key1")
+	has, err = kv.Has(context.TODO(), "key1")
 	assert.NoError(t, err)
 	assert.False(t, has)
 }
@@ -882,27 +883,27 @@ func TestHasPrefix(t *testing.T) {
 	assert.NoError(t, err)
 	rootPath := "/etcd/test/root/hasprefix"
 	kv := NewEtcdKV(etcdCli, rootPath)
-	err = kv.RemoveWithPrefix("")
+	err = kv.RemoveWithPrefix(context.TODO(), "")
 	require.NoError(t, err)
 
 	defer kv.Close()
-	defer kv.RemoveWithPrefix("")
+	defer kv.RemoveWithPrefix(context.TODO(), "")
 
-	has, err := kv.HasPrefix("key")
+	has, err := kv.HasPrefix(context.TODO(), "key")
 	assert.NoError(t, err)
 	assert.False(t, has)
 
-	err = kv.Save("key1", "value1")
+	err = kv.Save(context.TODO(), "key1", "value1")
 	assert.NoError(t, err)
 
-	has, err = kv.HasPrefix("key")
+	has, err = kv.HasPrefix(context.TODO(), "key")
 	assert.NoError(t, err)
 	assert.True(t, has)
 
-	err = kv.Remove("key1")
+	err = kv.Remove(context.TODO(), "key1")
 	assert.NoError(t, err)
 
-	has, err = kv.HasPrefix("key")
+	has, err = kv.HasPrefix(context.TODO(), "key")
 	assert.NoError(t, err)
 	assert.False(t, has)
 }
diff --git a/internal/kv/mem/mem_kv.go b/internal/kv/mem/mem_kv.go
index 0b95569ca74d7..4d341b97c6cf7 100644
--- a/internal/kv/mem/mem_kv.go
+++ b/internal/kv/mem/mem_kv.go
@@ -17,6 +17,8 @@
 package memkv
 
 import (
+	"context"
+	"github.com/milvus-io/milvus/pkg/kv"
 	"strings"
 	"sync"
 
@@ -26,6 +28,9 @@ import (
 	"github.com/milvus-io/milvus/pkg/util/merr"
 )
 
+// implementation assertion
+var _ kv.TxnKV = (*MemoryKV)(nil)
+
 // MemoryKV implements BaseKv interface and relies on underling btree.BTree.
 // As its name implies, all data is stored in memory.
 type MemoryKV struct {
@@ -78,7 +83,7 @@ func (s memoryKVItem) Less(than btree.Item) bool {
 }
 
 // Load loads an object with @key.
-func (kv *MemoryKV) Load(key string) (string, error) {
+func (kv *MemoryKV) Load(ctx context.Context, key string) (string, error) {
 	kv.RLock()
 	defer kv.RUnlock()
 	item := kv.tree.Get(memoryKVItem{key: key})
@@ -89,7 +94,7 @@ func (kv *MemoryKV) Load(key string) (string, error) {
 }
 
 // LoadBytes loads an object with @key.
-func (kv *MemoryKV) LoadBytes(key string) ([]byte, error) {
+func (kv *MemoryKV) LoadBytes(ctx context.Context, key string) ([]byte, error) {
 	kv.RLock()
 	defer kv.RUnlock()
 	item := kv.tree.Get(memoryKVItem{key: key})
@@ -100,7 +105,7 @@ func (kv *MemoryKV) LoadBytes(key string) ([]byte, error) {
 }
 
 // Get return value if key exists, or return empty string
-func (kv *MemoryKV) Get(key string) string {
+func (kv *MemoryKV) Get(ctx context.Context, key string) string {
 	kv.RLock()
 	defer kv.RUnlock()
 	item := kv.tree.Get(memoryKVItem{key: key})
@@ -111,7 +116,7 @@ func (kv *MemoryKV) Get(key string) string {
 }
 
 // LoadBytesWithDefault loads an object with @key. If the object does not exist, @defaultValue will be returned.
-func (kv *MemoryKV) LoadBytesWithDefault(key string, defaultValue []byte) []byte {
+func (kv *MemoryKV) LoadBytesWithDefault(ctx context.Context, key string, defaultValue []byte) []byte {
 	kv.RLock()
 	defer kv.RUnlock()
 	item := kv.tree.Get(memoryKVItem{key: key})
@@ -122,7 +127,7 @@ func (kv *MemoryKV) LoadBytesWithDefault(key string, defaultValue []byte) []byte
 }
 
 // LoadBytesRange loads objects with range @startKey to @endKey with @limit number of objects.
-func (kv *MemoryKV) LoadBytesRange(key, endKey string, limit int) ([]string, [][]byte, error) {
+func (kv *MemoryKV) LoadBytesRange(ctx context.Context, key, endKey string, limit int) ([]string, [][]byte, error) {
 	kv.RLock()
 	defer kv.RUnlock()
 	keys := make([]string, 0, limit)
@@ -139,7 +144,7 @@ func (kv *MemoryKV) LoadBytesRange(key, endKey string, limit int) ([]string, [][
 }
 
 // Save object with @key to btree. Object value is @value.
-func (kv *MemoryKV) Save(key, value string) error {
+func (kv *MemoryKV) Save(ctx context.Context, key, value string) error {
 	kv.Lock()
 	defer kv.Unlock()
 	kv.tree.ReplaceOrInsert(memoryKVItem{key, StringValue(value)})
@@ -147,7 +152,7 @@ func (kv *MemoryKV) Save(key, value string) error {
 }
 
 // SaveBytes object with @key to btree. Object value is @value.
-func (kv *MemoryKV) SaveBytes(key string, value []byte) error {
+func (kv *MemoryKV) SaveBytes(ctx context.Context, key string, value []byte) error {
 	kv.Lock()
 	defer kv.Unlock()
 	kv.tree.ReplaceOrInsert(memoryKVItem{key, ByteSliceValue(value)})
@@ -155,7 +160,7 @@ func (kv *MemoryKV) SaveBytes(key string, value []byte) error {
 }
 
 // Remove deletes an object with @key.
-func (kv *MemoryKV) Remove(key string) error {
+func (kv *MemoryKV) Remove(ctx context.Context, key string) error {
 	kv.Lock()
 	defer kv.Unlock()
 
@@ -164,7 +169,7 @@ func (kv *MemoryKV) Remove(key string) error {
 }
 
 // MultiLoad loads objects with multi @keys.
-func (kv *MemoryKV) MultiLoad(keys []string) ([]string, error) {
+func (kv *MemoryKV) MultiLoad(ctx context.Context, keys []string) ([]string, error) {
 	kv.RLock()
 	defer kv.RUnlock()
 	result := make([]string, 0, len(keys))
@@ -176,7 +181,7 @@ func (kv *MemoryKV) MultiLoad(keys []string) ([]string, error) {
 }
 
 // MultiLoadBytes loads objects with multi @keys.
-func (kv *MemoryKV) MultiLoadBytes(keys []string) ([][]byte, error) {
+func (kv *MemoryKV) MultiLoadBytes(ctx context.Context, keys []string) ([][]byte, error) {
 	kv.RLock()
 	defer kv.RUnlock()
 	result := make([][]byte, 0, len(keys))
@@ -188,7 +193,7 @@ func (kv *MemoryKV) MultiLoadBytes(keys []string) ([][]byte, error) {
 }
 
 // MultiSave saves given key-value pairs in MemoryKV atomicly.
-func (kv *MemoryKV) MultiSave(kvs map[string]string) error {
+func (kv *MemoryKV) MultiSave(ctx context.Context, kvs map[string]string) error {
 	kv.Lock()
 	defer kv.Unlock()
 	for key, value := range kvs {
@@ -198,7 +203,7 @@ func (kv *MemoryKV) MultiSave(kvs map[string]string) error {
 }
 
 // MultiSaveBytes saves given key-value pairs in MemoryKV atomicly.
-func (kv *MemoryKV) MultiSaveBytes(kvs map[string][]byte) error {
+func (kv *MemoryKV) MultiSaveBytes(ctx context.Context, kvs map[string][]byte) error {
 	kv.Lock()
 	defer kv.Unlock()
 	for key, value := range kvs {
@@ -208,7 +213,7 @@ func (kv *MemoryKV) MultiSaveBytes(kvs map[string][]byte) error {
 }
 
 // MultiRemove removes given @keys in MemoryKV atomicly.
-func (kv *MemoryKV) MultiRemove(keys []string) error {
+func (kv *MemoryKV) MultiRemove(ctx context.Context, keys []string) error {
 	kv.Lock()
 	defer kv.Unlock()
 	for _, key := range keys {
@@ -218,7 +223,7 @@ func (kv *MemoryKV) MultiRemove(keys []string) error {
 }
 
 // MultiSaveAndRemove saves and removes given key-value pairs in MemoryKV atomicly.
-func (kv *MemoryKV) MultiSaveAndRemove(saves map[string]string, removals []string, preds ...predicates.Predicate) error {
+func (kv *MemoryKV) MultiSaveAndRemove(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error {
 	if len(preds) > 0 {
 		return merr.WrapErrServiceUnavailable("predicates not supported")
 	}
@@ -234,7 +239,7 @@ func (kv *MemoryKV) MultiSaveAndRemove(saves map[string]string, removals []strin
 }
 
 // MultiSaveBytesAndRemove saves and removes given key-value pairs in MemoryKV atomicly.
-func (kv *MemoryKV) MultiSaveBytesAndRemove(saves map[string][]byte, removals []string) error {
+func (kv *MemoryKV) MultiSaveBytesAndRemove(ctx context.Context, saves map[string][]byte, removals []string) error {
 	kv.Lock()
 	defer kv.Unlock()
 	for key, value := range saves {
@@ -247,7 +252,7 @@ func (kv *MemoryKV) MultiSaveBytesAndRemove(saves map[string][]byte, removals []
 }
 
 // LoadWithPrefix returns all keys & values with given prefix.
-func (kv *MemoryKV) LoadWithPrefix(key string) ([]string, []string, error) {
+func (kv *MemoryKV) LoadWithPrefix(ctx context.Context, key string) ([]string, []string, error) {
 	kv.Lock()
 	defer kv.Unlock()
 
@@ -265,7 +270,7 @@ func (kv *MemoryKV) LoadWithPrefix(key string) ([]string, []string, error) {
 }
 
 // LoadBytesWithPrefix returns all keys & values with given prefix.
-func (kv *MemoryKV) LoadBytesWithPrefix(key string) ([]string, [][]byte, error) {
+func (kv *MemoryKV) LoadBytesWithPrefix(ctx context.Context, key string) ([]string, [][]byte, error) {
 	kv.Lock()
 	defer kv.Unlock()
 
@@ -287,7 +292,7 @@ func (kv *MemoryKV) Close() {
 }
 
 // MultiSaveAndRemoveWithPrefix saves key-value pairs in @saves, & remove key with prefix in @removals in MemoryKV atomically.
-func (kv *MemoryKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string, preds ...predicates.Predicate) error {
+func (kv *MemoryKV) MultiSaveAndRemoveWithPrefix(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error {
 	if len(preds) > 0 {
 		return merr.WrapErrServiceUnavailable("predicates not supported")
 	}
@@ -314,7 +319,7 @@ func (kv *MemoryKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, remova
 }
 
 // MultiSaveBytesAndRemoveWithPrefix saves key-value pairs in @saves, & remove key with prefix in @removals in MemoryKV atomically.
-func (kv *MemoryKV) MultiSaveBytesAndRemoveWithPrefix(saves map[string][]byte, removals []string) error {
+func (kv *MemoryKV) MultiSaveBytesAndRemoveWithPrefix(ctx context.Context, saves map[string][]byte, removals []string) error {
 	kv.Lock()
 	defer kv.Unlock()
 
@@ -338,7 +343,7 @@ func (kv *MemoryKV) MultiSaveBytesAndRemoveWithPrefix(saves map[string][]byte, r
 }
 
 // RemoveWithPrefix remove key of given prefix in MemoryKV atomicly.
-func (kv *MemoryKV) RemoveWithPrefix(key string) error {
+func (kv *MemoryKV) RemoveWithPrefix(ctx context.Context, key string) error {
 	kv.Lock()
 	defer kv.Unlock()
 
@@ -356,13 +361,13 @@ func (kv *MemoryKV) RemoveWithPrefix(key string) error {
 	return nil
 }
 
-func (kv *MemoryKV) Has(key string) (bool, error) {
+func (kv *MemoryKV) Has(ctx context.Context, key string) (bool, error) {
 	kv.Lock()
 	defer kv.Unlock()
 	return kv.tree.Has(memoryKVItem{key: key}), nil
 }
 
-func (kv *MemoryKV) HasPrefix(prefix string) (bool, error) {
+func (kv *MemoryKV) HasPrefix(ctx context.Context, prefix string) (bool, error) {
 	kv.Lock()
 	defer kv.Unlock()
 
diff --git a/internal/kv/mem/mem_kv_test.go b/internal/kv/mem/mem_kv_test.go
index 7fbb66bc13e28..4d81294c5e5ae 100644
--- a/internal/kv/mem/mem_kv_test.go
+++ b/internal/kv/mem/mem_kv_test.go
@@ -17,6 +17,7 @@
 package memkv
 
 import (
+	"context"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -30,15 +31,15 @@ func TestMemoryKV_SaveAndLoadBytes(t *testing.T) {
 
 	key := "key"
 	value := []byte("value")
-	err := mem.SaveBytes(key, value)
+	err := mem.SaveBytes(context.TODO(), key, value)
 	assert.NoError(t, err)
 
-	_value, err := mem.LoadBytes(key)
+	_value, err := mem.LoadBytes(context.TODO(), key)
 	assert.NoError(t, err)
 	assert.Equal(t, value, _value)
 
 	noKey := "no_key"
-	_value, err = mem.LoadBytes(noKey)
+	_value, err = mem.LoadBytes(context.TODO(), noKey)
 	assert.Error(t, err)
 	assert.Empty(t, _value)
 }
@@ -56,16 +57,16 @@ func TestMemoryKV_LoadBytesRange(t *testing.T) {
 
 	mem := NewMemoryKV()
 	for _, kv := range saveAndLoadBytesTests {
-		err := mem.SaveBytes(kv.key, kv.value)
+		err := mem.SaveBytes(context.TODO(), kv.key, kv.value)
 		assert.NoError(t, err)
 	}
 
-	keys, values, err := mem.LoadBytesRange("test1", "test1/b", 0)
+	keys, values, err := mem.LoadBytesRange(context.TODO(), "test1", "test1/b", 0)
 	assert.Equal(t, len(keys), 2)
 	assert.Equal(t, len(values), 2)
 	assert.NoError(t, err)
 
-	keys, values, err = mem.LoadBytesRange("test1", "test1/a", 2)
+	keys, values, err = mem.LoadBytesRange(context.TODO(), "test1", "test1/a", 2)
 	assert.Equal(t, len(keys), 1)
 	assert.Equal(t, len(values), 1)
 	assert.NoError(t, err)
@@ -76,15 +77,15 @@ func TestMemoryKV_LoadBytesWithDefault(t *testing.T) {
 
 	key := "key"
 	value := []byte("value")
-	err := mem.SaveBytes(key, value)
+	err := mem.SaveBytes(context.TODO(), key, value)
 	assert.NoError(t, err)
 
 	_default := []byte("default")
-	_value := mem.LoadBytesWithDefault(key, _default)
+	_value := mem.LoadBytesWithDefault(context.TODO(), key, _default)
 	assert.Equal(t, value, _value)
 
 	noKey := "no_key"
-	_value = mem.LoadBytesWithDefault(noKey, _default)
+	_value = mem.LoadBytesWithDefault(context.TODO(), noKey, _default)
 	assert.Equal(t, _value, _default)
 }
 
@@ -101,21 +102,21 @@ func TestMemoryKV_LoadBytesWithPrefix(t *testing.T) {
 
 	mem := NewMemoryKV()
 	for _, kv := range saveAndLoadBytesTests {
-		err := mem.SaveBytes(kv.key, kv.value)
+		err := mem.SaveBytes(context.TODO(), kv.key, kv.value)
 		assert.NoError(t, err)
 	}
 
-	keys, values, err := mem.LoadBytesWithPrefix("test1")
+	keys, values, err := mem.LoadBytesWithPrefix(context.TODO(), "test1")
 	assert.Equal(t, len(keys), 3)
 	assert.Equal(t, len(values), 3)
 	assert.NoError(t, err)
 
-	keys, values, err = mem.LoadBytesWithPrefix("test")
+	keys, values, err = mem.LoadBytesWithPrefix(context.TODO(), "test")
 	assert.Equal(t, len(keys), 4)
 	assert.Equal(t, len(values), 4)
 	assert.NoError(t, err)
 
-	keys, values, err = mem.LoadBytesWithPrefix("a")
+	keys, values, err = mem.LoadBytesWithPrefix(context.TODO(), "a")
 	assert.Equal(t, len(keys), 0)
 	assert.Equal(t, len(values), 0)
 	assert.NoError(t, err)
@@ -137,14 +138,14 @@ func TestMemoryKV_MultiSaveBytes(t *testing.T) {
 	}
 
 	mem := NewMemoryKV()
-	err := mem.MultiSaveBytes(saveAndLoadBytesTests)
+	err := mem.MultiSaveBytes(context.TODO(), saveAndLoadBytesTests)
 	assert.NoError(t, err)
 
-	_values, err := mem.MultiLoadBytes(keys)
+	_values, err := mem.MultiLoadBytes(context.TODO(), keys)
 	assert.Equal(t, values, _values)
 	assert.NoError(t, err)
 
-	_values, err = mem.MultiLoadBytes([]string{})
+	_values, err = mem.MultiLoadBytes(context.TODO(), []string{})
 	assert.Empty(t, _values)
 	assert.NoError(t, err)
 }
@@ -165,14 +166,14 @@ func TestMemoryKV_MultiSaveBytesAndRemove(t *testing.T) {
 	}
 
 	mem := NewMemoryKV()
-	err := mem.MultiSaveBytesAndRemove(saveAndLoadBytesTests, []string{keys[0]})
+	err := mem.MultiSaveBytesAndRemove(context.TODO(), saveAndLoadBytesTests, []string{keys[0]})
 	assert.NoError(t, err)
 
-	_value, err := mem.LoadBytes(keys[0])
+	_value, err := mem.LoadBytes(context.TODO(), keys[0])
 	assert.Empty(t, _value)
 	assert.Error(t, err)
 
-	_values, err := mem.MultiLoadBytes(keys[1:])
+	_values, err := mem.MultiLoadBytes(context.TODO(), keys[1:])
 	assert.Equal(t, values[1:], _values)
 	assert.NoError(t, err)
 }
@@ -193,10 +194,10 @@ func TestMemoryKV_MultiSaveBytesAndRemoveWithPrefix(t *testing.T) {
 	}
 
 	mem := NewMemoryKV()
-	err := mem.MultiSaveBytesAndRemoveWithPrefix(saveAndLoadBytesTests, []string{"test1"})
+	err := mem.MultiSaveBytesAndRemoveWithPrefix(context.TODO(), saveAndLoadBytesTests, []string{"test1"})
 	assert.NoError(t, err)
 
-	_keys, _values, err := mem.LoadBytesWithPrefix("test")
+	_keys, _values, err := mem.LoadBytesWithPrefix(context.TODO(), "test")
 	assert.ElementsMatch(t, keys, _keys)
 	assert.ElementsMatch(t, values, _values)
 	assert.NoError(t, err)
@@ -205,21 +206,21 @@ func TestMemoryKV_MultiSaveBytesAndRemoveWithPrefix(t *testing.T) {
 func TestHas(t *testing.T) {
 	kv := NewMemoryKV()
 
-	has, err := kv.Has("key1")
+	has, err := kv.Has(context.TODO(), "key1")
 	assert.NoError(t, err)
 	assert.False(t, has)
 
-	err = kv.Save("key1", "value1")
+	err = kv.Save(context.TODO(), "key1", "value1")
 	assert.NoError(t, err)
 
-	has, err = kv.Has("key1")
+	has, err = kv.Has(context.TODO(), "key1")
 	assert.NoError(t, err)
 	assert.True(t, has)
 
-	err = kv.Remove("key1")
+	err = kv.Remove(context.TODO(), "key1")
 	assert.NoError(t, err)
 
-	has, err = kv.Has("key1")
+	has, err = kv.Has(context.TODO(), "key1")
 	assert.NoError(t, err)
 	assert.False(t, has)
 }
@@ -227,21 +228,21 @@ func TestHas(t *testing.T) {
 func TestHasPrefix(t *testing.T) {
 	kv := NewMemoryKV()
 
-	has, err := kv.HasPrefix("key")
+	has, err := kv.HasPrefix(context.TODO(), "key")
 	assert.NoError(t, err)
 	assert.False(t, has)
 
-	err = kv.Save("key1", "value1")
+	err = kv.Save(context.TODO(), "key1", "value1")
 	assert.NoError(t, err)
 
-	has, err = kv.HasPrefix("key")
+	has, err = kv.HasPrefix(context.TODO(), "key")
 	assert.NoError(t, err)
 	assert.True(t, has)
 
-	err = kv.Remove("key1")
+	err = kv.Remove(context.TODO(), "key1")
 	assert.NoError(t, err)
 
-	has, err = kv.HasPrefix("key")
+	has, err = kv.HasPrefix(context.TODO(), "key")
 	assert.NoError(t, err)
 	assert.False(t, has)
 }
@@ -250,11 +251,11 @@ func TestPredicates(t *testing.T) {
 	kv := NewMemoryKV()
 
 	// predicates not supported for mem kv for now
-	err := kv.MultiSaveAndRemove(map[string]string{}, []string{}, predicates.ValueEqual("a", "b"))
+	err := kv.MultiSaveAndRemove(context.TODO(), map[string]string{}, []string{}, predicates.ValueEqual("a", "b"))
 	assert.Error(t, err)
 	assert.ErrorIs(t, err, merr.ErrServiceUnavailable)
 
-	err = kv.MultiSaveAndRemoveWithPrefix(map[string]string{}, []string{}, predicates.ValueEqual("a", "b"))
+	err = kv.MultiSaveAndRemoveWithPrefix(context.TODO(), map[string]string{}, []string{}, predicates.ValueEqual("a", "b"))
 	assert.Error(t, err)
 	assert.ErrorIs(t, err, merr.ErrServiceUnavailable)
 }
diff --git a/internal/kv/mock_snapshot_kv.go b/internal/kv/mock_snapshot_kv.go
index 4c5ac49a72083..fda86ef107b87 100644
--- a/internal/kv/mock_snapshot_kv.go
+++ b/internal/kv/mock_snapshot_kv.go
@@ -1,62 +1,65 @@
 package kv
 
 import (
+	"context"
 	"github.com/milvus-io/milvus/pkg/kv"
 	"github.com/milvus-io/milvus/pkg/util/typeutil"
 )
 
+var _ kv.SnapShotKV = &mockSnapshotKV{}
+
 type mockSnapshotKV struct {
 	kv.SnapShotKV
-	SaveFunc                         func(key string, value string, ts typeutil.Timestamp) error
-	LoadFunc                         func(key string, ts typeutil.Timestamp) (string, error)
-	MultiSaveFunc                    func(kvs map[string]string, ts typeutil.Timestamp) error
-	LoadWithPrefixFunc               func(key string, ts typeutil.Timestamp) ([]string, []string, error)
-	MultiSaveAndRemoveWithPrefixFunc func(saves map[string]string, removals []string, ts typeutil.Timestamp) error
-	MultiSaveAndRemoveFunc           func(saves map[string]string, removals []string, ts typeutil.Timestamp) error
+	SaveFunc                         func(ctx context.Context, key string, value string, ts typeutil.Timestamp) error
+	LoadFunc                         func(ctx context.Context, key string, ts typeutil.Timestamp) (string, error)
+	MultiSaveFunc                    func(ctx context.Context, kvs map[string]string, ts typeutil.Timestamp) error
+	LoadWithPrefixFunc               func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error)
+	MultiSaveAndRemoveWithPrefixFunc func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error
+	MultiSaveAndRemoveFunc           func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error
 }
 
 func NewMockSnapshotKV() *mockSnapshotKV {
 	return &mockSnapshotKV{}
 }
 
-func (m mockSnapshotKV) Save(key string, value string, ts typeutil.Timestamp) error {
+func (m mockSnapshotKV) Save(ctx context.Context, key string, value string, ts typeutil.Timestamp) error {
 	if m.SaveFunc != nil {
-		return m.SaveFunc(key, value, ts)
+		return m.SaveFunc(ctx, key, value, ts)
 	}
 	return nil
 }
 
-func (m mockSnapshotKV) Load(key string, ts typeutil.Timestamp) (string, error) {
+func (m mockSnapshotKV) Load(ctx context.Context, key string, ts typeutil.Timestamp) (string, error) {
 	if m.LoadFunc != nil {
-		return m.LoadFunc(key, ts)
+		return m.LoadFunc(ctx, key, ts)
 	}
 	return "", nil
 }
 
-func (m mockSnapshotKV) MultiSave(kvs map[string]string, ts typeutil.Timestamp) error {
+func (m mockSnapshotKV) MultiSave(ctx context.Context, kvs map[string]string, ts typeutil.Timestamp) error {
 	if m.MultiSaveFunc != nil {
-		return m.MultiSaveFunc(kvs, ts)
+		return m.MultiSaveFunc(ctx, kvs, ts)
 	}
 	return nil
 }
 
-func (m mockSnapshotKV) LoadWithPrefix(key string, ts typeutil.Timestamp) ([]string, []string, error) {
+func (m mockSnapshotKV) LoadWithPrefix(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
 	if m.LoadWithPrefixFunc != nil {
-		return m.LoadWithPrefixFunc(key, ts)
+		return m.LoadWithPrefixFunc(ctx, key, ts)
 	}
 	return nil, nil, nil
 }
 
-func (m mockSnapshotKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string, ts typeutil.Timestamp) error {
+func (m mockSnapshotKV) MultiSaveAndRemoveWithPrefix(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error {
 	if m.MultiSaveAndRemoveWithPrefixFunc != nil {
-		return m.MultiSaveAndRemoveWithPrefixFunc(saves, removals, ts)
+		return m.MultiSaveAndRemoveWithPrefixFunc(ctx, saves, removals, ts)
 	}
 	return nil
 }
 
-func (m mockSnapshotKV) MultiSaveAndRemove(saves map[string]string, removals []string, ts typeutil.Timestamp) error {
+func (m mockSnapshotKV) MultiSaveAndRemove(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error {
 	if m.MultiSaveAndRemoveFunc != nil {
-		return m.MultiSaveAndRemoveFunc(saves, removals, ts)
+		return m.MultiSaveAndRemoveFunc(ctx, saves, removals, ts)
 	}
 	return nil
 }
diff --git a/internal/kv/mock_snapshot_kv_test.go b/internal/kv/mock_snapshot_kv_test.go
index 0b2df70f9173f..19fabf54b0540 100644
--- a/internal/kv/mock_snapshot_kv_test.go
+++ b/internal/kv/mock_snapshot_kv_test.go
@@ -1,6 +1,7 @@
 package kv
 
 import (
+	"context"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -11,15 +12,15 @@ import (
 func Test_mockSnapshotKV_Save(t *testing.T) {
 	t.Run("func not set", func(t *testing.T) {
 		snapshot := NewMockSnapshotKV()
-		err := snapshot.Save("k", "v", 0)
+		err := snapshot.Save(context.TODO(), "k", "v", 0)
 		assert.NoError(t, err)
 	})
 	t.Run("func set", func(t *testing.T) {
 		snapshot := NewMockSnapshotKV()
-		snapshot.SaveFunc = func(key string, value string, ts typeutil.Timestamp) error {
+		snapshot.SaveFunc = func(ctx context.Context, key string, value string, ts typeutil.Timestamp) error {
 			return nil
 		}
-		err := snapshot.Save("k", "v", 0)
+		err := snapshot.Save(context.TODO(), "k", "v", 0)
 		assert.NoError(t, err)
 	})
 }
@@ -27,15 +28,15 @@ func Test_mockSnapshotKV_Save(t *testing.T) {
 func Test_mockSnapshotKV_Load(t *testing.T) {
 	t.Run("func not set", func(t *testing.T) {
 		snapshot := NewMockSnapshotKV()
-		_, err := snapshot.Load("k", 0)
+		_, err := snapshot.Load(context.TODO(), "k", 0)
 		assert.NoError(t, err)
 	})
 	t.Run("func set", func(t *testing.T) {
 		snapshot := NewMockSnapshotKV()
-		snapshot.LoadFunc = func(key string, ts typeutil.Timestamp) (string, error) {
+		snapshot.LoadFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) (string, error) {
 			return "", nil
 		}
-		_, err := snapshot.Load("k", 0)
+		_, err := snapshot.Load(context.TODO(), "k", 0)
 		assert.NoError(t, err)
 	})
 }
@@ -43,15 +44,15 @@ func Test_mockSnapshotKV_Load(t *testing.T) {
 func Test_mockSnapshotKV_MultiSave(t *testing.T) {
 	t.Run("func not set", func(t *testing.T) {
 		snapshot := NewMockSnapshotKV()
-		err := snapshot.MultiSave(nil, 0)
+		err := snapshot.MultiSave(context.TODO(), nil, 0)
 		assert.NoError(t, err)
 	})
 	t.Run("func set", func(t *testing.T) {
 		snapshot := NewMockSnapshotKV()
-		snapshot.MultiSaveFunc = func(kvs map[string]string, ts typeutil.Timestamp) error {
+		snapshot.MultiSaveFunc = func(ctx context.Context, kvs map[string]string, ts typeutil.Timestamp) error {
 			return nil
 		}
-		err := snapshot.MultiSave(nil, 0)
+		err := snapshot.MultiSave(context.TODO(), nil, 0)
 		assert.NoError(t, err)
 	})
 }
@@ -59,15 +60,15 @@ func Test_mockSnapshotKV_MultiSave(t *testing.T) {
 func Test_mockSnapshotKV_LoadWithPrefix(t *testing.T) {
 	t.Run("func not set", func(t *testing.T) {
 		snapshot := NewMockSnapshotKV()
-		_, _, err := snapshot.LoadWithPrefix("prefix", 0)
+		_, _, err := snapshot.LoadWithPrefix(context.TODO(), "prefix", 0)
 		assert.NoError(t, err)
 	})
 	t.Run("func set", func(t *testing.T) {
 		snapshot := NewMockSnapshotKV()
-		snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) {
+		snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
 			return nil, nil, nil
 		}
-		_, _, err := snapshot.LoadWithPrefix("prefix", 0)
+		_, _, err := snapshot.LoadWithPrefix(context.TODO(), "prefix", 0)
 		assert.NoError(t, err)
 	})
 }
@@ -75,15 +76,15 @@ func Test_mockSnapshotKV_LoadWithPrefix(t *testing.T) {
 func Test_mockSnapshotKV_MultiSaveAndRemoveWithPrefix(t *testing.T) {
 	t.Run("func not set", func(t *testing.T) {
 		snapshot := NewMockSnapshotKV()
-		err := snapshot.MultiSaveAndRemoveWithPrefix(nil, nil, 0)
+		err := snapshot.MultiSaveAndRemoveWithPrefix(context.TODO(), nil, nil, 0)
 		assert.NoError(t, err)
 	})
 	t.Run("func set", func(t *testing.T) {
 		snapshot := NewMockSnapshotKV()
-		snapshot.MultiSaveAndRemoveWithPrefixFunc = func(saves map[string]string, removals []string, ts typeutil.Timestamp) error {
+		snapshot.MultiSaveAndRemoveWithPrefixFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error {
 			return nil
 		}
-		err := snapshot.MultiSaveAndRemoveWithPrefix(nil, nil, 0)
+		err := snapshot.MultiSaveAndRemoveWithPrefix(context.TODO(), nil, nil, 0)
 		assert.NoError(t, err)
 	})
 }
@@ -91,15 +92,15 @@ func Test_mockSnapshotKV_MultiSaveAndRemoveWithPrefix(t *testing.T) {
 func Test_mockSnapshotKV_MultiSaveAndRemove(t *testing.T) {
 	t.Run("func not set", func(t *testing.T) {
 		snapshot := NewMockSnapshotKV()
-		err := snapshot.MultiSaveAndRemove(nil, nil, 0)
+		err := snapshot.MultiSaveAndRemove(context.TODO(), nil, nil, 0)
 		assert.NoError(t, err)
 	})
 	t.Run("func set", func(t *testing.T) {
 		snapshot := NewMockSnapshotKV()
-		snapshot.MultiSaveAndRemoveWithPrefixFunc = func(saves map[string]string, removals []string, ts typeutil.Timestamp) error {
+		snapshot.MultiSaveAndRemoveWithPrefixFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error {
 			return nil
 		}
-		err := snapshot.MultiSaveAndRemove(nil, nil, 0)
+		err := snapshot.MultiSaveAndRemove(context.TODO(), nil, nil, 0)
 		assert.NoError(t, err)
 	})
 }
diff --git a/internal/kv/mocks/meta_kv.go b/internal/kv/mocks/meta_kv.go
index 85094ae478440..82813bf3cb6df 100644
--- a/internal/kv/mocks/meta_kv.go
+++ b/internal/kv/mocks/meta_kv.go
@@ -3,8 +3,11 @@
 package mocks
 
 import (
-	predicates "github.com/milvus-io/milvus/pkg/kv/predicates"
+	context "context"
+
 	mock "github.com/stretchr/testify/mock"
+
+	predicates "github.com/milvus-io/milvus/pkg/kv/predicates"
 )
 
 // MetaKv is an autogenerated mock type for the MetaKv type
@@ -52,9 +55,9 @@ func (_c *MetaKv_Close_Call) RunAndReturn(run func()) *MetaKv_Close_Call {
 	return _c
 }
 
-// CompareVersionAndSwap provides a mock function with given fields: key, version, target
-func (_m *MetaKv) CompareVersionAndSwap(key string, version int64, target string) (bool, error) {
-	ret := _m.Called(key, version, target)
+// CompareVersionAndSwap provides a mock function with given fields: ctx, key, version, target
+func (_m *MetaKv) CompareVersionAndSwap(ctx context.Context, key string, version int64, target string) (bool, error) {
+	ret := _m.Called(ctx, key, version, target)
 
 	if len(ret) == 0 {
 		panic("no return value specified for CompareVersionAndSwap")
@@ -62,17 +65,17 @@ func (_m *MetaKv) CompareVersionAndSwap(key string, version int64, target string
 
 	var r0 bool
 	var r1 error
-	if rf, ok := ret.Get(0).(func(string, int64, string) (bool, error)); ok {
-		return rf(key, version, target)
+	if rf, ok := ret.Get(0).(func(context.Context, string, int64, string) (bool, error)); ok {
+		return rf(ctx, key, version, target)
 	}
-	if rf, ok := ret.Get(0).(func(string, int64, string) bool); ok {
-		r0 = rf(key, version, target)
+	if rf, ok := ret.Get(0).(func(context.Context, string, int64, string) bool); ok {
+		r0 = rf(ctx, key, version, target)
 	} else {
 		r0 = ret.Get(0).(bool)
 	}
 
-	if rf, ok := ret.Get(1).(func(string, int64, string) error); ok {
-		r1 = rf(key, version, target)
+	if rf, ok := ret.Get(1).(func(context.Context, string, int64, string) error); ok {
+		r1 = rf(ctx, key, version, target)
 	} else {
 		r1 = ret.Error(1)
 	}
@@ -86,16 +89,17 @@ type MetaKv_CompareVersionAndSwap_Call struct {
 }
 
 // CompareVersionAndSwap is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
 //   - version int64
 //   - target string
-func (_e *MetaKv_Expecter) CompareVersionAndSwap(key interface{}, version interface{}, target interface{}) *MetaKv_CompareVersionAndSwap_Call {
-	return &MetaKv_CompareVersionAndSwap_Call{Call: _e.mock.On("CompareVersionAndSwap", key, version, target)}
+func (_e *MetaKv_Expecter) CompareVersionAndSwap(ctx interface{}, key interface{}, version interface{}, target interface{}) *MetaKv_CompareVersionAndSwap_Call {
+	return &MetaKv_CompareVersionAndSwap_Call{Call: _e.mock.On("CompareVersionAndSwap", ctx, key, version, target)}
 }
 
-func (_c *MetaKv_CompareVersionAndSwap_Call) Run(run func(key string, version int64, target string)) *MetaKv_CompareVersionAndSwap_Call {
+func (_c *MetaKv_CompareVersionAndSwap_Call) Run(run func(ctx context.Context, key string, version int64, target string)) *MetaKv_CompareVersionAndSwap_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string), args[1].(int64), args[2].(string))
+		run(args[0].(context.Context), args[1].(string), args[2].(int64), args[3].(string))
 	})
 	return _c
 }
@@ -105,7 +109,7 @@ func (_c *MetaKv_CompareVersionAndSwap_Call) Return(_a0 bool, _a1 error) *MetaKv
 	return _c
 }
 
-func (_c *MetaKv_CompareVersionAndSwap_Call) RunAndReturn(run func(string, int64, string) (bool, error)) *MetaKv_CompareVersionAndSwap_Call {
+func (_c *MetaKv_CompareVersionAndSwap_Call) RunAndReturn(run func(context.Context, string, int64, string) (bool, error)) *MetaKv_CompareVersionAndSwap_Call {
 	_c.Call.Return(run)
 	return _c
 }
@@ -156,9 +160,9 @@ func (_c *MetaKv_GetPath_Call) RunAndReturn(run func(string) string) *MetaKv_Get
 	return _c
 }
 
-// Has provides a mock function with given fields: key
-func (_m *MetaKv) Has(key string) (bool, error) {
-	ret := _m.Called(key)
+// Has provides a mock function with given fields: ctx, key
+func (_m *MetaKv) Has(ctx context.Context, key string) (bool, error) {
+	ret := _m.Called(ctx, key)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Has")
@@ -166,17 +170,17 @@ func (_m *MetaKv) Has(key string) (bool, error) {
 
 	var r0 bool
 	var r1 error
-	if rf, ok := ret.Get(0).(func(string) (bool, error)); ok {
-		return rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) (bool, error)); ok {
+		return rf(ctx, key)
 	}
-	if rf, ok := ret.Get(0).(func(string) bool); ok {
-		r0 = rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) bool); ok {
+		r0 = rf(ctx, key)
 	} else {
 		r0 = ret.Get(0).(bool)
 	}
 
-	if rf, ok := ret.Get(1).(func(string) error); ok {
-		r1 = rf(key)
+	if rf, ok := ret.Get(1).(func(context.Context, string) error); ok {
+		r1 = rf(ctx, key)
 	} else {
 		r1 = ret.Error(1)
 	}
@@ -190,14 +194,15 @@ type MetaKv_Has_Call struct {
 }
 
 // Has is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
-func (_e *MetaKv_Expecter) Has(key interface{}) *MetaKv_Has_Call {
-	return &MetaKv_Has_Call{Call: _e.mock.On("Has", key)}
+func (_e *MetaKv_Expecter) Has(ctx interface{}, key interface{}) *MetaKv_Has_Call {
+	return &MetaKv_Has_Call{Call: _e.mock.On("Has", ctx, key)}
 }
 
-func (_c *MetaKv_Has_Call) Run(run func(key string)) *MetaKv_Has_Call {
+func (_c *MetaKv_Has_Call) Run(run func(ctx context.Context, key string)) *MetaKv_Has_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string))
+		run(args[0].(context.Context), args[1].(string))
 	})
 	return _c
 }
@@ -207,14 +212,14 @@ func (_c *MetaKv_Has_Call) Return(_a0 bool, _a1 error) *MetaKv_Has_Call {
 	return _c
 }
 
-func (_c *MetaKv_Has_Call) RunAndReturn(run func(string) (bool, error)) *MetaKv_Has_Call {
+func (_c *MetaKv_Has_Call) RunAndReturn(run func(context.Context, string) (bool, error)) *MetaKv_Has_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// HasPrefix provides a mock function with given fields: prefix
-func (_m *MetaKv) HasPrefix(prefix string) (bool, error) {
-	ret := _m.Called(prefix)
+// HasPrefix provides a mock function with given fields: ctx, prefix
+func (_m *MetaKv) HasPrefix(ctx context.Context, prefix string) (bool, error) {
+	ret := _m.Called(ctx, prefix)
 
 	if len(ret) == 0 {
 		panic("no return value specified for HasPrefix")
@@ -222,17 +227,17 @@ func (_m *MetaKv) HasPrefix(prefix string) (bool, error) {
 
 	var r0 bool
 	var r1 error
-	if rf, ok := ret.Get(0).(func(string) (bool, error)); ok {
-		return rf(prefix)
+	if rf, ok := ret.Get(0).(func(context.Context, string) (bool, error)); ok {
+		return rf(ctx, prefix)
 	}
-	if rf, ok := ret.Get(0).(func(string) bool); ok {
-		r0 = rf(prefix)
+	if rf, ok := ret.Get(0).(func(context.Context, string) bool); ok {
+		r0 = rf(ctx, prefix)
 	} else {
 		r0 = ret.Get(0).(bool)
 	}
 
-	if rf, ok := ret.Get(1).(func(string) error); ok {
-		r1 = rf(prefix)
+	if rf, ok := ret.Get(1).(func(context.Context, string) error); ok {
+		r1 = rf(ctx, prefix)
 	} else {
 		r1 = ret.Error(1)
 	}
@@ -246,14 +251,15 @@ type MetaKv_HasPrefix_Call struct {
 }
 
 // HasPrefix is a helper method to define mock.On call
+//   - ctx context.Context
 //   - prefix string
-func (_e *MetaKv_Expecter) HasPrefix(prefix interface{}) *MetaKv_HasPrefix_Call {
-	return &MetaKv_HasPrefix_Call{Call: _e.mock.On("HasPrefix", prefix)}
+func (_e *MetaKv_Expecter) HasPrefix(ctx interface{}, prefix interface{}) *MetaKv_HasPrefix_Call {
+	return &MetaKv_HasPrefix_Call{Call: _e.mock.On("HasPrefix", ctx, prefix)}
 }
 
-func (_c *MetaKv_HasPrefix_Call) Run(run func(prefix string)) *MetaKv_HasPrefix_Call {
+func (_c *MetaKv_HasPrefix_Call) Run(run func(ctx context.Context, prefix string)) *MetaKv_HasPrefix_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string))
+		run(args[0].(context.Context), args[1].(string))
 	})
 	return _c
 }
@@ -263,14 +269,14 @@ func (_c *MetaKv_HasPrefix_Call) Return(_a0 bool, _a1 error) *MetaKv_HasPrefix_C
 	return _c
 }
 
-func (_c *MetaKv_HasPrefix_Call) RunAndReturn(run func(string) (bool, error)) *MetaKv_HasPrefix_Call {
+func (_c *MetaKv_HasPrefix_Call) RunAndReturn(run func(context.Context, string) (bool, error)) *MetaKv_HasPrefix_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// Load provides a mock function with given fields: key
-func (_m *MetaKv) Load(key string) (string, error) {
-	ret := _m.Called(key)
+// Load provides a mock function with given fields: ctx, key
+func (_m *MetaKv) Load(ctx context.Context, key string) (string, error) {
+	ret := _m.Called(ctx, key)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Load")
@@ -278,17 +284,17 @@ func (_m *MetaKv) Load(key string) (string, error) {
 
 	var r0 string
 	var r1 error
-	if rf, ok := ret.Get(0).(func(string) (string, error)); ok {
-		return rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) (string, error)); ok {
+		return rf(ctx, key)
 	}
-	if rf, ok := ret.Get(0).(func(string) string); ok {
-		r0 = rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) string); ok {
+		r0 = rf(ctx, key)
 	} else {
 		r0 = ret.Get(0).(string)
 	}
 
-	if rf, ok := ret.Get(1).(func(string) error); ok {
-		r1 = rf(key)
+	if rf, ok := ret.Get(1).(func(context.Context, string) error); ok {
+		r1 = rf(ctx, key)
 	} else {
 		r1 = ret.Error(1)
 	}
@@ -302,14 +308,15 @@ type MetaKv_Load_Call struct {
 }
 
 // Load is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
-func (_e *MetaKv_Expecter) Load(key interface{}) *MetaKv_Load_Call {
-	return &MetaKv_Load_Call{Call: _e.mock.On("Load", key)}
+func (_e *MetaKv_Expecter) Load(ctx interface{}, key interface{}) *MetaKv_Load_Call {
+	return &MetaKv_Load_Call{Call: _e.mock.On("Load", ctx, key)}
 }
 
-func (_c *MetaKv_Load_Call) Run(run func(key string)) *MetaKv_Load_Call {
+func (_c *MetaKv_Load_Call) Run(run func(ctx context.Context, key string)) *MetaKv_Load_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string))
+		run(args[0].(context.Context), args[1].(string))
 	})
 	return _c
 }
@@ -319,14 +326,14 @@ func (_c *MetaKv_Load_Call) Return(_a0 string, _a1 error) *MetaKv_Load_Call {
 	return _c
 }
 
-func (_c *MetaKv_Load_Call) RunAndReturn(run func(string) (string, error)) *MetaKv_Load_Call {
+func (_c *MetaKv_Load_Call) RunAndReturn(run func(context.Context, string) (string, error)) *MetaKv_Load_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// LoadWithPrefix provides a mock function with given fields: key
-func (_m *MetaKv) LoadWithPrefix(key string) ([]string, []string, error) {
-	ret := _m.Called(key)
+// LoadWithPrefix provides a mock function with given fields: ctx, key
+func (_m *MetaKv) LoadWithPrefix(ctx context.Context, key string) ([]string, []string, error) {
+	ret := _m.Called(ctx, key)
 
 	if len(ret) == 0 {
 		panic("no return value specified for LoadWithPrefix")
@@ -335,27 +342,27 @@ func (_m *MetaKv) LoadWithPrefix(key string) ([]string, []string, error) {
 	var r0 []string
 	var r1 []string
 	var r2 error
-	if rf, ok := ret.Get(0).(func(string) ([]string, []string, error)); ok {
-		return rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) ([]string, []string, error)); ok {
+		return rf(ctx, key)
 	}
-	if rf, ok := ret.Get(0).(func(string) []string); ok {
-		r0 = rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) []string); ok {
+		r0 = rf(ctx, key)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).([]string)
 		}
 	}
 
-	if rf, ok := ret.Get(1).(func(string) []string); ok {
-		r1 = rf(key)
+	if rf, ok := ret.Get(1).(func(context.Context, string) []string); ok {
+		r1 = rf(ctx, key)
 	} else {
 		if ret.Get(1) != nil {
 			r1 = ret.Get(1).([]string)
 		}
 	}
 
-	if rf, ok := ret.Get(2).(func(string) error); ok {
-		r2 = rf(key)
+	if rf, ok := ret.Get(2).(func(context.Context, string) error); ok {
+		r2 = rf(ctx, key)
 	} else {
 		r2 = ret.Error(2)
 	}
@@ -369,14 +376,15 @@ type MetaKv_LoadWithPrefix_Call struct {
 }
 
 // LoadWithPrefix is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
-func (_e *MetaKv_Expecter) LoadWithPrefix(key interface{}) *MetaKv_LoadWithPrefix_Call {
-	return &MetaKv_LoadWithPrefix_Call{Call: _e.mock.On("LoadWithPrefix", key)}
+func (_e *MetaKv_Expecter) LoadWithPrefix(ctx interface{}, key interface{}) *MetaKv_LoadWithPrefix_Call {
+	return &MetaKv_LoadWithPrefix_Call{Call: _e.mock.On("LoadWithPrefix", ctx, key)}
 }
 
-func (_c *MetaKv_LoadWithPrefix_Call) Run(run func(key string)) *MetaKv_LoadWithPrefix_Call {
+func (_c *MetaKv_LoadWithPrefix_Call) Run(run func(ctx context.Context, key string)) *MetaKv_LoadWithPrefix_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string))
+		run(args[0].(context.Context), args[1].(string))
 	})
 	return _c
 }
@@ -386,14 +394,14 @@ func (_c *MetaKv_LoadWithPrefix_Call) Return(_a0 []string, _a1 []string, _a2 err
 	return _c
 }
 
-func (_c *MetaKv_LoadWithPrefix_Call) RunAndReturn(run func(string) ([]string, []string, error)) *MetaKv_LoadWithPrefix_Call {
+func (_c *MetaKv_LoadWithPrefix_Call) RunAndReturn(run func(context.Context, string) ([]string, []string, error)) *MetaKv_LoadWithPrefix_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// MultiLoad provides a mock function with given fields: keys
-func (_m *MetaKv) MultiLoad(keys []string) ([]string, error) {
-	ret := _m.Called(keys)
+// MultiLoad provides a mock function with given fields: ctx, keys
+func (_m *MetaKv) MultiLoad(ctx context.Context, keys []string) ([]string, error) {
+	ret := _m.Called(ctx, keys)
 
 	if len(ret) == 0 {
 		panic("no return value specified for MultiLoad")
@@ -401,19 +409,19 @@ func (_m *MetaKv) MultiLoad(keys []string) ([]string, error) {
 
 	var r0 []string
 	var r1 error
-	if rf, ok := ret.Get(0).(func([]string) ([]string, error)); ok {
-		return rf(keys)
+	if rf, ok := ret.Get(0).(func(context.Context, []string) ([]string, error)); ok {
+		return rf(ctx, keys)
 	}
-	if rf, ok := ret.Get(0).(func([]string) []string); ok {
-		r0 = rf(keys)
+	if rf, ok := ret.Get(0).(func(context.Context, []string) []string); ok {
+		r0 = rf(ctx, keys)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).([]string)
 		}
 	}
 
-	if rf, ok := ret.Get(1).(func([]string) error); ok {
-		r1 = rf(keys)
+	if rf, ok := ret.Get(1).(func(context.Context, []string) error); ok {
+		r1 = rf(ctx, keys)
 	} else {
 		r1 = ret.Error(1)
 	}
@@ -427,14 +435,15 @@ type MetaKv_MultiLoad_Call struct {
 }
 
 // MultiLoad is a helper method to define mock.On call
+//   - ctx context.Context
 //   - keys []string
-func (_e *MetaKv_Expecter) MultiLoad(keys interface{}) *MetaKv_MultiLoad_Call {
-	return &MetaKv_MultiLoad_Call{Call: _e.mock.On("MultiLoad", keys)}
+func (_e *MetaKv_Expecter) MultiLoad(ctx interface{}, keys interface{}) *MetaKv_MultiLoad_Call {
+	return &MetaKv_MultiLoad_Call{Call: _e.mock.On("MultiLoad", ctx, keys)}
 }
 
-func (_c *MetaKv_MultiLoad_Call) Run(run func(keys []string)) *MetaKv_MultiLoad_Call {
+func (_c *MetaKv_MultiLoad_Call) Run(run func(ctx context.Context, keys []string)) *MetaKv_MultiLoad_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].([]string))
+		run(args[0].(context.Context), args[1].([]string))
 	})
 	return _c
 }
@@ -444,22 +453,22 @@ func (_c *MetaKv_MultiLoad_Call) Return(_a0 []string, _a1 error) *MetaKv_MultiLo
 	return _c
 }
 
-func (_c *MetaKv_MultiLoad_Call) RunAndReturn(run func([]string) ([]string, error)) *MetaKv_MultiLoad_Call {
+func (_c *MetaKv_MultiLoad_Call) RunAndReturn(run func(context.Context, []string) ([]string, error)) *MetaKv_MultiLoad_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// MultiRemove provides a mock function with given fields: keys
-func (_m *MetaKv) MultiRemove(keys []string) error {
-	ret := _m.Called(keys)
+// MultiRemove provides a mock function with given fields: ctx, keys
+func (_m *MetaKv) MultiRemove(ctx context.Context, keys []string) error {
+	ret := _m.Called(ctx, keys)
 
 	if len(ret) == 0 {
 		panic("no return value specified for MultiRemove")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func([]string) error); ok {
-		r0 = rf(keys)
+	if rf, ok := ret.Get(0).(func(context.Context, []string) error); ok {
+		r0 = rf(ctx, keys)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -473,14 +482,15 @@ type MetaKv_MultiRemove_Call struct {
 }
 
 // MultiRemove is a helper method to define mock.On call
+//   - ctx context.Context
 //   - keys []string
-func (_e *MetaKv_Expecter) MultiRemove(keys interface{}) *MetaKv_MultiRemove_Call {
-	return &MetaKv_MultiRemove_Call{Call: _e.mock.On("MultiRemove", keys)}
+func (_e *MetaKv_Expecter) MultiRemove(ctx interface{}, keys interface{}) *MetaKv_MultiRemove_Call {
+	return &MetaKv_MultiRemove_Call{Call: _e.mock.On("MultiRemove", ctx, keys)}
 }
 
-func (_c *MetaKv_MultiRemove_Call) Run(run func(keys []string)) *MetaKv_MultiRemove_Call {
+func (_c *MetaKv_MultiRemove_Call) Run(run func(ctx context.Context, keys []string)) *MetaKv_MultiRemove_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].([]string))
+		run(args[0].(context.Context), args[1].([]string))
 	})
 	return _c
 }
@@ -490,22 +500,22 @@ func (_c *MetaKv_MultiRemove_Call) Return(_a0 error) *MetaKv_MultiRemove_Call {
 	return _c
 }
 
-func (_c *MetaKv_MultiRemove_Call) RunAndReturn(run func([]string) error) *MetaKv_MultiRemove_Call {
+func (_c *MetaKv_MultiRemove_Call) RunAndReturn(run func(context.Context, []string) error) *MetaKv_MultiRemove_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// MultiSave provides a mock function with given fields: kvs
-func (_m *MetaKv) MultiSave(kvs map[string]string) error {
-	ret := _m.Called(kvs)
+// MultiSave provides a mock function with given fields: ctx, kvs
+func (_m *MetaKv) MultiSave(ctx context.Context, kvs map[string]string) error {
+	ret := _m.Called(ctx, kvs)
 
 	if len(ret) == 0 {
 		panic("no return value specified for MultiSave")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(map[string]string) error); ok {
-		r0 = rf(kvs)
+	if rf, ok := ret.Get(0).(func(context.Context, map[string]string) error); ok {
+		r0 = rf(ctx, kvs)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -519,14 +529,15 @@ type MetaKv_MultiSave_Call struct {
 }
 
 // MultiSave is a helper method to define mock.On call
+//   - ctx context.Context
 //   - kvs map[string]string
-func (_e *MetaKv_Expecter) MultiSave(kvs interface{}) *MetaKv_MultiSave_Call {
-	return &MetaKv_MultiSave_Call{Call: _e.mock.On("MultiSave", kvs)}
+func (_e *MetaKv_Expecter) MultiSave(ctx interface{}, kvs interface{}) *MetaKv_MultiSave_Call {
+	return &MetaKv_MultiSave_Call{Call: _e.mock.On("MultiSave", ctx, kvs)}
 }
 
-func (_c *MetaKv_MultiSave_Call) Run(run func(kvs map[string]string)) *MetaKv_MultiSave_Call {
+func (_c *MetaKv_MultiSave_Call) Run(run func(ctx context.Context, kvs map[string]string)) *MetaKv_MultiSave_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(map[string]string))
+		run(args[0].(context.Context), args[1].(map[string]string))
 	})
 	return _c
 }
@@ -536,19 +547,19 @@ func (_c *MetaKv_MultiSave_Call) Return(_a0 error) *MetaKv_MultiSave_Call {
 	return _c
 }
 
-func (_c *MetaKv_MultiSave_Call) RunAndReturn(run func(map[string]string) error) *MetaKv_MultiSave_Call {
+func (_c *MetaKv_MultiSave_Call) RunAndReturn(run func(context.Context, map[string]string) error) *MetaKv_MultiSave_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// MultiSaveAndRemove provides a mock function with given fields: saves, removals, preds
-func (_m *MetaKv) MultiSaveAndRemove(saves map[string]string, removals []string, preds ...predicates.Predicate) error {
+// MultiSaveAndRemove provides a mock function with given fields: ctx, saves, removals, preds
+func (_m *MetaKv) MultiSaveAndRemove(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error {
 	_va := make([]interface{}, len(preds))
 	for _i := range preds {
 		_va[_i] = preds[_i]
 	}
 	var _ca []interface{}
-	_ca = append(_ca, saves, removals)
+	_ca = append(_ca, ctx, saves, removals)
 	_ca = append(_ca, _va...)
 	ret := _m.Called(_ca...)
 
@@ -557,8 +568,8 @@ func (_m *MetaKv) MultiSaveAndRemove(saves map[string]string, removals []string,
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(map[string]string, []string, ...predicates.Predicate) error); ok {
-		r0 = rf(saves, removals, preds...)
+	if rf, ok := ret.Get(0).(func(context.Context, map[string]string, []string, ...predicates.Predicate) error); ok {
+		r0 = rf(ctx, saves, removals, preds...)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -572,23 +583,24 @@ type MetaKv_MultiSaveAndRemove_Call struct {
 }
 
 // MultiSaveAndRemove is a helper method to define mock.On call
+//   - ctx context.Context
 //   - saves map[string]string
 //   - removals []string
 //   - preds ...predicates.Predicate
-func (_e *MetaKv_Expecter) MultiSaveAndRemove(saves interface{}, removals interface{}, preds ...interface{}) *MetaKv_MultiSaveAndRemove_Call {
+func (_e *MetaKv_Expecter) MultiSaveAndRemove(ctx interface{}, saves interface{}, removals interface{}, preds ...interface{}) *MetaKv_MultiSaveAndRemove_Call {
 	return &MetaKv_MultiSaveAndRemove_Call{Call: _e.mock.On("MultiSaveAndRemove",
-		append([]interface{}{saves, removals}, preds...)...)}
+		append([]interface{}{ctx, saves, removals}, preds...)...)}
 }
 
-func (_c *MetaKv_MultiSaveAndRemove_Call) Run(run func(saves map[string]string, removals []string, preds ...predicates.Predicate)) *MetaKv_MultiSaveAndRemove_Call {
+func (_c *MetaKv_MultiSaveAndRemove_Call) Run(run func(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate)) *MetaKv_MultiSaveAndRemove_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		variadicArgs := make([]predicates.Predicate, len(args)-2)
-		for i, a := range args[2:] {
+		variadicArgs := make([]predicates.Predicate, len(args)-3)
+		for i, a := range args[3:] {
 			if a != nil {
 				variadicArgs[i] = a.(predicates.Predicate)
 			}
 		}
-		run(args[0].(map[string]string), args[1].([]string), variadicArgs...)
+		run(args[0].(context.Context), args[1].(map[string]string), args[2].([]string), variadicArgs...)
 	})
 	return _c
 }
@@ -598,19 +610,19 @@ func (_c *MetaKv_MultiSaveAndRemove_Call) Return(_a0 error) *MetaKv_MultiSaveAnd
 	return _c
 }
 
-func (_c *MetaKv_MultiSaveAndRemove_Call) RunAndReturn(run func(map[string]string, []string, ...predicates.Predicate) error) *MetaKv_MultiSaveAndRemove_Call {
+func (_c *MetaKv_MultiSaveAndRemove_Call) RunAndReturn(run func(context.Context, map[string]string, []string, ...predicates.Predicate) error) *MetaKv_MultiSaveAndRemove_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// MultiSaveAndRemoveWithPrefix provides a mock function with given fields: saves, removals, preds
-func (_m *MetaKv) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string, preds ...predicates.Predicate) error {
+// MultiSaveAndRemoveWithPrefix provides a mock function with given fields: ctx, saves, removals, preds
+func (_m *MetaKv) MultiSaveAndRemoveWithPrefix(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error {
 	_va := make([]interface{}, len(preds))
 	for _i := range preds {
 		_va[_i] = preds[_i]
 	}
 	var _ca []interface{}
-	_ca = append(_ca, saves, removals)
+	_ca = append(_ca, ctx, saves, removals)
 	_ca = append(_ca, _va...)
 	ret := _m.Called(_ca...)
 
@@ -619,8 +631,8 @@ func (_m *MetaKv) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(map[string]string, []string, ...predicates.Predicate) error); ok {
-		r0 = rf(saves, removals, preds...)
+	if rf, ok := ret.Get(0).(func(context.Context, map[string]string, []string, ...predicates.Predicate) error); ok {
+		r0 = rf(ctx, saves, removals, preds...)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -634,23 +646,24 @@ type MetaKv_MultiSaveAndRemoveWithPrefix_Call struct {
 }
 
 // MultiSaveAndRemoveWithPrefix is a helper method to define mock.On call
+//   - ctx context.Context
 //   - saves map[string]string
 //   - removals []string
 //   - preds ...predicates.Predicate
-func (_e *MetaKv_Expecter) MultiSaveAndRemoveWithPrefix(saves interface{}, removals interface{}, preds ...interface{}) *MetaKv_MultiSaveAndRemoveWithPrefix_Call {
+func (_e *MetaKv_Expecter) MultiSaveAndRemoveWithPrefix(ctx interface{}, saves interface{}, removals interface{}, preds ...interface{}) *MetaKv_MultiSaveAndRemoveWithPrefix_Call {
 	return &MetaKv_MultiSaveAndRemoveWithPrefix_Call{Call: _e.mock.On("MultiSaveAndRemoveWithPrefix",
-		append([]interface{}{saves, removals}, preds...)...)}
+		append([]interface{}{ctx, saves, removals}, preds...)...)}
 }
 
-func (_c *MetaKv_MultiSaveAndRemoveWithPrefix_Call) Run(run func(saves map[string]string, removals []string, preds ...predicates.Predicate)) *MetaKv_MultiSaveAndRemoveWithPrefix_Call {
+func (_c *MetaKv_MultiSaveAndRemoveWithPrefix_Call) Run(run func(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate)) *MetaKv_MultiSaveAndRemoveWithPrefix_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		variadicArgs := make([]predicates.Predicate, len(args)-2)
-		for i, a := range args[2:] {
+		variadicArgs := make([]predicates.Predicate, len(args)-3)
+		for i, a := range args[3:] {
 			if a != nil {
 				variadicArgs[i] = a.(predicates.Predicate)
 			}
 		}
-		run(args[0].(map[string]string), args[1].([]string), variadicArgs...)
+		run(args[0].(context.Context), args[1].(map[string]string), args[2].([]string), variadicArgs...)
 	})
 	return _c
 }
@@ -660,22 +673,22 @@ func (_c *MetaKv_MultiSaveAndRemoveWithPrefix_Call) Return(_a0 error) *MetaKv_Mu
 	return _c
 }
 
-func (_c *MetaKv_MultiSaveAndRemoveWithPrefix_Call) RunAndReturn(run func(map[string]string, []string, ...predicates.Predicate) error) *MetaKv_MultiSaveAndRemoveWithPrefix_Call {
+func (_c *MetaKv_MultiSaveAndRemoveWithPrefix_Call) RunAndReturn(run func(context.Context, map[string]string, []string, ...predicates.Predicate) error) *MetaKv_MultiSaveAndRemoveWithPrefix_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// Remove provides a mock function with given fields: key
-func (_m *MetaKv) Remove(key string) error {
-	ret := _m.Called(key)
+// Remove provides a mock function with given fields: ctx, key
+func (_m *MetaKv) Remove(ctx context.Context, key string) error {
+	ret := _m.Called(ctx, key)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Remove")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(string) error); ok {
-		r0 = rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) error); ok {
+		r0 = rf(ctx, key)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -689,14 +702,15 @@ type MetaKv_Remove_Call struct {
 }
 
 // Remove is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
-func (_e *MetaKv_Expecter) Remove(key interface{}) *MetaKv_Remove_Call {
-	return &MetaKv_Remove_Call{Call: _e.mock.On("Remove", key)}
+func (_e *MetaKv_Expecter) Remove(ctx interface{}, key interface{}) *MetaKv_Remove_Call {
+	return &MetaKv_Remove_Call{Call: _e.mock.On("Remove", ctx, key)}
 }
 
-func (_c *MetaKv_Remove_Call) Run(run func(key string)) *MetaKv_Remove_Call {
+func (_c *MetaKv_Remove_Call) Run(run func(ctx context.Context, key string)) *MetaKv_Remove_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string))
+		run(args[0].(context.Context), args[1].(string))
 	})
 	return _c
 }
@@ -706,22 +720,22 @@ func (_c *MetaKv_Remove_Call) Return(_a0 error) *MetaKv_Remove_Call {
 	return _c
 }
 
-func (_c *MetaKv_Remove_Call) RunAndReturn(run func(string) error) *MetaKv_Remove_Call {
+func (_c *MetaKv_Remove_Call) RunAndReturn(run func(context.Context, string) error) *MetaKv_Remove_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// RemoveWithPrefix provides a mock function with given fields: key
-func (_m *MetaKv) RemoveWithPrefix(key string) error {
-	ret := _m.Called(key)
+// RemoveWithPrefix provides a mock function with given fields: ctx, key
+func (_m *MetaKv) RemoveWithPrefix(ctx context.Context, key string) error {
+	ret := _m.Called(ctx, key)
 
 	if len(ret) == 0 {
 		panic("no return value specified for RemoveWithPrefix")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(string) error); ok {
-		r0 = rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) error); ok {
+		r0 = rf(ctx, key)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -735,14 +749,15 @@ type MetaKv_RemoveWithPrefix_Call struct {
 }
 
 // RemoveWithPrefix is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
-func (_e *MetaKv_Expecter) RemoveWithPrefix(key interface{}) *MetaKv_RemoveWithPrefix_Call {
-	return &MetaKv_RemoveWithPrefix_Call{Call: _e.mock.On("RemoveWithPrefix", key)}
+func (_e *MetaKv_Expecter) RemoveWithPrefix(ctx interface{}, key interface{}) *MetaKv_RemoveWithPrefix_Call {
+	return &MetaKv_RemoveWithPrefix_Call{Call: _e.mock.On("RemoveWithPrefix", ctx, key)}
 }
 
-func (_c *MetaKv_RemoveWithPrefix_Call) Run(run func(key string)) *MetaKv_RemoveWithPrefix_Call {
+func (_c *MetaKv_RemoveWithPrefix_Call) Run(run func(ctx context.Context, key string)) *MetaKv_RemoveWithPrefix_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string))
+		run(args[0].(context.Context), args[1].(string))
 	})
 	return _c
 }
@@ -752,22 +767,22 @@ func (_c *MetaKv_RemoveWithPrefix_Call) Return(_a0 error) *MetaKv_RemoveWithPref
 	return _c
 }
 
-func (_c *MetaKv_RemoveWithPrefix_Call) RunAndReturn(run func(string) error) *MetaKv_RemoveWithPrefix_Call {
+func (_c *MetaKv_RemoveWithPrefix_Call) RunAndReturn(run func(context.Context, string) error) *MetaKv_RemoveWithPrefix_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// Save provides a mock function with given fields: key, value
-func (_m *MetaKv) Save(key string, value string) error {
-	ret := _m.Called(key, value)
+// Save provides a mock function with given fields: ctx, key, value
+func (_m *MetaKv) Save(ctx context.Context, key string, value string) error {
+	ret := _m.Called(ctx, key, value)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Save")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(string, string) error); ok {
-		r0 = rf(key, value)
+	if rf, ok := ret.Get(0).(func(context.Context, string, string) error); ok {
+		r0 = rf(ctx, key, value)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -781,15 +796,16 @@ type MetaKv_Save_Call struct {
 }
 
 // Save is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
 //   - value string
-func (_e *MetaKv_Expecter) Save(key interface{}, value interface{}) *MetaKv_Save_Call {
-	return &MetaKv_Save_Call{Call: _e.mock.On("Save", key, value)}
+func (_e *MetaKv_Expecter) Save(ctx interface{}, key interface{}, value interface{}) *MetaKv_Save_Call {
+	return &MetaKv_Save_Call{Call: _e.mock.On("Save", ctx, key, value)}
 }
 
-func (_c *MetaKv_Save_Call) Run(run func(key string, value string)) *MetaKv_Save_Call {
+func (_c *MetaKv_Save_Call) Run(run func(ctx context.Context, key string, value string)) *MetaKv_Save_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string), args[1].(string))
+		run(args[0].(context.Context), args[1].(string), args[2].(string))
 	})
 	return _c
 }
@@ -799,22 +815,22 @@ func (_c *MetaKv_Save_Call) Return(_a0 error) *MetaKv_Save_Call {
 	return _c
 }
 
-func (_c *MetaKv_Save_Call) RunAndReturn(run func(string, string) error) *MetaKv_Save_Call {
+func (_c *MetaKv_Save_Call) RunAndReturn(run func(context.Context, string, string) error) *MetaKv_Save_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// WalkWithPrefix provides a mock function with given fields: prefix, paginationSize, fn
-func (_m *MetaKv) WalkWithPrefix(prefix string, paginationSize int, fn func([]byte, []byte) error) error {
-	ret := _m.Called(prefix, paginationSize, fn)
+// WalkWithPrefix provides a mock function with given fields: ctx, prefix, paginationSize, fn
+func (_m *MetaKv) WalkWithPrefix(ctx context.Context, prefix string, paginationSize int, fn func([]byte, []byte) error) error {
+	ret := _m.Called(ctx, prefix, paginationSize, fn)
 
 	if len(ret) == 0 {
 		panic("no return value specified for WalkWithPrefix")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(string, int, func([]byte, []byte) error) error); ok {
-		r0 = rf(prefix, paginationSize, fn)
+	if rf, ok := ret.Get(0).(func(context.Context, string, int, func([]byte, []byte) error) error); ok {
+		r0 = rf(ctx, prefix, paginationSize, fn)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -828,16 +844,17 @@ type MetaKv_WalkWithPrefix_Call struct {
 }
 
 // WalkWithPrefix is a helper method to define mock.On call
+//   - ctx context.Context
 //   - prefix string
 //   - paginationSize int
 //   - fn func([]byte , []byte) error
-func (_e *MetaKv_Expecter) WalkWithPrefix(prefix interface{}, paginationSize interface{}, fn interface{}) *MetaKv_WalkWithPrefix_Call {
-	return &MetaKv_WalkWithPrefix_Call{Call: _e.mock.On("WalkWithPrefix", prefix, paginationSize, fn)}
+func (_e *MetaKv_Expecter) WalkWithPrefix(ctx interface{}, prefix interface{}, paginationSize interface{}, fn interface{}) *MetaKv_WalkWithPrefix_Call {
+	return &MetaKv_WalkWithPrefix_Call{Call: _e.mock.On("WalkWithPrefix", ctx, prefix, paginationSize, fn)}
 }
 
-func (_c *MetaKv_WalkWithPrefix_Call) Run(run func(prefix string, paginationSize int, fn func([]byte, []byte) error)) *MetaKv_WalkWithPrefix_Call {
+func (_c *MetaKv_WalkWithPrefix_Call) Run(run func(ctx context.Context, prefix string, paginationSize int, fn func([]byte, []byte) error)) *MetaKv_WalkWithPrefix_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string), args[1].(int), args[2].(func([]byte, []byte) error))
+		run(args[0].(context.Context), args[1].(string), args[2].(int), args[3].(func([]byte, []byte) error))
 	})
 	return _c
 }
@@ -847,7 +864,7 @@ func (_c *MetaKv_WalkWithPrefix_Call) Return(_a0 error) *MetaKv_WalkWithPrefix_C
 	return _c
 }
 
-func (_c *MetaKv_WalkWithPrefix_Call) RunAndReturn(run func(string, int, func([]byte, []byte) error) error) *MetaKv_WalkWithPrefix_Call {
+func (_c *MetaKv_WalkWithPrefix_Call) RunAndReturn(run func(context.Context, string, int, func([]byte, []byte) error) error) *MetaKv_WalkWithPrefix_Call {
 	_c.Call.Return(run)
 	return _c
 }
diff --git a/internal/kv/mocks/snapshot_kv.go b/internal/kv/mocks/snapshot_kv.go
index 42bc5d7ad9ed9..220385636b8d7 100644
--- a/internal/kv/mocks/snapshot_kv.go
+++ b/internal/kv/mocks/snapshot_kv.go
@@ -2,7 +2,11 @@
 
 package mocks
 
-import mock "github.com/stretchr/testify/mock"
+import (
+	context "context"
+
+	mock "github.com/stretchr/testify/mock"
+)
 
 // SnapShotKV is an autogenerated mock type for the SnapShotKV type
 type SnapShotKV struct {
@@ -17,9 +21,9 @@ func (_m *SnapShotKV) EXPECT() *SnapShotKV_Expecter {
 	return &SnapShotKV_Expecter{mock: &_m.Mock}
 }
 
-// Load provides a mock function with given fields: key, ts
-func (_m *SnapShotKV) Load(key string, ts uint64) (string, error) {
-	ret := _m.Called(key, ts)
+// Load provides a mock function with given fields: ctx, key, ts
+func (_m *SnapShotKV) Load(ctx context.Context, key string, ts uint64) (string, error) {
+	ret := _m.Called(ctx, key, ts)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Load")
@@ -27,17 +31,17 @@ func (_m *SnapShotKV) Load(key string, ts uint64) (string, error) {
 
 	var r0 string
 	var r1 error
-	if rf, ok := ret.Get(0).(func(string, uint64) (string, error)); ok {
-		return rf(key, ts)
+	if rf, ok := ret.Get(0).(func(context.Context, string, uint64) (string, error)); ok {
+		return rf(ctx, key, ts)
 	}
-	if rf, ok := ret.Get(0).(func(string, uint64) string); ok {
-		r0 = rf(key, ts)
+	if rf, ok := ret.Get(0).(func(context.Context, string, uint64) string); ok {
+		r0 = rf(ctx, key, ts)
 	} else {
 		r0 = ret.Get(0).(string)
 	}
 
-	if rf, ok := ret.Get(1).(func(string, uint64) error); ok {
-		r1 = rf(key, ts)
+	if rf, ok := ret.Get(1).(func(context.Context, string, uint64) error); ok {
+		r1 = rf(ctx, key, ts)
 	} else {
 		r1 = ret.Error(1)
 	}
@@ -51,15 +55,16 @@ type SnapShotKV_Load_Call struct {
 }
 
 // Load is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
 //   - ts uint64
-func (_e *SnapShotKV_Expecter) Load(key interface{}, ts interface{}) *SnapShotKV_Load_Call {
-	return &SnapShotKV_Load_Call{Call: _e.mock.On("Load", key, ts)}
+func (_e *SnapShotKV_Expecter) Load(ctx interface{}, key interface{}, ts interface{}) *SnapShotKV_Load_Call {
+	return &SnapShotKV_Load_Call{Call: _e.mock.On("Load", ctx, key, ts)}
 }
 
-func (_c *SnapShotKV_Load_Call) Run(run func(key string, ts uint64)) *SnapShotKV_Load_Call {
+func (_c *SnapShotKV_Load_Call) Run(run func(ctx context.Context, key string, ts uint64)) *SnapShotKV_Load_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string), args[1].(uint64))
+		run(args[0].(context.Context), args[1].(string), args[2].(uint64))
 	})
 	return _c
 }
@@ -69,14 +74,14 @@ func (_c *SnapShotKV_Load_Call) Return(_a0 string, _a1 error) *SnapShotKV_Load_C
 	return _c
 }
 
-func (_c *SnapShotKV_Load_Call) RunAndReturn(run func(string, uint64) (string, error)) *SnapShotKV_Load_Call {
+func (_c *SnapShotKV_Load_Call) RunAndReturn(run func(context.Context, string, uint64) (string, error)) *SnapShotKV_Load_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// LoadWithPrefix provides a mock function with given fields: key, ts
-func (_m *SnapShotKV) LoadWithPrefix(key string, ts uint64) ([]string, []string, error) {
-	ret := _m.Called(key, ts)
+// LoadWithPrefix provides a mock function with given fields: ctx, key, ts
+func (_m *SnapShotKV) LoadWithPrefix(ctx context.Context, key string, ts uint64) ([]string, []string, error) {
+	ret := _m.Called(ctx, key, ts)
 
 	if len(ret) == 0 {
 		panic("no return value specified for LoadWithPrefix")
@@ -85,27 +90,27 @@ func (_m *SnapShotKV) LoadWithPrefix(key string, ts uint64) ([]string, []string,
 	var r0 []string
 	var r1 []string
 	var r2 error
-	if rf, ok := ret.Get(0).(func(string, uint64) ([]string, []string, error)); ok {
-		return rf(key, ts)
+	if rf, ok := ret.Get(0).(func(context.Context, string, uint64) ([]string, []string, error)); ok {
+		return rf(ctx, key, ts)
 	}
-	if rf, ok := ret.Get(0).(func(string, uint64) []string); ok {
-		r0 = rf(key, ts)
+	if rf, ok := ret.Get(0).(func(context.Context, string, uint64) []string); ok {
+		r0 = rf(ctx, key, ts)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).([]string)
 		}
 	}
 
-	if rf, ok := ret.Get(1).(func(string, uint64) []string); ok {
-		r1 = rf(key, ts)
+	if rf, ok := ret.Get(1).(func(context.Context, string, uint64) []string); ok {
+		r1 = rf(ctx, key, ts)
 	} else {
 		if ret.Get(1) != nil {
 			r1 = ret.Get(1).([]string)
 		}
 	}
 
-	if rf, ok := ret.Get(2).(func(string, uint64) error); ok {
-		r2 = rf(key, ts)
+	if rf, ok := ret.Get(2).(func(context.Context, string, uint64) error); ok {
+		r2 = rf(ctx, key, ts)
 	} else {
 		r2 = ret.Error(2)
 	}
@@ -119,15 +124,16 @@ type SnapShotKV_LoadWithPrefix_Call struct {
 }
 
 // LoadWithPrefix is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
 //   - ts uint64
-func (_e *SnapShotKV_Expecter) LoadWithPrefix(key interface{}, ts interface{}) *SnapShotKV_LoadWithPrefix_Call {
-	return &SnapShotKV_LoadWithPrefix_Call{Call: _e.mock.On("LoadWithPrefix", key, ts)}
+func (_e *SnapShotKV_Expecter) LoadWithPrefix(ctx interface{}, key interface{}, ts interface{}) *SnapShotKV_LoadWithPrefix_Call {
+	return &SnapShotKV_LoadWithPrefix_Call{Call: _e.mock.On("LoadWithPrefix", ctx, key, ts)}
 }
 
-func (_c *SnapShotKV_LoadWithPrefix_Call) Run(run func(key string, ts uint64)) *SnapShotKV_LoadWithPrefix_Call {
+func (_c *SnapShotKV_LoadWithPrefix_Call) Run(run func(ctx context.Context, key string, ts uint64)) *SnapShotKV_LoadWithPrefix_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string), args[1].(uint64))
+		run(args[0].(context.Context), args[1].(string), args[2].(uint64))
 	})
 	return _c
 }
@@ -137,22 +143,22 @@ func (_c *SnapShotKV_LoadWithPrefix_Call) Return(_a0 []string, _a1 []string, _a2
 	return _c
 }
 
-func (_c *SnapShotKV_LoadWithPrefix_Call) RunAndReturn(run func(string, uint64) ([]string, []string, error)) *SnapShotKV_LoadWithPrefix_Call {
+func (_c *SnapShotKV_LoadWithPrefix_Call) RunAndReturn(run func(context.Context, string, uint64) ([]string, []string, error)) *SnapShotKV_LoadWithPrefix_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// MultiSave provides a mock function with given fields: kvs, ts
-func (_m *SnapShotKV) MultiSave(kvs map[string]string, ts uint64) error {
-	ret := _m.Called(kvs, ts)
+// MultiSave provides a mock function with given fields: ctx, kvs, ts
+func (_m *SnapShotKV) MultiSave(ctx context.Context, kvs map[string]string, ts uint64) error {
+	ret := _m.Called(ctx, kvs, ts)
 
 	if len(ret) == 0 {
 		panic("no return value specified for MultiSave")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(map[string]string, uint64) error); ok {
-		r0 = rf(kvs, ts)
+	if rf, ok := ret.Get(0).(func(context.Context, map[string]string, uint64) error); ok {
+		r0 = rf(ctx, kvs, ts)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -166,15 +172,16 @@ type SnapShotKV_MultiSave_Call struct {
 }
 
 // MultiSave is a helper method to define mock.On call
+//   - ctx context.Context
 //   - kvs map[string]string
 //   - ts uint64
-func (_e *SnapShotKV_Expecter) MultiSave(kvs interface{}, ts interface{}) *SnapShotKV_MultiSave_Call {
-	return &SnapShotKV_MultiSave_Call{Call: _e.mock.On("MultiSave", kvs, ts)}
+func (_e *SnapShotKV_Expecter) MultiSave(ctx interface{}, kvs interface{}, ts interface{}) *SnapShotKV_MultiSave_Call {
+	return &SnapShotKV_MultiSave_Call{Call: _e.mock.On("MultiSave", ctx, kvs, ts)}
 }
 
-func (_c *SnapShotKV_MultiSave_Call) Run(run func(kvs map[string]string, ts uint64)) *SnapShotKV_MultiSave_Call {
+func (_c *SnapShotKV_MultiSave_Call) Run(run func(ctx context.Context, kvs map[string]string, ts uint64)) *SnapShotKV_MultiSave_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(map[string]string), args[1].(uint64))
+		run(args[0].(context.Context), args[1].(map[string]string), args[2].(uint64))
 	})
 	return _c
 }
@@ -184,22 +191,22 @@ func (_c *SnapShotKV_MultiSave_Call) Return(_a0 error) *SnapShotKV_MultiSave_Cal
 	return _c
 }
 
-func (_c *SnapShotKV_MultiSave_Call) RunAndReturn(run func(map[string]string, uint64) error) *SnapShotKV_MultiSave_Call {
+func (_c *SnapShotKV_MultiSave_Call) RunAndReturn(run func(context.Context, map[string]string, uint64) error) *SnapShotKV_MultiSave_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// MultiSaveAndRemove provides a mock function with given fields: saves, removals, ts
-func (_m *SnapShotKV) MultiSaveAndRemove(saves map[string]string, removals []string, ts uint64) error {
-	ret := _m.Called(saves, removals, ts)
+// MultiSaveAndRemove provides a mock function with given fields: ctx, saves, removals, ts
+func (_m *SnapShotKV) MultiSaveAndRemove(ctx context.Context, saves map[string]string, removals []string, ts uint64) error {
+	ret := _m.Called(ctx, saves, removals, ts)
 
 	if len(ret) == 0 {
 		panic("no return value specified for MultiSaveAndRemove")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(map[string]string, []string, uint64) error); ok {
-		r0 = rf(saves, removals, ts)
+	if rf, ok := ret.Get(0).(func(context.Context, map[string]string, []string, uint64) error); ok {
+		r0 = rf(ctx, saves, removals, ts)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -213,16 +220,17 @@ type SnapShotKV_MultiSaveAndRemove_Call struct {
 }
 
 // MultiSaveAndRemove is a helper method to define mock.On call
+//   - ctx context.Context
 //   - saves map[string]string
 //   - removals []string
 //   - ts uint64
-func (_e *SnapShotKV_Expecter) MultiSaveAndRemove(saves interface{}, removals interface{}, ts interface{}) *SnapShotKV_MultiSaveAndRemove_Call {
-	return &SnapShotKV_MultiSaveAndRemove_Call{Call: _e.mock.On("MultiSaveAndRemove", saves, removals, ts)}
+func (_e *SnapShotKV_Expecter) MultiSaveAndRemove(ctx interface{}, saves interface{}, removals interface{}, ts interface{}) *SnapShotKV_MultiSaveAndRemove_Call {
+	return &SnapShotKV_MultiSaveAndRemove_Call{Call: _e.mock.On("MultiSaveAndRemove", ctx, saves, removals, ts)}
 }
 
-func (_c *SnapShotKV_MultiSaveAndRemove_Call) Run(run func(saves map[string]string, removals []string, ts uint64)) *SnapShotKV_MultiSaveAndRemove_Call {
+func (_c *SnapShotKV_MultiSaveAndRemove_Call) Run(run func(ctx context.Context, saves map[string]string, removals []string, ts uint64)) *SnapShotKV_MultiSaveAndRemove_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(map[string]string), args[1].([]string), args[2].(uint64))
+		run(args[0].(context.Context), args[1].(map[string]string), args[2].([]string), args[3].(uint64))
 	})
 	return _c
 }
@@ -232,22 +240,22 @@ func (_c *SnapShotKV_MultiSaveAndRemove_Call) Return(_a0 error) *SnapShotKV_Mult
 	return _c
 }
 
-func (_c *SnapShotKV_MultiSaveAndRemove_Call) RunAndReturn(run func(map[string]string, []string, uint64) error) *SnapShotKV_MultiSaveAndRemove_Call {
+func (_c *SnapShotKV_MultiSaveAndRemove_Call) RunAndReturn(run func(context.Context, map[string]string, []string, uint64) error) *SnapShotKV_MultiSaveAndRemove_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// MultiSaveAndRemoveWithPrefix provides a mock function with given fields: saves, removals, ts
-func (_m *SnapShotKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string, ts uint64) error {
-	ret := _m.Called(saves, removals, ts)
+// MultiSaveAndRemoveWithPrefix provides a mock function with given fields: ctx, saves, removals, ts
+func (_m *SnapShotKV) MultiSaveAndRemoveWithPrefix(ctx context.Context, saves map[string]string, removals []string, ts uint64) error {
+	ret := _m.Called(ctx, saves, removals, ts)
 
 	if len(ret) == 0 {
 		panic("no return value specified for MultiSaveAndRemoveWithPrefix")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(map[string]string, []string, uint64) error); ok {
-		r0 = rf(saves, removals, ts)
+	if rf, ok := ret.Get(0).(func(context.Context, map[string]string, []string, uint64) error); ok {
+		r0 = rf(ctx, saves, removals, ts)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -261,16 +269,17 @@ type SnapShotKV_MultiSaveAndRemoveWithPrefix_Call struct {
 }
 
 // MultiSaveAndRemoveWithPrefix is a helper method to define mock.On call
+//   - ctx context.Context
 //   - saves map[string]string
 //   - removals []string
 //   - ts uint64
-func (_e *SnapShotKV_Expecter) MultiSaveAndRemoveWithPrefix(saves interface{}, removals interface{}, ts interface{}) *SnapShotKV_MultiSaveAndRemoveWithPrefix_Call {
-	return &SnapShotKV_MultiSaveAndRemoveWithPrefix_Call{Call: _e.mock.On("MultiSaveAndRemoveWithPrefix", saves, removals, ts)}
+func (_e *SnapShotKV_Expecter) MultiSaveAndRemoveWithPrefix(ctx interface{}, saves interface{}, removals interface{}, ts interface{}) *SnapShotKV_MultiSaveAndRemoveWithPrefix_Call {
+	return &SnapShotKV_MultiSaveAndRemoveWithPrefix_Call{Call: _e.mock.On("MultiSaveAndRemoveWithPrefix", ctx, saves, removals, ts)}
 }
 
-func (_c *SnapShotKV_MultiSaveAndRemoveWithPrefix_Call) Run(run func(saves map[string]string, removals []string, ts uint64)) *SnapShotKV_MultiSaveAndRemoveWithPrefix_Call {
+func (_c *SnapShotKV_MultiSaveAndRemoveWithPrefix_Call) Run(run func(ctx context.Context, saves map[string]string, removals []string, ts uint64)) *SnapShotKV_MultiSaveAndRemoveWithPrefix_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(map[string]string), args[1].([]string), args[2].(uint64))
+		run(args[0].(context.Context), args[1].(map[string]string), args[2].([]string), args[3].(uint64))
 	})
 	return _c
 }
@@ -280,22 +289,22 @@ func (_c *SnapShotKV_MultiSaveAndRemoveWithPrefix_Call) Return(_a0 error) *SnapS
 	return _c
 }
 
-func (_c *SnapShotKV_MultiSaveAndRemoveWithPrefix_Call) RunAndReturn(run func(map[string]string, []string, uint64) error) *SnapShotKV_MultiSaveAndRemoveWithPrefix_Call {
+func (_c *SnapShotKV_MultiSaveAndRemoveWithPrefix_Call) RunAndReturn(run func(context.Context, map[string]string, []string, uint64) error) *SnapShotKV_MultiSaveAndRemoveWithPrefix_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// Save provides a mock function with given fields: key, value, ts
-func (_m *SnapShotKV) Save(key string, value string, ts uint64) error {
-	ret := _m.Called(key, value, ts)
+// Save provides a mock function with given fields: ctx, key, value, ts
+func (_m *SnapShotKV) Save(ctx context.Context, key string, value string, ts uint64) error {
+	ret := _m.Called(ctx, key, value, ts)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Save")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(string, string, uint64) error); ok {
-		r0 = rf(key, value, ts)
+	if rf, ok := ret.Get(0).(func(context.Context, string, string, uint64) error); ok {
+		r0 = rf(ctx, key, value, ts)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -309,16 +318,17 @@ type SnapShotKV_Save_Call struct {
 }
 
 // Save is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
 //   - value string
 //   - ts uint64
-func (_e *SnapShotKV_Expecter) Save(key interface{}, value interface{}, ts interface{}) *SnapShotKV_Save_Call {
-	return &SnapShotKV_Save_Call{Call: _e.mock.On("Save", key, value, ts)}
+func (_e *SnapShotKV_Expecter) Save(ctx interface{}, key interface{}, value interface{}, ts interface{}) *SnapShotKV_Save_Call {
+	return &SnapShotKV_Save_Call{Call: _e.mock.On("Save", ctx, key, value, ts)}
 }
 
-func (_c *SnapShotKV_Save_Call) Run(run func(key string, value string, ts uint64)) *SnapShotKV_Save_Call {
+func (_c *SnapShotKV_Save_Call) Run(run func(ctx context.Context, key string, value string, ts uint64)) *SnapShotKV_Save_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string), args[1].(string), args[2].(uint64))
+		run(args[0].(context.Context), args[1].(string), args[2].(string), args[3].(uint64))
 	})
 	return _c
 }
@@ -328,7 +338,7 @@ func (_c *SnapShotKV_Save_Call) Return(_a0 error) *SnapShotKV_Save_Call {
 	return _c
 }
 
-func (_c *SnapShotKV_Save_Call) RunAndReturn(run func(string, string, uint64) error) *SnapShotKV_Save_Call {
+func (_c *SnapShotKV_Save_Call) RunAndReturn(run func(context.Context, string, string, uint64) error) *SnapShotKV_Save_Call {
 	_c.Call.Return(run)
 	return _c
 }
diff --git a/internal/kv/mocks/txn_kv.go b/internal/kv/mocks/txn_kv.go
index 76811bbfb7a9f..7a149148fdc7e 100644
--- a/internal/kv/mocks/txn_kv.go
+++ b/internal/kv/mocks/txn_kv.go
@@ -3,8 +3,11 @@
 package mocks
 
 import (
-	predicates "github.com/milvus-io/milvus/pkg/kv/predicates"
+	context "context"
+
 	mock "github.com/stretchr/testify/mock"
+
+	predicates "github.com/milvus-io/milvus/pkg/kv/predicates"
 )
 
 // TxnKV is an autogenerated mock type for the TxnKV type
@@ -52,9 +55,9 @@ func (_c *TxnKV_Close_Call) RunAndReturn(run func()) *TxnKV_Close_Call {
 	return _c
 }
 
-// Has provides a mock function with given fields: key
-func (_m *TxnKV) Has(key string) (bool, error) {
-	ret := _m.Called(key)
+// Has provides a mock function with given fields: ctx, key
+func (_m *TxnKV) Has(ctx context.Context, key string) (bool, error) {
+	ret := _m.Called(ctx, key)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Has")
@@ -62,17 +65,17 @@ func (_m *TxnKV) Has(key string) (bool, error) {
 
 	var r0 bool
 	var r1 error
-	if rf, ok := ret.Get(0).(func(string) (bool, error)); ok {
-		return rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) (bool, error)); ok {
+		return rf(ctx, key)
 	}
-	if rf, ok := ret.Get(0).(func(string) bool); ok {
-		r0 = rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) bool); ok {
+		r0 = rf(ctx, key)
 	} else {
 		r0 = ret.Get(0).(bool)
 	}
 
-	if rf, ok := ret.Get(1).(func(string) error); ok {
-		r1 = rf(key)
+	if rf, ok := ret.Get(1).(func(context.Context, string) error); ok {
+		r1 = rf(ctx, key)
 	} else {
 		r1 = ret.Error(1)
 	}
@@ -86,14 +89,15 @@ type TxnKV_Has_Call struct {
 }
 
 // Has is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
-func (_e *TxnKV_Expecter) Has(key interface{}) *TxnKV_Has_Call {
-	return &TxnKV_Has_Call{Call: _e.mock.On("Has", key)}
+func (_e *TxnKV_Expecter) Has(ctx interface{}, key interface{}) *TxnKV_Has_Call {
+	return &TxnKV_Has_Call{Call: _e.mock.On("Has", ctx, key)}
 }
 
-func (_c *TxnKV_Has_Call) Run(run func(key string)) *TxnKV_Has_Call {
+func (_c *TxnKV_Has_Call) Run(run func(ctx context.Context, key string)) *TxnKV_Has_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string))
+		run(args[0].(context.Context), args[1].(string))
 	})
 	return _c
 }
@@ -103,14 +107,14 @@ func (_c *TxnKV_Has_Call) Return(_a0 bool, _a1 error) *TxnKV_Has_Call {
 	return _c
 }
 
-func (_c *TxnKV_Has_Call) RunAndReturn(run func(string) (bool, error)) *TxnKV_Has_Call {
+func (_c *TxnKV_Has_Call) RunAndReturn(run func(context.Context, string) (bool, error)) *TxnKV_Has_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// HasPrefix provides a mock function with given fields: prefix
-func (_m *TxnKV) HasPrefix(prefix string) (bool, error) {
-	ret := _m.Called(prefix)
+// HasPrefix provides a mock function with given fields: ctx, prefix
+func (_m *TxnKV) HasPrefix(ctx context.Context, prefix string) (bool, error) {
+	ret := _m.Called(ctx, prefix)
 
 	if len(ret) == 0 {
 		panic("no return value specified for HasPrefix")
@@ -118,17 +122,17 @@ func (_m *TxnKV) HasPrefix(prefix string) (bool, error) {
 
 	var r0 bool
 	var r1 error
-	if rf, ok := ret.Get(0).(func(string) (bool, error)); ok {
-		return rf(prefix)
+	if rf, ok := ret.Get(0).(func(context.Context, string) (bool, error)); ok {
+		return rf(ctx, prefix)
 	}
-	if rf, ok := ret.Get(0).(func(string) bool); ok {
-		r0 = rf(prefix)
+	if rf, ok := ret.Get(0).(func(context.Context, string) bool); ok {
+		r0 = rf(ctx, prefix)
 	} else {
 		r0 = ret.Get(0).(bool)
 	}
 
-	if rf, ok := ret.Get(1).(func(string) error); ok {
-		r1 = rf(prefix)
+	if rf, ok := ret.Get(1).(func(context.Context, string) error); ok {
+		r1 = rf(ctx, prefix)
 	} else {
 		r1 = ret.Error(1)
 	}
@@ -142,14 +146,15 @@ type TxnKV_HasPrefix_Call struct {
 }
 
 // HasPrefix is a helper method to define mock.On call
+//   - ctx context.Context
 //   - prefix string
-func (_e *TxnKV_Expecter) HasPrefix(prefix interface{}) *TxnKV_HasPrefix_Call {
-	return &TxnKV_HasPrefix_Call{Call: _e.mock.On("HasPrefix", prefix)}
+func (_e *TxnKV_Expecter) HasPrefix(ctx interface{}, prefix interface{}) *TxnKV_HasPrefix_Call {
+	return &TxnKV_HasPrefix_Call{Call: _e.mock.On("HasPrefix", ctx, prefix)}
 }
 
-func (_c *TxnKV_HasPrefix_Call) Run(run func(prefix string)) *TxnKV_HasPrefix_Call {
+func (_c *TxnKV_HasPrefix_Call) Run(run func(ctx context.Context, prefix string)) *TxnKV_HasPrefix_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string))
+		run(args[0].(context.Context), args[1].(string))
 	})
 	return _c
 }
@@ -159,14 +164,14 @@ func (_c *TxnKV_HasPrefix_Call) Return(_a0 bool, _a1 error) *TxnKV_HasPrefix_Cal
 	return _c
 }
 
-func (_c *TxnKV_HasPrefix_Call) RunAndReturn(run func(string) (bool, error)) *TxnKV_HasPrefix_Call {
+func (_c *TxnKV_HasPrefix_Call) RunAndReturn(run func(context.Context, string) (bool, error)) *TxnKV_HasPrefix_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// Load provides a mock function with given fields: key
-func (_m *TxnKV) Load(key string) (string, error) {
-	ret := _m.Called(key)
+// Load provides a mock function with given fields: ctx, key
+func (_m *TxnKV) Load(ctx context.Context, key string) (string, error) {
+	ret := _m.Called(ctx, key)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Load")
@@ -174,17 +179,17 @@ func (_m *TxnKV) Load(key string) (string, error) {
 
 	var r0 string
 	var r1 error
-	if rf, ok := ret.Get(0).(func(string) (string, error)); ok {
-		return rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) (string, error)); ok {
+		return rf(ctx, key)
 	}
-	if rf, ok := ret.Get(0).(func(string) string); ok {
-		r0 = rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) string); ok {
+		r0 = rf(ctx, key)
 	} else {
 		r0 = ret.Get(0).(string)
 	}
 
-	if rf, ok := ret.Get(1).(func(string) error); ok {
-		r1 = rf(key)
+	if rf, ok := ret.Get(1).(func(context.Context, string) error); ok {
+		r1 = rf(ctx, key)
 	} else {
 		r1 = ret.Error(1)
 	}
@@ -198,14 +203,15 @@ type TxnKV_Load_Call struct {
 }
 
 // Load is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
-func (_e *TxnKV_Expecter) Load(key interface{}) *TxnKV_Load_Call {
-	return &TxnKV_Load_Call{Call: _e.mock.On("Load", key)}
+func (_e *TxnKV_Expecter) Load(ctx interface{}, key interface{}) *TxnKV_Load_Call {
+	return &TxnKV_Load_Call{Call: _e.mock.On("Load", ctx, key)}
 }
 
-func (_c *TxnKV_Load_Call) Run(run func(key string)) *TxnKV_Load_Call {
+func (_c *TxnKV_Load_Call) Run(run func(ctx context.Context, key string)) *TxnKV_Load_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string))
+		run(args[0].(context.Context), args[1].(string))
 	})
 	return _c
 }
@@ -215,14 +221,14 @@ func (_c *TxnKV_Load_Call) Return(_a0 string, _a1 error) *TxnKV_Load_Call {
 	return _c
 }
 
-func (_c *TxnKV_Load_Call) RunAndReturn(run func(string) (string, error)) *TxnKV_Load_Call {
+func (_c *TxnKV_Load_Call) RunAndReturn(run func(context.Context, string) (string, error)) *TxnKV_Load_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// LoadWithPrefix provides a mock function with given fields: key
-func (_m *TxnKV) LoadWithPrefix(key string) ([]string, []string, error) {
-	ret := _m.Called(key)
+// LoadWithPrefix provides a mock function with given fields: ctx, key
+func (_m *TxnKV) LoadWithPrefix(ctx context.Context, key string) ([]string, []string, error) {
+	ret := _m.Called(ctx, key)
 
 	if len(ret) == 0 {
 		panic("no return value specified for LoadWithPrefix")
@@ -231,27 +237,27 @@ func (_m *TxnKV) LoadWithPrefix(key string) ([]string, []string, error) {
 	var r0 []string
 	var r1 []string
 	var r2 error
-	if rf, ok := ret.Get(0).(func(string) ([]string, []string, error)); ok {
-		return rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) ([]string, []string, error)); ok {
+		return rf(ctx, key)
 	}
-	if rf, ok := ret.Get(0).(func(string) []string); ok {
-		r0 = rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) []string); ok {
+		r0 = rf(ctx, key)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).([]string)
 		}
 	}
 
-	if rf, ok := ret.Get(1).(func(string) []string); ok {
-		r1 = rf(key)
+	if rf, ok := ret.Get(1).(func(context.Context, string) []string); ok {
+		r1 = rf(ctx, key)
 	} else {
 		if ret.Get(1) != nil {
 			r1 = ret.Get(1).([]string)
 		}
 	}
 
-	if rf, ok := ret.Get(2).(func(string) error); ok {
-		r2 = rf(key)
+	if rf, ok := ret.Get(2).(func(context.Context, string) error); ok {
+		r2 = rf(ctx, key)
 	} else {
 		r2 = ret.Error(2)
 	}
@@ -265,14 +271,15 @@ type TxnKV_LoadWithPrefix_Call struct {
 }
 
 // LoadWithPrefix is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
-func (_e *TxnKV_Expecter) LoadWithPrefix(key interface{}) *TxnKV_LoadWithPrefix_Call {
-	return &TxnKV_LoadWithPrefix_Call{Call: _e.mock.On("LoadWithPrefix", key)}
+func (_e *TxnKV_Expecter) LoadWithPrefix(ctx interface{}, key interface{}) *TxnKV_LoadWithPrefix_Call {
+	return &TxnKV_LoadWithPrefix_Call{Call: _e.mock.On("LoadWithPrefix", ctx, key)}
 }
 
-func (_c *TxnKV_LoadWithPrefix_Call) Run(run func(key string)) *TxnKV_LoadWithPrefix_Call {
+func (_c *TxnKV_LoadWithPrefix_Call) Run(run func(ctx context.Context, key string)) *TxnKV_LoadWithPrefix_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string))
+		run(args[0].(context.Context), args[1].(string))
 	})
 	return _c
 }
@@ -282,14 +289,14 @@ func (_c *TxnKV_LoadWithPrefix_Call) Return(_a0 []string, _a1 []string, _a2 erro
 	return _c
 }
 
-func (_c *TxnKV_LoadWithPrefix_Call) RunAndReturn(run func(string) ([]string, []string, error)) *TxnKV_LoadWithPrefix_Call {
+func (_c *TxnKV_LoadWithPrefix_Call) RunAndReturn(run func(context.Context, string) ([]string, []string, error)) *TxnKV_LoadWithPrefix_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// MultiLoad provides a mock function with given fields: keys
-func (_m *TxnKV) MultiLoad(keys []string) ([]string, error) {
-	ret := _m.Called(keys)
+// MultiLoad provides a mock function with given fields: ctx, keys
+func (_m *TxnKV) MultiLoad(ctx context.Context, keys []string) ([]string, error) {
+	ret := _m.Called(ctx, keys)
 
 	if len(ret) == 0 {
 		panic("no return value specified for MultiLoad")
@@ -297,19 +304,19 @@ func (_m *TxnKV) MultiLoad(keys []string) ([]string, error) {
 
 	var r0 []string
 	var r1 error
-	if rf, ok := ret.Get(0).(func([]string) ([]string, error)); ok {
-		return rf(keys)
+	if rf, ok := ret.Get(0).(func(context.Context, []string) ([]string, error)); ok {
+		return rf(ctx, keys)
 	}
-	if rf, ok := ret.Get(0).(func([]string) []string); ok {
-		r0 = rf(keys)
+	if rf, ok := ret.Get(0).(func(context.Context, []string) []string); ok {
+		r0 = rf(ctx, keys)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).([]string)
 		}
 	}
 
-	if rf, ok := ret.Get(1).(func([]string) error); ok {
-		r1 = rf(keys)
+	if rf, ok := ret.Get(1).(func(context.Context, []string) error); ok {
+		r1 = rf(ctx, keys)
 	} else {
 		r1 = ret.Error(1)
 	}
@@ -323,14 +330,15 @@ type TxnKV_MultiLoad_Call struct {
 }
 
 // MultiLoad is a helper method to define mock.On call
+//   - ctx context.Context
 //   - keys []string
-func (_e *TxnKV_Expecter) MultiLoad(keys interface{}) *TxnKV_MultiLoad_Call {
-	return &TxnKV_MultiLoad_Call{Call: _e.mock.On("MultiLoad", keys)}
+func (_e *TxnKV_Expecter) MultiLoad(ctx interface{}, keys interface{}) *TxnKV_MultiLoad_Call {
+	return &TxnKV_MultiLoad_Call{Call: _e.mock.On("MultiLoad", ctx, keys)}
 }
 
-func (_c *TxnKV_MultiLoad_Call) Run(run func(keys []string)) *TxnKV_MultiLoad_Call {
+func (_c *TxnKV_MultiLoad_Call) Run(run func(ctx context.Context, keys []string)) *TxnKV_MultiLoad_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].([]string))
+		run(args[0].(context.Context), args[1].([]string))
 	})
 	return _c
 }
@@ -340,22 +348,22 @@ func (_c *TxnKV_MultiLoad_Call) Return(_a0 []string, _a1 error) *TxnKV_MultiLoad
 	return _c
 }
 
-func (_c *TxnKV_MultiLoad_Call) RunAndReturn(run func([]string) ([]string, error)) *TxnKV_MultiLoad_Call {
+func (_c *TxnKV_MultiLoad_Call) RunAndReturn(run func(context.Context, []string) ([]string, error)) *TxnKV_MultiLoad_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// MultiRemove provides a mock function with given fields: keys
-func (_m *TxnKV) MultiRemove(keys []string) error {
-	ret := _m.Called(keys)
+// MultiRemove provides a mock function with given fields: ctx, keys
+func (_m *TxnKV) MultiRemove(ctx context.Context, keys []string) error {
+	ret := _m.Called(ctx, keys)
 
 	if len(ret) == 0 {
 		panic("no return value specified for MultiRemove")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func([]string) error); ok {
-		r0 = rf(keys)
+	if rf, ok := ret.Get(0).(func(context.Context, []string) error); ok {
+		r0 = rf(ctx, keys)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -369,14 +377,15 @@ type TxnKV_MultiRemove_Call struct {
 }
 
 // MultiRemove is a helper method to define mock.On call
+//   - ctx context.Context
 //   - keys []string
-func (_e *TxnKV_Expecter) MultiRemove(keys interface{}) *TxnKV_MultiRemove_Call {
-	return &TxnKV_MultiRemove_Call{Call: _e.mock.On("MultiRemove", keys)}
+func (_e *TxnKV_Expecter) MultiRemove(ctx interface{}, keys interface{}) *TxnKV_MultiRemove_Call {
+	return &TxnKV_MultiRemove_Call{Call: _e.mock.On("MultiRemove", ctx, keys)}
 }
 
-func (_c *TxnKV_MultiRemove_Call) Run(run func(keys []string)) *TxnKV_MultiRemove_Call {
+func (_c *TxnKV_MultiRemove_Call) Run(run func(ctx context.Context, keys []string)) *TxnKV_MultiRemove_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].([]string))
+		run(args[0].(context.Context), args[1].([]string))
 	})
 	return _c
 }
@@ -386,22 +395,22 @@ func (_c *TxnKV_MultiRemove_Call) Return(_a0 error) *TxnKV_MultiRemove_Call {
 	return _c
 }
 
-func (_c *TxnKV_MultiRemove_Call) RunAndReturn(run func([]string) error) *TxnKV_MultiRemove_Call {
+func (_c *TxnKV_MultiRemove_Call) RunAndReturn(run func(context.Context, []string) error) *TxnKV_MultiRemove_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// MultiSave provides a mock function with given fields: kvs
-func (_m *TxnKV) MultiSave(kvs map[string]string) error {
-	ret := _m.Called(kvs)
+// MultiSave provides a mock function with given fields: ctx, kvs
+func (_m *TxnKV) MultiSave(ctx context.Context, kvs map[string]string) error {
+	ret := _m.Called(ctx, kvs)
 
 	if len(ret) == 0 {
 		panic("no return value specified for MultiSave")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(map[string]string) error); ok {
-		r0 = rf(kvs)
+	if rf, ok := ret.Get(0).(func(context.Context, map[string]string) error); ok {
+		r0 = rf(ctx, kvs)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -415,14 +424,15 @@ type TxnKV_MultiSave_Call struct {
 }
 
 // MultiSave is a helper method to define mock.On call
+//   - ctx context.Context
 //   - kvs map[string]string
-func (_e *TxnKV_Expecter) MultiSave(kvs interface{}) *TxnKV_MultiSave_Call {
-	return &TxnKV_MultiSave_Call{Call: _e.mock.On("MultiSave", kvs)}
+func (_e *TxnKV_Expecter) MultiSave(ctx interface{}, kvs interface{}) *TxnKV_MultiSave_Call {
+	return &TxnKV_MultiSave_Call{Call: _e.mock.On("MultiSave", ctx, kvs)}
 }
 
-func (_c *TxnKV_MultiSave_Call) Run(run func(kvs map[string]string)) *TxnKV_MultiSave_Call {
+func (_c *TxnKV_MultiSave_Call) Run(run func(ctx context.Context, kvs map[string]string)) *TxnKV_MultiSave_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(map[string]string))
+		run(args[0].(context.Context), args[1].(map[string]string))
 	})
 	return _c
 }
@@ -432,19 +442,19 @@ func (_c *TxnKV_MultiSave_Call) Return(_a0 error) *TxnKV_MultiSave_Call {
 	return _c
 }
 
-func (_c *TxnKV_MultiSave_Call) RunAndReturn(run func(map[string]string) error) *TxnKV_MultiSave_Call {
+func (_c *TxnKV_MultiSave_Call) RunAndReturn(run func(context.Context, map[string]string) error) *TxnKV_MultiSave_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// MultiSaveAndRemove provides a mock function with given fields: saves, removals, preds
-func (_m *TxnKV) MultiSaveAndRemove(saves map[string]string, removals []string, preds ...predicates.Predicate) error {
+// MultiSaveAndRemove provides a mock function with given fields: ctx, saves, removals, preds
+func (_m *TxnKV) MultiSaveAndRemove(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error {
 	_va := make([]interface{}, len(preds))
 	for _i := range preds {
 		_va[_i] = preds[_i]
 	}
 	var _ca []interface{}
-	_ca = append(_ca, saves, removals)
+	_ca = append(_ca, ctx, saves, removals)
 	_ca = append(_ca, _va...)
 	ret := _m.Called(_ca...)
 
@@ -453,8 +463,8 @@ func (_m *TxnKV) MultiSaveAndRemove(saves map[string]string, removals []string,
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(map[string]string, []string, ...predicates.Predicate) error); ok {
-		r0 = rf(saves, removals, preds...)
+	if rf, ok := ret.Get(0).(func(context.Context, map[string]string, []string, ...predicates.Predicate) error); ok {
+		r0 = rf(ctx, saves, removals, preds...)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -468,23 +478,24 @@ type TxnKV_MultiSaveAndRemove_Call struct {
 }
 
 // MultiSaveAndRemove is a helper method to define mock.On call
+//   - ctx context.Context
 //   - saves map[string]string
 //   - removals []string
 //   - preds ...predicates.Predicate
-func (_e *TxnKV_Expecter) MultiSaveAndRemove(saves interface{}, removals interface{}, preds ...interface{}) *TxnKV_MultiSaveAndRemove_Call {
+func (_e *TxnKV_Expecter) MultiSaveAndRemove(ctx interface{}, saves interface{}, removals interface{}, preds ...interface{}) *TxnKV_MultiSaveAndRemove_Call {
 	return &TxnKV_MultiSaveAndRemove_Call{Call: _e.mock.On("MultiSaveAndRemove",
-		append([]interface{}{saves, removals}, preds...)...)}
+		append([]interface{}{ctx, saves, removals}, preds...)...)}
 }
 
-func (_c *TxnKV_MultiSaveAndRemove_Call) Run(run func(saves map[string]string, removals []string, preds ...predicates.Predicate)) *TxnKV_MultiSaveAndRemove_Call {
+func (_c *TxnKV_MultiSaveAndRemove_Call) Run(run func(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate)) *TxnKV_MultiSaveAndRemove_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		variadicArgs := make([]predicates.Predicate, len(args)-2)
-		for i, a := range args[2:] {
+		variadicArgs := make([]predicates.Predicate, len(args)-3)
+		for i, a := range args[3:] {
 			if a != nil {
 				variadicArgs[i] = a.(predicates.Predicate)
 			}
 		}
-		run(args[0].(map[string]string), args[1].([]string), variadicArgs...)
+		run(args[0].(context.Context), args[1].(map[string]string), args[2].([]string), variadicArgs...)
 	})
 	return _c
 }
@@ -494,19 +505,19 @@ func (_c *TxnKV_MultiSaveAndRemove_Call) Return(_a0 error) *TxnKV_MultiSaveAndRe
 	return _c
 }
 
-func (_c *TxnKV_MultiSaveAndRemove_Call) RunAndReturn(run func(map[string]string, []string, ...predicates.Predicate) error) *TxnKV_MultiSaveAndRemove_Call {
+func (_c *TxnKV_MultiSaveAndRemove_Call) RunAndReturn(run func(context.Context, map[string]string, []string, ...predicates.Predicate) error) *TxnKV_MultiSaveAndRemove_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// MultiSaveAndRemoveWithPrefix provides a mock function with given fields: saves, removals, preds
-func (_m *TxnKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string, preds ...predicates.Predicate) error {
+// MultiSaveAndRemoveWithPrefix provides a mock function with given fields: ctx, saves, removals, preds
+func (_m *TxnKV) MultiSaveAndRemoveWithPrefix(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error {
 	_va := make([]interface{}, len(preds))
 	for _i := range preds {
 		_va[_i] = preds[_i]
 	}
 	var _ca []interface{}
-	_ca = append(_ca, saves, removals)
+	_ca = append(_ca, ctx, saves, removals)
 	_ca = append(_ca, _va...)
 	ret := _m.Called(_ca...)
 
@@ -515,8 +526,8 @@ func (_m *TxnKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(map[string]string, []string, ...predicates.Predicate) error); ok {
-		r0 = rf(saves, removals, preds...)
+	if rf, ok := ret.Get(0).(func(context.Context, map[string]string, []string, ...predicates.Predicate) error); ok {
+		r0 = rf(ctx, saves, removals, preds...)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -530,23 +541,24 @@ type TxnKV_MultiSaveAndRemoveWithPrefix_Call struct {
 }
 
 // MultiSaveAndRemoveWithPrefix is a helper method to define mock.On call
+//   - ctx context.Context
 //   - saves map[string]string
 //   - removals []string
 //   - preds ...predicates.Predicate
-func (_e *TxnKV_Expecter) MultiSaveAndRemoveWithPrefix(saves interface{}, removals interface{}, preds ...interface{}) *TxnKV_MultiSaveAndRemoveWithPrefix_Call {
+func (_e *TxnKV_Expecter) MultiSaveAndRemoveWithPrefix(ctx interface{}, saves interface{}, removals interface{}, preds ...interface{}) *TxnKV_MultiSaveAndRemoveWithPrefix_Call {
 	return &TxnKV_MultiSaveAndRemoveWithPrefix_Call{Call: _e.mock.On("MultiSaveAndRemoveWithPrefix",
-		append([]interface{}{saves, removals}, preds...)...)}
+		append([]interface{}{ctx, saves, removals}, preds...)...)}
 }
 
-func (_c *TxnKV_MultiSaveAndRemoveWithPrefix_Call) Run(run func(saves map[string]string, removals []string, preds ...predicates.Predicate)) *TxnKV_MultiSaveAndRemoveWithPrefix_Call {
+func (_c *TxnKV_MultiSaveAndRemoveWithPrefix_Call) Run(run func(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate)) *TxnKV_MultiSaveAndRemoveWithPrefix_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		variadicArgs := make([]predicates.Predicate, len(args)-2)
-		for i, a := range args[2:] {
+		variadicArgs := make([]predicates.Predicate, len(args)-3)
+		for i, a := range args[3:] {
 			if a != nil {
 				variadicArgs[i] = a.(predicates.Predicate)
 			}
 		}
-		run(args[0].(map[string]string), args[1].([]string), variadicArgs...)
+		run(args[0].(context.Context), args[1].(map[string]string), args[2].([]string), variadicArgs...)
 	})
 	return _c
 }
@@ -556,22 +568,22 @@ func (_c *TxnKV_MultiSaveAndRemoveWithPrefix_Call) Return(_a0 error) *TxnKV_Mult
 	return _c
 }
 
-func (_c *TxnKV_MultiSaveAndRemoveWithPrefix_Call) RunAndReturn(run func(map[string]string, []string, ...predicates.Predicate) error) *TxnKV_MultiSaveAndRemoveWithPrefix_Call {
+func (_c *TxnKV_MultiSaveAndRemoveWithPrefix_Call) RunAndReturn(run func(context.Context, map[string]string, []string, ...predicates.Predicate) error) *TxnKV_MultiSaveAndRemoveWithPrefix_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// Remove provides a mock function with given fields: key
-func (_m *TxnKV) Remove(key string) error {
-	ret := _m.Called(key)
+// Remove provides a mock function with given fields: ctx, key
+func (_m *TxnKV) Remove(ctx context.Context, key string) error {
+	ret := _m.Called(ctx, key)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Remove")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(string) error); ok {
-		r0 = rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) error); ok {
+		r0 = rf(ctx, key)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -585,14 +597,15 @@ type TxnKV_Remove_Call struct {
 }
 
 // Remove is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
-func (_e *TxnKV_Expecter) Remove(key interface{}) *TxnKV_Remove_Call {
-	return &TxnKV_Remove_Call{Call: _e.mock.On("Remove", key)}
+func (_e *TxnKV_Expecter) Remove(ctx interface{}, key interface{}) *TxnKV_Remove_Call {
+	return &TxnKV_Remove_Call{Call: _e.mock.On("Remove", ctx, key)}
 }
 
-func (_c *TxnKV_Remove_Call) Run(run func(key string)) *TxnKV_Remove_Call {
+func (_c *TxnKV_Remove_Call) Run(run func(ctx context.Context, key string)) *TxnKV_Remove_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string))
+		run(args[0].(context.Context), args[1].(string))
 	})
 	return _c
 }
@@ -602,22 +615,22 @@ func (_c *TxnKV_Remove_Call) Return(_a0 error) *TxnKV_Remove_Call {
 	return _c
 }
 
-func (_c *TxnKV_Remove_Call) RunAndReturn(run func(string) error) *TxnKV_Remove_Call {
+func (_c *TxnKV_Remove_Call) RunAndReturn(run func(context.Context, string) error) *TxnKV_Remove_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// RemoveWithPrefix provides a mock function with given fields: key
-func (_m *TxnKV) RemoveWithPrefix(key string) error {
-	ret := _m.Called(key)
+// RemoveWithPrefix provides a mock function with given fields: ctx, key
+func (_m *TxnKV) RemoveWithPrefix(ctx context.Context, key string) error {
+	ret := _m.Called(ctx, key)
 
 	if len(ret) == 0 {
 		panic("no return value specified for RemoveWithPrefix")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(string) error); ok {
-		r0 = rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) error); ok {
+		r0 = rf(ctx, key)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -631,14 +644,15 @@ type TxnKV_RemoveWithPrefix_Call struct {
 }
 
 // RemoveWithPrefix is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
-func (_e *TxnKV_Expecter) RemoveWithPrefix(key interface{}) *TxnKV_RemoveWithPrefix_Call {
-	return &TxnKV_RemoveWithPrefix_Call{Call: _e.mock.On("RemoveWithPrefix", key)}
+func (_e *TxnKV_Expecter) RemoveWithPrefix(ctx interface{}, key interface{}) *TxnKV_RemoveWithPrefix_Call {
+	return &TxnKV_RemoveWithPrefix_Call{Call: _e.mock.On("RemoveWithPrefix", ctx, key)}
 }
 
-func (_c *TxnKV_RemoveWithPrefix_Call) Run(run func(key string)) *TxnKV_RemoveWithPrefix_Call {
+func (_c *TxnKV_RemoveWithPrefix_Call) Run(run func(ctx context.Context, key string)) *TxnKV_RemoveWithPrefix_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string))
+		run(args[0].(context.Context), args[1].(string))
 	})
 	return _c
 }
@@ -648,22 +662,22 @@ func (_c *TxnKV_RemoveWithPrefix_Call) Return(_a0 error) *TxnKV_RemoveWithPrefix
 	return _c
 }
 
-func (_c *TxnKV_RemoveWithPrefix_Call) RunAndReturn(run func(string) error) *TxnKV_RemoveWithPrefix_Call {
+func (_c *TxnKV_RemoveWithPrefix_Call) RunAndReturn(run func(context.Context, string) error) *TxnKV_RemoveWithPrefix_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// Save provides a mock function with given fields: key, value
-func (_m *TxnKV) Save(key string, value string) error {
-	ret := _m.Called(key, value)
+// Save provides a mock function with given fields: ctx, key, value
+func (_m *TxnKV) Save(ctx context.Context, key string, value string) error {
+	ret := _m.Called(ctx, key, value)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Save")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(string, string) error); ok {
-		r0 = rf(key, value)
+	if rf, ok := ret.Get(0).(func(context.Context, string, string) error); ok {
+		r0 = rf(ctx, key, value)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -677,15 +691,16 @@ type TxnKV_Save_Call struct {
 }
 
 // Save is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
 //   - value string
-func (_e *TxnKV_Expecter) Save(key interface{}, value interface{}) *TxnKV_Save_Call {
-	return &TxnKV_Save_Call{Call: _e.mock.On("Save", key, value)}
+func (_e *TxnKV_Expecter) Save(ctx interface{}, key interface{}, value interface{}) *TxnKV_Save_Call {
+	return &TxnKV_Save_Call{Call: _e.mock.On("Save", ctx, key, value)}
 }
 
-func (_c *TxnKV_Save_Call) Run(run func(key string, value string)) *TxnKV_Save_Call {
+func (_c *TxnKV_Save_Call) Run(run func(ctx context.Context, key string, value string)) *TxnKV_Save_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string), args[1].(string))
+		run(args[0].(context.Context), args[1].(string), args[2].(string))
 	})
 	return _c
 }
@@ -695,7 +710,7 @@ func (_c *TxnKV_Save_Call) Return(_a0 error) *TxnKV_Save_Call {
 	return _c
 }
 
-func (_c *TxnKV_Save_Call) RunAndReturn(run func(string, string) error) *TxnKV_Save_Call {
+func (_c *TxnKV_Save_Call) RunAndReturn(run func(context.Context, string, string) error) *TxnKV_Save_Call {
 	_c.Call.Return(run)
 	return _c
 }
diff --git a/internal/kv/mocks/watch_kv.go b/internal/kv/mocks/watch_kv.go
index af3328ef5dbe2..60a38a5f97bb0 100644
--- a/internal/kv/mocks/watch_kv.go
+++ b/internal/kv/mocks/watch_kv.go
@@ -3,6 +3,8 @@
 package mocks
 
 import (
+	context "context"
+
 	clientv3 "go.etcd.io/etcd/client/v3"
 
 	mock "github.com/stretchr/testify/mock"
@@ -55,9 +57,9 @@ func (_c *WatchKV_Close_Call) RunAndReturn(run func()) *WatchKV_Close_Call {
 	return _c
 }
 
-// CompareVersionAndSwap provides a mock function with given fields: key, version, target
-func (_m *WatchKV) CompareVersionAndSwap(key string, version int64, target string) (bool, error) {
-	ret := _m.Called(key, version, target)
+// CompareVersionAndSwap provides a mock function with given fields: ctx, key, version, target
+func (_m *WatchKV) CompareVersionAndSwap(ctx context.Context, key string, version int64, target string) (bool, error) {
+	ret := _m.Called(ctx, key, version, target)
 
 	if len(ret) == 0 {
 		panic("no return value specified for CompareVersionAndSwap")
@@ -65,17 +67,17 @@ func (_m *WatchKV) CompareVersionAndSwap(key string, version int64, target strin
 
 	var r0 bool
 	var r1 error
-	if rf, ok := ret.Get(0).(func(string, int64, string) (bool, error)); ok {
-		return rf(key, version, target)
+	if rf, ok := ret.Get(0).(func(context.Context, string, int64, string) (bool, error)); ok {
+		return rf(ctx, key, version, target)
 	}
-	if rf, ok := ret.Get(0).(func(string, int64, string) bool); ok {
-		r0 = rf(key, version, target)
+	if rf, ok := ret.Get(0).(func(context.Context, string, int64, string) bool); ok {
+		r0 = rf(ctx, key, version, target)
 	} else {
 		r0 = ret.Get(0).(bool)
 	}
 
-	if rf, ok := ret.Get(1).(func(string, int64, string) error); ok {
-		r1 = rf(key, version, target)
+	if rf, ok := ret.Get(1).(func(context.Context, string, int64, string) error); ok {
+		r1 = rf(ctx, key, version, target)
 	} else {
 		r1 = ret.Error(1)
 	}
@@ -89,16 +91,17 @@ type WatchKV_CompareVersionAndSwap_Call struct {
 }
 
 // CompareVersionAndSwap is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
 //   - version int64
 //   - target string
-func (_e *WatchKV_Expecter) CompareVersionAndSwap(key interface{}, version interface{}, target interface{}) *WatchKV_CompareVersionAndSwap_Call {
-	return &WatchKV_CompareVersionAndSwap_Call{Call: _e.mock.On("CompareVersionAndSwap", key, version, target)}
+func (_e *WatchKV_Expecter) CompareVersionAndSwap(ctx interface{}, key interface{}, version interface{}, target interface{}) *WatchKV_CompareVersionAndSwap_Call {
+	return &WatchKV_CompareVersionAndSwap_Call{Call: _e.mock.On("CompareVersionAndSwap", ctx, key, version, target)}
 }
 
-func (_c *WatchKV_CompareVersionAndSwap_Call) Run(run func(key string, version int64, target string)) *WatchKV_CompareVersionAndSwap_Call {
+func (_c *WatchKV_CompareVersionAndSwap_Call) Run(run func(ctx context.Context, key string, version int64, target string)) *WatchKV_CompareVersionAndSwap_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string), args[1].(int64), args[2].(string))
+		run(args[0].(context.Context), args[1].(string), args[2].(int64), args[3].(string))
 	})
 	return _c
 }
@@ -108,7 +111,7 @@ func (_c *WatchKV_CompareVersionAndSwap_Call) Return(_a0 bool, _a1 error) *Watch
 	return _c
 }
 
-func (_c *WatchKV_CompareVersionAndSwap_Call) RunAndReturn(run func(string, int64, string) (bool, error)) *WatchKV_CompareVersionAndSwap_Call {
+func (_c *WatchKV_CompareVersionAndSwap_Call) RunAndReturn(run func(context.Context, string, int64, string) (bool, error)) *WatchKV_CompareVersionAndSwap_Call {
 	_c.Call.Return(run)
 	return _c
 }
@@ -159,9 +162,9 @@ func (_c *WatchKV_GetPath_Call) RunAndReturn(run func(string) string) *WatchKV_G
 	return _c
 }
 
-// Has provides a mock function with given fields: key
-func (_m *WatchKV) Has(key string) (bool, error) {
-	ret := _m.Called(key)
+// Has provides a mock function with given fields: ctx, key
+func (_m *WatchKV) Has(ctx context.Context, key string) (bool, error) {
+	ret := _m.Called(ctx, key)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Has")
@@ -169,17 +172,17 @@ func (_m *WatchKV) Has(key string) (bool, error) {
 
 	var r0 bool
 	var r1 error
-	if rf, ok := ret.Get(0).(func(string) (bool, error)); ok {
-		return rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) (bool, error)); ok {
+		return rf(ctx, key)
 	}
-	if rf, ok := ret.Get(0).(func(string) bool); ok {
-		r0 = rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) bool); ok {
+		r0 = rf(ctx, key)
 	} else {
 		r0 = ret.Get(0).(bool)
 	}
 
-	if rf, ok := ret.Get(1).(func(string) error); ok {
-		r1 = rf(key)
+	if rf, ok := ret.Get(1).(func(context.Context, string) error); ok {
+		r1 = rf(ctx, key)
 	} else {
 		r1 = ret.Error(1)
 	}
@@ -193,14 +196,15 @@ type WatchKV_Has_Call struct {
 }
 
 // Has is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
-func (_e *WatchKV_Expecter) Has(key interface{}) *WatchKV_Has_Call {
-	return &WatchKV_Has_Call{Call: _e.mock.On("Has", key)}
+func (_e *WatchKV_Expecter) Has(ctx interface{}, key interface{}) *WatchKV_Has_Call {
+	return &WatchKV_Has_Call{Call: _e.mock.On("Has", ctx, key)}
 }
 
-func (_c *WatchKV_Has_Call) Run(run func(key string)) *WatchKV_Has_Call {
+func (_c *WatchKV_Has_Call) Run(run func(ctx context.Context, key string)) *WatchKV_Has_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string))
+		run(args[0].(context.Context), args[1].(string))
 	})
 	return _c
 }
@@ -210,14 +214,14 @@ func (_c *WatchKV_Has_Call) Return(_a0 bool, _a1 error) *WatchKV_Has_Call {
 	return _c
 }
 
-func (_c *WatchKV_Has_Call) RunAndReturn(run func(string) (bool, error)) *WatchKV_Has_Call {
+func (_c *WatchKV_Has_Call) RunAndReturn(run func(context.Context, string) (bool, error)) *WatchKV_Has_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// HasPrefix provides a mock function with given fields: prefix
-func (_m *WatchKV) HasPrefix(prefix string) (bool, error) {
-	ret := _m.Called(prefix)
+// HasPrefix provides a mock function with given fields: ctx, prefix
+func (_m *WatchKV) HasPrefix(ctx context.Context, prefix string) (bool, error) {
+	ret := _m.Called(ctx, prefix)
 
 	if len(ret) == 0 {
 		panic("no return value specified for HasPrefix")
@@ -225,17 +229,17 @@ func (_m *WatchKV) HasPrefix(prefix string) (bool, error) {
 
 	var r0 bool
 	var r1 error
-	if rf, ok := ret.Get(0).(func(string) (bool, error)); ok {
-		return rf(prefix)
+	if rf, ok := ret.Get(0).(func(context.Context, string) (bool, error)); ok {
+		return rf(ctx, prefix)
 	}
-	if rf, ok := ret.Get(0).(func(string) bool); ok {
-		r0 = rf(prefix)
+	if rf, ok := ret.Get(0).(func(context.Context, string) bool); ok {
+		r0 = rf(ctx, prefix)
 	} else {
 		r0 = ret.Get(0).(bool)
 	}
 
-	if rf, ok := ret.Get(1).(func(string) error); ok {
-		r1 = rf(prefix)
+	if rf, ok := ret.Get(1).(func(context.Context, string) error); ok {
+		r1 = rf(ctx, prefix)
 	} else {
 		r1 = ret.Error(1)
 	}
@@ -249,14 +253,15 @@ type WatchKV_HasPrefix_Call struct {
 }
 
 // HasPrefix is a helper method to define mock.On call
+//   - ctx context.Context
 //   - prefix string
-func (_e *WatchKV_Expecter) HasPrefix(prefix interface{}) *WatchKV_HasPrefix_Call {
-	return &WatchKV_HasPrefix_Call{Call: _e.mock.On("HasPrefix", prefix)}
+func (_e *WatchKV_Expecter) HasPrefix(ctx interface{}, prefix interface{}) *WatchKV_HasPrefix_Call {
+	return &WatchKV_HasPrefix_Call{Call: _e.mock.On("HasPrefix", ctx, prefix)}
 }
 
-func (_c *WatchKV_HasPrefix_Call) Run(run func(prefix string)) *WatchKV_HasPrefix_Call {
+func (_c *WatchKV_HasPrefix_Call) Run(run func(ctx context.Context, prefix string)) *WatchKV_HasPrefix_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string))
+		run(args[0].(context.Context), args[1].(string))
 	})
 	return _c
 }
@@ -266,14 +271,14 @@ func (_c *WatchKV_HasPrefix_Call) Return(_a0 bool, _a1 error) *WatchKV_HasPrefix
 	return _c
 }
 
-func (_c *WatchKV_HasPrefix_Call) RunAndReturn(run func(string) (bool, error)) *WatchKV_HasPrefix_Call {
+func (_c *WatchKV_HasPrefix_Call) RunAndReturn(run func(context.Context, string) (bool, error)) *WatchKV_HasPrefix_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// Load provides a mock function with given fields: key
-func (_m *WatchKV) Load(key string) (string, error) {
-	ret := _m.Called(key)
+// Load provides a mock function with given fields: ctx, key
+func (_m *WatchKV) Load(ctx context.Context, key string) (string, error) {
+	ret := _m.Called(ctx, key)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Load")
@@ -281,17 +286,17 @@ func (_m *WatchKV) Load(key string) (string, error) {
 
 	var r0 string
 	var r1 error
-	if rf, ok := ret.Get(0).(func(string) (string, error)); ok {
-		return rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) (string, error)); ok {
+		return rf(ctx, key)
 	}
-	if rf, ok := ret.Get(0).(func(string) string); ok {
-		r0 = rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) string); ok {
+		r0 = rf(ctx, key)
 	} else {
 		r0 = ret.Get(0).(string)
 	}
 
-	if rf, ok := ret.Get(1).(func(string) error); ok {
-		r1 = rf(key)
+	if rf, ok := ret.Get(1).(func(context.Context, string) error); ok {
+		r1 = rf(ctx, key)
 	} else {
 		r1 = ret.Error(1)
 	}
@@ -305,14 +310,15 @@ type WatchKV_Load_Call struct {
 }
 
 // Load is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
-func (_e *WatchKV_Expecter) Load(key interface{}) *WatchKV_Load_Call {
-	return &WatchKV_Load_Call{Call: _e.mock.On("Load", key)}
+func (_e *WatchKV_Expecter) Load(ctx interface{}, key interface{}) *WatchKV_Load_Call {
+	return &WatchKV_Load_Call{Call: _e.mock.On("Load", ctx, key)}
 }
 
-func (_c *WatchKV_Load_Call) Run(run func(key string)) *WatchKV_Load_Call {
+func (_c *WatchKV_Load_Call) Run(run func(ctx context.Context, key string)) *WatchKV_Load_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string))
+		run(args[0].(context.Context), args[1].(string))
 	})
 	return _c
 }
@@ -322,14 +328,14 @@ func (_c *WatchKV_Load_Call) Return(_a0 string, _a1 error) *WatchKV_Load_Call {
 	return _c
 }
 
-func (_c *WatchKV_Load_Call) RunAndReturn(run func(string) (string, error)) *WatchKV_Load_Call {
+func (_c *WatchKV_Load_Call) RunAndReturn(run func(context.Context, string) (string, error)) *WatchKV_Load_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// LoadWithPrefix provides a mock function with given fields: key
-func (_m *WatchKV) LoadWithPrefix(key string) ([]string, []string, error) {
-	ret := _m.Called(key)
+// LoadWithPrefix provides a mock function with given fields: ctx, key
+func (_m *WatchKV) LoadWithPrefix(ctx context.Context, key string) ([]string, []string, error) {
+	ret := _m.Called(ctx, key)
 
 	if len(ret) == 0 {
 		panic("no return value specified for LoadWithPrefix")
@@ -338,27 +344,27 @@ func (_m *WatchKV) LoadWithPrefix(key string) ([]string, []string, error) {
 	var r0 []string
 	var r1 []string
 	var r2 error
-	if rf, ok := ret.Get(0).(func(string) ([]string, []string, error)); ok {
-		return rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) ([]string, []string, error)); ok {
+		return rf(ctx, key)
 	}
-	if rf, ok := ret.Get(0).(func(string) []string); ok {
-		r0 = rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) []string); ok {
+		r0 = rf(ctx, key)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).([]string)
 		}
 	}
 
-	if rf, ok := ret.Get(1).(func(string) []string); ok {
-		r1 = rf(key)
+	if rf, ok := ret.Get(1).(func(context.Context, string) []string); ok {
+		r1 = rf(ctx, key)
 	} else {
 		if ret.Get(1) != nil {
 			r1 = ret.Get(1).([]string)
 		}
 	}
 
-	if rf, ok := ret.Get(2).(func(string) error); ok {
-		r2 = rf(key)
+	if rf, ok := ret.Get(2).(func(context.Context, string) error); ok {
+		r2 = rf(ctx, key)
 	} else {
 		r2 = ret.Error(2)
 	}
@@ -372,14 +378,15 @@ type WatchKV_LoadWithPrefix_Call struct {
 }
 
 // LoadWithPrefix is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
-func (_e *WatchKV_Expecter) LoadWithPrefix(key interface{}) *WatchKV_LoadWithPrefix_Call {
-	return &WatchKV_LoadWithPrefix_Call{Call: _e.mock.On("LoadWithPrefix", key)}
+func (_e *WatchKV_Expecter) LoadWithPrefix(ctx interface{}, key interface{}) *WatchKV_LoadWithPrefix_Call {
+	return &WatchKV_LoadWithPrefix_Call{Call: _e.mock.On("LoadWithPrefix", ctx, key)}
 }
 
-func (_c *WatchKV_LoadWithPrefix_Call) Run(run func(key string)) *WatchKV_LoadWithPrefix_Call {
+func (_c *WatchKV_LoadWithPrefix_Call) Run(run func(ctx context.Context, key string)) *WatchKV_LoadWithPrefix_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string))
+		run(args[0].(context.Context), args[1].(string))
 	})
 	return _c
 }
@@ -389,14 +396,14 @@ func (_c *WatchKV_LoadWithPrefix_Call) Return(_a0 []string, _a1 []string, _a2 er
 	return _c
 }
 
-func (_c *WatchKV_LoadWithPrefix_Call) RunAndReturn(run func(string) ([]string, []string, error)) *WatchKV_LoadWithPrefix_Call {
+func (_c *WatchKV_LoadWithPrefix_Call) RunAndReturn(run func(context.Context, string) ([]string, []string, error)) *WatchKV_LoadWithPrefix_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// MultiLoad provides a mock function with given fields: keys
-func (_m *WatchKV) MultiLoad(keys []string) ([]string, error) {
-	ret := _m.Called(keys)
+// MultiLoad provides a mock function with given fields: ctx, keys
+func (_m *WatchKV) MultiLoad(ctx context.Context, keys []string) ([]string, error) {
+	ret := _m.Called(ctx, keys)
 
 	if len(ret) == 0 {
 		panic("no return value specified for MultiLoad")
@@ -404,19 +411,19 @@ func (_m *WatchKV) MultiLoad(keys []string) ([]string, error) {
 
 	var r0 []string
 	var r1 error
-	if rf, ok := ret.Get(0).(func([]string) ([]string, error)); ok {
-		return rf(keys)
+	if rf, ok := ret.Get(0).(func(context.Context, []string) ([]string, error)); ok {
+		return rf(ctx, keys)
 	}
-	if rf, ok := ret.Get(0).(func([]string) []string); ok {
-		r0 = rf(keys)
+	if rf, ok := ret.Get(0).(func(context.Context, []string) []string); ok {
+		r0 = rf(ctx, keys)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).([]string)
 		}
 	}
 
-	if rf, ok := ret.Get(1).(func([]string) error); ok {
-		r1 = rf(keys)
+	if rf, ok := ret.Get(1).(func(context.Context, []string) error); ok {
+		r1 = rf(ctx, keys)
 	} else {
 		r1 = ret.Error(1)
 	}
@@ -430,14 +437,15 @@ type WatchKV_MultiLoad_Call struct {
 }
 
 // MultiLoad is a helper method to define mock.On call
+//   - ctx context.Context
 //   - keys []string
-func (_e *WatchKV_Expecter) MultiLoad(keys interface{}) *WatchKV_MultiLoad_Call {
-	return &WatchKV_MultiLoad_Call{Call: _e.mock.On("MultiLoad", keys)}
+func (_e *WatchKV_Expecter) MultiLoad(ctx interface{}, keys interface{}) *WatchKV_MultiLoad_Call {
+	return &WatchKV_MultiLoad_Call{Call: _e.mock.On("MultiLoad", ctx, keys)}
 }
 
-func (_c *WatchKV_MultiLoad_Call) Run(run func(keys []string)) *WatchKV_MultiLoad_Call {
+func (_c *WatchKV_MultiLoad_Call) Run(run func(ctx context.Context, keys []string)) *WatchKV_MultiLoad_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].([]string))
+		run(args[0].(context.Context), args[1].([]string))
 	})
 	return _c
 }
@@ -447,22 +455,22 @@ func (_c *WatchKV_MultiLoad_Call) Return(_a0 []string, _a1 error) *WatchKV_Multi
 	return _c
 }
 
-func (_c *WatchKV_MultiLoad_Call) RunAndReturn(run func([]string) ([]string, error)) *WatchKV_MultiLoad_Call {
+func (_c *WatchKV_MultiLoad_Call) RunAndReturn(run func(context.Context, []string) ([]string, error)) *WatchKV_MultiLoad_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// MultiRemove provides a mock function with given fields: keys
-func (_m *WatchKV) MultiRemove(keys []string) error {
-	ret := _m.Called(keys)
+// MultiRemove provides a mock function with given fields: ctx, keys
+func (_m *WatchKV) MultiRemove(ctx context.Context, keys []string) error {
+	ret := _m.Called(ctx, keys)
 
 	if len(ret) == 0 {
 		panic("no return value specified for MultiRemove")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func([]string) error); ok {
-		r0 = rf(keys)
+	if rf, ok := ret.Get(0).(func(context.Context, []string) error); ok {
+		r0 = rf(ctx, keys)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -476,14 +484,15 @@ type WatchKV_MultiRemove_Call struct {
 }
 
 // MultiRemove is a helper method to define mock.On call
+//   - ctx context.Context
 //   - keys []string
-func (_e *WatchKV_Expecter) MultiRemove(keys interface{}) *WatchKV_MultiRemove_Call {
-	return &WatchKV_MultiRemove_Call{Call: _e.mock.On("MultiRemove", keys)}
+func (_e *WatchKV_Expecter) MultiRemove(ctx interface{}, keys interface{}) *WatchKV_MultiRemove_Call {
+	return &WatchKV_MultiRemove_Call{Call: _e.mock.On("MultiRemove", ctx, keys)}
 }
 
-func (_c *WatchKV_MultiRemove_Call) Run(run func(keys []string)) *WatchKV_MultiRemove_Call {
+func (_c *WatchKV_MultiRemove_Call) Run(run func(ctx context.Context, keys []string)) *WatchKV_MultiRemove_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].([]string))
+		run(args[0].(context.Context), args[1].([]string))
 	})
 	return _c
 }
@@ -493,22 +502,22 @@ func (_c *WatchKV_MultiRemove_Call) Return(_a0 error) *WatchKV_MultiRemove_Call
 	return _c
 }
 
-func (_c *WatchKV_MultiRemove_Call) RunAndReturn(run func([]string) error) *WatchKV_MultiRemove_Call {
+func (_c *WatchKV_MultiRemove_Call) RunAndReturn(run func(context.Context, []string) error) *WatchKV_MultiRemove_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// MultiSave provides a mock function with given fields: kvs
-func (_m *WatchKV) MultiSave(kvs map[string]string) error {
-	ret := _m.Called(kvs)
+// MultiSave provides a mock function with given fields: ctx, kvs
+func (_m *WatchKV) MultiSave(ctx context.Context, kvs map[string]string) error {
+	ret := _m.Called(ctx, kvs)
 
 	if len(ret) == 0 {
 		panic("no return value specified for MultiSave")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(map[string]string) error); ok {
-		r0 = rf(kvs)
+	if rf, ok := ret.Get(0).(func(context.Context, map[string]string) error); ok {
+		r0 = rf(ctx, kvs)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -522,14 +531,15 @@ type WatchKV_MultiSave_Call struct {
 }
 
 // MultiSave is a helper method to define mock.On call
+//   - ctx context.Context
 //   - kvs map[string]string
-func (_e *WatchKV_Expecter) MultiSave(kvs interface{}) *WatchKV_MultiSave_Call {
-	return &WatchKV_MultiSave_Call{Call: _e.mock.On("MultiSave", kvs)}
+func (_e *WatchKV_Expecter) MultiSave(ctx interface{}, kvs interface{}) *WatchKV_MultiSave_Call {
+	return &WatchKV_MultiSave_Call{Call: _e.mock.On("MultiSave", ctx, kvs)}
 }
 
-func (_c *WatchKV_MultiSave_Call) Run(run func(kvs map[string]string)) *WatchKV_MultiSave_Call {
+func (_c *WatchKV_MultiSave_Call) Run(run func(ctx context.Context, kvs map[string]string)) *WatchKV_MultiSave_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(map[string]string))
+		run(args[0].(context.Context), args[1].(map[string]string))
 	})
 	return _c
 }
@@ -539,19 +549,19 @@ func (_c *WatchKV_MultiSave_Call) Return(_a0 error) *WatchKV_MultiSave_Call {
 	return _c
 }
 
-func (_c *WatchKV_MultiSave_Call) RunAndReturn(run func(map[string]string) error) *WatchKV_MultiSave_Call {
+func (_c *WatchKV_MultiSave_Call) RunAndReturn(run func(context.Context, map[string]string) error) *WatchKV_MultiSave_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// MultiSaveAndRemove provides a mock function with given fields: saves, removals, preds
-func (_m *WatchKV) MultiSaveAndRemove(saves map[string]string, removals []string, preds ...predicates.Predicate) error {
+// MultiSaveAndRemove provides a mock function with given fields: ctx, saves, removals, preds
+func (_m *WatchKV) MultiSaveAndRemove(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error {
 	_va := make([]interface{}, len(preds))
 	for _i := range preds {
 		_va[_i] = preds[_i]
 	}
 	var _ca []interface{}
-	_ca = append(_ca, saves, removals)
+	_ca = append(_ca, ctx, saves, removals)
 	_ca = append(_ca, _va...)
 	ret := _m.Called(_ca...)
 
@@ -560,8 +570,8 @@ func (_m *WatchKV) MultiSaveAndRemove(saves map[string]string, removals []string
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(map[string]string, []string, ...predicates.Predicate) error); ok {
-		r0 = rf(saves, removals, preds...)
+	if rf, ok := ret.Get(0).(func(context.Context, map[string]string, []string, ...predicates.Predicate) error); ok {
+		r0 = rf(ctx, saves, removals, preds...)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -575,23 +585,24 @@ type WatchKV_MultiSaveAndRemove_Call struct {
 }
 
 // MultiSaveAndRemove is a helper method to define mock.On call
+//   - ctx context.Context
 //   - saves map[string]string
 //   - removals []string
 //   - preds ...predicates.Predicate
-func (_e *WatchKV_Expecter) MultiSaveAndRemove(saves interface{}, removals interface{}, preds ...interface{}) *WatchKV_MultiSaveAndRemove_Call {
+func (_e *WatchKV_Expecter) MultiSaveAndRemove(ctx interface{}, saves interface{}, removals interface{}, preds ...interface{}) *WatchKV_MultiSaveAndRemove_Call {
 	return &WatchKV_MultiSaveAndRemove_Call{Call: _e.mock.On("MultiSaveAndRemove",
-		append([]interface{}{saves, removals}, preds...)...)}
+		append([]interface{}{ctx, saves, removals}, preds...)...)}
 }
 
-func (_c *WatchKV_MultiSaveAndRemove_Call) Run(run func(saves map[string]string, removals []string, preds ...predicates.Predicate)) *WatchKV_MultiSaveAndRemove_Call {
+func (_c *WatchKV_MultiSaveAndRemove_Call) Run(run func(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate)) *WatchKV_MultiSaveAndRemove_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		variadicArgs := make([]predicates.Predicate, len(args)-2)
-		for i, a := range args[2:] {
+		variadicArgs := make([]predicates.Predicate, len(args)-3)
+		for i, a := range args[3:] {
 			if a != nil {
 				variadicArgs[i] = a.(predicates.Predicate)
 			}
 		}
-		run(args[0].(map[string]string), args[1].([]string), variadicArgs...)
+		run(args[0].(context.Context), args[1].(map[string]string), args[2].([]string), variadicArgs...)
 	})
 	return _c
 }
@@ -601,19 +612,19 @@ func (_c *WatchKV_MultiSaveAndRemove_Call) Return(_a0 error) *WatchKV_MultiSaveA
 	return _c
 }
 
-func (_c *WatchKV_MultiSaveAndRemove_Call) RunAndReturn(run func(map[string]string, []string, ...predicates.Predicate) error) *WatchKV_MultiSaveAndRemove_Call {
+func (_c *WatchKV_MultiSaveAndRemove_Call) RunAndReturn(run func(context.Context, map[string]string, []string, ...predicates.Predicate) error) *WatchKV_MultiSaveAndRemove_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// MultiSaveAndRemoveWithPrefix provides a mock function with given fields: saves, removals, preds
-func (_m *WatchKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string, preds ...predicates.Predicate) error {
+// MultiSaveAndRemoveWithPrefix provides a mock function with given fields: ctx, saves, removals, preds
+func (_m *WatchKV) MultiSaveAndRemoveWithPrefix(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error {
 	_va := make([]interface{}, len(preds))
 	for _i := range preds {
 		_va[_i] = preds[_i]
 	}
 	var _ca []interface{}
-	_ca = append(_ca, saves, removals)
+	_ca = append(_ca, ctx, saves, removals)
 	_ca = append(_ca, _va...)
 	ret := _m.Called(_ca...)
 
@@ -622,8 +633,8 @@ func (_m *WatchKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, removal
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(map[string]string, []string, ...predicates.Predicate) error); ok {
-		r0 = rf(saves, removals, preds...)
+	if rf, ok := ret.Get(0).(func(context.Context, map[string]string, []string, ...predicates.Predicate) error); ok {
+		r0 = rf(ctx, saves, removals, preds...)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -637,23 +648,24 @@ type WatchKV_MultiSaveAndRemoveWithPrefix_Call struct {
 }
 
 // MultiSaveAndRemoveWithPrefix is a helper method to define mock.On call
+//   - ctx context.Context
 //   - saves map[string]string
 //   - removals []string
 //   - preds ...predicates.Predicate
-func (_e *WatchKV_Expecter) MultiSaveAndRemoveWithPrefix(saves interface{}, removals interface{}, preds ...interface{}) *WatchKV_MultiSaveAndRemoveWithPrefix_Call {
+func (_e *WatchKV_Expecter) MultiSaveAndRemoveWithPrefix(ctx interface{}, saves interface{}, removals interface{}, preds ...interface{}) *WatchKV_MultiSaveAndRemoveWithPrefix_Call {
 	return &WatchKV_MultiSaveAndRemoveWithPrefix_Call{Call: _e.mock.On("MultiSaveAndRemoveWithPrefix",
-		append([]interface{}{saves, removals}, preds...)...)}
+		append([]interface{}{ctx, saves, removals}, preds...)...)}
 }
 
-func (_c *WatchKV_MultiSaveAndRemoveWithPrefix_Call) Run(run func(saves map[string]string, removals []string, preds ...predicates.Predicate)) *WatchKV_MultiSaveAndRemoveWithPrefix_Call {
+func (_c *WatchKV_MultiSaveAndRemoveWithPrefix_Call) Run(run func(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate)) *WatchKV_MultiSaveAndRemoveWithPrefix_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		variadicArgs := make([]predicates.Predicate, len(args)-2)
-		for i, a := range args[2:] {
+		variadicArgs := make([]predicates.Predicate, len(args)-3)
+		for i, a := range args[3:] {
 			if a != nil {
 				variadicArgs[i] = a.(predicates.Predicate)
 			}
 		}
-		run(args[0].(map[string]string), args[1].([]string), variadicArgs...)
+		run(args[0].(context.Context), args[1].(map[string]string), args[2].([]string), variadicArgs...)
 	})
 	return _c
 }
@@ -663,22 +675,22 @@ func (_c *WatchKV_MultiSaveAndRemoveWithPrefix_Call) Return(_a0 error) *WatchKV_
 	return _c
 }
 
-func (_c *WatchKV_MultiSaveAndRemoveWithPrefix_Call) RunAndReturn(run func(map[string]string, []string, ...predicates.Predicate) error) *WatchKV_MultiSaveAndRemoveWithPrefix_Call {
+func (_c *WatchKV_MultiSaveAndRemoveWithPrefix_Call) RunAndReturn(run func(context.Context, map[string]string, []string, ...predicates.Predicate) error) *WatchKV_MultiSaveAndRemoveWithPrefix_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// Remove provides a mock function with given fields: key
-func (_m *WatchKV) Remove(key string) error {
-	ret := _m.Called(key)
+// Remove provides a mock function with given fields: ctx, key
+func (_m *WatchKV) Remove(ctx context.Context, key string) error {
+	ret := _m.Called(ctx, key)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Remove")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(string) error); ok {
-		r0 = rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) error); ok {
+		r0 = rf(ctx, key)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -692,14 +704,15 @@ type WatchKV_Remove_Call struct {
 }
 
 // Remove is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
-func (_e *WatchKV_Expecter) Remove(key interface{}) *WatchKV_Remove_Call {
-	return &WatchKV_Remove_Call{Call: _e.mock.On("Remove", key)}
+func (_e *WatchKV_Expecter) Remove(ctx interface{}, key interface{}) *WatchKV_Remove_Call {
+	return &WatchKV_Remove_Call{Call: _e.mock.On("Remove", ctx, key)}
 }
 
-func (_c *WatchKV_Remove_Call) Run(run func(key string)) *WatchKV_Remove_Call {
+func (_c *WatchKV_Remove_Call) Run(run func(ctx context.Context, key string)) *WatchKV_Remove_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string))
+		run(args[0].(context.Context), args[1].(string))
 	})
 	return _c
 }
@@ -709,22 +722,22 @@ func (_c *WatchKV_Remove_Call) Return(_a0 error) *WatchKV_Remove_Call {
 	return _c
 }
 
-func (_c *WatchKV_Remove_Call) RunAndReturn(run func(string) error) *WatchKV_Remove_Call {
+func (_c *WatchKV_Remove_Call) RunAndReturn(run func(context.Context, string) error) *WatchKV_Remove_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// RemoveWithPrefix provides a mock function with given fields: key
-func (_m *WatchKV) RemoveWithPrefix(key string) error {
-	ret := _m.Called(key)
+// RemoveWithPrefix provides a mock function with given fields: ctx, key
+func (_m *WatchKV) RemoveWithPrefix(ctx context.Context, key string) error {
+	ret := _m.Called(ctx, key)
 
 	if len(ret) == 0 {
 		panic("no return value specified for RemoveWithPrefix")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(string) error); ok {
-		r0 = rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) error); ok {
+		r0 = rf(ctx, key)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -738,14 +751,15 @@ type WatchKV_RemoveWithPrefix_Call struct {
 }
 
 // RemoveWithPrefix is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
-func (_e *WatchKV_Expecter) RemoveWithPrefix(key interface{}) *WatchKV_RemoveWithPrefix_Call {
-	return &WatchKV_RemoveWithPrefix_Call{Call: _e.mock.On("RemoveWithPrefix", key)}
+func (_e *WatchKV_Expecter) RemoveWithPrefix(ctx interface{}, key interface{}) *WatchKV_RemoveWithPrefix_Call {
+	return &WatchKV_RemoveWithPrefix_Call{Call: _e.mock.On("RemoveWithPrefix", ctx, key)}
 }
 
-func (_c *WatchKV_RemoveWithPrefix_Call) Run(run func(key string)) *WatchKV_RemoveWithPrefix_Call {
+func (_c *WatchKV_RemoveWithPrefix_Call) Run(run func(ctx context.Context, key string)) *WatchKV_RemoveWithPrefix_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string))
+		run(args[0].(context.Context), args[1].(string))
 	})
 	return _c
 }
@@ -755,22 +769,22 @@ func (_c *WatchKV_RemoveWithPrefix_Call) Return(_a0 error) *WatchKV_RemoveWithPr
 	return _c
 }
 
-func (_c *WatchKV_RemoveWithPrefix_Call) RunAndReturn(run func(string) error) *WatchKV_RemoveWithPrefix_Call {
+func (_c *WatchKV_RemoveWithPrefix_Call) RunAndReturn(run func(context.Context, string) error) *WatchKV_RemoveWithPrefix_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// Save provides a mock function with given fields: key, value
-func (_m *WatchKV) Save(key string, value string) error {
-	ret := _m.Called(key, value)
+// Save provides a mock function with given fields: ctx, key, value
+func (_m *WatchKV) Save(ctx context.Context, key string, value string) error {
+	ret := _m.Called(ctx, key, value)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Save")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(string, string) error); ok {
-		r0 = rf(key, value)
+	if rf, ok := ret.Get(0).(func(context.Context, string, string) error); ok {
+		r0 = rf(ctx, key, value)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -784,15 +798,16 @@ type WatchKV_Save_Call struct {
 }
 
 // Save is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
 //   - value string
-func (_e *WatchKV_Expecter) Save(key interface{}, value interface{}) *WatchKV_Save_Call {
-	return &WatchKV_Save_Call{Call: _e.mock.On("Save", key, value)}
+func (_e *WatchKV_Expecter) Save(ctx interface{}, key interface{}, value interface{}) *WatchKV_Save_Call {
+	return &WatchKV_Save_Call{Call: _e.mock.On("Save", ctx, key, value)}
 }
 
-func (_c *WatchKV_Save_Call) Run(run func(key string, value string)) *WatchKV_Save_Call {
+func (_c *WatchKV_Save_Call) Run(run func(ctx context.Context, key string, value string)) *WatchKV_Save_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string), args[1].(string))
+		run(args[0].(context.Context), args[1].(string), args[2].(string))
 	})
 	return _c
 }
@@ -802,22 +817,22 @@ func (_c *WatchKV_Save_Call) Return(_a0 error) *WatchKV_Save_Call {
 	return _c
 }
 
-func (_c *WatchKV_Save_Call) RunAndReturn(run func(string, string) error) *WatchKV_Save_Call {
+func (_c *WatchKV_Save_Call) RunAndReturn(run func(context.Context, string, string) error) *WatchKV_Save_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// WalkWithPrefix provides a mock function with given fields: prefix, paginationSize, fn
-func (_m *WatchKV) WalkWithPrefix(prefix string, paginationSize int, fn func([]byte, []byte) error) error {
-	ret := _m.Called(prefix, paginationSize, fn)
+// WalkWithPrefix provides a mock function with given fields: ctx, prefix, paginationSize, fn
+func (_m *WatchKV) WalkWithPrefix(ctx context.Context, prefix string, paginationSize int, fn func([]byte, []byte) error) error {
+	ret := _m.Called(ctx, prefix, paginationSize, fn)
 
 	if len(ret) == 0 {
 		panic("no return value specified for WalkWithPrefix")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(string, int, func([]byte, []byte) error) error); ok {
-		r0 = rf(prefix, paginationSize, fn)
+	if rf, ok := ret.Get(0).(func(context.Context, string, int, func([]byte, []byte) error) error); ok {
+		r0 = rf(ctx, prefix, paginationSize, fn)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -831,16 +846,17 @@ type WatchKV_WalkWithPrefix_Call struct {
 }
 
 // WalkWithPrefix is a helper method to define mock.On call
+//   - ctx context.Context
 //   - prefix string
 //   - paginationSize int
 //   - fn func([]byte , []byte) error
-func (_e *WatchKV_Expecter) WalkWithPrefix(prefix interface{}, paginationSize interface{}, fn interface{}) *WatchKV_WalkWithPrefix_Call {
-	return &WatchKV_WalkWithPrefix_Call{Call: _e.mock.On("WalkWithPrefix", prefix, paginationSize, fn)}
+func (_e *WatchKV_Expecter) WalkWithPrefix(ctx interface{}, prefix interface{}, paginationSize interface{}, fn interface{}) *WatchKV_WalkWithPrefix_Call {
+	return &WatchKV_WalkWithPrefix_Call{Call: _e.mock.On("WalkWithPrefix", ctx, prefix, paginationSize, fn)}
 }
 
-func (_c *WatchKV_WalkWithPrefix_Call) Run(run func(prefix string, paginationSize int, fn func([]byte, []byte) error)) *WatchKV_WalkWithPrefix_Call {
+func (_c *WatchKV_WalkWithPrefix_Call) Run(run func(ctx context.Context, prefix string, paginationSize int, fn func([]byte, []byte) error)) *WatchKV_WalkWithPrefix_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string), args[1].(int), args[2].(func([]byte, []byte) error))
+		run(args[0].(context.Context), args[1].(string), args[2].(int), args[3].(func([]byte, []byte) error))
 	})
 	return _c
 }
@@ -850,22 +866,22 @@ func (_c *WatchKV_WalkWithPrefix_Call) Return(_a0 error) *WatchKV_WalkWithPrefix
 	return _c
 }
 
-func (_c *WatchKV_WalkWithPrefix_Call) RunAndReturn(run func(string, int, func([]byte, []byte) error) error) *WatchKV_WalkWithPrefix_Call {
+func (_c *WatchKV_WalkWithPrefix_Call) RunAndReturn(run func(context.Context, string, int, func([]byte, []byte) error) error) *WatchKV_WalkWithPrefix_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// Watch provides a mock function with given fields: key
-func (_m *WatchKV) Watch(key string) clientv3.WatchChan {
-	ret := _m.Called(key)
+// Watch provides a mock function with given fields: ctx, key
+func (_m *WatchKV) Watch(ctx context.Context, key string) clientv3.WatchChan {
+	ret := _m.Called(ctx, key)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Watch")
 	}
 
 	var r0 clientv3.WatchChan
-	if rf, ok := ret.Get(0).(func(string) clientv3.WatchChan); ok {
-		r0 = rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) clientv3.WatchChan); ok {
+		r0 = rf(ctx, key)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).(clientv3.WatchChan)
@@ -881,14 +897,15 @@ type WatchKV_Watch_Call struct {
 }
 
 // Watch is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
-func (_e *WatchKV_Expecter) Watch(key interface{}) *WatchKV_Watch_Call {
-	return &WatchKV_Watch_Call{Call: _e.mock.On("Watch", key)}
+func (_e *WatchKV_Expecter) Watch(ctx interface{}, key interface{}) *WatchKV_Watch_Call {
+	return &WatchKV_Watch_Call{Call: _e.mock.On("Watch", ctx, key)}
 }
 
-func (_c *WatchKV_Watch_Call) Run(run func(key string)) *WatchKV_Watch_Call {
+func (_c *WatchKV_Watch_Call) Run(run func(ctx context.Context, key string)) *WatchKV_Watch_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string))
+		run(args[0].(context.Context), args[1].(string))
 	})
 	return _c
 }
@@ -898,22 +915,22 @@ func (_c *WatchKV_Watch_Call) Return(_a0 clientv3.WatchChan) *WatchKV_Watch_Call
 	return _c
 }
 
-func (_c *WatchKV_Watch_Call) RunAndReturn(run func(string) clientv3.WatchChan) *WatchKV_Watch_Call {
+func (_c *WatchKV_Watch_Call) RunAndReturn(run func(context.Context, string) clientv3.WatchChan) *WatchKV_Watch_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// WatchWithPrefix provides a mock function with given fields: key
-func (_m *WatchKV) WatchWithPrefix(key string) clientv3.WatchChan {
-	ret := _m.Called(key)
+// WatchWithPrefix provides a mock function with given fields: ctx, key
+func (_m *WatchKV) WatchWithPrefix(ctx context.Context, key string) clientv3.WatchChan {
+	ret := _m.Called(ctx, key)
 
 	if len(ret) == 0 {
 		panic("no return value specified for WatchWithPrefix")
 	}
 
 	var r0 clientv3.WatchChan
-	if rf, ok := ret.Get(0).(func(string) clientv3.WatchChan); ok {
-		r0 = rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) clientv3.WatchChan); ok {
+		r0 = rf(ctx, key)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).(clientv3.WatchChan)
@@ -929,14 +946,15 @@ type WatchKV_WatchWithPrefix_Call struct {
 }
 
 // WatchWithPrefix is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
-func (_e *WatchKV_Expecter) WatchWithPrefix(key interface{}) *WatchKV_WatchWithPrefix_Call {
-	return &WatchKV_WatchWithPrefix_Call{Call: _e.mock.On("WatchWithPrefix", key)}
+func (_e *WatchKV_Expecter) WatchWithPrefix(ctx interface{}, key interface{}) *WatchKV_WatchWithPrefix_Call {
+	return &WatchKV_WatchWithPrefix_Call{Call: _e.mock.On("WatchWithPrefix", ctx, key)}
 }
 
-func (_c *WatchKV_WatchWithPrefix_Call) Run(run func(key string)) *WatchKV_WatchWithPrefix_Call {
+func (_c *WatchKV_WatchWithPrefix_Call) Run(run func(ctx context.Context, key string)) *WatchKV_WatchWithPrefix_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string))
+		run(args[0].(context.Context), args[1].(string))
 	})
 	return _c
 }
@@ -946,22 +964,22 @@ func (_c *WatchKV_WatchWithPrefix_Call) Return(_a0 clientv3.WatchChan) *WatchKV_
 	return _c
 }
 
-func (_c *WatchKV_WatchWithPrefix_Call) RunAndReturn(run func(string) clientv3.WatchChan) *WatchKV_WatchWithPrefix_Call {
+func (_c *WatchKV_WatchWithPrefix_Call) RunAndReturn(run func(context.Context, string) clientv3.WatchChan) *WatchKV_WatchWithPrefix_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// WatchWithRevision provides a mock function with given fields: key, revision
-func (_m *WatchKV) WatchWithRevision(key string, revision int64) clientv3.WatchChan {
-	ret := _m.Called(key, revision)
+// WatchWithRevision provides a mock function with given fields: ctx, key, revision
+func (_m *WatchKV) WatchWithRevision(ctx context.Context, key string, revision int64) clientv3.WatchChan {
+	ret := _m.Called(ctx, key, revision)
 
 	if len(ret) == 0 {
 		panic("no return value specified for WatchWithRevision")
 	}
 
 	var r0 clientv3.WatchChan
-	if rf, ok := ret.Get(0).(func(string, int64) clientv3.WatchChan); ok {
-		r0 = rf(key, revision)
+	if rf, ok := ret.Get(0).(func(context.Context, string, int64) clientv3.WatchChan); ok {
+		r0 = rf(ctx, key, revision)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).(clientv3.WatchChan)
@@ -977,15 +995,16 @@ type WatchKV_WatchWithRevision_Call struct {
 }
 
 // WatchWithRevision is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
 //   - revision int64
-func (_e *WatchKV_Expecter) WatchWithRevision(key interface{}, revision interface{}) *WatchKV_WatchWithRevision_Call {
-	return &WatchKV_WatchWithRevision_Call{Call: _e.mock.On("WatchWithRevision", key, revision)}
+func (_e *WatchKV_Expecter) WatchWithRevision(ctx interface{}, key interface{}, revision interface{}) *WatchKV_WatchWithRevision_Call {
+	return &WatchKV_WatchWithRevision_Call{Call: _e.mock.On("WatchWithRevision", ctx, key, revision)}
 }
 
-func (_c *WatchKV_WatchWithRevision_Call) Run(run func(key string, revision int64)) *WatchKV_WatchWithRevision_Call {
+func (_c *WatchKV_WatchWithRevision_Call) Run(run func(ctx context.Context, key string, revision int64)) *WatchKV_WatchWithRevision_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string), args[1].(int64))
+		run(args[0].(context.Context), args[1].(string), args[2].(int64))
 	})
 	return _c
 }
@@ -995,7 +1014,7 @@ func (_c *WatchKV_WatchWithRevision_Call) Return(_a0 clientv3.WatchChan) *WatchK
 	return _c
 }
 
-func (_c *WatchKV_WatchWithRevision_Call) RunAndReturn(run func(string, int64) clientv3.WatchChan) *WatchKV_WatchWithRevision_Call {
+func (_c *WatchKV_WatchWithRevision_Call) RunAndReturn(run func(context.Context, string, int64) clientv3.WatchChan) *WatchKV_WatchWithRevision_Call {
 	_c.Call.Return(run)
 	return _c
 }
diff --git a/internal/kv/tikv/txn_tikv.go b/internal/kv/tikv/txn_tikv.go
index f7fcd7d444a7e..cfb2b2e8374a8 100644
--- a/internal/kv/tikv/txn_tikv.go
+++ b/internal/kv/tikv/txn_tikv.go
@@ -138,7 +138,7 @@ func logWarnOnFailure(err *error, msg string, fields ...zap.Field) {
 }
 
 // Has returns if a key exists.
-func (kv *txnTiKV) Has(key string) (bool, error) {
+func (kv *txnTiKV) Has(ctx context.Context, key string) (bool, error) {
 	start := time.Now()
 	key = path.Join(kv.rootPath, key)
 	ctx, cancel := context.WithTimeout(context.Background(), kv.requestTimeout)
@@ -167,7 +167,7 @@ func rollbackOnFailure(err *error, txn *transaction.KVTxn) {
 }
 
 // HasPrefix returns if a key prefix exists.
-func (kv *txnTiKV) HasPrefix(prefix string) (bool, error) {
+func (kv *txnTiKV) HasPrefix(ctx context.Context, prefix string) (bool, error) {
 	start := time.Now()
 	prefix = path.Join(kv.rootPath, prefix)
 
@@ -197,7 +197,7 @@ func (kv *txnTiKV) HasPrefix(prefix string) (bool, error) {
 }
 
 // Load returns value of the key.
-func (kv *txnTiKV) Load(key string) (string, error) {
+func (kv *txnTiKV) Load(ctx context.Context, key string) (string, error) {
 	start := time.Now()
 	key = path.Join(kv.rootPath, key)
 	ctx, cancel := context.WithTimeout(context.Background(), kv.requestTimeout)
@@ -229,7 +229,7 @@ func batchConvertFromString(prefix string, keys []string) [][]byte {
 }
 
 // MultiLoad gets the values of input keys in a transaction.
-func (kv *txnTiKV) MultiLoad(keys []string) ([]string, error) {
+func (kv *txnTiKV) MultiLoad(ctx context.Context, keys []string) ([]string, error) {
 	start := time.Now()
 	ctx, cancel := context.WithTimeout(context.Background(), kv.requestTimeout)
 	defer cancel()
@@ -271,7 +271,7 @@ func (kv *txnTiKV) MultiLoad(keys []string) ([]string, error) {
 }
 
 // LoadWithPrefix returns all the keys and values for the given key prefix.
-func (kv *txnTiKV) LoadWithPrefix(prefix string) ([]string, []string, error) {
+func (kv *txnTiKV) LoadWithPrefix(ctx context.Context, prefix string) ([]string, []string, error) {
 	start := time.Now()
 	prefix = path.Join(kv.rootPath, prefix)
 
@@ -311,7 +311,7 @@ func (kv *txnTiKV) LoadWithPrefix(prefix string) ([]string, []string, error) {
 }
 
 // Save saves the input key-value pair.
-func (kv *txnTiKV) Save(key, value string) error {
+func (kv *txnTiKV) Save(ctx context.Context, key, value string) error {
 	key = path.Join(kv.rootPath, key)
 	ctx, cancel := context.WithTimeout(context.Background(), kv.requestTimeout)
 	defer cancel()
@@ -324,7 +324,7 @@ func (kv *txnTiKV) Save(key, value string) error {
 }
 
 // MultiSave saves the input key-value pairs in transaction manner.
-func (kv *txnTiKV) MultiSave(kvs map[string]string) error {
+func (kv *txnTiKV) MultiSave(ctx context.Context, kvs map[string]string) error {
 	start := time.Now()
 	ctx, cancel := context.WithTimeout(context.Background(), kv.requestTimeout)
 	defer cancel()
@@ -366,7 +366,7 @@ func (kv *txnTiKV) MultiSave(kvs map[string]string) error {
 }
 
 // Remove removes the input key.
-func (kv *txnTiKV) Remove(key string) error {
+func (kv *txnTiKV) Remove(ctx context.Context, key string) error {
 	key = path.Join(kv.rootPath, key)
 	ctx, cancel := context.WithTimeout(context.Background(), kv.requestTimeout)
 	defer cancel()
@@ -379,7 +379,7 @@ func (kv *txnTiKV) Remove(key string) error {
 }
 
 // MultiRemove removes the input keys in transaction manner.
-func (kv *txnTiKV) MultiRemove(keys []string) error {
+func (kv *txnTiKV) MultiRemove(ctx context.Context, keys []string) error {
 	start := time.Now()
 	ctx, cancel := context.WithTimeout(context.Background(), kv.requestTimeout)
 	defer cancel()
@@ -415,7 +415,7 @@ func (kv *txnTiKV) MultiRemove(keys []string) error {
 }
 
 // RemoveWithPrefix removes the keys for the given prefix.
-func (kv *txnTiKV) RemoveWithPrefix(prefix string) error {
+func (kv *txnTiKV) RemoveWithPrefix(ctx context.Context, prefix string) error {
 	start := time.Now()
 	prefix = path.Join(kv.rootPath, prefix)
 	ctx, cancel := context.WithTimeout(context.Background(), kv.requestTimeout)
@@ -436,7 +436,7 @@ func (kv *txnTiKV) RemoveWithPrefix(prefix string) error {
 }
 
 // MultiSaveAndRemove saves the key-value pairs and removes the keys in a transaction.
-func (kv *txnTiKV) MultiSaveAndRemove(saves map[string]string, removals []string, preds ...predicates.Predicate) error {
+func (kv *txnTiKV) MultiSaveAndRemove(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error {
 	start := time.Now()
 	ctx, cancel := context.WithTimeout(context.Background(), kv.requestTimeout)
 	defer cancel()
@@ -499,7 +499,7 @@ func (kv *txnTiKV) MultiSaveAndRemove(saves map[string]string, removals []string
 }
 
 // MultiSaveAndRemoveWithPrefix saves kv in @saves and removes the keys with given prefix in @removals.
-func (kv *txnTiKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string, preds ...predicates.Predicate) error {
+func (kv *txnTiKV) MultiSaveAndRemoveWithPrefix(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error {
 	start := time.Now()
 	ctx, cancel := context.WithTimeout(context.Background(), kv.requestTimeout)
 	defer cancel()
@@ -585,7 +585,7 @@ func (kv *txnTiKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, removal
 }
 
 // WalkWithPrefix visits each kv with input prefix and apply given fn to it.
-func (kv *txnTiKV) WalkWithPrefix(prefix string, paginationSize int, fn func([]byte, []byte) error) error {
+func (kv *txnTiKV) WalkWithPrefix(ctx context.Context, prefix string, paginationSize int, fn func([]byte, []byte) error) error {
 	start := time.Now()
 	prefix = path.Join(kv.rootPath, prefix)
 
@@ -746,7 +746,7 @@ func (kv *txnTiKV) removeTiKVMeta(ctx context.Context, key string) error {
 	return err
 }
 
-func (kv *txnTiKV) CompareVersionAndSwap(key string, version int64, target string) (bool, error) {
+func (kv *txnTiKV) CompareVersionAndSwap(ctx context.Context, key string, version int64, target string) (bool, error) {
 	err := fmt.Errorf("Unimplemented! CompareVersionAndSwap is under deprecation")
 	logWarnOnFailure(&err, "Unimplemented")
 	return false, err
diff --git a/internal/kv/tikv/txn_tikv_test.go b/internal/kv/tikv/txn_tikv_test.go
index eeff2338ffb9f..fd62f4ffc1249 100644
--- a/internal/kv/tikv/txn_tikv_test.go
+++ b/internal/kv/tikv/txn_tikv_test.go
@@ -37,11 +37,11 @@ func TestTiKVLoad(te *testing.T) {
 	te.Run("kv SaveAndLoad", func(t *testing.T) {
 		rootPath := "/tikv/test/root/saveandload"
 		kv := NewTiKV(txnClient, rootPath)
-		err := kv.RemoveWithPrefix("")
+		err := kv.RemoveWithPrefix(context.TODO(), "")
 		require.NoError(t, err)
 
 		defer kv.Close()
-		defer kv.RemoveWithPrefix("")
+		defer kv.RemoveWithPrefix(context.TODO(), "")
 
 		saveAndLoadTests := []struct {
 			key   string
@@ -55,11 +55,11 @@ func TestTiKVLoad(te *testing.T) {
 
 		for i, test := range saveAndLoadTests {
 			if i < 4 {
-				err = kv.Save(test.key, test.value)
+				err = kv.Save(context.TODO(), test.key, test.value)
 				assert.NoError(t, err)
 			}
 
-			val, err := kv.Load(test.key)
+			val, err := kv.Load(context.TODO(), test.key)
 			assert.NoError(t, err)
 			assert.Equal(t, test.value, val)
 		}
@@ -73,7 +73,7 @@ func TestTiKVLoad(te *testing.T) {
 		}
 
 		for _, test := range invalidLoadTests {
-			val, err := kv.Load(test.invalidKey)
+			val, err := kv.Load(context.TODO(), test.invalidKey)
 			assert.Error(t, err)
 			assert.Zero(t, val)
 		}
@@ -110,7 +110,7 @@ func TestTiKVLoad(te *testing.T) {
 		}
 
 		for _, test := range loadPrefixTests {
-			actualKeys, actualValues, err := kv.LoadWithPrefix(test.prefix)
+			actualKeys, actualValues, err := kv.LoadWithPrefix(context.TODO(), test.prefix)
 			assert.ElementsMatch(t, test.expectedKeys, actualKeys)
 			assert.ElementsMatch(t, test.expectedValues, actualValues)
 			assert.Equal(t, test.expectedError, err)
@@ -127,15 +127,15 @@ func TestTiKVLoad(te *testing.T) {
 		}
 
 		for _, test := range removeTests {
-			err = kv.Remove(test.validKey)
+			err = kv.Remove(context.TODO(), test.validKey)
 			assert.NoError(t, err)
 
-			_, err = kv.Load(test.validKey)
+			_, err = kv.Load(context.TODO(), test.validKey)
 			assert.Error(t, err)
 
-			err = kv.Remove(test.validKey)
+			err = kv.Remove(context.TODO(), test.validKey)
 			assert.NoError(t, err)
-			err = kv.Remove(test.invalidKey)
+			err = kv.Remove(context.TODO(), test.invalidKey)
 			assert.NoError(t, err)
 		}
 	})
@@ -145,7 +145,7 @@ func TestTiKVLoad(te *testing.T) {
 		kv := NewTiKV(txnClient, rootPath)
 
 		defer kv.Close()
-		defer kv.RemoveWithPrefix("")
+		defer kv.RemoveWithPrefix(context.TODO(), "")
 
 		multiSaveTests := map[string]string{
 			"key_1":      "value_1",
@@ -156,10 +156,10 @@ func TestTiKVLoad(te *testing.T) {
 			"_":          "other",
 		}
 
-		err := kv.MultiSave(multiSaveTests)
+		err := kv.MultiSave(context.TODO(), multiSaveTests)
 		assert.NoError(t, err)
 		for k, v := range multiSaveTests {
-			actualV, err := kv.Load(k)
+			actualV, err := kv.Load(context.TODO(), k)
 			assert.NoError(t, err)
 			assert.Equal(t, v, actualV)
 		}
@@ -175,7 +175,7 @@ func TestTiKVLoad(te *testing.T) {
 		}
 
 		for _, test := range multiLoadTests {
-			vs, err := kv.MultiLoad(test.inputKeys)
+			vs, err := kv.MultiLoad(context.TODO(), test.inputKeys)
 			assert.NoError(t, err)
 			assert.Equal(t, test.expectedValues, vs)
 		}
@@ -191,7 +191,7 @@ func TestTiKVLoad(te *testing.T) {
 		}
 
 		for _, test := range invalidMultiLoad {
-			vs, err := kv.MultiLoad(test.invalidKeys)
+			vs, err := kv.MultiLoad(context.TODO(), test.invalidKeys)
 			assert.Error(t, err)
 			assert.Equal(t, test.expectedValues, vs)
 		}
@@ -202,10 +202,10 @@ func TestTiKVLoad(te *testing.T) {
 		}
 
 		for _, k := range removeWithPrefixTests {
-			err = kv.RemoveWithPrefix(k)
+			err = kv.RemoveWithPrefix(context.TODO(), k)
 			assert.NoError(t, err)
 
-			ks, vs, err := kv.LoadWithPrefix(k)
+			ks, vs, err := kv.LoadWithPrefix(context.TODO(), k)
 			assert.Empty(t, ks)
 			assert.Empty(t, vs)
 			assert.NoError(t, err)
@@ -218,10 +218,10 @@ func TestTiKVLoad(te *testing.T) {
 			"_",
 		}
 
-		err = kv.MultiRemove(multiRemoveTests)
+		err = kv.MultiRemove(context.TODO(), multiRemoveTests)
 		assert.NoError(t, err)
 
-		ks, vs, err := kv.LoadWithPrefix("")
+		ks, vs, err := kv.LoadWithPrefix(context.TODO(), "")
 		assert.NoError(t, err)
 		assert.Empty(t, ks)
 		assert.Empty(t, vs)
@@ -238,11 +238,11 @@ func TestTiKVLoad(te *testing.T) {
 			{make(map[string]string), []string{"multikey_2"}},
 		}
 		for _, test := range multiSaveAndRemoveTests {
-			err = kv.MultiSaveAndRemove(test.multiSaves, test.multiRemoves)
+			err = kv.MultiSaveAndRemove(context.TODO(), test.multiSaves, test.multiRemoves)
 			assert.NoError(t, err)
 		}
 
-		ks, vs, err = kv.LoadWithPrefix("")
+		ks, vs, err = kv.LoadWithPrefix(context.TODO(), "")
 		assert.NoError(t, err)
 		assert.Empty(t, ks)
 		assert.Empty(t, vs)
@@ -252,7 +252,7 @@ func TestTiKVLoad(te *testing.T) {
 		rootPath := "/tikv/test/root/multi_remove_with_prefix"
 		kv := NewTiKV(txnClient, rootPath)
 		defer kv.Close()
-		defer kv.RemoveWithPrefix("")
+		defer kv.RemoveWithPrefix(context.TODO(), "")
 
 		prepareTests := map[string]string{
 			"x/abc/1": "1",
@@ -264,7 +264,7 @@ func TestTiKVLoad(te *testing.T) {
 		}
 
 		// MultiSaveAndRemoveWithPrefix
-		err := kv.MultiSave(prepareTests)
+		err := kv.MultiSave(context.TODO(), prepareTests)
 		require.NoError(t, err)
 		multiSaveAndRemoveWithPrefixTests := []struct {
 			multiSave map[string]string
@@ -282,14 +282,14 @@ func TestTiKVLoad(te *testing.T) {
 		}
 
 		for _, test := range multiSaveAndRemoveWithPrefixTests {
-			k, _, err := kv.LoadWithPrefix(test.loadPrefix)
+			k, _, err := kv.LoadWithPrefix(context.TODO(), test.loadPrefix)
 			assert.NoError(t, err)
 			assert.Equal(t, test.lengthBeforeRemove, len(k))
 
-			err = kv.MultiSaveAndRemoveWithPrefix(test.multiSave, test.prefix)
+			err = kv.MultiSaveAndRemoveWithPrefix(context.TODO(), test.multiSave, test.prefix)
 			assert.NoError(t, err)
 
-			k, _, err = kv.LoadWithPrefix(test.loadPrefix)
+			k, _, err = kv.LoadWithPrefix(context.TODO(), test.loadPrefix)
 			assert.NoError(t, err)
 			assert.Equal(t, test.lengthAfterRemove, len(k))
 		}
@@ -306,17 +306,17 @@ func TestTiKVLoad(te *testing.T) {
 		defer func() {
 			beginTxn = tiTxnBegin
 		}()
-		err := kv.Save("key1", "v1")
+		err := kv.Save(context.TODO(), "key1", "v1")
 		assert.Error(t, err)
-		err = kv.MultiSave(map[string]string{"A/100": "v1"})
+		err = kv.MultiSave(context.TODO(), map[string]string{"A/100": "v1"})
 		assert.Error(t, err)
-		err = kv.Remove("key1")
+		err = kv.Remove(context.TODO(), "key1")
 		assert.Error(t, err)
-		err = kv.MultiRemove([]string{"key_1", "key_2"})
+		err = kv.MultiRemove(context.TODO(), []string{"key_1", "key_2"})
 		assert.Error(t, err)
-		err = kv.MultiSaveAndRemove(map[string]string{"key_1": "value_1"}, []string{})
+		err = kv.MultiSaveAndRemove(context.TODO(), map[string]string{"key_1": "value_1"}, []string{})
 		assert.Error(t, err)
-		err = kv.MultiSaveAndRemoveWithPrefix(map[string]string{"y/c": "vvv"}, []string{"/"})
+		err = kv.MultiSaveAndRemoveWithPrefix(context.TODO(), map[string]string{"y/c": "vvv"}, []string{"/"})
 		assert.Error(t, err)
 	})
 
@@ -332,17 +332,17 @@ func TestTiKVLoad(te *testing.T) {
 			commitTxn = tiTxnCommit
 		}()
 		var err error
-		err = kv.Save("key1", "v1")
+		err = kv.Save(context.TODO(), "key1", "v1")
 		assert.Error(t, err)
-		err = kv.MultiSave(map[string]string{"A/100": "v1"})
+		err = kv.MultiSave(context.TODO(), map[string]string{"A/100": "v1"})
 		assert.Error(t, err)
-		err = kv.Remove("key1")
+		err = kv.Remove(context.TODO(), "key1")
 		assert.Error(t, err)
-		err = kv.MultiRemove([]string{"key_1", "key_2"})
+		err = kv.MultiRemove(context.TODO(), []string{"key_1", "key_2"})
 		assert.Error(t, err)
-		err = kv.MultiSaveAndRemove(map[string]string{"key_1": "value_1"}, []string{})
+		err = kv.MultiSaveAndRemove(context.TODO(), map[string]string{"key_1": "value_1"}, []string{})
 		assert.Error(t, err)
-		err = kv.MultiSaveAndRemoveWithPrefix(map[string]string{"y/c": "vvv"}, []string{"/"})
+		err = kv.MultiSaveAndRemoveWithPrefix(context.TODO(), map[string]string{"y/c": "vvv"}, []string{"/"})
 		assert.Error(t, err)
 	})
 }
@@ -352,7 +352,7 @@ func TestWalkWithPagination(t *testing.T) {
 	kv := NewTiKV(txnClient, rootPath)
 
 	defer kv.Close()
-	defer kv.RemoveWithPrefix("")
+	defer kv.RemoveWithPrefix(context.TODO(), "")
 
 	kvs := map[string]string{
 		"A/100":    "v1",
@@ -362,23 +362,23 @@ func TestWalkWithPagination(t *testing.T) {
 		"B/100":    "v5",
 	}
 
-	err := kv.MultiSave(kvs)
+	err := kv.MultiSave(context.TODO(), kvs)
 	assert.NoError(t, err)
 	for k, v := range kvs {
-		actualV, err := kv.Load(k)
+		actualV, err := kv.Load(context.TODO(), k)
 		assert.NoError(t, err)
 		assert.Equal(t, v, actualV)
 	}
 
 	t.Run("apply function error ", func(t *testing.T) {
-		err = kv.WalkWithPrefix("A", 5, func(key []byte, value []byte) error {
+		err = kv.WalkWithPrefix(context.TODO(), "A", 5, func(key []byte, value []byte) error {
 			return errors.New("error")
 		})
 		assert.Error(t, err)
 	})
 
 	t.Run("get with non-exist prefix ", func(t *testing.T) {
-		err = kv.WalkWithPrefix("non-exist-prefix", 5, func(key []byte, value []byte) error {
+		err = kv.WalkWithPrefix(context.TODO(), "non-exist-prefix", 5, func(key []byte, value []byte) error {
 			return nil
 		})
 		assert.NoError(t, err)
@@ -399,7 +399,7 @@ func TestWalkWithPagination(t *testing.T) {
 			ret := make(map[string]string)
 			actualKeys := make([]string, 0)
 
-			err = kv.WalkWithPrefix("A", pagination, func(key []byte, value []byte) error {
+			err = kv.WalkWithPrefix(context.TODO(), "A", pagination, func(key []byte, value []byte) error {
 				k := string(key)
 				k = k[len(rootPath)+1:]
 				ret[k] = string(value)
@@ -434,30 +434,30 @@ func TestElapse(t *testing.T) {
 func TestHas(t *testing.T) {
 	rootPath := "/tikv/test/root/pagination"
 	kv := NewTiKV(txnClient, rootPath)
-	err := kv.RemoveWithPrefix("")
+	err := kv.RemoveWithPrefix(context.TODO(), "")
 	require.NoError(t, err)
 
 	defer kv.Close()
-	defer kv.RemoveWithPrefix("")
+	defer kv.RemoveWithPrefix(context.TODO(), "")
 
-	has, err := kv.Has("key1")
+	has, err := kv.Has(context.TODO(), "key1")
 	assert.NoError(t, err)
 	assert.False(t, has)
 
-	err = kv.Save("key1", "value1")
+	err = kv.Save(context.TODO(), "key1", "value1")
 	assert.NoError(t, err)
 
-	err = kv.Save("key1", EmptyValueString)
+	err = kv.Save(context.TODO(), "key1", EmptyValueString)
 	assert.Error(t, err)
 
-	has, err = kv.Has("key1")
+	has, err = kv.Has(context.TODO(), "key1")
 	assert.NoError(t, err)
 	assert.True(t, has)
 
-	err = kv.Remove("key1")
+	err = kv.Remove(context.TODO(), "key1")
 	assert.NoError(t, err)
 
-	has, err = kv.Has("key1")
+	has, err = kv.Has(context.TODO(), "key1")
 	assert.NoError(t, err)
 	assert.False(t, has)
 }
@@ -465,27 +465,27 @@ func TestHas(t *testing.T) {
 func TestHasPrefix(t *testing.T) {
 	rootPath := "/etcd/test/root/hasprefix"
 	kv := NewTiKV(txnClient, rootPath)
-	err := kv.RemoveWithPrefix("")
+	err := kv.RemoveWithPrefix(context.TODO(), "")
 	require.NoError(t, err)
 
 	defer kv.Close()
-	defer kv.RemoveWithPrefix("")
+	defer kv.RemoveWithPrefix(context.TODO(), "")
 
-	has, err := kv.HasPrefix("key")
+	has, err := kv.HasPrefix(context.TODO(), "key")
 	assert.NoError(t, err)
 	assert.False(t, has)
 
-	err = kv.Save("key1", "value1")
+	err = kv.Save(context.TODO(), "key1", "value1")
 	assert.NoError(t, err)
 
-	has, err = kv.HasPrefix("key")
+	has, err = kv.HasPrefix(context.TODO(), "key")
 	assert.NoError(t, err)
 	assert.True(t, has)
 
-	err = kv.Remove("key1")
+	err = kv.Remove(context.TODO(), "key1")
 	assert.NoError(t, err)
 
-	has, err = kv.HasPrefix("key")
+	has, err = kv.HasPrefix(context.TODO(), "key")
 	assert.NoError(t, err)
 	assert.False(t, has)
 }
@@ -493,32 +493,32 @@ func TestHasPrefix(t *testing.T) {
 func TestEmptyKey(t *testing.T) {
 	rootPath := "/etcd/test/root/loadempty"
 	kv := NewTiKV(txnClient, rootPath)
-	err := kv.RemoveWithPrefix("")
+	err := kv.RemoveWithPrefix(context.TODO(), "")
 	require.NoError(t, err)
 
 	defer kv.Close()
-	defer kv.RemoveWithPrefix("")
+	defer kv.RemoveWithPrefix(context.TODO(), "")
 
-	has, err := kv.HasPrefix("key")
+	has, err := kv.HasPrefix(context.TODO(), "key")
 	assert.NoError(t, err)
 	assert.False(t, has)
 
-	err = kv.Save("key", "")
+	err = kv.Save(context.TODO(), "key", "")
 	assert.NoError(t, err)
 
-	has, err = kv.HasPrefix("key")
+	has, err = kv.HasPrefix(context.TODO(), "key")
 	assert.NoError(t, err)
 	assert.True(t, has)
 
-	val, err := kv.Load("key")
+	val, err := kv.Load(context.TODO(), "key")
 	assert.NoError(t, err)
 	assert.Equal(t, val, "")
 
-	_, vals, err := kv.LoadWithPrefix("key")
+	_, vals, err := kv.LoadWithPrefix(context.TODO(), "key")
 	assert.NoError(t, err)
 	assert.Equal(t, vals[0], "")
 
-	vals, err = kv.MultiLoad([]string{"key"})
+	vals, err = kv.MultiLoad(context.TODO(), []string{"key"})
 	assert.NoError(t, err)
 	assert.Equal(t, vals[0], "")
 
@@ -528,34 +528,34 @@ func TestEmptyKey(t *testing.T) {
 		return nil
 	}
 
-	err = kv.WalkWithPrefix("", 1, nothing)
+	err = kv.WalkWithPrefix(context.TODO(), "", 1, nothing)
 	assert.NoError(t, err)
 	assert.Equal(t, res, "")
 
 	multiSaveTests := map[string]string{
 		"key1": "",
 	}
-	err = kv.MultiSave(multiSaveTests)
+	err = kv.MultiSave(context.TODO(), multiSaveTests)
 	assert.NoError(t, err)
-	val, err = kv.Load("key1")
+	val, err = kv.Load(context.TODO(), "key1")
 	assert.NoError(t, err)
 	assert.Equal(t, val, "")
 
 	multiSaveTests = map[string]string{
 		"key2": "",
 	}
-	err = kv.MultiSaveAndRemove(multiSaveTests, nil)
+	err = kv.MultiSaveAndRemove(context.TODO(), multiSaveTests, nil)
 	assert.NoError(t, err)
-	val, err = kv.Load("key2")
+	val, err = kv.Load(context.TODO(), "key2")
 	assert.NoError(t, err)
 	assert.Equal(t, val, "")
 
 	multiSaveTests = map[string]string{
 		"key3": "",
 	}
-	err = kv.MultiSaveAndRemoveWithPrefix(multiSaveTests, nil)
+	err = kv.MultiSaveAndRemoveWithPrefix(context.TODO(), multiSaveTests, nil)
 	assert.NoError(t, err)
-	val, err = kv.Load("key3")
+	val, err = kv.Load(context.TODO(), "key3")
 	assert.NoError(t, err)
 	assert.Equal(t, val, "")
 }
@@ -563,11 +563,11 @@ func TestEmptyKey(t *testing.T) {
 func TestScanSize(t *testing.T) {
 	scanSize := SnapshotScanSize
 	kv := NewTiKV(txnClient, "/")
-	err := kv.RemoveWithPrefix("")
+	err := kv.RemoveWithPrefix(context.TODO(), "")
 	require.NoError(t, err)
 
 	defer kv.Close()
-	defer kv.RemoveWithPrefix("")
+	defer kv.RemoveWithPrefix(context.TODO(), "")
 
 	// Test total > scansize
 	keyMap := map[string]string{}
@@ -576,32 +576,32 @@ func TestScanSize(t *testing.T) {
 		keyMap[a] = a
 	}
 
-	err = kv.MultiSave(keyMap)
+	err = kv.MultiSave(context.TODO(), keyMap)
 	assert.NoError(t, err)
 
-	keys, _, err := kv.LoadWithPrefix("")
+	keys, _, err := kv.LoadWithPrefix(context.TODO(), "")
 	assert.NoError(t, err)
 	assert.Equal(t, len(keys), scanSize+100)
 
-	err = kv.RemoveWithPrefix("")
+	err = kv.RemoveWithPrefix(context.TODO(), "")
 	require.NoError(t, err)
 }
 
 func TestTiKVUnimplemented(t *testing.T) {
 	kv := NewTiKV(txnClient, "/")
-	err := kv.RemoveWithPrefix("")
+	err := kv.RemoveWithPrefix(context.TODO(), "")
 	require.NoError(t, err)
 
 	defer kv.Close()
-	defer kv.RemoveWithPrefix("")
+	defer kv.RemoveWithPrefix(context.TODO(), "")
 
-	_, err = kv.CompareVersionAndSwap("k", 1, "target")
+	_, err = kv.CompareVersionAndSwap(context.TODO(), "k", 1, "target")
 	assert.Error(t, err)
 }
 
 func TestTxnWithPredicates(t *testing.T) {
 	kv := NewTiKV(txnClient, "/")
-	err := kv.RemoveWithPrefix("")
+	err := kv.RemoveWithPrefix(context.TODO(), "")
 	require.NoError(t, err)
 
 	prepareKV := map[string]string{
@@ -609,7 +609,7 @@ func TestTxnWithPredicates(t *testing.T) {
 		"lease2": "2",
 	}
 
-	err = kv.MultiSave(prepareKV)
+	err = kv.MultiSave(context.TODO(), prepareKV)
 	require.NoError(t, err)
 
 	multiSaveAndRemovePredTests := []struct {
@@ -624,14 +624,14 @@ func TestTxnWithPredicates(t *testing.T) {
 
 	for _, test := range multiSaveAndRemovePredTests {
 		t.Run(test.tag, func(t *testing.T) {
-			err := kv.MultiSaveAndRemove(test.multiSave, nil, test.preds...)
+			err := kv.MultiSaveAndRemove(context.TODO(), test.multiSave, nil, test.preds...)
 			t.Log(err)
 			if test.expectSuccess {
 				assert.NoError(t, err)
 			} else {
 				assert.Error(t, err)
 			}
-			err = kv.MultiSaveAndRemoveWithPrefix(test.multiSave, nil, test.preds...)
+			err = kv.MultiSaveAndRemoveWithPrefix(context.TODO(), test.multiSave, nil, test.preds...)
 			if test.expectSuccess {
 				assert.NoError(t, err)
 			} else {
diff --git a/internal/metastore/kv/datacoord/kv_catalog.go b/internal/metastore/kv/datacoord/kv_catalog.go
index 2233029ee30f6..7bea3819433ca 100644
--- a/internal/metastore/kv/datacoord/kv_catalog.go
+++ b/internal/metastore/kv/datacoord/kv_catalog.go
@@ -69,7 +69,7 @@ func (kc *Catalog) ListSegments(ctx context.Context) ([]*datapb.SegmentInfo, err
 
 	executeFn := func(binlogType storage.BinlogType, result map[typeutil.UniqueID][]*datapb.FieldBinlog) {
 		group.Go(func() error {
-			ret, err := kc.listBinlogs(binlogType)
+			ret, err := kc.listBinlogs(ctx, binlogType)
 			if err != nil {
 				return err
 			}
@@ -85,7 +85,7 @@ func (kc *Catalog) ListSegments(ctx context.Context) ([]*datapb.SegmentInfo, err
 	executeFn(storage.StatsBinlog, statsLogs)
 	executeFn(storage.BM25Binlog, bm25Logs)
 	group.Go(func() error {
-		ret, err := kc.listSegments()
+		ret, err := kc.listSegments(ctx)
 		if err != nil {
 			return err
 		}
@@ -105,7 +105,7 @@ func (kc *Catalog) ListSegments(ctx context.Context) ([]*datapb.SegmentInfo, err
 	return segments, nil
 }
 
-func (kc *Catalog) listSegments() ([]*datapb.SegmentInfo, error) {
+func (kc *Catalog) listSegments(ctx context.Context) ([]*datapb.SegmentInfo, error) {
 	segments := make([]*datapb.SegmentInfo, 0)
 
 	applyFn := func(key []byte, value []byte) error {
@@ -130,7 +130,7 @@ func (kc *Catalog) listSegments() ([]*datapb.SegmentInfo, error) {
 		return nil
 	}
 
-	err := kc.MetaKv.WalkWithPrefix(SegmentPrefix+"/", paginationSize, applyFn)
+	err := kc.MetaKv.WalkWithPrefix(ctx, SegmentPrefix+"/", paginationSize, applyFn)
 	if err != nil {
 		return nil, err
 	}
@@ -163,7 +163,7 @@ func (kc *Catalog) parseBinlogKey(key string, prefixIdx int) (int64, int64, int6
 	return collectionID, partitionID, segmentID, nil
 }
 
-func (kc *Catalog) listBinlogs(binlogType storage.BinlogType) (map[typeutil.UniqueID][]*datapb.FieldBinlog, error) {
+func (kc *Catalog) listBinlogs(ctx context.Context, binlogType storage.BinlogType) (map[typeutil.UniqueID][]*datapb.FieldBinlog, error) {
 	ret := make(map[typeutil.UniqueID][]*datapb.FieldBinlog)
 
 	var err error
@@ -215,7 +215,7 @@ func (kc *Catalog) listBinlogs(binlogType storage.BinlogType) (map[typeutil.Uniq
 		return nil
 	}
 
-	err = kc.MetaKv.WalkWithPrefix(logPathPrefix, paginationSize, applyFn)
+	err = kc.MetaKv.WalkWithPrefix(ctx, logPathPrefix, paginationSize, applyFn)
 	if err != nil {
 		return nil, err
 	}
@@ -263,13 +263,13 @@ func (kc *Catalog) AddSegment(ctx context.Context, segment *datapb.SegmentInfo)
 	if err != nil {
 		return err
 	}
-	return kc.MetaKv.MultiSave(kvs)
+	return kc.MetaKv.MultiSave(ctx, kvs)
 }
 
 // LoadFromSegmentPath loads segment info from persistent storage by given segment path.
 // # TESTING ONLY #
-func (kc *Catalog) LoadFromSegmentPath(colID, partID, segID typeutil.UniqueID) (*datapb.SegmentInfo, error) {
-	v, err := kc.MetaKv.Load(buildSegmentPath(colID, partID, segID))
+func (kc *Catalog) LoadFromSegmentPath(ctx context.Context, colID, partID, segID typeutil.UniqueID) (*datapb.SegmentInfo, error) {
+	v, err := kc.MetaKv.Load(ctx, buildSegmentPath(colID, partID, segID))
 	if err != nil {
 		log.Error("(testing only) failed to load segment info by segment path")
 		return nil, err
@@ -301,7 +301,7 @@ func (kc *Catalog) AlterSegments(ctx context.Context, segments []*datapb.Segment
 		}
 
 		if segment.GetState() == commonpb.SegmentState_Dropped {
-			binlogs, err := kc.handleDroppedSegment(segment)
+			binlogs, err := kc.handleDroppedSegment(ctx, segment)
 			if err != nil {
 				return err
 			}
@@ -327,12 +327,12 @@ func (kc *Catalog) AlterSegments(ctx context.Context, segments []*datapb.Segment
 		maps.Copy(kvs, binlogKvs)
 	}
 
-	return kc.SaveByBatch(kvs)
+	return kc.SaveByBatch(ctx, kvs)
 }
 
-func (kc *Catalog) handleDroppedSegment(segment *datapb.SegmentInfo) (kvs map[string]string, err error) {
+func (kc *Catalog) handleDroppedSegment(ctx context.Context, segment *datapb.SegmentInfo) (kvs map[string]string, err error) {
 	var has bool
-	has, err = kc.hasBinlogPrefix(segment)
+	has, err = kc.hasBinlogPrefix(ctx, segment)
 	if err != nil {
 		return
 	}
@@ -346,9 +346,9 @@ func (kc *Catalog) handleDroppedSegment(segment *datapb.SegmentInfo) (kvs map[st
 	return
 }
 
-func (kc *Catalog) SaveByBatch(kvs map[string]string) error {
+func (kc *Catalog) SaveByBatch(ctx context.Context, kvs map[string]string) error {
 	saveFn := func(partialKvs map[string]string) error {
-		return kc.MetaKv.MultiSave(partialKvs)
+		return kc.MetaKv.MultiSave(ctx, partialKvs)
 	}
 	err := etcd.SaveByBatchWithLimit(kvs, util.MaxEtcdTxnNum, saveFn)
 	if err != nil {
@@ -358,22 +358,22 @@ func (kc *Catalog) SaveByBatch(kvs map[string]string) error {
 	return nil
 }
 
-func (kc *Catalog) hasBinlogPrefix(segment *datapb.SegmentInfo) (bool, error) {
+func (kc *Catalog) hasBinlogPrefix(ctx context.Context, segment *datapb.SegmentInfo) (bool, error) {
 	collectionID, partitionID, segmentID := segment.GetCollectionID(), segment.GetPartitionID(), segment.GetID()
 	prefix := buildFieldBinlogPathPrefix(collectionID, partitionID, segmentID)
-	hasBinlogPrefix, err := kc.MetaKv.HasPrefix(prefix)
+	hasBinlogPrefix, err := kc.MetaKv.HasPrefix(ctx, prefix)
 	if err != nil {
 		return false, err
 	}
 
 	prefix = buildFieldDeltalogPathPrefix(collectionID, partitionID, segmentID)
-	hasDeltaPrefix, err := kc.MetaKv.HasPrefix(prefix)
+	hasDeltaPrefix, err := kc.MetaKv.HasPrefix(ctx, prefix)
 	if err != nil {
 		return false, err
 	}
 
 	prefix = buildFieldStatslogPathPrefix(collectionID, partitionID, segmentID)
-	hasStatsPrefix, err := kc.MetaKv.HasPrefix(prefix)
+	hasStatsPrefix, err := kc.MetaKv.HasPrefix(ctx, prefix)
 	if err != nil {
 		return false, err
 	}
@@ -400,7 +400,7 @@ func (kc *Catalog) SaveDroppedSegmentsInBatch(ctx context.Context, segments []*d
 	}
 
 	saveFn := func(partialKvs map[string]string) error {
-		return kc.MetaKv.MultiSave(partialKvs)
+		return kc.MetaKv.MultiSave(ctx, partialKvs)
 	}
 	if err := etcd.SaveByBatchWithLimit(kvs, util.MaxEtcdTxnNum, saveFn); err != nil {
 		return err
@@ -417,7 +417,7 @@ func (kc *Catalog) DropSegment(ctx context.Context, segment *datapb.SegmentInfo)
 	bm25logPrefix := fmt.Sprintf("%s/%d/%d/%d", SegmentBM25logPathPrefix, segment.GetCollectionID(), segment.GetPartitionID(), segment.GetID())
 
 	keys := []string{segKey, binlogPreix, deltalogPreix, statelogPreix, bm25logPrefix}
-	if err := kc.MetaKv.MultiSaveAndRemoveWithPrefix(nil, keys); err != nil {
+	if err := kc.MetaKv.MultiSaveAndRemoveWithPrefix(ctx, nil, keys); err != nil {
 		return err
 	}
 
@@ -426,7 +426,7 @@ func (kc *Catalog) DropSegment(ctx context.Context, segment *datapb.SegmentInfo)
 
 func (kc *Catalog) MarkChannelAdded(ctx context.Context, channel string) error {
 	key := buildChannelRemovePath(channel)
-	err := kc.MetaKv.Save(key, NonRemoveFlagTomestone)
+	err := kc.MetaKv.Save(ctx, key, NonRemoveFlagTomestone)
 	if err != nil {
 		log.Error("failed to mark channel added", zap.String("channel", channel), zap.Error(err))
 		return err
@@ -437,7 +437,7 @@ func (kc *Catalog) MarkChannelAdded(ctx context.Context, channel string) error {
 
 func (kc *Catalog) MarkChannelDeleted(ctx context.Context, channel string) error {
 	key := buildChannelRemovePath(channel)
-	err := kc.MetaKv.Save(key, RemoveFlagTomestone)
+	err := kc.MetaKv.Save(ctx, key, RemoveFlagTomestone)
 	if err != nil {
 		log.Error("Failed to mark channel dropped", zap.String("channel", channel), zap.Error(err))
 		return err
@@ -448,7 +448,7 @@ func (kc *Catalog) MarkChannelDeleted(ctx context.Context, channel string) error
 
 func (kc *Catalog) ShouldDropChannel(ctx context.Context, channel string) bool {
 	key := buildChannelRemovePath(channel)
-	v, err := kc.MetaKv.Load(key)
+	v, err := kc.MetaKv.Load(ctx, key)
 	if err != nil || v != RemoveFlagTomestone {
 		return false
 	}
@@ -457,7 +457,7 @@ func (kc *Catalog) ShouldDropChannel(ctx context.Context, channel string) bool {
 
 func (kc *Catalog) ChannelExists(ctx context.Context, channel string) bool {
 	key := buildChannelRemovePath(channel)
-	v, err := kc.MetaKv.Load(key)
+	v, err := kc.MetaKv.Load(ctx, key)
 	return err == nil && v == NonRemoveFlagTomestone
 }
 
@@ -465,7 +465,7 @@ func (kc *Catalog) ChannelExists(ctx context.Context, channel string) bool {
 func (kc *Catalog) DropChannel(ctx context.Context, channel string) error {
 	key := buildChannelRemovePath(channel)
 	log.Info("removing channel remove path", zap.String("channel", channel))
-	return kc.MetaKv.Remove(key)
+	return kc.MetaKv.Remove(ctx, key)
 }
 
 func (kc *Catalog) ListChannelCheckpoint(ctx context.Context) (map[string]*msgpb.MsgPosition, error) {
@@ -483,7 +483,7 @@ func (kc *Catalog) ListChannelCheckpoint(ctx context.Context) (map[string]*msgpb
 		return nil
 	}
 
-	err := kc.MetaKv.WalkWithPrefix(ChannelCheckpointPrefix, paginationSize, applyFn)
+	err := kc.MetaKv.WalkWithPrefix(ctx, ChannelCheckpointPrefix, paginationSize, applyFn)
 	if err != nil {
 		return nil, err
 	}
@@ -497,7 +497,7 @@ func (kc *Catalog) SaveChannelCheckpoint(ctx context.Context, vChannel string, p
 	if err != nil {
 		return err
 	}
-	return kc.MetaKv.Save(k, string(v))
+	return kc.MetaKv.Save(ctx, k, string(v))
 }
 
 func (kc *Catalog) SaveChannelCheckpoints(ctx context.Context, positions []*msgpb.MsgPosition) error {
@@ -510,15 +510,15 @@ func (kc *Catalog) SaveChannelCheckpoints(ctx context.Context, positions []*msgp
 		}
 		kvs[k] = string(v)
 	}
-	return kc.SaveByBatch(kvs)
+	return kc.SaveByBatch(ctx, kvs)
 }
 
 func (kc *Catalog) DropChannelCheckpoint(ctx context.Context, vChannel string) error {
 	k := buildChannelCPKey(vChannel)
-	return kc.MetaKv.Remove(k)
+	return kc.MetaKv.Remove(ctx, k)
 }
 
-func (kc *Catalog) getBinlogsWithPrefix(binlogType storage.BinlogType, collectionID, partitionID,
+func (kc *Catalog) getBinlogsWithPrefix(ctx context.Context, binlogType storage.BinlogType, collectionID, partitionID,
 	segmentID typeutil.UniqueID,
 ) ([]string, []string, error) {
 	var binlogPrefix string
@@ -532,7 +532,7 @@ func (kc *Catalog) getBinlogsWithPrefix(binlogType storage.BinlogType, collectio
 	default:
 		return nil, nil, fmt.Errorf("invalid binlog type: %d", binlogType)
 	}
-	keys, values, err := kc.MetaKv.LoadWithPrefix(binlogPrefix)
+	keys, values, err := kc.MetaKv.LoadWithPrefix(ctx, binlogPrefix)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -547,7 +547,7 @@ func (kc *Catalog) CreateIndex(ctx context.Context, index *model.Index) error {
 		return err
 	}
 
-	err = kc.MetaKv.Save(key, string(value))
+	err = kc.MetaKv.Save(ctx, key, string(value))
 	if err != nil {
 		return err
 	}
@@ -568,7 +568,7 @@ func (kc *Catalog) ListIndexes(ctx context.Context) ([]*model.Index, error) {
 		return nil
 	}
 
-	err := kc.MetaKv.WalkWithPrefix(util.FieldIndexPrefix, paginationSize, applyFn)
+	err := kc.MetaKv.WalkWithPrefix(ctx, util.FieldIndexPrefix, paginationSize, applyFn)
 	if err != nil {
 		return nil, err
 	}
@@ -588,7 +588,7 @@ func (kc *Catalog) AlterIndexes(ctx context.Context, indexes []*model.Index) err
 		kvs[key] = string(value)
 		// TODO when we have better txn kv we should make this as a transaction
 		if len(kvs) >= 64 {
-			err = kc.MetaKv.MultiSave(kvs)
+			err = kc.MetaKv.MultiSave(ctx, kvs)
 			if err != nil {
 				return err
 			}
@@ -596,7 +596,7 @@ func (kc *Catalog) AlterIndexes(ctx context.Context, indexes []*model.Index) err
 		}
 	}
 	if len(kvs) != 0 {
-		return kc.MetaKv.MultiSave(kvs)
+		return kc.MetaKv.MultiSave(ctx, kvs)
 	}
 	return nil
 }
@@ -604,7 +604,7 @@ func (kc *Catalog) AlterIndexes(ctx context.Context, indexes []*model.Index) err
 func (kc *Catalog) DropIndex(ctx context.Context, collID typeutil.UniqueID, dropIdxID typeutil.UniqueID) error {
 	key := BuildIndexKey(collID, dropIdxID)
 
-	err := kc.MetaKv.Remove(key)
+	err := kc.MetaKv.Remove(ctx, key)
 	if err != nil {
 		log.Error("drop collection index meta fail", zap.Int64("collectionID", collID),
 			zap.Int64("indexID", dropIdxID), zap.Error(err))
@@ -621,7 +621,7 @@ func (kc *Catalog) CreateSegmentIndex(ctx context.Context, segIdx *model.Segment
 	if err != nil {
 		return err
 	}
-	err = kc.MetaKv.Save(key, string(value))
+	err = kc.MetaKv.Save(ctx, key, string(value))
 	if err != nil {
 		log.Error("failed to save segment index meta in etcd", zap.Int64("buildID", segIdx.BuildID),
 			zap.Int64("segmentID", segIdx.SegmentID), zap.Error(err))
@@ -644,7 +644,7 @@ func (kc *Catalog) ListSegmentIndexes(ctx context.Context) ([]*model.SegmentInde
 		return nil
 	}
 
-	err := kc.MetaKv.WalkWithPrefix(util.SegmentIndexPrefix, paginationSize, applyFn)
+	err := kc.MetaKv.WalkWithPrefix(ctx, util.SegmentIndexPrefix, paginationSize, applyFn)
 	if err != nil {
 		return nil, err
 	}
@@ -662,13 +662,13 @@ func (kc *Catalog) AlterSegmentIndexes(ctx context.Context, segIdxes []*model.Se
 		}
 		kvs[key] = string(value)
 	}
-	return kc.MetaKv.MultiSave(kvs)
+	return kc.MetaKv.MultiSave(ctx, kvs)
 }
 
 func (kc *Catalog) DropSegmentIndex(ctx context.Context, collID, partID, segID, buildID typeutil.UniqueID) error {
 	key := BuildSegmentIndexKey(collID, partID, segID, buildID)
 
-	err := kc.MetaKv.Remove(key)
+	err := kc.MetaKv.Remove(ctx, key)
 	if err != nil {
 		log.Error("drop segment index meta fail", zap.Int64("buildID", buildID), zap.Error(err))
 		return err
@@ -683,7 +683,7 @@ func (kc *Catalog) SaveImportJob(ctx context.Context, job *datapb.ImportJob) err
 	if err != nil {
 		return err
 	}
-	return kc.MetaKv.Save(key, string(value))
+	return kc.MetaKv.Save(ctx, key, string(value))
 }
 
 func (kc *Catalog) ListImportJobs(ctx context.Context) ([]*datapb.ImportJob, error) {
@@ -698,7 +698,7 @@ func (kc *Catalog) ListImportJobs(ctx context.Context) ([]*datapb.ImportJob, err
 		return nil
 	}
 
-	err := kc.MetaKv.WalkWithPrefix(ImportJobPrefix, paginationSize, applyFn)
+	err := kc.MetaKv.WalkWithPrefix(ctx, ImportJobPrefix, paginationSize, applyFn)
 	if err != nil {
 		return nil, err
 	}
@@ -707,7 +707,7 @@ func (kc *Catalog) ListImportJobs(ctx context.Context) ([]*datapb.ImportJob, err
 
 func (kc *Catalog) DropImportJob(ctx context.Context, jobID int64) error {
 	key := buildImportJobKey(jobID)
-	return kc.MetaKv.Remove(key)
+	return kc.MetaKv.Remove(ctx, key)
 }
 
 func (kc *Catalog) SavePreImportTask(ctx context.Context, task *datapb.PreImportTask) error {
@@ -716,7 +716,7 @@ func (kc *Catalog) SavePreImportTask(ctx context.Context, task *datapb.PreImport
 	if err != nil {
 		return err
 	}
-	return kc.MetaKv.Save(key, string(value))
+	return kc.MetaKv.Save(ctx, key, string(value))
 }
 
 func (kc *Catalog) ListPreImportTasks(ctx context.Context) ([]*datapb.PreImportTask, error) {
@@ -732,7 +732,7 @@ func (kc *Catalog) ListPreImportTasks(ctx context.Context) ([]*datapb.PreImportT
 		return nil
 	}
 
-	err := kc.MetaKv.WalkWithPrefix(PreImportTaskPrefix, paginationSize, applyFn)
+	err := kc.MetaKv.WalkWithPrefix(ctx, PreImportTaskPrefix, paginationSize, applyFn)
 	if err != nil {
 		return nil, err
 	}
@@ -741,7 +741,7 @@ func (kc *Catalog) ListPreImportTasks(ctx context.Context) ([]*datapb.PreImportT
 
 func (kc *Catalog) DropPreImportTask(ctx context.Context, taskID int64) error {
 	key := buildPreImportTaskKey(taskID)
-	return kc.MetaKv.Remove(key)
+	return kc.MetaKv.Remove(ctx, key)
 }
 
 func (kc *Catalog) SaveImportTask(ctx context.Context, task *datapb.ImportTaskV2) error {
@@ -750,7 +750,7 @@ func (kc *Catalog) SaveImportTask(ctx context.Context, task *datapb.ImportTaskV2
 	if err != nil {
 		return err
 	}
-	return kc.MetaKv.Save(key, string(value))
+	return kc.MetaKv.Save(ctx, key, string(value))
 }
 
 func (kc *Catalog) ListImportTasks(ctx context.Context) ([]*datapb.ImportTaskV2, error) {
@@ -766,7 +766,7 @@ func (kc *Catalog) ListImportTasks(ctx context.Context) ([]*datapb.ImportTaskV2,
 		return nil
 	}
 
-	err := kc.MetaKv.WalkWithPrefix(ImportTaskPrefix, paginationSize, applyFn)
+	err := kc.MetaKv.WalkWithPrefix(ctx, ImportTaskPrefix, paginationSize, applyFn)
 	if err != nil {
 		return nil, err
 	}
@@ -775,7 +775,7 @@ func (kc *Catalog) ListImportTasks(ctx context.Context) ([]*datapb.ImportTaskV2,
 
 func (kc *Catalog) DropImportTask(ctx context.Context, taskID int64) error {
 	key := buildImportTaskKey(taskID)
-	return kc.MetaKv.Remove(key)
+	return kc.MetaKv.Remove(ctx, key)
 }
 
 // GcConfirm returns true if related collection/partition is not found.
@@ -785,7 +785,7 @@ func (kc *Catalog) GcConfirm(ctx context.Context, collectionID, partitionID type
 	if partitionID != common.AllPartitionsID {
 		prefix = buildPartitionPrefix(collectionID, partitionID)
 	}
-	keys, values, err := kc.MetaKv.LoadWithPrefix(prefix)
+	keys, values, err := kc.MetaKv.LoadWithPrefix(ctx, prefix)
 	if err != nil {
 		// error case can be regarded as not finished.
 		return false
@@ -806,7 +806,7 @@ func (kc *Catalog) ListCompactionTask(ctx context.Context) ([]*datapb.Compaction
 		return nil
 	}
 
-	err := kc.MetaKv.WalkWithPrefix(CompactionTaskPrefix, paginationSize, applyFn)
+	err := kc.MetaKv.WalkWithPrefix(ctx, CompactionTaskPrefix, paginationSize, applyFn)
 	if err != nil {
 		return nil, err
 	}
@@ -824,12 +824,12 @@ func (kc *Catalog) SaveCompactionTask(ctx context.Context, coll *datapb.Compacti
 	}
 	kvs := make(map[string]string)
 	kvs[k] = v
-	return kc.SaveByBatch(kvs)
+	return kc.SaveByBatch(ctx, kvs)
 }
 
 func (kc *Catalog) DropCompactionTask(ctx context.Context, task *datapb.CompactionTask) error {
 	key := buildCompactionTaskPath(task)
-	return kc.MetaKv.Remove(key)
+	return kc.MetaKv.Remove(ctx, key)
 }
 
 func (kc *Catalog) ListAnalyzeTasks(ctx context.Context) ([]*indexpb.AnalyzeTask, error) {
@@ -845,7 +845,7 @@ func (kc *Catalog) ListAnalyzeTasks(ctx context.Context) ([]*indexpb.AnalyzeTask
 		return nil
 	}
 
-	err := kc.MetaKv.WalkWithPrefix(AnalyzeTaskPrefix, paginationSize, applyFn)
+	err := kc.MetaKv.WalkWithPrefix(ctx, AnalyzeTaskPrefix, paginationSize, applyFn)
 	if err != nil {
 		return nil, err
 	}
@@ -860,7 +860,7 @@ func (kc *Catalog) SaveAnalyzeTask(ctx context.Context, task *indexpb.AnalyzeTas
 		return err
 	}
 
-	err = kc.MetaKv.Save(key, string(value))
+	err = kc.MetaKv.Save(ctx, key, string(value))
 	if err != nil {
 		return err
 	}
@@ -869,7 +869,7 @@ func (kc *Catalog) SaveAnalyzeTask(ctx context.Context, task *indexpb.AnalyzeTas
 
 func (kc *Catalog) DropAnalyzeTask(ctx context.Context, taskID typeutil.UniqueID) error {
 	key := buildAnalyzeTaskKey(taskID)
-	return kc.MetaKv.Remove(key)
+	return kc.MetaKv.Remove(ctx, key)
 }
 
 func (kc *Catalog) ListPartitionStatsInfos(ctx context.Context) ([]*datapb.PartitionStatsInfo, error) {
@@ -885,7 +885,7 @@ func (kc *Catalog) ListPartitionStatsInfos(ctx context.Context) ([]*datapb.Parti
 		return nil
 	}
 
-	err := kc.MetaKv.WalkWithPrefix(PartitionStatsInfoPrefix, paginationSize, applyFn)
+	err := kc.MetaKv.WalkWithPrefix(ctx, PartitionStatsInfoPrefix, paginationSize, applyFn)
 	if err != nil {
 		return nil, err
 	}
@@ -903,23 +903,23 @@ func (kc *Catalog) SavePartitionStatsInfo(ctx context.Context, coll *datapb.Part
 	}
 	kvs := make(map[string]string)
 	kvs[k] = v
-	return kc.SaveByBatch(kvs)
+	return kc.SaveByBatch(ctx, kvs)
 }
 
 func (kc *Catalog) DropPartitionStatsInfo(ctx context.Context, info *datapb.PartitionStatsInfo) error {
 	key := buildPartitionStatsInfoPath(info)
-	return kc.MetaKv.Remove(key)
+	return kc.MetaKv.Remove(ctx, key)
 }
 
 func (kc *Catalog) SaveCurrentPartitionStatsVersion(ctx context.Context, collID, partID int64, vChannel string, currentVersion int64) error {
 	key := buildCurrentPartitionStatsVersionPath(collID, partID, vChannel)
 	value := strconv.FormatInt(currentVersion, 10)
-	return kc.MetaKv.Save(key, value)
+	return kc.MetaKv.Save(ctx, key, value)
 }
 
 func (kc *Catalog) GetCurrentPartitionStatsVersion(ctx context.Context, collID, partID int64, vChannel string) (int64, error) {
 	key := buildCurrentPartitionStatsVersionPath(collID, partID, vChannel)
-	valueStr, err := kc.MetaKv.Load(key)
+	valueStr, err := kc.MetaKv.Load(ctx, key)
 	if err != nil {
 		if errors.Is(err, merr.ErrIoKeyNotFound) {
 			return 0, nil
@@ -932,7 +932,7 @@ func (kc *Catalog) GetCurrentPartitionStatsVersion(ctx context.Context, collID,
 
 func (kc *Catalog) DropCurrentPartitionStatsVersion(ctx context.Context, collID, partID int64, vChannel string) error {
 	key := buildCurrentPartitionStatsVersionPath(collID, partID, vChannel)
-	return kc.MetaKv.Remove(key)
+	return kc.MetaKv.Remove(ctx, key)
 }
 
 func (kc *Catalog) ListStatsTasks(ctx context.Context) ([]*indexpb.StatsTask, error) {
@@ -948,7 +948,7 @@ func (kc *Catalog) ListStatsTasks(ctx context.Context) ([]*indexpb.StatsTask, er
 		return nil
 	}
 
-	err := kc.MetaKv.WalkWithPrefix(StatsTaskPrefix, paginationSize, applyFn)
+	err := kc.MetaKv.WalkWithPrefix(ctx, StatsTaskPrefix, paginationSize, applyFn)
 	if err != nil {
 		return nil, err
 	}
@@ -962,7 +962,7 @@ func (kc *Catalog) SaveStatsTask(ctx context.Context, task *indexpb.StatsTask) e
 		return err
 	}
 
-	err = kc.MetaKv.Save(key, string(value))
+	err = kc.MetaKv.Save(ctx, key, string(value))
 	if err != nil {
 		return err
 	}
@@ -971,5 +971,5 @@ func (kc *Catalog) SaveStatsTask(ctx context.Context, task *indexpb.StatsTask) e
 
 func (kc *Catalog) DropStatsTask(ctx context.Context, taskID typeutil.UniqueID) error {
 	key := buildStatsTaskKey(taskID)
-	return kc.MetaKv.Remove(key)
+	return kc.MetaKv.Remove(ctx, key)
 }
diff --git a/internal/metastore/kv/datacoord/kv_catalog_test.go b/internal/metastore/kv/datacoord/kv_catalog_test.go
index baeb8bedd385c..992aeb0830aed 100644
--- a/internal/metastore/kv/datacoord/kv_catalog_test.go
+++ b/internal/metastore/kv/datacoord/kv_catalog_test.go
@@ -148,7 +148,7 @@ var (
 func Test_ListSegments(t *testing.T) {
 	t.Run("load failed", func(t *testing.T) {
 		metakv := mocks.NewMetaKv(t)
-		metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(errors.New("error"))
+		metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(errors.New("error"))
 
 		catalog := NewCatalog(metakv, rootPath, "")
 		ret, err := catalog.ListSegments(context.TODO())
@@ -190,7 +190,7 @@ func Test_ListSegments(t *testing.T) {
 		assert.NoError(t, err)
 
 		metakv := mocks.NewMetaKv(t)
-		metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(s string, i int, f func([]byte, []byte) error) error {
+		metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, s string, i int, f func([]byte, []byte) error) error {
 			if strings.HasPrefix(k5, s) {
 				return f([]byte(k5), segBytes)
 			}
@@ -207,7 +207,7 @@ func Test_ListSegments(t *testing.T) {
 
 	t.Run("test compatibility with stats task", func(t *testing.T) {
 		metakv := mocks.NewMetaKv(t)
-		metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(s string, i int, f func([]byte, []byte) error) error {
+		metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, s string, i int, f func([]byte, []byte) error) error {
 			if strings.HasPrefix(k9, path.Join(s)) {
 				return f([]byte(k9), nil)
 			}
@@ -225,7 +225,7 @@ func Test_ListSegments(t *testing.T) {
 		var savedKvs map[string]string
 
 		metakv := mocks.NewMetaKv(t)
-		metakv.EXPECT().MultiSave(mock.Anything).RunAndReturn(func(m map[string]string) error {
+		metakv.EXPECT().MultiSave(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, m map[string]string) error {
 			savedKvs = m
 			return nil
 		})
@@ -234,7 +234,7 @@ func Test_ListSegments(t *testing.T) {
 		err := catalog.AddSegment(context.TODO(), segment1)
 		assert.NoError(t, err)
 
-		metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(s string, i int, f func([]byte, []byte) error) error {
+		metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, s string, i int, f func([]byte, []byte) error) error {
 			if strings.HasPrefix(k5, s) {
 				return f([]byte(k5), []byte(savedKvs[k5]))
 			}
@@ -267,7 +267,7 @@ func Test_ListSegments(t *testing.T) {
 func Test_AddSegments(t *testing.T) {
 	t.Run("generate binlog kvs failed", func(t *testing.T) {
 		metakv := mocks.NewMetaKv(t)
-		metakv.EXPECT().MultiSave(mock.Anything).Return(errors.New("error")).Maybe()
+		metakv.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(errors.New("error")).Maybe()
 
 		catalog := NewCatalog(metakv, rootPath, "")
 
@@ -277,7 +277,7 @@ func Test_AddSegments(t *testing.T) {
 
 	t.Run("save error", func(t *testing.T) {
 		metakv := mocks.NewMetaKv(t)
-		metakv.EXPECT().MultiSave(mock.Anything).Return(errors.New("error"))
+		metakv.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(errors.New("error"))
 
 		catalog := NewCatalog(metakv, rootPath, "")
 		err := catalog.AddSegment(context.TODO(), segment1)
@@ -287,11 +287,11 @@ func Test_AddSegments(t *testing.T) {
 	t.Run("save successfully", func(t *testing.T) {
 		savedKvs := make(map[string]string)
 		metakv := mocks.NewMetaKv(t)
-		metakv.EXPECT().MultiSave(mock.Anything).RunAndReturn(func(m map[string]string) error {
+		metakv.EXPECT().MultiSave(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, m map[string]string) error {
 			savedKvs = m
 			return nil
 		})
-		metakv.EXPECT().Load(mock.Anything).RunAndReturn(func(s string) (string, error) {
+		metakv.EXPECT().Load(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, s string) (string, error) {
 			if v, ok := savedKvs[s]; ok {
 				return v, nil
 			}
@@ -301,7 +301,7 @@ func Test_AddSegments(t *testing.T) {
 		catalog := NewCatalog(metakv, rootPath, "")
 		err := catalog.AddSegment(context.TODO(), segment1)
 		assert.NoError(t, err)
-		adjustedSeg, err := catalog.LoadFromSegmentPath(segment1.CollectionID, segment1.PartitionID, segment1.ID)
+		adjustedSeg, err := catalog.LoadFromSegmentPath(context.TODO(), segment1.CollectionID, segment1.PartitionID, segment1.ID)
 		assert.NoError(t, err)
 		// Check that num of rows is corrected from 100 to 5.
 		assert.Equal(t, int64(100), segment1.GetNumOfRows())
@@ -364,7 +364,7 @@ func Test_AddSegments(t *testing.T) {
 func Test_AlterSegments(t *testing.T) {
 	t.Run("generate binlog kvs failed", func(t *testing.T) {
 		metakv := mocks.NewMetaKv(t)
-		metakv.EXPECT().MultiSave(mock.Anything).Return(errors.New("error")).Maybe()
+		metakv.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(errors.New("error")).Maybe()
 
 		catalog := NewCatalog(metakv, rootPath, "")
 		err := catalog.AlterSegments(context.TODO(), []*datapb.SegmentInfo{invalidSegment}, metastore.BinlogsIncrement{
@@ -375,7 +375,7 @@ func Test_AlterSegments(t *testing.T) {
 
 	t.Run("save error", func(t *testing.T) {
 		metakv := mocks.NewMetaKv(t)
-		metakv.EXPECT().MultiSave(mock.Anything).Return(errors.New("error"))
+		metakv.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(errors.New("error"))
 
 		catalog := NewCatalog(metakv, rootPath, "")
 		err := catalog.AlterSegments(context.TODO(), []*datapb.SegmentInfo{segment1})
@@ -386,7 +386,7 @@ func Test_AlterSegments(t *testing.T) {
 	t.Run("save successfully", func(t *testing.T) {
 		var savedKvs map[string]string
 		metakv := mocks.NewMetaKv(t)
-		metakv.EXPECT().MultiSave(mock.Anything).RunAndReturn(func(m map[string]string) error {
+		metakv.EXPECT().MultiSave(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, m map[string]string) error {
 			savedKvs = m
 			return nil
 		})
@@ -410,12 +410,12 @@ func Test_AlterSegments(t *testing.T) {
 		savedKvs := make(map[string]string)
 		opGroupCount := 0
 		metakv := mocks.NewMetaKv(t)
-		metakv.EXPECT().MultiSave(mock.Anything).RunAndReturn(func(m map[string]string) error {
+		metakv.EXPECT().MultiSave(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, m map[string]string) error {
 			maps.Copy(savedKvs, m)
 			opGroupCount++
 			return nil
 		})
-		metakv.EXPECT().Load(mock.Anything).RunAndReturn(func(s string) (string, error) {
+		metakv.EXPECT().Load(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, s string) (string, error) {
 			if v, ok := savedKvs[s]; ok {
 				return v, nil
 			}
@@ -458,7 +458,7 @@ func Test_AlterSegments(t *testing.T) {
 		assert.Equal(t, 255+3, len(savedKvs))
 		assert.Equal(t, 3, opGroupCount)
 
-		adjustedSeg, err := catalog.LoadFromSegmentPath(segmentXL.CollectionID, segmentXL.PartitionID, segmentXL.ID)
+		adjustedSeg, err := catalog.LoadFromSegmentPath(context.TODO(), segmentXL.CollectionID, segmentXL.PartitionID, segmentXL.ID)
 		assert.NoError(t, err)
 		// Check that num of rows is corrected from 100 to 1275.
 		assert.Equal(t, int64(100), segmentXL.GetNumOfRows())
@@ -507,7 +507,7 @@ func Test_AlterSegments(t *testing.T) {
 func Test_DropSegment(t *testing.T) {
 	t.Run("remove failed", func(t *testing.T) {
 		metakv := mocks.NewMetaKv(t)
-		metakv.EXPECT().MultiSaveAndRemoveWithPrefix(mock.Anything, mock.Anything).Return(errors.New("error"))
+		metakv.EXPECT().MultiSaveAndRemoveWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(errors.New("error"))
 
 		catalog := NewCatalog(metakv, rootPath, "")
 		err := catalog.DropSegment(context.TODO(), segment1)
@@ -517,7 +517,7 @@ func Test_DropSegment(t *testing.T) {
 	t.Run("remove successfully", func(t *testing.T) {
 		removedKvs := make(map[string]struct{}, 0)
 		metakv := mocks.NewMetaKv(t)
-		metakv.EXPECT().MultiSaveAndRemoveWithPrefix(mock.Anything, mock.Anything).RunAndReturn(func(m map[string]string, s []string, p ...predicates.Predicate) error {
+		metakv.EXPECT().MultiSaveAndRemoveWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, m map[string]string, s []string, p ...predicates.Predicate) error {
 			for _, key := range s {
 				removedKvs[key] = struct{}{}
 			}
@@ -543,7 +543,7 @@ func Test_DropSegment(t *testing.T) {
 
 func Test_SaveDroppedSegmentsInBatch_SaveError(t *testing.T) {
 	txn := mocks.NewMetaKv(t)
-	txn.EXPECT().MultiSave(mock.Anything).Return(errors.New("mock error"))
+	txn.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(errors.New("mock error"))
 
 	catalog := NewCatalog(txn, rootPath, "")
 	segments := []*datapb.SegmentInfo{
@@ -563,8 +563,8 @@ func Test_SaveDroppedSegmentsInBatch_MultiSave(t *testing.T) {
 	)
 	txn := mocks.NewMetaKv(t)
 	txn.EXPECT().
-		MultiSave(mock.Anything).
-		Run(func(kvs map[string]string) {
+		MultiSave(mock.Anything, mock.Anything).
+		Run(func(ctx context.Context, kvs map[string]string) {
 			count++
 			kvSize += len(kvs)
 		}).
@@ -632,12 +632,12 @@ func TestChannelCP(t *testing.T) {
 
 	t.Run("ListChannelCheckpoint", func(t *testing.T) {
 		txn := mocks.NewMetaKv(t)
-		txn.EXPECT().Save(mock.Anything, mock.Anything).Return(nil)
+		txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(nil)
 		catalog := NewCatalog(txn, rootPath, "")
 		err := catalog.SaveChannelCheckpoint(context.TODO(), mockVChannel, pos)
 		assert.NoError(t, err)
 
-		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error {
+		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error {
 			return f([]byte(k), v)
 		})
 		res, err := catalog.ListChannelCheckpoint(context.TODO())
@@ -648,14 +648,14 @@ func TestChannelCP(t *testing.T) {
 	t.Run("ListChannelCheckpoint failed", func(t *testing.T) {
 		txn := mocks.NewMetaKv(t)
 		catalog := NewCatalog(txn, rootPath, "")
-		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(errors.New("mock error"))
+		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(errors.New("mock error"))
 		_, err = catalog.ListChannelCheckpoint(context.TODO())
 		assert.Error(t, err)
 	})
 
 	t.Run("SaveChannelCheckpoint", func(t *testing.T) {
 		txn := mocks.NewMetaKv(t)
-		txn.EXPECT().Save(mock.Anything, mock.Anything).Return(nil)
+		txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(nil)
 		catalog := NewCatalog(txn, rootPath, "")
 		err := catalog.SaveChannelCheckpoint(context.TODO(), mockVChannel, pos)
 		assert.NoError(t, err)
@@ -664,14 +664,14 @@ func TestChannelCP(t *testing.T) {
 	t.Run("SaveChannelCheckpoint failed", func(t *testing.T) {
 		txn := mocks.NewMetaKv(t)
 		catalog := NewCatalog(txn, rootPath, "")
-		txn.EXPECT().Save(mock.Anything, mock.Anything).Return(errors.New("mock error"))
+		txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(errors.New("mock error"))
 		err = catalog.SaveChannelCheckpoint(context.TODO(), mockVChannel, &msgpb.MsgPosition{})
 		assert.Error(t, err)
 	})
 
 	t.Run("SaveChannelCheckpoints", func(t *testing.T) {
 		txn := mocks.NewMetaKv(t)
-		txn.EXPECT().MultiSave(mock.Anything).Return(nil)
+		txn.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(nil)
 		catalog := NewCatalog(txn, rootPath, "")
 		err := catalog.SaveChannelCheckpoints(context.TODO(), []*msgpb.MsgPosition{pos})
 		assert.NoError(t, err)
@@ -680,20 +680,20 @@ func TestChannelCP(t *testing.T) {
 	t.Run("SaveChannelCheckpoints failed", func(t *testing.T) {
 		txn := mocks.NewMetaKv(t)
 		catalog := NewCatalog(txn, rootPath, "")
-		txn.EXPECT().MultiSave(mock.Anything).Return(errors.New("mock error"))
+		txn.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(errors.New("mock error"))
 		err = catalog.SaveChannelCheckpoints(context.TODO(), []*msgpb.MsgPosition{pos})
 		assert.Error(t, err)
 	})
 
 	t.Run("DropChannelCheckpoint", func(t *testing.T) {
 		txn := mocks.NewMetaKv(t)
-		txn.EXPECT().Save(mock.Anything, mock.Anything).Return(nil)
+		txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(nil)
 		catalog := NewCatalog(txn, rootPath, "")
 		err := catalog.SaveChannelCheckpoint(context.TODO(), mockVChannel, pos)
 		assert.NoError(t, err)
 
-		txn.EXPECT().Remove(mock.Anything).Return(nil)
-		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(nil)
+		txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(nil)
+		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
 		err = catalog.DropChannelCheckpoint(context.TODO(), mockVChannel)
 		assert.NoError(t, err)
 		res, err := catalog.ListChannelCheckpoint(context.TODO())
@@ -704,7 +704,7 @@ func TestChannelCP(t *testing.T) {
 	t.Run("DropChannelCheckpoint failed", func(t *testing.T) {
 		txn := mocks.NewMetaKv(t)
 		catalog := NewCatalog(txn, rootPath, "")
-		txn.EXPECT().Remove(mock.Anything).Return(errors.New("mock error"))
+		txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(errors.New("mock error"))
 		err = catalog.DropChannelCheckpoint(context.TODO(), mockVChannel)
 		assert.Error(t, err)
 	})
@@ -713,7 +713,7 @@ func TestChannelCP(t *testing.T) {
 func Test_MarkChannelDeleted_SaveError(t *testing.T) {
 	txn := mocks.NewMetaKv(t)
 	txn.EXPECT().
-		Save(mock.Anything, mock.Anything).
+		Save(mock.Anything, mock.Anything, mock.Anything).
 		Return(errors.New("mock error"))
 
 	catalog := NewCatalog(txn, rootPath, "")
@@ -724,7 +724,7 @@ func Test_MarkChannelDeleted_SaveError(t *testing.T) {
 func Test_MarkChannelAdded_SaveError(t *testing.T) {
 	txn := mocks.NewMetaKv(t)
 	txn.EXPECT().
-		Save(mock.Anything, mock.Anything).
+		Save(mock.Anything, mock.Anything, mock.Anything).
 		Return(errors.New("mock error"))
 
 	catalog := NewCatalog(txn, rootPath, "")
@@ -735,7 +735,7 @@ func Test_MarkChannelAdded_SaveError(t *testing.T) {
 func Test_ChannelExists_SaveError(t *testing.T) {
 	txn := mocks.NewMetaKv(t)
 	txn.EXPECT().
-		Load(mock.Anything).
+		Load(mock.Anything, mock.Anything).
 		Return("", errors.New("mock error"))
 
 	catalog := NewCatalog(txn, rootPath, "")
@@ -854,7 +854,7 @@ func verifySavedKvsForDroppedSegment(t *testing.T, savedKvs map[string]string) {
 func TestCatalog_CreateIndex(t *testing.T) {
 	t.Run("success", func(t *testing.T) {
 		metakv := mocks.NewMetaKv(t)
-		metakv.EXPECT().Save(mock.Anything, mock.Anything).Return(nil)
+		metakv.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(nil)
 
 		catalog := &Catalog{
 			MetaKv: metakv,
@@ -866,7 +866,7 @@ func TestCatalog_CreateIndex(t *testing.T) {
 
 	t.Run("failed", func(t *testing.T) {
 		metakv := mocks.NewMetaKv(t)
-		metakv.EXPECT().Save(mock.Anything, mock.Anything).Return(errors.New("error"))
+		metakv.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(errors.New("error"))
 
 		catalog := &Catalog{
 			MetaKv: metakv,
@@ -880,7 +880,7 @@ func TestCatalog_CreateIndex(t *testing.T) {
 func TestCatalog_ListIndexes(t *testing.T) {
 	t.Run("success", func(t *testing.T) {
 		metakv := mocks.NewMetaKv(t)
-		metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error {
+		metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error {
 			i := &indexpb.FieldIndex{
 				IndexInfo: &indexpb.IndexInfo{
 					CollectionID: 0,
@@ -908,7 +908,7 @@ func TestCatalog_ListIndexes(t *testing.T) {
 
 	t.Run("failed", func(t *testing.T) {
 		txn := mocks.NewMetaKv(t)
-		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(errors.New("error"))
+		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(errors.New("error"))
 		catalog := &Catalog{
 			MetaKv: txn,
 		}
@@ -918,7 +918,7 @@ func TestCatalog_ListIndexes(t *testing.T) {
 
 	t.Run("unmarshal failed", func(t *testing.T) {
 		txn := mocks.NewMetaKv(t)
-		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error {
+		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error {
 			return f([]byte("1"), []byte("invalid"))
 		})
 
@@ -943,7 +943,7 @@ func TestCatalog_AlterIndexes(t *testing.T) {
 	}
 
 	metakv := mocks.NewMetaKv(t)
-	metakv.EXPECT().MultiSave(mock.Anything).Return(nil)
+	metakv.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(nil)
 	catalog := &Catalog{
 		MetaKv: metakv,
 	}
@@ -970,14 +970,14 @@ func TestCatalog_AlterSegmentMultiIndexes(t *testing.T) {
 	t.Run("batchAdd", func(t *testing.T) {
 		metakv := mocks.NewMetaKv(t)
 
-		metakv.EXPECT().MultiSave(mock.Anything).RunAndReturn(func(m map[string]string) error {
+		metakv.EXPECT().MultiSave(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, m map[string]string) error {
 			if len(m) > 64 {
 				return errors.New("fail")
 			}
 			return nil
 		})
 
-		metakv.EXPECT().MultiSave(mock.Anything).Return(nil)
+		metakv.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(nil)
 		catalog := &Catalog{
 			MetaKv: metakv,
 		}
@@ -990,7 +990,7 @@ func TestCatalog_AlterSegmentMultiIndexes(t *testing.T) {
 func TestCatalog_DropIndex(t *testing.T) {
 	t.Run("success", func(t *testing.T) {
 		metakv := mocks.NewMetaKv(t)
-		metakv.EXPECT().Remove(mock.Anything).Return(nil)
+		metakv.EXPECT().Remove(mock.Anything, mock.Anything).Return(nil)
 		catalog := &Catalog{
 			MetaKv: metakv,
 		}
@@ -1001,7 +1001,7 @@ func TestCatalog_DropIndex(t *testing.T) {
 
 	t.Run("failed", func(t *testing.T) {
 		metakv := mocks.NewMetaKv(t)
-		metakv.EXPECT().Remove(mock.Anything).Return(errors.New("error"))
+		metakv.EXPECT().Remove(mock.Anything, mock.Anything).Return(errors.New("error"))
 		catalog := &Catalog{
 			MetaKv: metakv,
 		}
@@ -1031,7 +1031,7 @@ func TestCatalog_CreateSegmentIndex(t *testing.T) {
 
 	t.Run("success", func(t *testing.T) {
 		metakv := mocks.NewMetaKv(t)
-		metakv.EXPECT().Save(mock.Anything, mock.Anything).Return(nil)
+		metakv.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(nil)
 		catalog := &Catalog{
 			MetaKv: metakv,
 		}
@@ -1042,7 +1042,7 @@ func TestCatalog_CreateSegmentIndex(t *testing.T) {
 
 	t.Run("failed", func(t *testing.T) {
 		metakv := mocks.NewMetaKv(t)
-		metakv.EXPECT().Save(mock.Anything, mock.Anything).Return(errors.New("error"))
+		metakv.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(errors.New("error"))
 		catalog := &Catalog{
 			MetaKv: metakv,
 		}
@@ -1074,7 +1074,7 @@ func TestCatalog_ListSegmentIndexes(t *testing.T) {
 		assert.NoError(t, err)
 
 		metakv := mocks.NewMetaKv(t)
-		metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error {
+		metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error {
 			return f([]byte("key"), v)
 		})
 		catalog := &Catalog{
@@ -1088,7 +1088,7 @@ func TestCatalog_ListSegmentIndexes(t *testing.T) {
 
 	t.Run("failed", func(t *testing.T) {
 		metakv := mocks.NewMetaKv(t)
-		metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(errors.New("error"))
+		metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(errors.New("error"))
 		catalog := &Catalog{
 			MetaKv: metakv,
 		}
@@ -1099,7 +1099,7 @@ func TestCatalog_ListSegmentIndexes(t *testing.T) {
 
 	t.Run("unmarshal failed", func(t *testing.T) {
 		metakv := mocks.NewMetaKv(t)
-		metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error {
+		metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error {
 			return f([]byte("key"), []byte("invalid"))
 		})
 		catalog := &Catalog{
@@ -1131,7 +1131,7 @@ func TestCatalog_AlterSegmentIndexes(t *testing.T) {
 
 	t.Run("add", func(t *testing.T) {
 		metakv := mocks.NewMetaKv(t)
-		metakv.EXPECT().MultiSave(mock.Anything).Return(nil)
+		metakv.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(nil)
 		catalog := &Catalog{
 			MetaKv: metakv,
 		}
@@ -1144,7 +1144,7 @@ func TestCatalog_AlterSegmentIndexes(t *testing.T) {
 func TestCatalog_DropSegmentIndex(t *testing.T) {
 	t.Run("success", func(t *testing.T) {
 		metakv := mocks.NewMetaKv(t)
-		metakv.EXPECT().Remove(mock.Anything).Return(nil)
+		metakv.EXPECT().Remove(mock.Anything, mock.Anything).Return(nil)
 		catalog := &Catalog{
 			MetaKv: metakv,
 		}
@@ -1155,7 +1155,7 @@ func TestCatalog_DropSegmentIndex(t *testing.T) {
 
 	t.Run("fail", func(t *testing.T) {
 		metakv := mocks.NewMetaKv(t)
-		metakv.EXPECT().Remove(mock.Anything).Return(errors.New("error"))
+		metakv.EXPECT().Remove(mock.Anything, mock.Anything).Return(errors.New("error"))
 		catalog := &Catalog{
 			MetaKv: metakv,
 		}
@@ -1336,12 +1336,14 @@ func TestCatalog_GcConfirm(t *testing.T) {
 	kc.MetaKv = txn
 
 	txn.On("LoadWithPrefix",
+		mock.Anything,
 		mock.AnythingOfType("string")).
 		Return(nil, nil, errors.New("error mock LoadWithPrefix")).
 		Once()
 	assert.False(t, kc.GcConfirm(context.TODO(), 100, 10000))
 
 	txn.On("LoadWithPrefix",
+		mock.Anything,
 		mock.AnythingOfType("string")).
 		Return(nil, nil, nil)
 	assert.True(t, kc.GcConfirm(context.TODO(), 100, 10000))
@@ -1365,7 +1367,7 @@ func TestCatalog_Import(t *testing.T) {
 
 	t.Run("SaveImportJob", func(t *testing.T) {
 		txn := mocks.NewMetaKv(t)
-		txn.EXPECT().Save(mock.Anything, mock.Anything).Return(nil)
+		txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(nil)
 		kc.MetaKv = txn
 		err := kc.SaveImportJob(context.TODO(), job)
 		assert.NoError(t, err)
@@ -1374,7 +1376,7 @@ func TestCatalog_Import(t *testing.T) {
 		assert.NoError(t, err)
 
 		txn = mocks.NewMetaKv(t)
-		txn.EXPECT().Save(mock.Anything, mock.Anything).Return(mockErr)
+		txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(mockErr)
 		kc.MetaKv = txn
 		err = kc.SaveImportJob(context.TODO(), job)
 		assert.Error(t, err)
@@ -1384,7 +1386,7 @@ func TestCatalog_Import(t *testing.T) {
 		txn := mocks.NewMetaKv(t)
 		value, err := proto.Marshal(job)
 		assert.NoError(t, err)
-		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error {
+		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error {
 			return f(nil, value)
 		})
 		kc.MetaKv = txn
@@ -1393,7 +1395,7 @@ func TestCatalog_Import(t *testing.T) {
 		assert.Equal(t, 1, len(jobs))
 
 		txn = mocks.NewMetaKv(t)
-		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error {
+		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error {
 			return f(nil, []byte("@#%#^#"))
 		})
 		kc.MetaKv = txn
@@ -1401,7 +1403,7 @@ func TestCatalog_Import(t *testing.T) {
 		assert.Error(t, err)
 
 		txn = mocks.NewMetaKv(t)
-		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(mockErr)
+		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(mockErr)
 		kc.MetaKv = txn
 		_, err = kc.ListImportJobs(context.TODO())
 		assert.Error(t, err)
@@ -1409,13 +1411,13 @@ func TestCatalog_Import(t *testing.T) {
 
 	t.Run("DropImportJob", func(t *testing.T) {
 		txn := mocks.NewMetaKv(t)
-		txn.EXPECT().Remove(mock.Anything).Return(nil)
+		txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(nil)
 		kc.MetaKv = txn
 		err := kc.DropImportJob(context.TODO(), job.GetJobID())
 		assert.NoError(t, err)
 
 		txn = mocks.NewMetaKv(t)
-		txn.EXPECT().Remove(mock.Anything).Return(mockErr)
+		txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(mockErr)
 		kc.MetaKv = txn
 		err = kc.DropImportJob(context.TODO(), job.GetJobID())
 		assert.Error(t, err)
@@ -1423,7 +1425,7 @@ func TestCatalog_Import(t *testing.T) {
 
 	t.Run("SavePreImportTask", func(t *testing.T) {
 		txn := mocks.NewMetaKv(t)
-		txn.EXPECT().Save(mock.Anything, mock.Anything).Return(nil)
+		txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(nil)
 		kc.MetaKv = txn
 		err := kc.SavePreImportTask(context.TODO(), pit)
 		assert.NoError(t, err)
@@ -1432,7 +1434,7 @@ func TestCatalog_Import(t *testing.T) {
 		assert.NoError(t, err)
 
 		txn = mocks.NewMetaKv(t)
-		txn.EXPECT().Save(mock.Anything, mock.Anything).Return(mockErr)
+		txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(mockErr)
 		kc.MetaKv = txn
 		err = kc.SavePreImportTask(context.TODO(), pit)
 		assert.Error(t, err)
@@ -1442,7 +1444,7 @@ func TestCatalog_Import(t *testing.T) {
 		txn := mocks.NewMetaKv(t)
 		value, err := proto.Marshal(pit)
 		assert.NoError(t, err)
-		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error {
+		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error {
 			return f(nil, value)
 		})
 		kc.MetaKv = txn
@@ -1451,7 +1453,7 @@ func TestCatalog_Import(t *testing.T) {
 		assert.Equal(t, 1, len(tasks))
 
 		txn = mocks.NewMetaKv(t)
-		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error {
+		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error {
 			return f(nil, []byte("@#%#^#"))
 		})
 		kc.MetaKv = txn
@@ -1459,7 +1461,7 @@ func TestCatalog_Import(t *testing.T) {
 		assert.Error(t, err)
 
 		txn = mocks.NewMetaKv(t)
-		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(mockErr)
+		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(mockErr)
 		kc.MetaKv = txn
 		_, err = kc.ListPreImportTasks(context.TODO())
 		assert.Error(t, err)
@@ -1467,13 +1469,13 @@ func TestCatalog_Import(t *testing.T) {
 
 	t.Run("DropPreImportTask", func(t *testing.T) {
 		txn := mocks.NewMetaKv(t)
-		txn.EXPECT().Remove(mock.Anything).Return(nil)
+		txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(nil)
 		kc.MetaKv = txn
 		err := kc.DropPreImportTask(context.TODO(), pit.GetTaskID())
 		assert.NoError(t, err)
 
 		txn = mocks.NewMetaKv(t)
-		txn.EXPECT().Remove(mock.Anything).Return(mockErr)
+		txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(mockErr)
 		kc.MetaKv = txn
 		err = kc.DropPreImportTask(context.TODO(), pit.GetTaskID())
 		assert.Error(t, err)
@@ -1481,7 +1483,7 @@ func TestCatalog_Import(t *testing.T) {
 
 	t.Run("SaveImportTask", func(t *testing.T) {
 		txn := mocks.NewMetaKv(t)
-		txn.EXPECT().Save(mock.Anything, mock.Anything).Return(nil)
+		txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(nil)
 		kc.MetaKv = txn
 		err := kc.SaveImportTask(context.TODO(), it)
 		assert.NoError(t, err)
@@ -1490,7 +1492,7 @@ func TestCatalog_Import(t *testing.T) {
 		assert.NoError(t, err)
 
 		txn = mocks.NewMetaKv(t)
-		txn.EXPECT().Save(mock.Anything, mock.Anything).Return(mockErr)
+		txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(mockErr)
 		kc.MetaKv = txn
 		err = kc.SaveImportTask(context.TODO(), it)
 		assert.Error(t, err)
@@ -1500,7 +1502,7 @@ func TestCatalog_Import(t *testing.T) {
 		txn := mocks.NewMetaKv(t)
 		value, err := proto.Marshal(it)
 		assert.NoError(t, err)
-		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error {
+		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error {
 			return f(nil, value)
 		})
 		kc.MetaKv = txn
@@ -1509,7 +1511,7 @@ func TestCatalog_Import(t *testing.T) {
 		assert.Equal(t, 1, len(tasks))
 
 		txn = mocks.NewMetaKv(t)
-		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error {
+		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error {
 			return f(nil, []byte("@#%#^#"))
 		})
 		kc.MetaKv = txn
@@ -1517,7 +1519,7 @@ func TestCatalog_Import(t *testing.T) {
 		assert.Error(t, err)
 
 		txn = mocks.NewMetaKv(t)
-		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(mockErr)
+		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(mockErr)
 		kc.MetaKv = txn
 		_, err = kc.ListImportTasks(context.TODO())
 		assert.Error(t, err)
@@ -1525,13 +1527,13 @@ func TestCatalog_Import(t *testing.T) {
 
 	t.Run("DropImportTask", func(t *testing.T) {
 		txn := mocks.NewMetaKv(t)
-		txn.EXPECT().Remove(mock.Anything).Return(nil)
+		txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(nil)
 		kc.MetaKv = txn
 		err := kc.DropImportTask(context.TODO(), it.GetTaskID())
 		assert.NoError(t, err)
 
 		txn = mocks.NewMetaKv(t)
-		txn.EXPECT().Remove(mock.Anything).Return(mockErr)
+		txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(mockErr)
 		kc.MetaKv = txn
 		err = kc.DropImportTask(context.TODO(), it.GetTaskID())
 		assert.Error(t, err)
@@ -1544,7 +1546,7 @@ func TestCatalog_AnalyzeTask(t *testing.T) {
 
 	t.Run("ListAnalyzeTasks", func(t *testing.T) {
 		txn := mocks.NewMetaKv(t)
-		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(mockErr)
+		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(mockErr)
 		kc.MetaKv = txn
 
 		tasks, err := kc.ListAnalyzeTasks(context.Background())
@@ -1570,7 +1572,7 @@ func TestCatalog_AnalyzeTask(t *testing.T) {
 		assert.NoError(t, err)
 
 		txn = mocks.NewMetaKv(t)
-		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error {
+		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error {
 			return f([]byte("key1"), value)
 		})
 		kc.MetaKv = txn
@@ -1580,7 +1582,7 @@ func TestCatalog_AnalyzeTask(t *testing.T) {
 		assert.Equal(t, 1, len(tasks))
 
 		txn = mocks.NewMetaKv(t)
-		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error {
+		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error {
 			return f([]byte("key1"), []byte("1234"))
 		})
 		kc.MetaKv = txn
@@ -1608,14 +1610,14 @@ func TestCatalog_AnalyzeTask(t *testing.T) {
 		}
 
 		txn := mocks.NewMetaKv(t)
-		txn.EXPECT().Save(mock.Anything, mock.Anything).Return(nil)
+		txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(nil)
 		kc.MetaKv = txn
 
 		err := kc.SaveAnalyzeTask(context.Background(), task)
 		assert.NoError(t, err)
 
 		txn = mocks.NewMetaKv(t)
-		txn.EXPECT().Save(mock.Anything, mock.Anything).Return(mockErr)
+		txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(mockErr)
 		kc.MetaKv = txn
 
 		err = kc.SaveAnalyzeTask(context.Background(), task)
@@ -1624,14 +1626,14 @@ func TestCatalog_AnalyzeTask(t *testing.T) {
 
 	t.Run("DropAnalyzeTask", func(t *testing.T) {
 		txn := mocks.NewMetaKv(t)
-		txn.EXPECT().Remove(mock.Anything).Return(nil)
+		txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(nil)
 		kc.MetaKv = txn
 
 		err := kc.DropAnalyzeTask(context.Background(), 1)
 		assert.NoError(t, err)
 
 		txn = mocks.NewMetaKv(t)
-		txn.EXPECT().Remove(mock.Anything).Return(mockErr)
+		txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(mockErr)
 		kc.MetaKv = txn
 
 		err = kc.DropAnalyzeTask(context.Background(), 1)
@@ -1645,7 +1647,7 @@ func Test_PartitionStatsInfo(t *testing.T) {
 
 	t.Run("ListPartitionStatsInfo", func(t *testing.T) {
 		txn := mocks.NewMetaKv(t)
-		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(mockErr)
+		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(mockErr)
 		kc.MetaKv = txn
 
 		infos, err := kc.ListPartitionStatsInfos(context.Background())
@@ -1665,7 +1667,7 @@ func Test_PartitionStatsInfo(t *testing.T) {
 		assert.NoError(t, err)
 
 		txn = mocks.NewMetaKv(t)
-		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error {
+		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error {
 			return f([]byte("key1"), value)
 		})
 		kc.MetaKv = txn
@@ -1675,7 +1677,7 @@ func Test_PartitionStatsInfo(t *testing.T) {
 		assert.Equal(t, 1, len(infos))
 
 		txn = mocks.NewMetaKv(t)
-		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error {
+		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error {
 			return f([]byte("key1"), []byte("1234"))
 		})
 		kc.MetaKv = txn
@@ -1687,7 +1689,7 @@ func Test_PartitionStatsInfo(t *testing.T) {
 
 	t.Run("SavePartitionStatsInfo", func(t *testing.T) {
 		txn := mocks.NewMetaKv(t)
-		txn.EXPECT().MultiSave(mock.Anything).Return(mockErr)
+		txn.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(mockErr)
 		kc.MetaKv = txn
 
 		info := &datapb.PartitionStatsInfo{
@@ -1704,7 +1706,7 @@ func Test_PartitionStatsInfo(t *testing.T) {
 		assert.Error(t, err)
 
 		txn = mocks.NewMetaKv(t)
-		txn.EXPECT().MultiSave(mock.Anything).Return(nil)
+		txn.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(nil)
 		kc.MetaKv = txn
 
 		err = kc.SavePartitionStatsInfo(context.Background(), info)
@@ -1713,7 +1715,7 @@ func Test_PartitionStatsInfo(t *testing.T) {
 
 	t.Run("DropPartitionStatsInfo", func(t *testing.T) {
 		txn := mocks.NewMetaKv(t)
-		txn.EXPECT().Remove(mock.Anything).Return(mockErr)
+		txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(mockErr)
 		kc.MetaKv = txn
 
 		info := &datapb.PartitionStatsInfo{
@@ -1730,7 +1732,7 @@ func Test_PartitionStatsInfo(t *testing.T) {
 		assert.Error(t, err)
 
 		txn = mocks.NewMetaKv(t)
-		txn.EXPECT().Remove(mock.Anything).Return(nil)
+		txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(nil)
 		kc.MetaKv = txn
 
 		err = kc.DropPartitionStatsInfo(context.Background(), info)
@@ -1748,14 +1750,14 @@ func Test_CurrentPartitionStatsVersion(t *testing.T) {
 
 	t.Run("SaveCurrentPartitionStatsVersion", func(t *testing.T) {
 		txn := mocks.NewMetaKv(t)
-		txn.EXPECT().Save(mock.Anything, mock.Anything).Return(mockErr)
+		txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(mockErr)
 		kc.MetaKv = txn
 
 		err := kc.SaveCurrentPartitionStatsVersion(context.Background(), collID, partID, vChannel, currentVersion)
 		assert.Error(t, err)
 
 		txn = mocks.NewMetaKv(t)
-		txn.EXPECT().Save(mock.Anything, mock.Anything).Return(nil)
+		txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(nil)
 		kc.MetaKv = txn
 
 		err = kc.SaveCurrentPartitionStatsVersion(context.Background(), collID, partID, vChannel, currentVersion)
@@ -1764,7 +1766,7 @@ func Test_CurrentPartitionStatsVersion(t *testing.T) {
 
 	t.Run("GetCurrentPartitionStatsVersion", func(t *testing.T) {
 		txn := mocks.NewMetaKv(t)
-		txn.EXPECT().Load(mock.Anything).Return("", mockErr)
+		txn.EXPECT().Load(mock.Anything, mock.Anything).Return("", mockErr)
 		kc.MetaKv = txn
 
 		version, err := kc.GetCurrentPartitionStatsVersion(context.Background(), collID, partID, vChannel)
@@ -1772,7 +1774,7 @@ func Test_CurrentPartitionStatsVersion(t *testing.T) {
 		assert.Equal(t, int64(0), version)
 
 		txn = mocks.NewMetaKv(t)
-		txn.EXPECT().Load(mock.Anything).Return("1", nil)
+		txn.EXPECT().Load(mock.Anything, mock.Anything).Return("1", nil)
 		kc.MetaKv = txn
 
 		version, err = kc.GetCurrentPartitionStatsVersion(context.Background(), collID, partID, vChannel)
@@ -1782,14 +1784,14 @@ func Test_CurrentPartitionStatsVersion(t *testing.T) {
 
 	t.Run("DropCurrentPartitionStatsVersion", func(t *testing.T) {
 		txn := mocks.NewMetaKv(t)
-		txn.EXPECT().Remove(mock.Anything).Return(mockErr)
+		txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(mockErr)
 		kc.MetaKv = txn
 
 		err := kc.DropCurrentPartitionStatsVersion(context.Background(), collID, partID, vChannel)
 		assert.Error(t, err)
 
 		txn = mocks.NewMetaKv(t)
-		txn.EXPECT().Remove(mock.Anything).Return(nil)
+		txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(nil)
 		kc.MetaKv = txn
 
 		err = kc.DropCurrentPartitionStatsVersion(context.Background(), collID, partID, vChannel)
@@ -1803,7 +1805,7 @@ func Test_StatsTasks(t *testing.T) {
 
 	t.Run("ListStatsTasks", func(t *testing.T) {
 		txn := mocks.NewMetaKv(t)
-		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(mockErr)
+		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(mockErr)
 		kc.MetaKv = txn
 
 		tasks, err := kc.ListStatsTasks(context.Background())
@@ -1825,7 +1827,7 @@ func Test_StatsTasks(t *testing.T) {
 		assert.NoError(t, err)
 
 		txn = mocks.NewMetaKv(t)
-		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error {
+		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error {
 			return f([]byte("key1"), value)
 		})
 		kc.MetaKv = txn
@@ -1835,7 +1837,7 @@ func Test_StatsTasks(t *testing.T) {
 		assert.Equal(t, 1, len(tasks))
 
 		txn = mocks.NewMetaKv(t)
-		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ string, _ int, f func([]byte, []byte) error) error {
+		txn.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, _ string, _ int, f func([]byte, []byte) error) error {
 			return f([]byte("key1"), []byte("1234"))
 		})
 		kc.MetaKv = txn
@@ -1847,7 +1849,7 @@ func Test_StatsTasks(t *testing.T) {
 
 	t.Run("SaveStatsTask", func(t *testing.T) {
 		txn := mocks.NewMetaKv(t)
-		txn.EXPECT().Save(mock.Anything, mock.Anything).Return(mockErr)
+		txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(mockErr)
 		kc.MetaKv = txn
 
 		task := &indexpb.StatsTask{
@@ -1866,7 +1868,7 @@ func Test_StatsTasks(t *testing.T) {
 		assert.Error(t, err)
 
 		txn = mocks.NewMetaKv(t)
-		txn.EXPECT().Save(mock.Anything, mock.Anything).Return(nil)
+		txn.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(nil)
 		kc.MetaKv = txn
 
 		err = kc.SaveStatsTask(context.Background(), task)
@@ -1875,14 +1877,14 @@ func Test_StatsTasks(t *testing.T) {
 
 	t.Run("DropStatsTask", func(t *testing.T) {
 		txn := mocks.NewMetaKv(t)
-		txn.EXPECT().Remove(mock.Anything).Return(mockErr)
+		txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(mockErr)
 		kc.MetaKv = txn
 
 		err := kc.DropStatsTask(context.Background(), 1)
 		assert.Error(t, err)
 
 		txn = mocks.NewMetaKv(t)
-		txn.EXPECT().Remove(mock.Anything).Return(nil)
+		txn.EXPECT().Remove(mock.Anything, mock.Anything).Return(nil)
 		kc.MetaKv = txn
 
 		err = kc.DropStatsTask(context.Background(), 1)
diff --git a/internal/metastore/kv/querycoord/kv_catalog.go b/internal/metastore/kv/querycoord/kv_catalog.go
index f121462e47848..fd0c1a8476c1b 100644
--- a/internal/metastore/kv/querycoord/kv_catalog.go
+++ b/internal/metastore/kv/querycoord/kv_catalog.go
@@ -51,7 +51,7 @@ func (s Catalog) SaveCollection(ctx context.Context, collection *querypb.Collect
 	if err != nil {
 		return err
 	}
-	err = s.cli.Save(k, string(v))
+	err = s.cli.Save(ctx, k, string(v))
 	if err != nil {
 		return err
 	}
@@ -65,7 +65,7 @@ func (s Catalog) SavePartition(ctx context.Context, info ...*querypb.PartitionLo
 		if err != nil {
 			return err
 		}
-		err = s.cli.Save(k, string(v))
+		err = s.cli.Save(ctx, k, string(v))
 		if err != nil {
 			return err
 		}
@@ -83,7 +83,7 @@ func (s Catalog) SaveReplica(ctx context.Context, replicas ...*querypb.Replica)
 		}
 		kvs[key] = string(value)
 	}
-	return s.cli.MultiSave(kvs)
+	return s.cli.MultiSave(ctx, kvs)
 }
 
 func (s Catalog) SaveResourceGroup(ctx context.Context, rgs ...*querypb.ResourceGroup) error {
@@ -98,12 +98,12 @@ func (s Catalog) SaveResourceGroup(ctx context.Context, rgs ...*querypb.Resource
 		ret[key] = string(value)
 	}
 
-	return s.cli.MultiSave(ret)
+	return s.cli.MultiSave(ctx, ret)
 }
 
 func (s Catalog) RemoveResourceGroup(ctx context.Context, rgName string) error {
 	key := encodeResourceGroupKey(rgName)
-	return s.cli.Remove(key)
+	return s.cli.Remove(ctx, key)
 }
 
 func (s Catalog) GetCollections(ctx context.Context) ([]*querypb.CollectionLoadInfo, error) {
@@ -117,7 +117,7 @@ func (s Catalog) GetCollections(ctx context.Context) ([]*querypb.CollectionLoadI
 		return nil
 	}
 
-	err := s.cli.WalkWithPrefix(CollectionLoadInfoPrefix, paginationSize, applyFn)
+	err := s.cli.WalkWithPrefix(ctx, CollectionLoadInfoPrefix, paginationSize, applyFn)
 	if err != nil {
 		return nil, err
 	}
@@ -136,7 +136,7 @@ func (s Catalog) GetPartitions(ctx context.Context) (map[int64][]*querypb.Partit
 		return nil
 	}
 
-	err := s.cli.WalkWithPrefix(PartitionLoadInfoPrefix, paginationSize, applyFn)
+	err := s.cli.WalkWithPrefix(ctx, PartitionLoadInfoPrefix, paginationSize, applyFn)
 	if err != nil {
 		return nil, err
 	}
@@ -155,7 +155,7 @@ func (s Catalog) GetReplicas(ctx context.Context) ([]*querypb.Replica, error) {
 		return nil
 	}
 
-	err := s.cli.WalkWithPrefix(ReplicaPrefix, paginationSize, applyFn)
+	err := s.cli.WalkWithPrefix(ctx, ReplicaPrefix, paginationSize, applyFn)
 	if err != nil {
 		return nil, err
 	}
@@ -170,7 +170,7 @@ func (s Catalog) GetReplicas(ctx context.Context) ([]*querypb.Replica, error) {
 }
 
 func (s Catalog) getReplicasFromV1(ctx context.Context) ([]*querypb.Replica, error) {
-	_, replicaValues, err := s.cli.LoadWithPrefix(ReplicaMetaPrefixV1)
+	_, replicaValues, err := s.cli.LoadWithPrefix(ctx, ReplicaMetaPrefixV1)
 	if err != nil {
 		return nil, err
 	}
@@ -193,7 +193,7 @@ func (s Catalog) getReplicasFromV1(ctx context.Context) ([]*querypb.Replica, err
 }
 
 func (s Catalog) GetResourceGroups(ctx context.Context) ([]*querypb.ResourceGroup, error) {
-	_, rgs, err := s.cli.LoadWithPrefix(ResourceGroupPrefix)
+	_, rgs, err := s.cli.LoadWithPrefix(ctx, ResourceGroupPrefix)
 	if err != nil {
 		return nil, err
 	}
@@ -214,12 +214,12 @@ func (s Catalog) GetResourceGroups(ctx context.Context) ([]*querypb.ResourceGrou
 func (s Catalog) ReleaseCollection(ctx context.Context, collection int64) error {
 	// remove collection and obtained partitions
 	collectionKey := EncodeCollectionLoadInfoKey(collection)
-	err := s.cli.Remove(collectionKey)
+	err := s.cli.Remove(ctx, collectionKey)
 	if err != nil {
 		return err
 	}
 	partitionsPrefix := fmt.Sprintf("%s/%d", PartitionLoadInfoPrefix, collection)
-	return s.cli.RemoveWithPrefix(partitionsPrefix)
+	return s.cli.RemoveWithPrefix(ctx, partitionsPrefix)
 }
 
 func (s Catalog) ReleasePartition(ctx context.Context, collection int64, partitions ...int64) error {
@@ -233,7 +233,7 @@ func (s Catalog) ReleasePartition(ctx context.Context, collection int64, partiti
 			if endIndex > len(partitions) {
 				endIndex = len(partitions)
 			}
-			err := s.cli.MultiRemove(keys[index:endIndex])
+			err := s.cli.MultiRemove(ctx, keys[index:endIndex])
 			if err != nil {
 				return err
 			}
@@ -241,12 +241,12 @@ func (s Catalog) ReleasePartition(ctx context.Context, collection int64, partiti
 		}
 		return nil
 	}
-	return s.cli.MultiRemove(keys)
+	return s.cli.MultiRemove(ctx, keys)
 }
 
 func (s Catalog) ReleaseReplicas(ctx context.Context, collectionID int64) error {
 	key := encodeCollectionReplicaKey(collectionID)
-	return s.cli.RemoveWithPrefix(key)
+	return s.cli.RemoveWithPrefix(ctx, key)
 }
 
 func (s Catalog) ReleaseReplica(ctx context.Context, collection int64, replicas ...int64) error {
@@ -260,7 +260,7 @@ func (s Catalog) ReleaseReplica(ctx context.Context, collection int64, replicas
 			if endIndex > len(replicas) {
 				endIndex = len(replicas)
 			}
-			err := s.cli.MultiRemove(keys[index:endIndex])
+			err := s.cli.MultiRemove(ctx, keys[index:endIndex])
 			if err != nil {
 				return err
 			}
@@ -268,7 +268,7 @@ func (s Catalog) ReleaseReplica(ctx context.Context, collection int64, replicas
 		}
 		return nil
 	}
-	return s.cli.MultiRemove(keys)
+	return s.cli.MultiRemove(ctx, keys)
 }
 
 func (s Catalog) SaveCollectionTargets(ctx context.Context, targets ...*querypb.CollectionTarget) error {
@@ -285,7 +285,7 @@ func (s Catalog) SaveCollectionTargets(ctx context.Context, targets ...*querypb.
 	}
 
 	// to reduce the target size, we do compress before write to etcd
-	err := s.cli.MultiSave(kvs)
+	err := s.cli.MultiSave(ctx, kvs)
 	if err != nil {
 		return err
 	}
@@ -294,7 +294,7 @@ func (s Catalog) SaveCollectionTargets(ctx context.Context, targets ...*querypb.
 
 func (s Catalog) RemoveCollectionTarget(ctx context.Context, collectionID int64) error {
 	k := encodeCollectionTargetKey(collectionID)
-	return s.cli.Remove(k)
+	return s.cli.Remove(ctx, k)
 }
 
 func (s Catalog) GetCollectionTargets(ctx context.Context) (map[int64]*querypb.CollectionTarget, error) {
@@ -312,7 +312,7 @@ func (s Catalog) GetCollectionTargets(ctx context.Context) (map[int64]*querypb.C
 		return nil
 	}
 
-	err := s.cli.WalkWithPrefix(CollectionTargetPrefix, paginationSize, applyFn)
+	err := s.cli.WalkWithPrefix(ctx, CollectionTargetPrefix, paginationSize, applyFn)
 	if err != nil {
 		return nil, err
 	}
diff --git a/internal/metastore/kv/querycoord/kv_catalog_test.go b/internal/metastore/kv/querycoord/kv_catalog_test.go
index 638ee61b524be..1c119e838f611 100644
--- a/internal/metastore/kv/querycoord/kv_catalog_test.go
+++ b/internal/metastore/kv/querycoord/kv_catalog_test.go
@@ -238,8 +238,8 @@ func (suite *CatalogTestSuite) TestCollectionTarget() {
 	// test access meta store failed
 	mockStore := mocks.NewMetaKv(suite.T())
 	mockErr := errors.New("failed to access etcd")
-	mockStore.EXPECT().MultiSave(mock.Anything).Return(mockErr)
-	mockStore.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(mockErr)
+	mockStore.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(mockErr)
+	mockStore.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(mockErr)
 
 	suite.catalog.cli = mockStore
 	err = suite.catalog.SaveCollectionTargets(ctx, &querypb.CollectionTarget{})
diff --git a/internal/metastore/kv/rootcoord/kv_catalog.go b/internal/metastore/kv/rootcoord/kv_catalog.go
index 7bc1e4e170496..e2f4c1bb704da 100644
--- a/internal/metastore/kv/rootcoord/kv_catalog.go
+++ b/internal/metastore/kv/rootcoord/kv_catalog.go
@@ -96,16 +96,16 @@ func BuildAliasPrefixWithDB(dbID int64) string {
 
 // since SnapshotKV may save both snapshot key and the original key if the original key is newest
 // MaxEtcdTxnNum need to divided by 2
-func batchMultiSaveAndRemove(snapshot kv.SnapShotKV, limit int, saves map[string]string, removals []string, ts typeutil.Timestamp) error {
+func batchMultiSaveAndRemove(ctx context.Context, snapshot kv.SnapShotKV, limit int, saves map[string]string, removals []string, ts typeutil.Timestamp) error {
 	saveFn := func(partialKvs map[string]string) error {
-		return snapshot.MultiSave(partialKvs, ts)
+		return snapshot.MultiSave(ctx, partialKvs, ts)
 	}
 	if err := etcd.SaveByBatchWithLimit(saves, limit, saveFn); err != nil {
 		return err
 	}
 
 	removeFn := func(partialKeys []string) error {
-		return snapshot.MultiSaveAndRemove(nil, partialKeys, ts)
+		return snapshot.MultiSaveAndRemove(ctx, nil, partialKeys, ts)
 	}
 	return etcd.RemoveByBatchWithLimit(removals, limit, removeFn)
 }
@@ -117,7 +117,7 @@ func (kc *Catalog) CreateDatabase(ctx context.Context, db *model.Database, ts ty
 	if err != nil {
 		return err
 	}
-	return kc.Snapshot.Save(key, string(v), ts)
+	return kc.Snapshot.Save(ctx, key, string(v), ts)
 }
 
 func (kc *Catalog) AlterDatabase(ctx context.Context, newColl *model.Database, ts typeutil.Timestamp) error {
@@ -127,16 +127,16 @@ func (kc *Catalog) AlterDatabase(ctx context.Context, newColl *model.Database, t
 	if err != nil {
 		return err
 	}
-	return kc.Snapshot.Save(key, string(v), ts)
+	return kc.Snapshot.Save(ctx, key, string(v), ts)
 }
 
 func (kc *Catalog) DropDatabase(ctx context.Context, dbID int64, ts typeutil.Timestamp) error {
 	key := BuildDatabaseKey(dbID)
-	return kc.Snapshot.MultiSaveAndRemove(nil, []string{key}, ts)
+	return kc.Snapshot.MultiSaveAndRemove(ctx, nil, []string{key}, ts)
 }
 
 func (kc *Catalog) ListDatabases(ctx context.Context, ts typeutil.Timestamp) ([]*model.Database, error) {
-	_, vals, err := kc.Snapshot.LoadWithPrefix(DBInfoMetaPrefix, ts)
+	_, vals, err := kc.Snapshot.LoadWithPrefix(ctx, DBInfoMetaPrefix, ts)
 	if err != nil {
 		return nil, err
 	}
@@ -171,7 +171,7 @@ func (kc *Catalog) CreateCollection(ctx context.Context, coll *model.Collection,
 	// If we succeeded to save collection but failed to save other related keys, the garbage meta can be removed
 	// outside and the collection won't be seen by any others (since it's of creating state).
 	// However, if we save other keys first, there is no chance to remove the intermediate meta.
-	if err := kc.Snapshot.Save(k1, string(v1), ts); err != nil {
+	if err := kc.Snapshot.Save(ctx, k1, string(v1), ts); err != nil {
 		return err
 	}
 
@@ -216,13 +216,13 @@ func (kc *Catalog) CreateCollection(ctx context.Context, coll *model.Collection,
 	// since SnapshotKV may save both snapshot key and the original key if the original key is newest
 	// MaxEtcdTxnNum need to divided by 2
 	return etcd.SaveByBatchWithLimit(kvs, util.MaxEtcdTxnNum/2, func(partialKvs map[string]string) error {
-		return kc.Snapshot.MultiSave(partialKvs, ts)
+		return kc.Snapshot.MultiSave(ctx, partialKvs, ts)
 	})
 }
 
 func (kc *Catalog) loadCollectionFromDb(ctx context.Context, dbID int64, collectionID typeutil.UniqueID, ts typeutil.Timestamp) (*pb.CollectionInfo, error) {
 	collKey := BuildCollectionKey(dbID, collectionID)
-	collVal, err := kc.Snapshot.Load(collKey, ts)
+	collVal, err := kc.Snapshot.Load(ctx, collKey, ts)
 	if err != nil {
 		return nil, merr.WrapErrCollectionNotFound(collectionID, err.Error())
 	}
@@ -280,7 +280,7 @@ func (kc *Catalog) CreatePartition(ctx context.Context, dbID int64, partition *m
 		if err != nil {
 			return err
 		}
-		return kc.Snapshot.Save(k, string(v), ts)
+		return kc.Snapshot.Save(ctx, k, string(v), ts)
 	}
 
 	if partitionExistByID(collMeta, partition.PartitionID) {
@@ -302,7 +302,7 @@ func (kc *Catalog) CreatePartition(ctx context.Context, dbID int64, partition *m
 	if err != nil {
 		return err
 	}
-	return kc.Snapshot.Save(k, string(v), ts)
+	return kc.Snapshot.Save(ctx, k, string(v), ts)
 }
 
 func (kc *Catalog) CreateAlias(ctx context.Context, alias *model.Alias, ts typeutil.Timestamp) error {
@@ -315,7 +315,7 @@ func (kc *Catalog) CreateAlias(ctx context.Context, alias *model.Alias, ts typeu
 		return err
 	}
 	kvs := map[string]string{k: string(v)}
-	return kc.Snapshot.MultiSaveAndRemove(kvs, []string{oldKBefore210, oldKeyWithoutDb}, ts)
+	return kc.Snapshot.MultiSaveAndRemove(ctx, kvs, []string{oldKBefore210, oldKeyWithoutDb}, ts)
 }
 
 func (kc *Catalog) CreateCredential(ctx context.Context, credential *model.Credential) error {
@@ -326,7 +326,7 @@ func (kc *Catalog) CreateCredential(ctx context.Context, credential *model.Crede
 		return err
 	}
 
-	err = kc.Txn.Save(k, string(v))
+	err = kc.Txn.Save(ctx, k, string(v))
 	if err != nil {
 		log.Error("create credential persist meta fail", zap.String("key", k), zap.Error(err))
 		return err
@@ -341,7 +341,7 @@ func (kc *Catalog) AlterCredential(ctx context.Context, credential *model.Creden
 
 func (kc *Catalog) listPartitionsAfter210(ctx context.Context, collectionID typeutil.UniqueID, ts typeutil.Timestamp) ([]*model.Partition, error) {
 	prefix := BuildPartitionPrefix(collectionID)
-	_, values, err := kc.Snapshot.LoadWithPrefix(prefix, ts)
+	_, values, err := kc.Snapshot.LoadWithPrefix(ctx, prefix, ts)
 	if err != nil {
 		return nil, err
 	}
@@ -358,7 +358,7 @@ func (kc *Catalog) listPartitionsAfter210(ctx context.Context, collectionID type
 }
 
 func (kc *Catalog) batchListPartitionsAfter210(ctx context.Context, ts typeutil.Timestamp) (map[int64][]*model.Partition, error) {
-	_, values, err := kc.Snapshot.LoadWithPrefix(PartitionMetaPrefix, ts)
+	_, values, err := kc.Snapshot.LoadWithPrefix(ctx, PartitionMetaPrefix, ts)
 	if err != nil {
 		return nil, err
 	}
@@ -385,7 +385,7 @@ func fieldVersionAfter210(collMeta *pb.CollectionInfo) bool {
 
 func (kc *Catalog) listFieldsAfter210(ctx context.Context, collectionID typeutil.UniqueID, ts typeutil.Timestamp) ([]*model.Field, error) {
 	prefix := BuildFieldPrefix(collectionID)
-	_, values, err := kc.Snapshot.LoadWithPrefix(prefix, ts)
+	_, values, err := kc.Snapshot.LoadWithPrefix(ctx, prefix, ts)
 	if err != nil {
 		return nil, err
 	}
@@ -402,7 +402,7 @@ func (kc *Catalog) listFieldsAfter210(ctx context.Context, collectionID typeutil
 }
 
 func (kc *Catalog) batchListFieldsAfter210(ctx context.Context, ts typeutil.Timestamp) (map[int64][]*model.Field, error) {
-	keys, values, err := kc.Snapshot.LoadWithPrefix(FieldMetaPrefix, ts)
+	keys, values, err := kc.Snapshot.LoadWithPrefix(ctx, FieldMetaPrefix, ts)
 	if err != nil {
 		return nil, err
 	}
@@ -427,9 +427,9 @@ func (kc *Catalog) batchListFieldsAfter210(ctx context.Context, ts typeutil.Time
 	return ret, nil
 }
 
-func (kc *Catalog) listFunctions(collectionID typeutil.UniqueID, ts typeutil.Timestamp) ([]*model.Function, error) {
+func (kc *Catalog) listFunctions(ctx context.Context, collectionID typeutil.UniqueID, ts typeutil.Timestamp) ([]*model.Function, error) {
 	prefix := BuildFunctionPrefix(collectionID)
-	_, values, err := kc.Snapshot.LoadWithPrefix(prefix, ts)
+	_, values, err := kc.Snapshot.LoadWithPrefix(ctx, prefix, ts)
 	if err != nil {
 		return nil, err
 	}
@@ -445,8 +445,8 @@ func (kc *Catalog) listFunctions(collectionID typeutil.UniqueID, ts typeutil.Tim
 	return functions, nil
 }
 
-func (kc *Catalog) batchListFunctions(ts typeutil.Timestamp) (map[int64][]*model.Function, error) {
-	keys, values, err := kc.Snapshot.LoadWithPrefix(FunctionMetaPrefix, ts)
+func (kc *Catalog) batchListFunctions(ctx context.Context, ts typeutil.Timestamp) (map[int64][]*model.Function, error) {
+	keys, values, err := kc.Snapshot.LoadWithPrefix(ctx, FunctionMetaPrefix, ts)
 	if err != nil {
 		return nil, err
 	}
@@ -490,7 +490,7 @@ func (kc *Catalog) appendPartitionAndFieldsInfo(ctx context.Context, collMeta *p
 	}
 	collection.Fields = fields
 
-	functions, err := kc.listFunctions(collection.CollectionID, ts)
+	functions, err := kc.listFunctions(ctx, collection.CollectionID, ts)
 	if err != nil {
 		return nil, err
 	}
@@ -520,7 +520,7 @@ func (kc *Catalog) batchAppendPartitionAndFieldsInfo(ctx context.Context, collMe
 					return nil, err
 				}
 
-				functionMetaMap, err = kc.batchListFunctions(ts)
+				functionMetaMap, err = kc.batchListFunctions(ctx, ts)
 				if err != nil {
 					return nil, err
 				}
@@ -558,7 +558,7 @@ func (kc *Catalog) CollectionExists(ctx context.Context, dbID int64, collectionI
 
 func (kc *Catalog) GetCredential(ctx context.Context, username string) (*model.Credential, error) {
 	k := fmt.Sprintf("%s/%s", CredentialPrefix, username)
-	v, err := kc.Txn.Load(k)
+	v, err := kc.Txn.Load(ctx, k)
 	if err != nil {
 		if errors.Is(err, merr.ErrIoKeyNotFound) {
 			log.Debug("not found the user", zap.String("key", k))
@@ -612,15 +612,15 @@ func (kc *Catalog) DropCollection(ctx context.Context, collectionInfo *model.Col
 	// However, if we remove collection first, we cannot remove other metas.
 	// since SnapshotKV may save both snapshot key and the original key if the original key is newest
 	// MaxEtcdTxnNum need to divided by 2
-	if err := batchMultiSaveAndRemove(kc.Snapshot, util.MaxEtcdTxnNum/2, nil, delMetakeysSnap, ts); err != nil {
+	if err := batchMultiSaveAndRemove(ctx, kc.Snapshot, util.MaxEtcdTxnNum/2, nil, delMetakeysSnap, ts); err != nil {
 		return err
 	}
 
 	// if we found collection dropping, we should try removing related resources.
-	return kc.Snapshot.MultiSaveAndRemove(nil, collectionKeys, ts)
+	return kc.Snapshot.MultiSaveAndRemove(ctx, nil, collectionKeys, ts)
 }
 
-func (kc *Catalog) alterModifyCollection(oldColl *model.Collection, newColl *model.Collection, ts typeutil.Timestamp) error {
+func (kc *Catalog) alterModifyCollection(ctx context.Context, oldColl *model.Collection, newColl *model.Collection, ts typeutil.Timestamp) error {
 	if oldColl.TenantID != newColl.TenantID || oldColl.CollectionID != newColl.CollectionID {
 		return fmt.Errorf("altering tenant id or collection id is forbidden")
 	}
@@ -646,19 +646,19 @@ func (kc *Catalog) alterModifyCollection(oldColl *model.Collection, newColl *mod
 	}
 	saves := map[string]string{newKey: string(value)}
 	if oldKey == newKey {
-		return kc.Snapshot.Save(newKey, string(value), ts)
+		return kc.Snapshot.Save(ctx, newKey, string(value), ts)
 	}
-	return kc.Snapshot.MultiSaveAndRemove(saves, []string{oldKey}, ts)
+	return kc.Snapshot.MultiSaveAndRemove(ctx, saves, []string{oldKey}, ts)
 }
 
 func (kc *Catalog) AlterCollection(ctx context.Context, oldColl *model.Collection, newColl *model.Collection, alterType metastore.AlterType, ts typeutil.Timestamp) error {
 	if alterType == metastore.MODIFY {
-		return kc.alterModifyCollection(oldColl, newColl, ts)
+		return kc.alterModifyCollection(ctx, oldColl, newColl, ts)
 	}
 	return fmt.Errorf("altering collection doesn't support %s", alterType.String())
 }
 
-func (kc *Catalog) alterModifyPartition(oldPart *model.Partition, newPart *model.Partition, ts typeutil.Timestamp) error {
+func (kc *Catalog) alterModifyPartition(ctx context.Context, oldPart *model.Partition, newPart *model.Partition, ts typeutil.Timestamp) error {
 	if oldPart.CollectionID != newPart.CollectionID || oldPart.PartitionID != newPart.PartitionID {
 		return fmt.Errorf("altering collection id or partition id is forbidden")
 	}
@@ -672,12 +672,12 @@ func (kc *Catalog) alterModifyPartition(oldPart *model.Partition, newPart *model
 	if err != nil {
 		return err
 	}
-	return kc.Snapshot.Save(key, string(value), ts)
+	return kc.Snapshot.Save(ctx, key, string(value), ts)
 }
 
 func (kc *Catalog) AlterPartition(ctx context.Context, dbID int64, oldPart *model.Partition, newPart *model.Partition, alterType metastore.AlterType, ts typeutil.Timestamp) error {
 	if alterType == metastore.MODIFY {
-		return kc.alterModifyPartition(oldPart, newPart, ts)
+		return kc.alterModifyPartition(ctx, oldPart, newPart, ts)
 	}
 	return fmt.Errorf("altering partition doesn't support %s", alterType.String())
 }
@@ -716,7 +716,7 @@ func (kc *Catalog) DropPartition(ctx context.Context, dbID int64, collectionID t
 
 	if partitionVersionAfter210(collMeta) {
 		k := BuildPartitionKey(collectionID, partitionID)
-		return kc.Snapshot.MultiSaveAndRemove(nil, []string{k}, ts)
+		return kc.Snapshot.MultiSaveAndRemove(ctx, nil, []string{k}, ts)
 	}
 
 	k := BuildCollectionKey(util.NonDBID, collectionID)
@@ -725,7 +725,7 @@ func (kc *Catalog) DropPartition(ctx context.Context, dbID int64, collectionID t
 	if err != nil {
 		return err
 	}
-	return kc.Snapshot.Save(k, string(v), ts)
+	return kc.Snapshot.Save(ctx, k, string(v), ts)
 }
 
 func (kc *Catalog) DropCredential(ctx context.Context, username string) error {
@@ -745,7 +745,7 @@ func (kc *Catalog) DropCredential(ctx context.Context, username string) error {
 			}
 		}
 	}
-	err = kc.Txn.MultiRemove(deleteKeys)
+	err = kc.Txn.MultiRemove(ctx, deleteKeys)
 	if err != nil {
 		log.Warn("fail to drop credential", zap.String("key", k), zap.Error(err))
 		return err
@@ -758,12 +758,12 @@ func (kc *Catalog) DropAlias(ctx context.Context, dbID int64, alias string, ts t
 	oldKBefore210 := BuildAliasKey210(alias)
 	oldKeyWithoutDb := BuildAliasKey(alias)
 	k := BuildAliasKeyWithDB(dbID, alias)
-	return kc.Snapshot.MultiSaveAndRemove(nil, []string{k, oldKeyWithoutDb, oldKBefore210}, ts)
+	return kc.Snapshot.MultiSaveAndRemove(ctx, nil, []string{k, oldKeyWithoutDb, oldKBefore210}, ts)
 }
 
 func (kc *Catalog) GetCollectionByName(ctx context.Context, dbID int64, collectionName string, ts typeutil.Timestamp) (*model.Collection, error) {
 	prefix := getDatabasePrefix(dbID)
-	_, vals, err := kc.Snapshot.LoadWithPrefix(prefix, ts)
+	_, vals, err := kc.Snapshot.LoadWithPrefix(ctx, prefix, ts)
 	if err != nil {
 		log.Warn("get collection meta fail", zap.String("collectionName", collectionName), zap.Error(err))
 		return nil, err
@@ -787,7 +787,7 @@ func (kc *Catalog) GetCollectionByName(ctx context.Context, dbID int64, collecti
 
 func (kc *Catalog) ListCollections(ctx context.Context, dbID int64, ts typeutil.Timestamp) ([]*model.Collection, error) {
 	prefix := getDatabasePrefix(dbID)
-	_, vals, err := kc.Snapshot.LoadWithPrefix(prefix, ts)
+	_, vals, err := kc.Snapshot.LoadWithPrefix(ctx, prefix, ts)
 	if err != nil {
 		log.Error("get collections meta fail",
 			zap.String("prefix", prefix),
@@ -825,19 +825,19 @@ func (kc *Catalog) ListCollections(ctx context.Context, dbID int64, ts typeutil.
 // all collections in default database should be marked with dbID 1.
 // this method also update dbid in meta store when dbid is 0
 // see also: https://github.com/milvus-io/milvus/issues/33608
-func (kv *Catalog) fixDefaultDBIDConsistency(_ context.Context, collMeta *pb.CollectionInfo, ts typeutil.Timestamp) {
+func (kv *Catalog) fixDefaultDBIDConsistency(ctx context.Context, collMeta *pb.CollectionInfo, ts typeutil.Timestamp) {
 	if collMeta.DbId == util.NonDBID {
 		coll := model.UnmarshalCollectionModel(collMeta)
 		cloned := coll.Clone()
 		cloned.DBID = util.DefaultDBID
-		kv.alterModifyCollection(coll, cloned, ts)
+		kv.alterModifyCollection(ctx, coll, cloned, ts)
 
 		collMeta.DbId = util.DefaultDBID
 	}
 }
 
 func (kc *Catalog) listAliasesBefore210(ctx context.Context, ts typeutil.Timestamp) ([]*model.Alias, error) {
-	_, values, err := kc.Snapshot.LoadWithPrefix(CollectionAliasMetaPrefix210, ts)
+	_, values, err := kc.Snapshot.LoadWithPrefix(ctx, CollectionAliasMetaPrefix210, ts)
 	if err != nil {
 		return nil, err
 	}
@@ -861,7 +861,7 @@ func (kc *Catalog) listAliasesBefore210(ctx context.Context, ts typeutil.Timesta
 
 func (kc *Catalog) listAliasesAfter210WithDb(ctx context.Context, dbID int64, ts typeutil.Timestamp) ([]*model.Alias, error) {
 	prefix := BuildAliasPrefixWithDB(dbID)
-	_, values, err := kc.Snapshot.LoadWithPrefix(prefix, ts)
+	_, values, err := kc.Snapshot.LoadWithPrefix(ctx, prefix, ts)
 	if err != nil {
 		return nil, err
 	}
@@ -917,7 +917,7 @@ func (kc *Catalog) ListCredentials(ctx context.Context) ([]string, error) {
 }
 
 func (kc *Catalog) ListCredentialsWithPasswd(ctx context.Context) (map[string]string, error) {
-	keys, values, err := kc.Txn.LoadWithPrefix(CredentialPrefix)
+	keys, values, err := kc.Txn.LoadWithPrefix(ctx, CredentialPrefix)
 	if err != nil {
 		log.Error("list all credential usernames fail", zap.String("prefix", CredentialPrefix), zap.Error(err))
 		return nil, err
@@ -942,32 +942,32 @@ func (kc *Catalog) ListCredentialsWithPasswd(ctx context.Context) (map[string]st
 	return users, nil
 }
 
-func (kc *Catalog) save(k string) error {
+func (kc *Catalog) save(ctx context.Context, k string) error {
 	var err error
-	if _, err = kc.Txn.Load(k); err != nil && !errors.Is(err, merr.ErrIoKeyNotFound) {
+	if _, err = kc.Txn.Load(ctx, k); err != nil && !errors.Is(err, merr.ErrIoKeyNotFound) {
 		return err
 	}
 	if err == nil {
 		log.Debug("the key has existed", zap.String("key", k))
 		return common.NewIgnorableError(fmt.Errorf("the key[%s] has existed", k))
 	}
-	return kc.Txn.Save(k, "")
+	return kc.Txn.Save(ctx, k, "")
 }
 
-func (kc *Catalog) remove(k string) error {
+func (kc *Catalog) remove(ctx context.Context, k string) error {
 	var err error
-	if _, err = kc.Txn.Load(k); err != nil && !errors.Is(err, merr.ErrIoKeyNotFound) {
+	if _, err = kc.Txn.Load(ctx, k); err != nil && !errors.Is(err, merr.ErrIoKeyNotFound) {
 		return err
 	}
 	if err != nil && errors.Is(err, merr.ErrIoKeyNotFound) {
 		return common.NewIgnorableError(fmt.Errorf("the key[%s] isn't existed", k))
 	}
-	return kc.Txn.Remove(k)
+	return kc.Txn.Remove(ctx, k)
 }
 
 func (kc *Catalog) CreateRole(ctx context.Context, tenant string, entity *milvuspb.RoleEntity) error {
 	k := funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, entity.Name)
-	err := kc.save(k)
+	err := kc.save(ctx, k)
 	if err != nil && !common.IsIgnorableError(err) {
 		log.Warn("fail to save the role", zap.String("key", k), zap.Error(err))
 	}
@@ -993,7 +993,7 @@ func (kc *Catalog) DropRole(ctx context.Context, tenant string, roleName string)
 		}
 	}
 
-	err = kc.Txn.MultiRemove(deleteKeys)
+	err = kc.Txn.MultiRemove(ctx, deleteKeys)
 	if err != nil {
 		log.Warn("fail to drop role", zap.String("key", k), zap.Error(err))
 		return err
@@ -1005,12 +1005,12 @@ func (kc *Catalog) AlterUserRole(ctx context.Context, tenant string, userEntity
 	k := funcutil.HandleTenantForEtcdKey(RoleMappingPrefix, tenant, fmt.Sprintf("%s/%s", userEntity.Name, roleEntity.Name))
 	var err error
 	if operateType == milvuspb.OperateUserRoleType_AddUserToRole {
-		err = kc.save(k)
+		err = kc.save(ctx, k)
 		if err != nil {
 			log.Error("fail to save the user-role", zap.String("key", k), zap.Error(err))
 		}
 	} else if operateType == milvuspb.OperateUserRoleType_RemoveUserFromRole {
-		err = kc.remove(k)
+		err = kc.remove(ctx, k)
 		if err != nil {
 			log.Error("fail to remove the user-role", zap.String("key", k), zap.Error(err))
 		}
@@ -1026,7 +1026,7 @@ func (kc *Catalog) ListRole(ctx context.Context, tenant string, entity *milvuspb
 	roleToUsers := make(map[string][]string)
 	if includeUserInfo {
 		roleMappingKey := funcutil.HandleTenantForEtcdKey(RoleMappingPrefix, tenant, "")
-		keys, _, err := kc.Txn.LoadWithPrefix(roleMappingKey)
+		keys, _, err := kc.Txn.LoadWithPrefix(ctx, roleMappingKey)
 		if err != nil {
 			log.Error("fail to load role mappings", zap.String("key", roleMappingKey), zap.Error(err))
 			return results, err
@@ -1057,7 +1057,7 @@ func (kc *Catalog) ListRole(ctx context.Context, tenant string, entity *milvuspb
 
 	if entity == nil {
 		roleKey := funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, "")
-		keys, _, err := kc.Txn.LoadWithPrefix(roleKey)
+		keys, _, err := kc.Txn.LoadWithPrefix(ctx, roleKey)
 		if err != nil {
 			log.Error("fail to load roles", zap.String("key", roleKey), zap.Error(err))
 			return results, err
@@ -1075,7 +1075,7 @@ func (kc *Catalog) ListRole(ctx context.Context, tenant string, entity *milvuspb
 			return results, fmt.Errorf("role name in the role entity is empty")
 		}
 		roleKey := funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, entity.Name)
-		_, err := kc.Txn.Load(roleKey)
+		_, err := kc.Txn.Load(ctx, roleKey)
 		if err != nil {
 			log.Warn("fail to load a role", zap.String("key", roleKey), zap.Error(err))
 			return results, err
@@ -1086,10 +1086,10 @@ func (kc *Catalog) ListRole(ctx context.Context, tenant string, entity *milvuspb
 	return results, nil
 }
 
-func (kc *Catalog) getRolesByUsername(tenant string, username string) ([]string, error) {
+func (kc *Catalog) getRolesByUsername(ctx context.Context, tenant string, username string) ([]string, error) {
 	var roles []string
 	k := funcutil.HandleTenantForEtcdKey(RoleMappingPrefix, tenant, username)
-	keys, _, err := kc.Txn.LoadWithPrefix(k)
+	keys, _, err := kc.Txn.LoadWithPrefix(ctx, k)
 	if err != nil {
 		log.Error("fail to load role mappings by the username", zap.String("key", k), zap.Error(err))
 		return roles, err
@@ -1106,12 +1106,12 @@ func (kc *Catalog) getRolesByUsername(tenant string, username string) ([]string,
 }
 
 // getUserResult get the user result by the username. And never return the error because the error means the user isn't added to a role.
-func (kc *Catalog) getUserResult(tenant string, username string, includeRoleInfo bool) (*milvuspb.UserResult, error) {
+func (kc *Catalog) getUserResult(ctx context.Context, tenant string, username string, includeRoleInfo bool) (*milvuspb.UserResult, error) {
 	result := &milvuspb.UserResult{User: &milvuspb.UserEntity{Name: username}}
 	if !includeRoleInfo {
 		return result, nil
 	}
-	roleNames, err := kc.getRolesByUsername(tenant, username)
+	roleNames, err := kc.getRolesByUsername(ctx, tenant, username)
 	if err != nil {
 		log.Warn("fail to get roles by the username", zap.Error(err))
 		return result, err
@@ -1132,7 +1132,7 @@ func (kc *Catalog) ListUser(ctx context.Context, tenant string, entity *milvuspb
 	)
 
 	appendUserResult := func(username string) error {
-		result, err := kc.getUserResult(tenant, username, includeRoleInfo)
+		result, err := kc.getUserResult(ctx, tenant, username, includeRoleInfo)
 		if err != nil {
 			return err
 		}
@@ -1175,13 +1175,13 @@ func (kc *Catalog) AlterGrant(ctx context.Context, tenant string, entity *milvus
 
 	// Compatible with logic without db
 	if entity.DbName == util.DefaultDBName {
-		v, err = kc.Txn.Load(funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, fmt.Sprintf("%s/%s/%s", entity.Role.Name, entity.Object.Name, entity.ObjectName)))
+		v, err = kc.Txn.Load(ctx, funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, fmt.Sprintf("%s/%s/%s", entity.Role.Name, entity.Object.Name, entity.ObjectName)))
 		if err == nil {
 			idStr = v
 		}
 	}
 	if idStr == "" {
-		if v, err = kc.Txn.Load(k); err == nil {
+		if v, err = kc.Txn.Load(ctx, k); err == nil {
 			idStr = v
 		} else {
 			log.Warn("fail to load grant privilege entity", zap.String("key", k), zap.Any("type", operateType), zap.Error(err))
@@ -1196,7 +1196,7 @@ func (kc *Catalog) AlterGrant(ctx context.Context, tenant string, entity *milvus
 			}
 
 			idStr = crypto.MD5(k)
-			err = kc.Txn.Save(k, idStr)
+			err = kc.Txn.Save(ctx, k, idStr)
 			if err != nil {
 				log.Error("fail to allocate id when altering the grant", zap.Error(err))
 				return err
@@ -1204,7 +1204,7 @@ func (kc *Catalog) AlterGrant(ctx context.Context, tenant string, entity *milvus
 		}
 	}
 	k = funcutil.HandleTenantForEtcdKey(GranteeIDPrefix, tenant, fmt.Sprintf("%s/%s", idStr, privilegeName))
-	_, err = kc.Txn.Load(k)
+	_, err = kc.Txn.Load(ctx, k)
 	if err != nil {
 		log.Warn("fail to load the grantee id", zap.String("key", k), zap.Error(err))
 		if !errors.Is(err, merr.ErrIoKeyNotFound) {
@@ -1216,7 +1216,7 @@ func (kc *Catalog) AlterGrant(ctx context.Context, tenant string, entity *milvus
 			return common.NewIgnorableError(fmt.Errorf("the grantee-id[%s] isn't existed", k))
 		}
 		if funcutil.IsGrant(operateType) {
-			if err = kc.Txn.Save(k, entity.Grantor.User.Name); err != nil {
+			if err = kc.Txn.Save(ctx, k, entity.Grantor.User.Name); err != nil {
 				log.Error("fail to save the grantee id", zap.String("key", k), zap.Error(err))
 			}
 			return err
@@ -1224,7 +1224,7 @@ func (kc *Catalog) AlterGrant(ctx context.Context, tenant string, entity *milvus
 		return nil
 	}
 	if funcutil.IsRevoke(operateType) {
-		if err = kc.Txn.Remove(k); err != nil {
+		if err = kc.Txn.Remove(ctx, k); err != nil {
 			log.Error("fail to remove the grantee id", zap.String("key", k), zap.Error(err))
 			return err
 		}
@@ -1244,7 +1244,7 @@ func (kc *Catalog) ListGrant(ctx context.Context, tenant string, entity *milvusp
 			return nil
 		}
 		granteeIDKey := funcutil.HandleTenantForEtcdKey(GranteeIDPrefix, tenant, v)
-		keys, values, err := kc.Txn.LoadWithPrefix(granteeIDKey)
+		keys, values, err := kc.Txn.LoadWithPrefix(ctx, granteeIDKey)
 		if err != nil {
 			log.Error("fail to load the grantee ids", zap.String("key", granteeIDKey), zap.Error(err))
 			return err
@@ -1276,7 +1276,7 @@ func (kc *Catalog) ListGrant(ctx context.Context, tenant string, entity *milvusp
 	if !funcutil.IsEmptyString(entity.ObjectName) && entity.Object != nil && !funcutil.IsEmptyString(entity.Object.Name) {
 		if entity.DbName == util.DefaultDBName {
 			granteeKey = funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, fmt.Sprintf("%s/%s/%s", entity.Role.Name, entity.Object.Name, entity.ObjectName))
-			v, err := kc.Txn.Load(granteeKey)
+			v, err := kc.Txn.Load(ctx, granteeKey)
 			if err == nil {
 				err = appendGrantEntity(v, entity.Object.Name, entity.ObjectName)
 				if err == nil {
@@ -1287,14 +1287,14 @@ func (kc *Catalog) ListGrant(ctx context.Context, tenant string, entity *milvusp
 
 		if entity.DbName != util.AnyWord {
 			granteeKey = funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, fmt.Sprintf("%s/%s/%s", entity.Role.Name, entity.Object.Name, funcutil.CombineObjectName(util.AnyWord, entity.ObjectName)))
-			v, err := kc.Txn.Load(granteeKey)
+			v, err := kc.Txn.Load(ctx, granteeKey)
 			if err == nil {
 				_ = appendGrantEntity(v, entity.Object.Name, funcutil.CombineObjectName(util.AnyWord, entity.ObjectName))
 			}
 		}
 
 		granteeKey = funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, fmt.Sprintf("%s/%s/%s", entity.Role.Name, entity.Object.Name, funcutil.CombineObjectName(entity.DbName, entity.ObjectName)))
-		v, err := kc.Txn.Load(granteeKey)
+		v, err := kc.Txn.Load(ctx, granteeKey)
 		if err != nil {
 			log.Error("fail to load the grant privilege entity", zap.String("key", granteeKey), zap.Error(err))
 			return entities, err
@@ -1305,7 +1305,7 @@ func (kc *Catalog) ListGrant(ctx context.Context, tenant string, entity *milvusp
 		}
 	} else {
 		granteeKey = funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, entity.Role.Name)
-		keys, values, err := kc.Txn.LoadWithPrefix(granteeKey)
+		keys, values, err := kc.Txn.LoadWithPrefix(ctx, granteeKey)
 		if err != nil {
 			log.Error("fail to load grant privilege entities", zap.String("key", granteeKey), zap.Error(err))
 			return entities, err
@@ -1336,7 +1336,7 @@ func (kc *Catalog) DeleteGrant(ctx context.Context, tenant string, role *milvusp
 	removeKeys = append(removeKeys, k)
 
 	// the values are the grantee id list
-	_, values, err := kc.Txn.LoadWithPrefix(k)
+	_, values, err := kc.Txn.LoadWithPrefix(ctx, k)
 	if err != nil {
 		log.Warn("fail to load grant privilege entities", zap.String("key", k), zap.Error(err))
 		return err
@@ -1346,7 +1346,7 @@ func (kc *Catalog) DeleteGrant(ctx context.Context, tenant string, role *milvusp
 		removeKeys = append(removeKeys, granteeIDKey)
 	}
 
-	if err = kc.Txn.MultiSaveAndRemoveWithPrefix(nil, removeKeys); err != nil {
+	if err = kc.Txn.MultiSaveAndRemoveWithPrefix(ctx, nil, removeKeys); err != nil {
 		log.Error("fail to remove with the prefix", zap.String("key", k), zap.Error(err))
 	}
 	return err
@@ -1355,7 +1355,7 @@ func (kc *Catalog) DeleteGrant(ctx context.Context, tenant string, role *milvusp
 func (kc *Catalog) ListPolicy(ctx context.Context, tenant string) ([]string, error) {
 	var grantInfoStrs []string
 	granteeKey := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, "")
-	keys, values, err := kc.Txn.LoadWithPrefix(granteeKey)
+	keys, values, err := kc.Txn.LoadWithPrefix(ctx, granteeKey)
 	if err != nil {
 		log.Error("fail to load all grant privilege entities", zap.String("key", granteeKey), zap.Error(err))
 		return []string{}, err
@@ -1368,7 +1368,7 @@ func (kc *Catalog) ListPolicy(ctx context.Context, tenant string) ([]string, err
 			continue
 		}
 		granteeIDKey := funcutil.HandleTenantForEtcdKey(GranteeIDPrefix, tenant, values[i])
-		idKeys, _, err := kc.Txn.LoadWithPrefix(granteeIDKey)
+		idKeys, _, err := kc.Txn.LoadWithPrefix(ctx, granteeIDKey)
 		if err != nil {
 			log.Error("fail to load the grantee ids", zap.String("key", granteeIDKey), zap.Error(err))
 			return []string{}, err
@@ -1390,7 +1390,7 @@ func (kc *Catalog) ListPolicy(ctx context.Context, tenant string) ([]string, err
 func (kc *Catalog) ListUserRole(ctx context.Context, tenant string) ([]string, error) {
 	var userRoles []string
 	k := funcutil.HandleTenantForEtcdKey(RoleMappingPrefix, tenant, "")
-	keys, _, err := kc.Txn.LoadWithPrefix(k)
+	keys, _, err := kc.Txn.LoadWithPrefix(ctx, k)
 	if err != nil {
 		log.Error("fail to load all user-role mappings", zap.String("key", k), zap.Error(err))
 		return []string{}, err
@@ -1613,7 +1613,7 @@ func (kc *Catalog) RestoreRBAC(ctx context.Context, tenant string, meta *milvusp
 
 func (kc *Catalog) GetPrivilegeGroup(ctx context.Context, groupName string) (*milvuspb.PrivilegeGroupInfo, error) {
 	k := BuildPrivilegeGroupkey(groupName)
-	val, err := kc.Txn.Load(k)
+	val, err := kc.Txn.Load(ctx, k)
 	if err != nil {
 		if errors.Is(err, merr.ErrIoKeyNotFound) {
 			return nil, fmt.Errorf("privilege group [%s] does not exist", groupName)
@@ -1632,7 +1632,7 @@ func (kc *Catalog) GetPrivilegeGroup(ctx context.Context, groupName string) (*mi
 
 func (kc *Catalog) DropPrivilegeGroup(ctx context.Context, groupName string) error {
 	k := BuildPrivilegeGroupkey(groupName)
-	err := kc.Txn.Remove(k)
+	err := kc.Txn.Remove(ctx, k)
 	if err != nil {
 		log.Warn("fail to drop privilege group", zap.String("key", k), zap.Error(err))
 		return err
@@ -1651,7 +1651,7 @@ func (kc *Catalog) SavePrivilegeGroup(ctx context.Context, data *milvuspb.Privil
 		log.Error("failed to marshal privilege group info", zap.Error(err))
 		return err
 	}
-	if err = kc.Txn.Save(k, string(v)); err != nil {
+	if err = kc.Txn.Save(ctx, k, string(v)); err != nil {
 		log.Warn("fail to put privilege group", zap.String("key", k), zap.Error(err))
 		return err
 	}
@@ -1659,7 +1659,7 @@ func (kc *Catalog) SavePrivilegeGroup(ctx context.Context, data *milvuspb.Privil
 }
 
 func (kc *Catalog) ListPrivilegeGroups(ctx context.Context) ([]*milvuspb.PrivilegeGroupInfo, error) {
-	_, vals, err := kc.Txn.LoadWithPrefix(PrivilegeGroupPrefix)
+	_, vals, err := kc.Txn.LoadWithPrefix(ctx, PrivilegeGroupPrefix)
 	if err != nil {
 		log.Error("failed to list privilege groups", zap.String("prefix", PrivilegeGroupPrefix), zap.Error(err))
 		return nil, err
diff --git a/internal/metastore/kv/rootcoord/kv_catalog_test.go b/internal/metastore/kv/rootcoord/kv_catalog_test.go
index c0dcc6368b880..9cf8dfcfcfd38 100644
--- a/internal/metastore/kv/rootcoord/kv_catalog_test.go
+++ b/internal/metastore/kv/rootcoord/kv_catalog_test.go
@@ -102,7 +102,7 @@ func TestCatalog_ListCollections(t *testing.T) {
 	t.Run("load collection with prefix fail", func(t *testing.T) {
 		kv := mocks.NewSnapShotKV(t)
 		ts := uint64(1)
-		kv.On("LoadWithPrefix", CollectionMetaPrefix, ts).
+		kv.On("LoadWithPrefix", mock.Anything, CollectionMetaPrefix, ts).
 			Return(nil, nil, targetErr)
 
 		kc := Catalog{Snapshot: kv}
@@ -117,9 +117,9 @@ func TestCatalog_ListCollections(t *testing.T) {
 
 		bColl, err := proto.Marshal(coll2)
 		assert.NoError(t, err)
-		kv.On("LoadWithPrefix", CollectionMetaPrefix, ts).
+		kv.On("LoadWithPrefix", mock.Anything, CollectionMetaPrefix, ts).
 			Return([]string{"key"}, []string{string(bColl)}, nil)
-		kv.On("LoadWithPrefix", mock.MatchedBy(
+		kv.On("LoadWithPrefix", mock.Anything, mock.MatchedBy(
 			func(prefix string) bool {
 				return strings.HasPrefix(prefix, PartitionMetaPrefix)
 			}), ts).
@@ -137,20 +137,20 @@ func TestCatalog_ListCollections(t *testing.T) {
 
 		bColl, err := proto.Marshal(coll2)
 		assert.NoError(t, err)
-		kv.On("LoadWithPrefix", CollectionMetaPrefix, ts).
+		kv.On("LoadWithPrefix", mock.Anything, CollectionMetaPrefix, ts).
 			Return([]string{"key"}, []string{string(bColl)}, nil)
 
 		partitionMeta := &pb.PartitionInfo{}
 		pm, err := proto.Marshal(partitionMeta)
 		assert.NoError(t, err)
 
-		kv.On("LoadWithPrefix", mock.MatchedBy(
+		kv.On("LoadWithPrefix", mock.Anything, mock.MatchedBy(
 			func(prefix string) bool {
 				return strings.HasPrefix(prefix, PartitionMetaPrefix)
 			}), ts).
 			Return([]string{"key"}, []string{string(pm)}, nil)
 
-		kv.On("LoadWithPrefix", mock.MatchedBy(
+		kv.On("LoadWithPrefix", mock.Anything, mock.MatchedBy(
 			func(prefix string) bool {
 				return strings.HasPrefix(prefix, FieldMetaPrefix)
 			}), ts).
@@ -168,9 +168,9 @@ func TestCatalog_ListCollections(t *testing.T) {
 
 		bColl, err := proto.Marshal(coll1)
 		assert.NoError(t, err)
-		kv.On("LoadWithPrefix", CollectionMetaPrefix, ts).
+		kv.On("LoadWithPrefix", mock.Anything, CollectionMetaPrefix, ts).
 			Return([]string{"key"}, []string{string(bColl)}, nil)
-		kv.On("MultiSaveAndRemove", mock.Anything, mock.Anything, ts).Return(nil)
+		kv.On("MultiSaveAndRemove", mock.Anything, mock.Anything, mock.Anything, ts).Return(nil)
 		kc := Catalog{Snapshot: kv}
 
 		ret, err := kc.ListCollections(ctx, util.NonDBID, ts)
@@ -186,14 +186,14 @@ func TestCatalog_ListCollections(t *testing.T) {
 
 		bColl, err := proto.Marshal(coll2)
 		assert.NoError(t, err)
-		kv.On("LoadWithPrefix", BuildDatabasePrefixWithDBID(testDb), ts).
+		kv.On("LoadWithPrefix", mock.Anything, BuildDatabasePrefixWithDBID(testDb), ts).
 			Return([]string{"key"}, []string{string(bColl)}, nil)
 
 		partitionMeta := &pb.PartitionInfo{}
 		pm, err := proto.Marshal(partitionMeta)
 		assert.NoError(t, err)
 
-		kv.On("LoadWithPrefix", mock.MatchedBy(
+		kv.On("LoadWithPrefix", mock.Anything, mock.MatchedBy(
 			func(prefix string) bool {
 				return strings.HasPrefix(prefix, PartitionMetaPrefix)
 			}), ts).
@@ -203,7 +203,7 @@ func TestCatalog_ListCollections(t *testing.T) {
 		fm, err := proto.Marshal(fieldMeta)
 		assert.NoError(t, err)
 
-		kv.On("LoadWithPrefix", mock.MatchedBy(
+		kv.On("LoadWithPrefix", mock.Anything, mock.MatchedBy(
 			func(prefix string) bool {
 				return strings.HasPrefix(prefix, FieldMetaPrefix)
 			}), ts).
@@ -212,7 +212,7 @@ func TestCatalog_ListCollections(t *testing.T) {
 		functionMeta := &schemapb.FunctionSchema{}
 		fcm, err := proto.Marshal(functionMeta)
 		assert.NoError(t, err)
-		kv.On("LoadWithPrefix", mock.MatchedBy(
+		kv.On("LoadWithPrefix", mock.Anything, mock.MatchedBy(
 			func(prefix string) bool {
 				return strings.HasPrefix(prefix, FunctionMetaPrefix)
 			}), ts).
@@ -236,14 +236,14 @@ func TestCatalog_ListCollections(t *testing.T) {
 		aColl, err := proto.Marshal(coll3)
 		assert.NoError(t, err)
 
-		kv.On("LoadWithPrefix", CollectionMetaPrefix, ts).
+		kv.On("LoadWithPrefix", mock.Anything, CollectionMetaPrefix, ts).
 			Return([]string{"key", "key2"}, []string{string(bColl), string(aColl)}, nil)
 
 		partitionMeta := &pb.PartitionInfo{}
 		pm, err := proto.Marshal(partitionMeta)
 		assert.NoError(t, err)
 
-		kv.On("LoadWithPrefix", mock.MatchedBy(
+		kv.On("LoadWithPrefix", mock.Anything, mock.MatchedBy(
 			func(prefix string) bool {
 				return strings.HasPrefix(prefix, PartitionMetaPrefix)
 			}), ts).
@@ -253,7 +253,7 @@ func TestCatalog_ListCollections(t *testing.T) {
 		fm, err := proto.Marshal(fieldMeta)
 		assert.NoError(t, err)
 
-		kv.On("LoadWithPrefix", mock.MatchedBy(
+		kv.On("LoadWithPrefix", mock.Anything, mock.MatchedBy(
 			func(prefix string) bool {
 				return strings.HasPrefix(prefix, FieldMetaPrefix)
 			}), ts).
@@ -262,13 +262,13 @@ func TestCatalog_ListCollections(t *testing.T) {
 		functionMeta := &schemapb.FunctionSchema{}
 		fcm, err := proto.Marshal(functionMeta)
 		assert.NoError(t, err)
-		kv.On("LoadWithPrefix", mock.MatchedBy(
+		kv.On("LoadWithPrefix", mock.Anything, mock.MatchedBy(
 			func(prefix string) bool {
 				return strings.HasPrefix(prefix, FunctionMetaPrefix)
 			}), ts).
 			Return([]string{"rootcoord/functions/1/1"}, []string{string(fcm)}, nil)
 
-		kv.On("MultiSaveAndRemove", mock.Anything, mock.Anything, ts).Return(nil)
+		kv.On("MultiSaveAndRemove", mock.Anything, mock.Anything, mock.Anything, ts).Return(nil)
 		kc := Catalog{Snapshot: kv}
 
 		ret, err := kc.ListCollections(ctx, util.NonDBID, ts)
@@ -284,7 +284,7 @@ func TestCatalog_loadCollection(t *testing.T) {
 	t.Run("load failed", func(t *testing.T) {
 		ctx := context.Background()
 		kv := mocks.NewSnapShotKV(t)
-		kv.EXPECT().Load(mock.Anything, mock.Anything).Return("", errors.New("mock"))
+		kv.EXPECT().Load(mock.Anything, mock.Anything, mock.Anything).Return("", errors.New("mock"))
 		kc := Catalog{Snapshot: kv}
 		_, err := kc.loadCollection(ctx, testDb, 1, 0)
 		assert.Error(t, err)
@@ -293,7 +293,7 @@ func TestCatalog_loadCollection(t *testing.T) {
 	t.Run("load, not collection info", func(t *testing.T) {
 		ctx := context.Background()
 		kv := mocks.NewSnapShotKV(t)
-		kv.EXPECT().Load(mock.Anything, mock.Anything).Return("not in pb format", nil)
+		kv.EXPECT().Load(mock.Anything, mock.Anything, mock.Anything).Return("not in pb format", nil)
 		kc := Catalog{Snapshot: kv}
 		_, err := kc.loadCollection(ctx, testDb, 1, 0)
 		assert.Error(t, err)
@@ -305,7 +305,7 @@ func TestCatalog_loadCollection(t *testing.T) {
 		value, err := proto.Marshal(coll)
 		assert.NoError(t, err)
 		kv := mocks.NewSnapShotKV(t)
-		kv.EXPECT().Load(mock.Anything, mock.Anything).Return(string(value), nil)
+		kv.EXPECT().Load(mock.Anything, mock.Anything, mock.Anything).Return(string(value), nil)
 		kc := Catalog{Snapshot: kv}
 		got, err := kc.loadCollection(ctx, util.DefaultDBID, 1, 0)
 		assert.NoError(t, err)
@@ -322,8 +322,8 @@ func TestCatalog_loadCollection(t *testing.T) {
 		value, err := proto.Marshal(coll)
 		assert.NoError(t, err)
 		kv := mocks.NewSnapShotKV(t)
-		kv.EXPECT().Load(mock.Anything, mock.Anything).Return(string(value), nil)
-		kv.EXPECT().MultiSaveAndRemove(mock.Anything, mock.Anything, mock.Anything).Return(nil)
+		kv.EXPECT().Load(mock.Anything, mock.Anything, mock.Anything).Return(string(value), nil)
+		kv.EXPECT().MultiSaveAndRemove(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
 		kc := Catalog{Snapshot: kv}
 		got, err := kc.loadCollection(ctx, util.NonDBID, 1, 0)
 		assert.NoError(t, err)
@@ -380,12 +380,12 @@ func TestCatalog_GetCollectionByID(t *testing.T) {
 	ss := mocks.NewSnapShotKV(t)
 	c := Catalog{Snapshot: ss}
 
-	ss.EXPECT().Load(mock.Anything, mock.Anything).Return("", errors.New("load error")).Twice()
+	ss.EXPECT().Load(mock.Anything, mock.Anything, mock.Anything).Return("", errors.New("load error")).Twice()
 	coll, err := c.GetCollectionByID(ctx, 0, 1, 1)
 	assert.Error(t, err)
 	assert.Nil(t, coll)
 
-	ss.EXPECT().Load(mock.Anything, mock.Anything).RunAndReturn(func(key string, ts uint64) (string, error) {
+	ss.EXPECT().Load(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, key string, ts uint64) (string, error) {
 		collByte, err := proto.Marshal(&pb.CollectionInfo{
 			ID: 1,
 			Schema: &schemapb.CollectionSchema{
@@ -400,7 +400,7 @@ func TestCatalog_GetCollectionByID(t *testing.T) {
 		require.NoError(t, err)
 		return string(collByte), nil
 	}).Once()
-	ss.EXPECT().MultiSaveAndRemove(mock.Anything, mock.Anything, mock.Anything).Return(nil)
+	ss.EXPECT().MultiSaveAndRemove(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
 
 	coll, err = c.GetCollectionByID(ctx, 0, 10000, 1)
 	assert.NoError(t, err)
@@ -412,7 +412,7 @@ func TestCatalog_CreatePartitionV2(t *testing.T) {
 	t.Run("collection not exist", func(t *testing.T) {
 		ctx := context.Background()
 		snapshot := kv.NewMockSnapshotKV()
-		snapshot.LoadFunc = func(key string, ts typeutil.Timestamp) (string, error) {
+		snapshot.LoadFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) (string, error) {
 			return "", errors.New("mock")
 		}
 		kc := Catalog{Snapshot: snapshot}
@@ -430,10 +430,10 @@ func TestCatalog_CreatePartitionV2(t *testing.T) {
 		assert.NoError(t, err)
 
 		snapshot := kv.NewMockSnapshotKV()
-		snapshot.LoadFunc = func(key string, ts typeutil.Timestamp) (string, error) {
+		snapshot.LoadFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) (string, error) {
 			return string(value), nil
 		}
-		snapshot.SaveFunc = func(key string, value string, ts typeutil.Timestamp) error {
+		snapshot.SaveFunc = func(ctx context.Context, key string, value string, ts typeutil.Timestamp) error {
 			return errors.New("mock")
 		}
 
@@ -442,7 +442,7 @@ func TestCatalog_CreatePartitionV2(t *testing.T) {
 		err = kc.CreatePartition(ctx, 0, &model.Partition{}, 0)
 		assert.Error(t, err)
 
-		snapshot.SaveFunc = func(key string, value string, ts typeutil.Timestamp) error {
+		snapshot.SaveFunc = func(ctx context.Context, key string, value string, ts typeutil.Timestamp) error {
 			return nil
 		}
 		err = kc.CreatePartition(ctx, 0, &model.Partition{}, 0)
@@ -458,7 +458,7 @@ func TestCatalog_CreatePartitionV2(t *testing.T) {
 		assert.NoError(t, err)
 
 		snapshot := kv.NewMockSnapshotKV()
-		snapshot.LoadFunc = func(key string, ts typeutil.Timestamp) (string, error) {
+		snapshot.LoadFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) (string, error) {
 			return string(value), nil
 		}
 
@@ -477,7 +477,7 @@ func TestCatalog_CreatePartitionV2(t *testing.T) {
 		assert.NoError(t, err)
 
 		snapshot := kv.NewMockSnapshotKV()
-		snapshot.LoadFunc = func(key string, ts typeutil.Timestamp) (string, error) {
+		snapshot.LoadFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) (string, error) {
 			return string(value), nil
 		}
 
@@ -500,10 +500,10 @@ func TestCatalog_CreatePartitionV2(t *testing.T) {
 		assert.NoError(t, err)
 
 		snapshot := kv.NewMockSnapshotKV()
-		snapshot.LoadFunc = func(key string, ts typeutil.Timestamp) (string, error) {
+		snapshot.LoadFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) (string, error) {
 			return string(value), nil
 		}
-		snapshot.SaveFunc = func(key string, value string, ts typeutil.Timestamp) error {
+		snapshot.SaveFunc = func(ctx context.Context, key string, value string, ts typeutil.Timestamp) error {
 			return errors.New("mock")
 		}
 
@@ -512,7 +512,7 @@ func TestCatalog_CreatePartitionV2(t *testing.T) {
 		err = kc.CreatePartition(ctx, 0, &model.Partition{}, 0)
 		assert.Error(t, err)
 
-		snapshot.SaveFunc = func(key string, value string, ts typeutil.Timestamp) error {
+		snapshot.SaveFunc = func(ctx context.Context, key string, value string, ts typeutil.Timestamp) error {
 			return nil
 		}
 		err = kc.CreatePartition(ctx, 0, &model.Partition{}, 0)
@@ -524,7 +524,7 @@ func TestCatalog_CreateAliasV2(t *testing.T) {
 	ctx := context.Background()
 
 	snapshot := kv.NewMockSnapshotKV()
-	snapshot.MultiSaveAndRemoveFunc = func(saves map[string]string, removals []string, ts typeutil.Timestamp) error {
+	snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error {
 		return errors.New("mock")
 	}
 
@@ -533,7 +533,7 @@ func TestCatalog_CreateAliasV2(t *testing.T) {
 	err := kc.CreateAlias(ctx, &model.Alias{}, 0)
 	assert.Error(t, err)
 
-	snapshot.MultiSaveAndRemoveFunc = func(saves map[string]string, removals []string, ts typeutil.Timestamp) error {
+	snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error {
 		return nil
 	}
 	err = kc.CreateAlias(ctx, &model.Alias{}, 0)
@@ -545,7 +545,7 @@ func TestCatalog_listPartitionsAfter210(t *testing.T) {
 		ctx := context.Background()
 
 		snapshot := kv.NewMockSnapshotKV()
-		snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) {
+		snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
 			return nil, nil, errors.New("mock")
 		}
 
@@ -559,7 +559,7 @@ func TestCatalog_listPartitionsAfter210(t *testing.T) {
 		ctx := context.Background()
 
 		snapshot := kv.NewMockSnapshotKV()
-		snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) {
+		snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
 			return []string{"key"}, []string{"not in pb format"}, nil
 		}
 
@@ -577,7 +577,7 @@ func TestCatalog_listPartitionsAfter210(t *testing.T) {
 		assert.NoError(t, err)
 
 		snapshot := kv.NewMockSnapshotKV()
-		snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) {
+		snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
 			return []string{"key"}, []string{string(value)}, nil
 		}
 
@@ -603,7 +603,7 @@ func TestCatalog_listFieldsAfter210(t *testing.T) {
 		ctx := context.Background()
 
 		snapshot := kv.NewMockSnapshotKV()
-		snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) {
+		snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
 			return nil, nil, errors.New("mock")
 		}
 
@@ -617,7 +617,7 @@ func TestCatalog_listFieldsAfter210(t *testing.T) {
 		ctx := context.Background()
 
 		snapshot := kv.NewMockSnapshotKV()
-		snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) {
+		snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
 			return []string{"key"}, []string{"not in pb format"}, nil
 		}
 
@@ -635,7 +635,7 @@ func TestCatalog_listFieldsAfter210(t *testing.T) {
 		assert.NoError(t, err)
 
 		snapshot := kv.NewMockSnapshotKV()
-		snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) {
+		snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
 			return []string{"key"}, []string{string(value)}, nil
 		}
 
@@ -652,7 +652,7 @@ func TestCatalog_AlterAliasV2(t *testing.T) {
 	ctx := context.Background()
 
 	snapshot := kv.NewMockSnapshotKV()
-	snapshot.MultiSaveAndRemoveFunc = func(saves map[string]string, removals []string, ts typeutil.Timestamp) error {
+	snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error {
 		return errors.New("mock")
 	}
 
@@ -661,7 +661,7 @@ func TestCatalog_AlterAliasV2(t *testing.T) {
 	err := kc.AlterAlias(ctx, &model.Alias{}, 0)
 	assert.Error(t, err)
 
-	snapshot.MultiSaveAndRemoveFunc = func(saves map[string]string, removals []string, ts typeutil.Timestamp) error {
+	snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error {
 		return nil
 	}
 	err = kc.AlterAlias(ctx, &model.Alias{}, 0)
@@ -703,7 +703,7 @@ func TestCatalog_DropPartitionV2(t *testing.T) {
 
 		snapshot := mocks.NewSnapShotKV(t)
 		snapshot.On("Load",
-			mock.Anything, mock.Anything).Return("not in codec format", nil)
+			mock.Anything, mock.Anything, mock.Anything).Return("not in codec format", nil)
 
 		kc := Catalog{Snapshot: snapshot}
 
@@ -716,7 +716,7 @@ func TestCatalog_DropPartitionV2(t *testing.T) {
 
 		snapshot := mocks.NewSnapShotKV(t)
 		snapshot.On("Load",
-			mock.Anything, mock.Anything).Return("", merr.WrapErrIoKeyNotFound("partition"))
+			mock.Anything, mock.Anything, mock.Anything).Return("", merr.WrapErrIoKeyNotFound("partition"))
 
 		kc := Catalog{Snapshot: snapshot}
 
@@ -732,10 +732,10 @@ func TestCatalog_DropPartitionV2(t *testing.T) {
 		assert.NoError(t, err)
 
 		snapshot := kv.NewMockSnapshotKV()
-		snapshot.LoadFunc = func(key string, ts typeutil.Timestamp) (string, error) {
+		snapshot.LoadFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) (string, error) {
 			return string(value), nil
 		}
-		snapshot.MultiSaveAndRemoveFunc = func(saves map[string]string, removals []string, ts typeutil.Timestamp) error {
+		snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error {
 			return errors.New("mock")
 		}
 
@@ -744,7 +744,7 @@ func TestCatalog_DropPartitionV2(t *testing.T) {
 		err = kc.DropPartition(ctx, 0, 100, 101, 0)
 		assert.Error(t, err)
 
-		snapshot.MultiSaveAndRemoveFunc = func(saves map[string]string, removals []string, ts typeutil.Timestamp) error {
+		snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error {
 			return nil
 		}
 		err = kc.DropPartition(ctx, 0, 100, 101, 0)
@@ -764,10 +764,10 @@ func TestCatalog_DropPartitionV2(t *testing.T) {
 		assert.NoError(t, err)
 
 		snapshot := kv.NewMockSnapshotKV()
-		snapshot.LoadFunc = func(key string, ts typeutil.Timestamp) (string, error) {
+		snapshot.LoadFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) (string, error) {
 			return string(value), nil
 		}
-		snapshot.SaveFunc = func(key string, value string, ts typeutil.Timestamp) error {
+		snapshot.SaveFunc = func(ctx context.Context, key string, value string, ts typeutil.Timestamp) error {
 			return errors.New("mock")
 		}
 
@@ -776,7 +776,7 @@ func TestCatalog_DropPartitionV2(t *testing.T) {
 		err = kc.DropPartition(ctx, 0, 100, 101, 0)
 		assert.Error(t, err)
 
-		snapshot.SaveFunc = func(key string, value string, ts typeutil.Timestamp) error {
+		snapshot.SaveFunc = func(ctx context.Context, key string, value string, ts typeutil.Timestamp) error {
 			return nil
 		}
 		err = kc.DropPartition(ctx, 0, 100, 102, 0)
@@ -788,7 +788,7 @@ func TestCatalog_DropAliasV2(t *testing.T) {
 	ctx := context.Background()
 
 	snapshot := kv.NewMockSnapshotKV()
-	snapshot.MultiSaveAndRemoveFunc = func(saves map[string]string, removals []string, ts typeutil.Timestamp) error {
+	snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error {
 		return errors.New("mock")
 	}
 
@@ -797,7 +797,7 @@ func TestCatalog_DropAliasV2(t *testing.T) {
 	err := kc.DropAlias(ctx, testDb, "alias", 0)
 	assert.Error(t, err)
 
-	snapshot.MultiSaveAndRemoveFunc = func(saves map[string]string, removals []string, ts typeutil.Timestamp) error {
+	snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error {
 		return nil
 	}
 	err = kc.DropAlias(ctx, testDb, "alias", 0)
@@ -809,7 +809,7 @@ func TestCatalog_listAliasesBefore210(t *testing.T) {
 		ctx := context.Background()
 
 		snapshot := kv.NewMockSnapshotKV()
-		snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) {
+		snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
 			return nil, nil, errors.New("mock")
 		}
 
@@ -823,7 +823,7 @@ func TestCatalog_listAliasesBefore210(t *testing.T) {
 		ctx := context.Background()
 
 		snapshot := kv.NewMockSnapshotKV()
-		snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) {
+		snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
 			return []string{"key"}, []string{"not in pb format"}, nil
 		}
 
@@ -841,7 +841,7 @@ func TestCatalog_listAliasesBefore210(t *testing.T) {
 		assert.NoError(t, err)
 
 		snapshot := kv.NewMockSnapshotKV()
-		snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) {
+		snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
 			return []string{"key"}, []string{string(value)}, nil
 		}
 
@@ -859,7 +859,7 @@ func TestCatalog_listAliasesAfter210(t *testing.T) {
 		ctx := context.Background()
 
 		snapshot := kv.NewMockSnapshotKV()
-		snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) {
+		snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
 			return nil, nil, errors.New("mock")
 		}
 
@@ -873,7 +873,7 @@ func TestCatalog_listAliasesAfter210(t *testing.T) {
 		ctx := context.Background()
 
 		snapshot := kv.NewMockSnapshotKV()
-		snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) {
+		snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
 			return []string{"key"}, []string{"not in pb format"}, nil
 		}
 
@@ -891,7 +891,7 @@ func TestCatalog_listAliasesAfter210(t *testing.T) {
 		assert.NoError(t, err)
 
 		snapshot := kv.NewMockSnapshotKV()
-		snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) {
+		snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
 			return []string{"key"}, []string{string(value)}, nil
 		}
 
@@ -909,7 +909,7 @@ func TestCatalog_ListAliasesV2(t *testing.T) {
 		ctx := context.Background()
 
 		snapshot := kv.NewMockSnapshotKV()
-		snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) {
+		snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
 			return []string{"key"}, []string{"not in pb format"}, nil
 		}
 
@@ -927,7 +927,7 @@ func TestCatalog_ListAliasesV2(t *testing.T) {
 		assert.NoError(t, err)
 
 		snapshot := kv.NewMockSnapshotKV()
-		snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) {
+		snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
 			if key == AliasMetaPrefix {
 				return nil, nil, errors.New("mock")
 			}
@@ -955,7 +955,7 @@ func TestCatalog_ListAliasesV2(t *testing.T) {
 		assert.NoError(t, err)
 
 		snapshot := kv.NewMockSnapshotKV()
-		snapshot.LoadWithPrefixFunc = func(key string, ts typeutil.Timestamp) ([]string, []string, error) {
+		snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
 			dbStr := fmt.Sprintf("%d", testDb)
 			if strings.Contains(key, dbStr) && strings.Contains(key, Aliases) {
 				return []string{"key1"}, []string{string(value2)}, nil
@@ -975,33 +975,33 @@ func TestCatalog_ListAliasesV2(t *testing.T) {
 func Test_batchMultiSaveAndRemove(t *testing.T) {
 	t.Run("failed to save", func(t *testing.T) {
 		snapshot := kv.NewMockSnapshotKV()
-		snapshot.MultiSaveFunc = func(kvs map[string]string, ts typeutil.Timestamp) error {
+		snapshot.MultiSaveFunc = func(ctx context.Context, kvs map[string]string, ts typeutil.Timestamp) error {
 			return errors.New("error mock MultiSave")
 		}
 		saves := map[string]string{"k": "v"}
-		err := batchMultiSaveAndRemove(snapshot, util.MaxEtcdTxnNum/2, saves, []string{}, 0)
+		err := batchMultiSaveAndRemove(context.TODO(), snapshot, util.MaxEtcdTxnNum/2, saves, []string{}, 0)
 		assert.Error(t, err)
 	})
 	t.Run("failed to remove", func(t *testing.T) {
 		snapshot := kv.NewMockSnapshotKV()
-		snapshot.MultiSaveFunc = func(kvs map[string]string, ts typeutil.Timestamp) error {
+		snapshot.MultiSaveFunc = func(ctx context.Context, kvs map[string]string, ts typeutil.Timestamp) error {
 			return nil
 		}
-		snapshot.MultiSaveAndRemoveFunc = func(saves map[string]string, removals []string, ts typeutil.Timestamp) error {
+		snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error {
 			return errors.New("error mock MultiSaveAndRemove")
 		}
 		saves := map[string]string{"k": "v"}
 		removals := []string{"prefix1", "prefix2"}
-		err := batchMultiSaveAndRemove(snapshot, util.MaxEtcdTxnNum/2, saves, removals, 0)
+		err := batchMultiSaveAndRemove(context.TODO(), snapshot, util.MaxEtcdTxnNum/2, saves, removals, 0)
 		assert.Error(t, err)
 	})
 	t.Run("normal case", func(t *testing.T) {
 		snapshot := kv.NewMockSnapshotKV()
-		snapshot.MultiSaveFunc = func(kvs map[string]string, ts typeutil.Timestamp) error {
+		snapshot.MultiSaveFunc = func(ctx context.Context, kvs map[string]string, ts typeutil.Timestamp) error {
 			log.Info("multi save", zap.Any("len", len(kvs)), zap.Any("saves", kvs))
 			return nil
 		}
-		snapshot.MultiSaveAndRemoveFunc = func(saves map[string]string, removals []string, ts typeutil.Timestamp) error {
+		snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error {
 			log.Info("multi save and remove with prefix", zap.Any("len of saves", len(saves)), zap.Any("len of removals", len(removals)),
 				zap.Any("saves", saves), zap.Any("removals", removals))
 			return nil
@@ -1013,7 +1013,7 @@ func Test_batchMultiSaveAndRemove(t *testing.T) {
 			saves[fmt.Sprintf("k%d", i)] = fmt.Sprintf("v%d", i)
 			removals = append(removals, fmt.Sprintf("k%d", i))
 		}
-		err := batchMultiSaveAndRemove(snapshot, util.MaxEtcdTxnNum/2, saves, removals, 0)
+		err := batchMultiSaveAndRemove(context.TODO(), snapshot, util.MaxEtcdTxnNum/2, saves, removals, 0)
 		assert.NoError(t, err)
 	})
 }
@@ -1036,7 +1036,7 @@ func TestCatalog_AlterCollection(t *testing.T) {
 	t.Run("modify", func(t *testing.T) {
 		snapshot := kv.NewMockSnapshotKV()
 		kvs := map[string]string{}
-		snapshot.SaveFunc = func(key string, value string, ts typeutil.Timestamp) error {
+		snapshot.SaveFunc = func(ctx context.Context, key string, value string, ts typeutil.Timestamp) error {
 			kvs[key] = value
 			return nil
 		}
@@ -1070,7 +1070,7 @@ func TestCatalog_AlterCollection(t *testing.T) {
 	t.Run("modify db name", func(t *testing.T) {
 		var collectionID int64 = 1
 		snapshot := kv.NewMockSnapshotKV()
-		snapshot.MultiSaveAndRemoveFunc = func(saves map[string]string, removals []string, ts typeutil.Timestamp) error {
+		snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error {
 			assert.ElementsMatch(t, []string{BuildCollectionKey(0, collectionID)}, removals)
 			assert.Equal(t, len(saves), 1)
 			assert.Contains(t, maps.Keys(saves), BuildCollectionKey(1, collectionID))
@@ -1104,7 +1104,7 @@ func TestCatalog_AlterPartition(t *testing.T) {
 	t.Run("modify", func(t *testing.T) {
 		snapshot := kv.NewMockSnapshotKV()
 		kvs := map[string]string{}
-		snapshot.SaveFunc = func(key string, value string, ts typeutil.Timestamp) error {
+		snapshot.SaveFunc = func(ctx context.Context, key string, value string, ts typeutil.Timestamp) error {
 			kvs[key] = value
 			return nil
 		}
@@ -1151,6 +1151,7 @@ func withMockSave(saveErr error) mockSnapshotOpt {
 	return func(ss *mocks.SnapShotKV) {
 		ss.On(
 			"Save",
+			mock.Anything,
 			mock.AnythingOfType("string"),
 			mock.AnythingOfType("string"),
 			mock.AnythingOfType("uint64")).
@@ -1162,6 +1163,7 @@ func withMockMultiSave(multiSaveErr error) mockSnapshotOpt {
 	return func(ss *mocks.SnapShotKV) {
 		ss.On(
 			"MultiSave",
+			mock.Anything,
 			mock.AnythingOfType("map[string]string"),
 			mock.AnythingOfType("uint64")).
 			Return(multiSaveErr)
@@ -1172,6 +1174,7 @@ func withMockMultiSaveAndRemoveWithPrefix(err error) mockSnapshotOpt {
 	return func(ss *mocks.SnapShotKV) {
 		ss.On(
 			"MultiSaveAndRemoveWithPrefix",
+			mock.Anything,
 			mock.AnythingOfType("map[string]string"),
 			mock.AnythingOfType("[]string"),
 			mock.AnythingOfType("uint64")).
@@ -1183,6 +1186,7 @@ func withMockMultiSaveAndRemove(err error) mockSnapshotOpt {
 	return func(ss *mocks.SnapShotKV) {
 		ss.On(
 			"MultiSaveAndRemove",
+			mock.Anything,
 			mock.AnythingOfType("map[string]string"),
 			mock.AnythingOfType("[]string"),
 			mock.AnythingOfType("uint64")).
@@ -1296,19 +1300,21 @@ func TestCatalog_DropCollection(t *testing.T) {
 		removeCollectionCalled := false
 		mockSnapshot.On(
 			"MultiSaveAndRemove",
+			mock.Anything,
 			mock.AnythingOfType("map[string]string"),
 			mock.AnythingOfType("[]string"),
 			mock.AnythingOfType("uint64")).
-			Return(func(map[string]string, []string, typeutil.Timestamp) error {
+			Return(func(context.Context, map[string]string, []string, typeutil.Timestamp) error {
 				removeOtherCalled = true
 				return nil
 			}).Once()
 		mockSnapshot.On(
 			"MultiSaveAndRemove",
+			mock.Anything,
 			mock.AnythingOfType("map[string]string"),
 			mock.AnythingOfType("[]string"),
 			mock.AnythingOfType("uint64")).
-			Return(func(map[string]string, []string, typeutil.Timestamp) error {
+			Return(func(context.Context, map[string]string, []string, typeutil.Timestamp) error {
 				removeCollectionCalled = true
 				return errors.New("error mock MultiSaveAndRemove")
 			}).Once()
@@ -1399,10 +1405,10 @@ func TestRBAC_Credential(t *testing.T) {
 			marshalFailkey  = fmt.Sprintf("%s/%s", CredentialPrefix, marshalFailName)
 		)
 
-		kvmock.EXPECT().Load(loadFailKey).Return("", errors.New("Mock invalid load"))
-		kvmock.EXPECT().Load(marshalFailkey).Return("random", nil)
-		kvmock.EXPECT().Load(mock.Anything).Call.Return(
-			func(key string) string {
+		kvmock.EXPECT().Load(mock.Anything, loadFailKey).Return("", errors.New("Mock invalid load"))
+		kvmock.EXPECT().Load(mock.Anything, marshalFailkey).Return("random", nil)
+		kvmock.EXPECT().Load(mock.Anything, mock.Anything).Call.Return(
+			func(ctx context.Context, key string) string {
 				v, err := json.Marshal(&internalpb.CredentialInfo{EncryptedPassword: key})
 				require.NoError(t, err)
 				return string(v)
@@ -1447,8 +1453,8 @@ func TestRBAC_Credential(t *testing.T) {
 		)
 
 		mockFailPath := fmt.Sprintf("%s/%s", CredentialPrefix, invalidName)
-		kvmock.EXPECT().Save(mockFailPath, mock.Anything).Return(errors.New("Mock invalid save"))
-		kvmock.EXPECT().Save(mock.Anything, mock.Anything).Return(nil)
+		kvmock.EXPECT().Save(mock.Anything, mockFailPath, mock.Anything).Return(errors.New("Mock invalid save"))
+		kvmock.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(nil)
 
 		tests := []struct {
 			description string
@@ -1502,26 +1508,27 @@ func TestRBAC_Credential(t *testing.T) {
 			getFailName           = "get-fail"
 		)
 
-		kvmock.EXPECT().MultiRemove([]string{fmt.Sprintf("%s/%s", CredentialPrefix, dropFailName)}).Return(errors.New("Mock drop fail"))
+		kvmock.EXPECT().MultiRemove(mock.Anything, []string{fmt.Sprintf("%s/%s", CredentialPrefix, dropFailName)}).Return(errors.New("Mock drop fail"))
 		kvmock.EXPECT().MultiRemove(
+			mock.Anything,
 			[]string{
 				fmt.Sprintf("%s/%s", CredentialPrefix, validName),
 				validUserRoleKeyPrefix + "/role1",
 				validUserRoleKeyPrefix + "/role2",
 			},
 		).Return(nil)
-		kvmock.EXPECT().MultiRemove(mock.Anything).Return(errors.New("Mock invalid multi remove"))
+		kvmock.EXPECT().MultiRemove(mock.Anything, mock.Anything).Return(errors.New("Mock invalid multi remove"))
 
-		kvmock.EXPECT().Load(fmt.Sprintf("%s/%s", CredentialPrefix, getFailName)).Return("", errors.New("Mock invalid load"))
-		kvmock.EXPECT().Load(fmt.Sprintf("%s/%s", CredentialPrefix, validName)).Return(getUserInfoMetaString(validName), nil)
-		kvmock.EXPECT().Load(fmt.Sprintf("%s/%s", CredentialPrefix, dropFailName)).Return(getUserInfoMetaString(dropFailName), nil)
+		kvmock.EXPECT().Load(mock.Anything, fmt.Sprintf("%s/%s", CredentialPrefix, getFailName)).Return("", errors.New("Mock invalid load"))
+		kvmock.EXPECT().Load(mock.Anything, fmt.Sprintf("%s/%s", CredentialPrefix, validName)).Return(getUserInfoMetaString(validName), nil)
+		kvmock.EXPECT().Load(mock.Anything, fmt.Sprintf("%s/%s", CredentialPrefix, dropFailName)).Return(getUserInfoMetaString(dropFailName), nil)
 
-		kvmock.EXPECT().LoadWithPrefix(validUserRoleKeyPrefix).Return(
+		kvmock.EXPECT().LoadWithPrefix(mock.Anything, validUserRoleKeyPrefix).Return(
 			[]string{validUserRoleKeyPrefix + "/role1", validUserRoleKeyPrefix + "/role2"},
 			[]string{"", ""},
 			nil,
 		)
-		kvmock.EXPECT().LoadWithPrefix(dropUserRoleKeyPrefix).Return([]string{}, []string{}, nil)
+		kvmock.EXPECT().LoadWithPrefix(mock.Anything, dropUserRoleKeyPrefix).Return([]string{}, []string{}, nil)
 
 		tests := []struct {
 			description string
@@ -1557,8 +1564,8 @@ func TestRBAC_Credential(t *testing.T) {
 
 		// Return valid keys if count==0
 		// return error if count!=0
-		kvmock.EXPECT().LoadWithPrefix(mock.Anything).Call.Return(
-			func(key string) []string {
+		kvmock.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Call.Return(
+			func(ctx context.Context, key string) []string {
 				cmu.RLock()
 				defer cmu.RUnlock()
 				if count == 0 {
@@ -1571,7 +1578,7 @@ func TestRBAC_Credential(t *testing.T) {
 				}
 				return nil
 			},
-			func(key string) []string {
+			func(ctx context.Context, key string) []string {
 				cmu.RLock()
 				defer cmu.RUnlock()
 				passwd, _ := json.Marshal(&model.Credential{EncryptedPassword: crypto.Base64Encode("passwd")})
@@ -1585,7 +1592,7 @@ func TestRBAC_Credential(t *testing.T) {
 				}
 				return nil
 			},
-			func(key string) error {
+			func(ctx context.Context, key string) error {
 				cmu.RLock()
 				defer cmu.RUnlock()
 				if count == 0 {
@@ -1638,10 +1645,10 @@ func TestRBAC_Role(t *testing.T) {
 			otherError  = fmt.Errorf("mock load error")
 		)
 
-		kvmock.EXPECT().Load(notExistKey).Return("", merr.WrapErrIoKeyNotFound(notExistKey)).Once()
-		kvmock.EXPECT().Load(errorKey).Return("", otherError).Once()
-		kvmock.EXPECT().Load(mock.Anything).Return("", nil).Once()
-		kvmock.EXPECT().Remove(mock.Anything).Call.Return(nil).Once()
+		kvmock.EXPECT().Load(mock.Anything, notExistKey).Return("", merr.WrapErrIoKeyNotFound(notExistKey)).Once()
+		kvmock.EXPECT().Load(mock.Anything, errorKey).Return("", otherError).Once()
+		kvmock.EXPECT().Load(mock.Anything, mock.Anything).Return("", nil).Once()
+		kvmock.EXPECT().Remove(mock.Anything, mock.Anything).Call.Return(nil).Once()
 		tests := []struct {
 			description string
 
@@ -1657,7 +1664,7 @@ func TestRBAC_Role(t *testing.T) {
 		}
 		for _, test := range tests {
 			t.Run(test.description, func(t *testing.T) {
-				err := c.remove(test.key)
+				err := c.remove(context.TODO(), test.key)
 				if test.isValid {
 					assert.NoError(t, err)
 				} else {
@@ -1682,10 +1689,10 @@ func TestRBAC_Role(t *testing.T) {
 			otherError  = fmt.Errorf("mock load error")
 		)
 
-		kvmock.EXPECT().Load(notExistKey).Return("", merr.WrapErrIoKeyNotFound(notExistKey)).Once()
-		kvmock.EXPECT().Load(errorKey).Return("", otherError).Once()
-		kvmock.EXPECT().Load(mock.Anything).Return("", nil).Once()
-		kvmock.EXPECT().Save(mock.Anything, mock.Anything).Call.Return(nil).Once()
+		kvmock.EXPECT().Load(mock.Anything, notExistKey).Return("", merr.WrapErrIoKeyNotFound(notExistKey)).Once()
+		kvmock.EXPECT().Load(mock.Anything, errorKey).Return("", otherError).Once()
+		kvmock.EXPECT().Load(mock.Anything, mock.Anything).Return("", nil).Once()
+		kvmock.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Call.Return(nil).Once()
 
 		tests := []struct {
 			description string
@@ -1703,7 +1710,7 @@ func TestRBAC_Role(t *testing.T) {
 
 		for _, test := range tests {
 			t.Run(test.description, func(t *testing.T) {
-				err := c.save(test.key)
+				err := c.save(context.TODO(), test.key)
 				if test.isValid {
 					assert.NoError(t, err)
 				} else {
@@ -1729,10 +1736,10 @@ func TestRBAC_Role(t *testing.T) {
 			otherError   = fmt.Errorf("mock load error")
 		)
 
-		kvmock.EXPECT().Load(notExistPath).Return("", merr.WrapErrIoKeyNotFound(notExistName)).Once()
-		kvmock.EXPECT().Load(errorPath).Return("", otherError).Once()
-		kvmock.EXPECT().Load(mock.Anything).Return("", nil).Once()
-		kvmock.EXPECT().Save(mock.Anything, mock.Anything).Call.Return(nil).Once()
+		kvmock.EXPECT().Load(mock.Anything, notExistPath).Return("", merr.WrapErrIoKeyNotFound(notExistName)).Once()
+		kvmock.EXPECT().Load(mock.Anything, errorPath).Return("", otherError).Once()
+		kvmock.EXPECT().Load(mock.Anything, mock.Anything).Return("", nil).Once()
+		kvmock.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Call.Return(nil).Once()
 
 		tests := []struct {
 			description string
@@ -1770,27 +1777,27 @@ func TestRBAC_Role(t *testing.T) {
 			getFailName = "get-fail"
 		)
 
-		kvmock.EXPECT().MultiRemove([]string{funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, errorName)}).Return(errors.New("remove error"))
-		kvmock.EXPECT().MultiRemove([]string{
+		kvmock.EXPECT().MultiRemove(mock.Anything, []string{funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, errorName)}).Return(errors.New("remove error"))
+		kvmock.EXPECT().MultiRemove(mock.Anything, []string{
 			funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, validName),
 			funcutil.HandleTenantForEtcdKey(RoleMappingPrefix, tenant, fmt.Sprintf("%s/%s", "user1", validName)),
 			funcutil.HandleTenantForEtcdKey(RoleMappingPrefix, tenant, fmt.Sprintf("%s/%s", "user2", validName)),
 		}).Return(nil)
-		kvmock.EXPECT().MultiRemove(mock.Anything).Return(errors.New("mock multi remove error"))
+		kvmock.EXPECT().MultiRemove(mock.Anything, mock.Anything).Return(errors.New("mock multi remove error"))
 
 		getRoleMappingKey := func(username, rolename string) string {
 			return funcutil.HandleTenantForEtcdKey(RoleMappingPrefix, tenant, fmt.Sprintf("%s/%s", username, rolename))
 		}
 
-		kvmock.EXPECT().LoadWithPrefix(funcutil.HandleTenantForEtcdKey(RoleMappingPrefix, tenant, "")).Return(
+		kvmock.EXPECT().LoadWithPrefix(mock.Anything, funcutil.HandleTenantForEtcdKey(RoleMappingPrefix, tenant, "")).Return(
 			[]string{getRoleMappingKey("user1", validName), getRoleMappingKey("user2", validName), getRoleMappingKey("user3", "role3")},
 			[]string{},
 			nil,
 		)
 
-		kvmock.EXPECT().Load(funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, getFailName)).Return("", errors.New("mock load error"))
-		kvmock.EXPECT().Load(funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, validName)).Return("", nil)
-		kvmock.EXPECT().Load(funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, errorName)).Return("", nil)
+		kvmock.EXPECT().Load(mock.Anything, funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, getFailName)).Return("", errors.New("mock load error"))
+		kvmock.EXPECT().Load(mock.Anything, funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, validName)).Return("", nil)
+		kvmock.EXPECT().Load(mock.Anything, funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, errorName)).Return("", nil)
 
 		tests := []struct {
 			description string
@@ -1829,20 +1836,20 @@ func TestRBAC_Role(t *testing.T) {
 			errorRoleRemove     = "error-role-remove"
 			errorRoleRemovepath = funcutil.HandleTenantForEtcdKey(RoleMappingPrefix, tenant, fmt.Sprintf("%s/%s", user, errorRoleRemove))
 		)
-		kvmock.EXPECT().Save(mock.Anything, mock.Anything).Return(nil)
-		kvmock.EXPECT().Remove(mock.Anything).Return(nil)
+		kvmock.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(nil)
+		kvmock.EXPECT().Remove(mock.Anything, mock.Anything).Return(nil)
 
 		// Catalog.save() returns error
-		kvmock.EXPECT().Load(errorRoleSavepath).Return("", nil)
+		kvmock.EXPECT().Load(mock.Anything, errorRoleSavepath).Return("", nil)
 
 		// Catalog.save() returns nil
-		kvmock.EXPECT().Load(noErrorRoleSavepath).Return("", merr.WrapErrIoKeyNotFound(noErrorRoleSavepath))
+		kvmock.EXPECT().Load(mock.Anything, noErrorRoleSavepath).Return("", merr.WrapErrIoKeyNotFound(noErrorRoleSavepath))
 
 		// Catalog.remove() returns error
-		kvmock.EXPECT().Load(errorRoleRemovepath).Return("", errors.New("not exists"))
+		kvmock.EXPECT().Load(mock.Anything, errorRoleRemovepath).Return("", errors.New("not exists"))
 
 		// Catalog.remove() returns nil
-		kvmock.EXPECT().Load(mock.Anything).Return("", nil)
+		kvmock.EXPECT().Load(mock.Anything, mock.Anything).Return("", nil)
 
 		tests := []struct {
 			description string
@@ -1885,13 +1892,13 @@ func TestRBAC_Role(t *testing.T) {
 				errorLoadPath = funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, errorLoad)
 			)
 
-			kvmock.EXPECT().Load(errorLoadPath).Call.Return("", errors.New("mock load error"))
-			kvmock.EXPECT().Load(mock.Anything).Call.Return("", nil)
+			kvmock.EXPECT().Load(mock.Anything, errorLoadPath).Call.Return("", errors.New("mock load error"))
+			kvmock.EXPECT().Load(mock.Anything, mock.Anything).Call.Return("", nil)
 
 			// Return valid keys if loadWithPrefixReturn == True
 			// return error if loadWithPrefixReturn == False
-			kvmock.EXPECT().LoadWithPrefix(mock.Anything).Call.Return(
-				func(key string) []string {
+			kvmock.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Call.Return(
+				func(ctx context.Context, key string) []string {
 					if loadWithPrefixReturn.Load() {
 						return []string{
 							fmt.Sprintf("%s/%s/%s", RoleMappingPrefix, tenant, "user1/role1"),
@@ -1903,7 +1910,7 @@ func TestRBAC_Role(t *testing.T) {
 					return nil
 				},
 				nil,
-				func(key string) error {
+				func(ctx context.Context, key string) error {
 					if loadWithPrefixReturn.Load() {
 						return nil
 					}
@@ -1961,8 +1968,8 @@ func TestRBAC_Role(t *testing.T) {
 			// Return valid keys if loadWithPrefixReturn == True
 			// return error if loadWithPrefixReturn == False
 			// Mocking the return of kv_catalog.go:L699
-			kvmock.EXPECT().LoadWithPrefix(mock.Anything).Call.Return(
-				func(key string) []string {
+			kvmock.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Call.Return(
+				func(ctx context.Context, key string) []string {
 					if loadWithPrefixReturn.Load() {
 						return []string{
 							fmt.Sprintf("%s/%s/%s", RolePrefix, tenant, "role1"),
@@ -1974,7 +1981,7 @@ func TestRBAC_Role(t *testing.T) {
 					return nil
 				},
 				nil,
-				func(key string) error {
+				func(ctx context.Context, key string) error {
 					if loadWithPrefixReturn.Load() {
 						return nil
 					}
@@ -2018,13 +2025,13 @@ func TestRBAC_Role(t *testing.T) {
 			invalidUserKey = funcutil.HandleTenantForEtcdKey(RoleMappingPrefix, tenant, invalidUser)
 		)
 		// returns error for invalidUserKey
-		kvmock.EXPECT().LoadWithPrefix(invalidUserKey).Call.Return(
+		kvmock.EXPECT().LoadWithPrefix(mock.Anything, invalidUserKey).Call.Return(
 			nil, nil, errors.New("Mock load with prefix wrong"))
 
 		// Returns keys for RoleMappingPrefix/tenant/user1
 		user1Key := fmt.Sprintf("%s/%s/%s", RoleMappingPrefix, tenant, "user1")
-		kvmock.EXPECT().LoadWithPrefix(user1Key).Call.Return(
-			func(key string) []string {
+		kvmock.EXPECT().LoadWithPrefix(mock.Anything, user1Key).Call.Return(
+			func(ctx context.Context, key string) []string {
 				return []string{
 					fmt.Sprintf("%s/%s/%s", RoleMappingPrefix, tenant, "user1/role1"),
 					fmt.Sprintf("%s/%s/%s", RoleMappingPrefix, tenant, "user1/role2"),
@@ -2034,8 +2041,8 @@ func TestRBAC_Role(t *testing.T) {
 
 		// Returns keys for CredentialPrefix
 		var loadCredentialPrefixReturn atomic.Bool
-		kvmock.EXPECT().LoadWithPrefix(CredentialPrefix).Call.Return(
-			func(key string) []string {
+		kvmock.EXPECT().LoadWithPrefix(mock.Anything, CredentialPrefix).Call.Return(
+			func(ctx context.Context, key string) []string {
 				if loadCredentialPrefixReturn.Load() {
 					return []string{
 						fmt.Sprintf("%s/%s/%s", CredentialPrefix, UserSubPrefix, "user1"),
@@ -2043,14 +2050,14 @@ func TestRBAC_Role(t *testing.T) {
 				}
 				return nil
 			},
-			func(key string) []string {
+			func(ctx context.Context, key string) []string {
 				if loadCredentialPrefixReturn.Load() {
 					passwd, _ := json.Marshal(&model.Credential{EncryptedPassword: crypto.Base64Encode("passwd")})
 					return []string{string(passwd)}
 				}
 				return nil
 			},
-			func(key string) error {
+			func(ctx context.Context, key string) error {
 				if loadCredentialPrefixReturn.Load() {
 					return nil
 				}
@@ -2074,7 +2081,7 @@ func TestRBAC_Role(t *testing.T) {
 
 			for _, test := range tests {
 				t.Run(test.description, func(t *testing.T) {
-					res, err := c.getUserResult(tenant, test.user, test.includeRoleInfo)
+					res, err := c.getUserResult(context.TODO(), tenant, test.user, test.includeRoleInfo)
 
 					assert.Equal(t, test.user, res.GetUser().GetName())
 
@@ -2101,9 +2108,9 @@ func TestRBAC_Role(t *testing.T) {
 				invalidUserLoadKey = fmt.Sprintf("%s/%s", CredentialPrefix, invalidUserLoad)
 			)
 			// Returns error for invalidUserLoadKey
-			kvmock.EXPECT().Load(invalidUserLoadKey).Call.Return("", errors.New("Mock load wrong"))
-			kvmock.EXPECT().Load(mock.Anything).Call.Return(
-				func(key string) string {
+			kvmock.EXPECT().Load(mock.Anything, invalidUserLoadKey).Call.Return("", errors.New("Mock load wrong"))
+			kvmock.EXPECT().Load(mock.Anything, mock.Anything).Call.Return(
+				func(ctx context.Context, key string) string {
 					v, err := json.Marshal(&internalpb.CredentialInfo{EncryptedPassword: key})
 					require.NoError(t, err)
 					return string(v)
@@ -2160,8 +2167,8 @@ func TestRBAC_Role(t *testing.T) {
 		// Return valid keys if loadWithPrefixReturn == True
 		// return error if loadWithPrefixReturn == False
 		// Mocking the return of kv_catalog.go:ListUserRole:L982
-		kvmock.EXPECT().LoadWithPrefix(mock.Anything).Call.Return(
-			func(key string) []string {
+		kvmock.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Call.Return(
+			func(ctx context.Context, key string) []string {
 				if loadWithPrefixReturn.Load() {
 					return []string{
 						fmt.Sprintf("%s/%s/%s", RoleMappingPrefix, tenant, "user1/role1"),
@@ -2173,7 +2180,7 @@ func TestRBAC_Role(t *testing.T) {
 				return nil
 			},
 			nil,
-			func(key string) error {
+			func(ctx context.Context, key string) error {
 				if loadWithPrefixReturn.Load() {
 					return nil
 				}
@@ -2247,35 +2254,35 @@ func TestRBAC_Grant(t *testing.T) {
 		errorSaveRoleKeyWithDb := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, fmt.Sprintf("%s/%s/%s", errorSaveRole, object, funcutil.CombineObjectName(util.DefaultDBName, objName)))
 
 		// Mock return in kv_catalog.go:AlterGrant:L815
-		kvmock.EXPECT().Load(validRoleKey).Call.
-			Return(func(key string) string { return validRoleValue }, nil)
+		kvmock.EXPECT().Load(mock.Anything, validRoleKey).Call.
+			Return(func(ctx context.Context, key string) string { return validRoleValue }, nil)
 
-		kvmock.EXPECT().Load(invalidRoleKey).Call.
-			Return("", func(key string) error {
+		kvmock.EXPECT().Load(mock.Anything, invalidRoleKey).Call.
+			Return("", func(ctx context.Context, key string) error {
 				return fmt.Errorf("mock load error, key=%s", key)
 			})
-		kvmock.EXPECT().Load(invalidRoleKeyWithDb).Call.
-			Return("", func(key string) error {
+		kvmock.EXPECT().Load(mock.Anything, invalidRoleKeyWithDb).Call.
+			Return("", func(ctx context.Context, key string) error {
 				return fmt.Errorf("mock load error, key=%s", key)
 			})
-		kvmock.EXPECT().Load(keyNotExistRoleKey).Call.
-			Return("", func(key string) error {
+		kvmock.EXPECT().Load(mock.Anything, keyNotExistRoleKey).Call.
+			Return("", func(ctx context.Context, key string) error {
 				return merr.WrapErrIoKeyNotFound(key)
 			})
-		kvmock.EXPECT().Load(keyNotExistRoleKeyWithDb).Call.
-			Return("", func(key string) error {
+		kvmock.EXPECT().Load(mock.Anything, keyNotExistRoleKeyWithDb).Call.
+			Return("", func(ctx context.Context, key string) error {
 				return merr.WrapErrIoKeyNotFound(key)
 			})
-		kvmock.EXPECT().Load(errorSaveRoleKey).Call.
-			Return("", func(key string) error {
+		kvmock.EXPECT().Load(mock.Anything, errorSaveRoleKey).Call.
+			Return("", func(ctx context.Context, key string) error {
 				return merr.WrapErrIoKeyNotFound(key)
 			})
-		kvmock.EXPECT().Load(errorSaveRoleKeyWithDb).Call.
-			Return("", func(key string) error {
+		kvmock.EXPECT().Load(mock.Anything, errorSaveRoleKeyWithDb).Call.
+			Return("", func(ctx context.Context, key string) error {
 				return merr.WrapErrIoKeyNotFound(key)
 			})
-		kvmock.EXPECT().Save(keyNotExistRoleKeyWithDb, mock.Anything).Return(nil)
-		kvmock.EXPECT().Save(errorSaveRoleKeyWithDb, mock.Anything).Return(errors.New("mock save error role"))
+		kvmock.EXPECT().Save(mock.Anything, keyNotExistRoleKeyWithDb, mock.Anything).Return(nil)
+		kvmock.EXPECT().Save(mock.Anything, errorSaveRoleKeyWithDb, mock.Anything).Return(errors.New("mock save error role"))
 
 		validPrivilegeKey := funcutil.HandleTenantForEtcdKey(GranteeIDPrefix, tenant, fmt.Sprintf("%s/%s", validRoleValue, validPrivilege))
 		invalidPrivilegeKey := funcutil.HandleTenantForEtcdKey(GranteeIDPrefix, tenant, fmt.Sprintf("%s/%s", validRoleValue, invalidPrivilege))
@@ -2283,24 +2290,24 @@ func TestRBAC_Grant(t *testing.T) {
 		keyNotExistPrivilegeKey2WithDb := funcutil.HandleTenantForEtcdKey(GranteeIDPrefix, tenant, fmt.Sprintf("%s/%s", keyNotExistRoleValueWithDb, keyNotExistPrivilege2))
 
 		// Mock return in kv_catalog.go:AlterGrant:L838
-		kvmock.EXPECT().Load(validPrivilegeKey).Call.Return("", nil)
-		kvmock.EXPECT().Load(invalidPrivilegeKey).Call.
-			Return("", func(key string) error {
+		kvmock.EXPECT().Load(mock.Anything, validPrivilegeKey).Call.Return("", nil)
+		kvmock.EXPECT().Load(mock.Anything, invalidPrivilegeKey).Call.
+			Return("", func(ctx context.Context, key string) error {
 				return fmt.Errorf("mock load error, key=%s", key)
 			})
-		kvmock.EXPECT().Load(keyNotExistPrivilegeKey).Call.
-			Return("", func(key string) error {
+		kvmock.EXPECT().Load(mock.Anything, keyNotExistPrivilegeKey).Call.
+			Return("", func(ctx context.Context, key string) error {
 				return merr.WrapErrIoKeyNotFound(key)
 			})
-		kvmock.EXPECT().Load(keyNotExistPrivilegeKey2WithDb).Call.
-			Return("", func(key string) error {
+		kvmock.EXPECT().Load(mock.Anything, keyNotExistPrivilegeKey2WithDb).Call.
+			Return("", func(ctx context.Context, key string) error {
 				return merr.WrapErrIoKeyNotFound(key)
 			})
-		kvmock.EXPECT().Load(mock.Anything).Call.Return("", nil)
+		kvmock.EXPECT().Load(mock.Anything, mock.Anything).Call.Return("", nil)
 
 		t.Run("test Grant", func(t *testing.T) {
-			kvmock.EXPECT().Save(mock.Anything, validUser).Return(nil)
-			kvmock.EXPECT().Save(mock.Anything, invalidUser).Return(errors.New("mock save invalid user"))
+			kvmock.EXPECT().Save(mock.Anything, mock.Anything, validUser).Return(nil)
+			kvmock.EXPECT().Save(mock.Anything, mock.Anything, invalidUser).Return(errors.New("mock save invalid user"))
 
 			tests := []struct {
 				isValid bool
@@ -2359,9 +2366,9 @@ func TestRBAC_Grant(t *testing.T) {
 			invalidPrivilegeRemove := "p-remove"
 			invalidPrivilegeRemoveKey := funcutil.HandleTenantForEtcdKey(GranteeIDPrefix, tenant, fmt.Sprintf("%s/%s", validRoleValue, invalidPrivilegeRemove))
 
-			kvmock.EXPECT().Load(invalidPrivilegeRemoveKey).Call.Return("", nil)
-			kvmock.EXPECT().Remove(invalidPrivilegeRemoveKey).Return(errors.New("mock remove error"))
-			kvmock.EXPECT().Remove(mock.Anything).Return(nil)
+			kvmock.EXPECT().Load(mock.Anything, invalidPrivilegeRemoveKey).Call.Return("", nil)
+			kvmock.EXPECT().Remove(mock.Anything, invalidPrivilegeRemoveKey).Return(errors.New("mock remove error"))
+			kvmock.EXPECT().Remove(mock.Anything, mock.Anything).Return(nil)
 			tests := []struct {
 				isValid bool
 
@@ -2426,10 +2433,10 @@ func TestRBAC_Grant(t *testing.T) {
 			granteePrefix       = funcutil.HandleTenantForEtcdKey(GranteeIDPrefix, tenant, granteeID+"/")
 		)
 
-		kvmock.EXPECT().LoadWithPrefix(loadErrorRolePrefix).Call.Return(nil, nil, errors.New("mock loadWithPrefix error"))
-		kvmock.EXPECT().LoadWithPrefix(mock.Anything).Call.Return(nil, []string{granteeID}, nil)
-		kvmock.EXPECT().MultiSaveAndRemoveWithPrefix(mock.Anything, []string{errorRolePrefix, granteePrefix}, mock.Anything).Call.Return(errors.New("mock removeWithPrefix error"))
-		kvmock.EXPECT().MultiSaveAndRemoveWithPrefix(mock.Anything, mock.Anything, mock.Anything).Call.Return(nil)
+		kvmock.EXPECT().LoadWithPrefix(mock.Anything, loadErrorRolePrefix).Call.Return(nil, nil, errors.New("mock loadWithPrefix error"))
+		kvmock.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Call.Return(nil, []string{granteeID}, nil)
+		kvmock.EXPECT().MultiSaveAndRemoveWithPrefix(mock.Anything, mock.Anything, []string{errorRolePrefix, granteePrefix}, mock.Anything).Call.Return(errors.New("mock removeWithPrefix error"))
+		kvmock.EXPECT().MultiSaveAndRemoveWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Call.Return(nil)
 
 		tests := []struct {
 			isValid bool
@@ -2462,24 +2469,24 @@ func TestRBAC_Grant(t *testing.T) {
 		// Mock Load in kv_catalog.go:L901
 		validGranteeKey := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant,
 			fmt.Sprintf("%s/%s/%s", "role1", "obj1", "obj_name1"))
-		kvmock.EXPECT().Load(validGranteeKey).Call.
-			Return(func(key string) string { return crypto.MD5(key) }, nil)
+		kvmock.EXPECT().Load(mock.Anything, validGranteeKey).Call.
+			Return(func(ctx context.Context, key string) string { return crypto.MD5(key) }, nil)
 		validGranteeKey2 := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant,
 			fmt.Sprintf("%s/%s/%s", "role1", "obj2", "foo.obj_name2"))
-		kvmock.EXPECT().Load(validGranteeKey2).Call.
-			Return(func(key string) string { return crypto.MD5(key) }, nil)
+		kvmock.EXPECT().Load(mock.Anything, validGranteeKey2).Call.
+			Return(func(ctx context.Context, key string) string { return crypto.MD5(key) }, nil)
 		validGranteeKey3 := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant,
 			fmt.Sprintf("%s/%s/%s", "role1", "obj3", "*.obj_name3"))
-		kvmock.EXPECT().Load(validGranteeKey3).Call.
-			Return(func(key string) string { return crypto.MD5(key) }, nil)
+		kvmock.EXPECT().Load(mock.Anything, validGranteeKey3).Call.
+			Return(func(ctx context.Context, key string) string { return crypto.MD5(key) }, nil)
 
-		kvmock.EXPECT().Load(mock.Anything).Call.
+		kvmock.EXPECT().Load(mock.Anything, mock.Anything).Call.
 			Return("", errors.New("mock Load error"))
 
 		invalidRoleKey := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, invalidRole)
-		kvmock.EXPECT().LoadWithPrefix(invalidRoleKey).Call.Return(nil, nil, errors.New("mock loadWithPrefix error"))
-		kvmock.EXPECT().LoadWithPrefix(mock.Anything).Call.Return(
-			func(key string) []string {
+		kvmock.EXPECT().LoadWithPrefix(mock.Anything, invalidRoleKey).Call.Return(nil, nil, errors.New("mock loadWithPrefix error"))
+		kvmock.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Call.Return(
+			func(ctx context.Context, key string) []string {
 				// Mock kv_catalog.go:ListGrant:L871
 				if strings.Contains(key, GranteeIDPrefix) {
 					return []string{
@@ -2495,7 +2502,7 @@ func TestRBAC_Grant(t *testing.T) {
 					fmt.Sprintf("%s/%s", key, "obj3/*.obj_name3"),
 				}
 			},
-			func(key string) []string {
+			func(ctx context.Context, key string) []string {
 				if strings.Contains(key, GranteeIDPrefix) {
 					return []string{"user1", "user2"}
 				}
@@ -2582,8 +2589,8 @@ func TestRBAC_Grant(t *testing.T) {
 			secondLoadWithPrefixReturn atomic.Bool
 		)
 
-		kvmock.EXPECT().LoadWithPrefix(mock.Anything).Call.Return(
-			func(key string) []string {
+		kvmock.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Call.Return(
+			func(ctx context.Context, key string) []string {
 				contains := strings.Contains(key, GranteeIDPrefix)
 				if contains {
 					if secondLoadWithPrefixReturn.Load() {
@@ -2606,7 +2613,7 @@ func TestRBAC_Grant(t *testing.T) {
 				return nil
 			},
 
-			func(key string) []string {
+			func(ctx context.Context, key string) []string {
 				if firstLoadWithPrefixReturn.Load() {
 					return []string{
 						crypto.MD5(fmt.Sprintf("%s/%s", key, "obj1/obj_name1")),
@@ -2616,7 +2623,7 @@ func TestRBAC_Grant(t *testing.T) {
 				return nil
 			},
 
-			func(key string) error {
+			func(ctx context.Context, key string) error {
 				contains := strings.Contains(key, GranteeIDPrefix)
 				if contains {
 					if secondLoadWithPrefixReturn.Load() {
@@ -2680,7 +2687,7 @@ func TestRBAC_Backup(t *testing.T) {
 		Params.EtcdCfg.EtcdTLSMinVersion.GetValue())
 	rootPath := "/test/rbac"
 	metaKV := etcdkv.NewEtcdKV(etcdCli, rootPath)
-	defer metaKV.RemoveWithPrefix("")
+	defer metaKV.RemoveWithPrefix(context.TODO(), "")
 	defer metaKV.Close()
 	c := &Catalog{Txn: metaKV}
 
@@ -2733,7 +2740,7 @@ func TestRBAC_Restore(t *testing.T) {
 		Params.EtcdCfg.EtcdTLSMinVersion.GetValue())
 	rootPath := "/test/rbac"
 	metaKV := etcdkv.NewEtcdKV(etcdCli, rootPath)
-	defer metaKV.RemoveWithPrefix("")
+	defer metaKV.RemoveWithPrefix(context.TODO(), "")
 	defer metaKV.Close()
 	c := &Catalog{Txn: metaKV}
 
@@ -2901,8 +2908,8 @@ func TestRBAC_PrivilegeGroup(t *testing.T) {
 			kvmock = mocks.NewTxnKV(t)
 			c      = &Catalog{Txn: kvmock}
 		)
-		kvmock.EXPECT().Load(key1).Return(string(v1), nil)
-		kvmock.EXPECT().Load(key2).Return("", merr.ErrIoKeyNotFound)
+		kvmock.EXPECT().Load(mock.Anything, key1).Return(string(v1), nil)
+		kvmock.EXPECT().Load(mock.Anything, key2).Return("", merr.ErrIoKeyNotFound)
 
 		tests := []struct {
 			description        string
@@ -2932,8 +2939,8 @@ func TestRBAC_PrivilegeGroup(t *testing.T) {
 			c      = &Catalog{Txn: kvmock}
 		)
 
-		kvmock.EXPECT().Remove(key1).Return(nil)
-		kvmock.EXPECT().Remove(key2).Return(errors.New("Mock remove failure"))
+		kvmock.EXPECT().Remove(mock.Anything, key1).Return(nil)
+		kvmock.EXPECT().Remove(mock.Anything, key2).Return(errors.New("Mock remove failure"))
 
 		tests := []struct {
 			description string
@@ -2962,8 +2969,8 @@ func TestRBAC_PrivilegeGroup(t *testing.T) {
 			c      = &Catalog{Txn: kvmock}
 		)
 
-		kvmock.EXPECT().Save(key1, mock.Anything).Return(nil)
-		kvmock.EXPECT().Save(key2, mock.Anything).Return(nil)
+		kvmock.EXPECT().Save(mock.Anything, key1, mock.Anything).Return(nil)
+		kvmock.EXPECT().Save(mock.Anything, key2, mock.Anything).Return(nil)
 
 		tests := []struct {
 			description string
@@ -2992,7 +2999,7 @@ func TestRBAC_PrivilegeGroup(t *testing.T) {
 			c      = &Catalog{Txn: kvmock}
 		)
 
-		kvmock.EXPECT().LoadWithPrefix(PrivilegeGroupPrefix).Return(
+		kvmock.EXPECT().LoadWithPrefix(mock.Anything, PrivilegeGroupPrefix).Return(
 			[]string{key1, key2},
 			[]string{string(v1), string(v2)},
 			nil,
@@ -3022,7 +3029,7 @@ func TestCatalog_AlterDatabase(t *testing.T) {
 	c := &Catalog{Snapshot: kvmock}
 	db := model.NewDatabase(1, "db", pb.DatabaseState_DatabaseCreated, nil)
 
-	kvmock.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(nil)
+	kvmock.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
 	ctx := context.Background()
 
 	// test alter database success
@@ -3039,7 +3046,7 @@ func TestCatalog_AlterDatabase(t *testing.T) {
 	// test alter database fail
 	mockErr := errors.New("access kv store error")
 	kvmock.ExpectedCalls = nil
-	kvmock.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(mockErr)
+	kvmock.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(mockErr)
 	err = c.AlterDatabase(ctx, newDB, typeutil.ZeroTimestamp)
 	assert.ErrorIs(t, err, mockErr)
 }
@@ -3047,11 +3054,11 @@ func TestCatalog_AlterDatabase(t *testing.T) {
 func TestCatalog_listFunctionError(t *testing.T) {
 	mockSnapshot := newMockSnapshot(t)
 	kc := &Catalog{Snapshot: mockSnapshot}
-	mockSnapshot.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Return(nil, nil, fmt.Errorf("mock error"))
-	_, err := kc.listFunctions(1, 1)
+	mockSnapshot.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(nil, nil, fmt.Errorf("mock error"))
+	_, err := kc.listFunctions(context.TODO(), 1, 1)
 	assert.Error(t, err)
 
-	mockSnapshot.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Return([]string{"test-key"}, []string{"invalid bytes"}, nil)
-	_, err = kc.listFunctions(1, 1)
+	mockSnapshot.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return([]string{"test-key"}, []string{"invalid bytes"}, nil)
+	_, err = kc.listFunctions(context.TODO(), 1, 1)
 	assert.Error(t, err)
 }
diff --git a/internal/metastore/kv/rootcoord/suffix_snapshot.go b/internal/metastore/kv/rootcoord/suffix_snapshot.go
index 5a625d1347450..0457f631fcace 100644
--- a/internal/metastore/kv/rootcoord/suffix_snapshot.go
+++ b/internal/metastore/kv/rootcoord/suffix_snapshot.go
@@ -18,6 +18,7 @@ package rootcoord
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"path"
 	"sort"
@@ -119,7 +120,7 @@ func NewSuffixSnapshot(metaKV kv.MetaKv, sep, root, snapshot string) (*SuffixSna
 		rootLen:        rootLen,
 		closeGC:        make(chan struct{}, 1),
 	}
-	go ss.startBackgroundGC()
+	go ss.startBackgroundGC(context.TODO())
 	return ss, nil
 }
 
@@ -196,10 +197,10 @@ func (ss *SuffixSnapshot) isTSOfKey(key string, groupKey string) (typeutil.Times
 
 // checkKeyTS checks provided key's latest ts is before provided ts
 // lock is needed
-func (ss *SuffixSnapshot) checkKeyTS(key string, ts typeutil.Timestamp) (bool, error) {
+func (ss *SuffixSnapshot) checkKeyTS(ctx context.Context, key string, ts typeutil.Timestamp) (bool, error) {
 	latest, has := ss.lastestTS[key]
 	if !has {
-		err := ss.loadLatestTS(key)
+		err := ss.loadLatestTS(ctx, key)
 		if err != nil {
 			return false, err
 		}
@@ -209,9 +210,9 @@ func (ss *SuffixSnapshot) checkKeyTS(key string, ts typeutil.Timestamp) (bool, e
 }
 
 // loadLatestTS load the loatest ts for specified key
-func (ss *SuffixSnapshot) loadLatestTS(key string) error {
+func (ss *SuffixSnapshot) loadLatestTS(ctx context.Context, key string) error {
 	prefix := ss.composeSnapshotPrefix(key)
-	keys, _, err := ss.MetaKv.LoadWithPrefix(prefix)
+	keys, _, err := ss.MetaKv.LoadWithPrefix(ctx, prefix)
 	if err != nil {
 		log.Warn("SuffixSnapshot MetaKv LoadWithPrefix failed", zap.String("key", key),
 			zap.Error(err))
@@ -273,11 +274,11 @@ func binarySearchRecords(records []tsv, ts typeutil.Timestamp) (string, bool) {
 // if ts == 0, SuffixSnapshot works as a MetaKv
 // otherwise, SuffixSnapshot will store a ts-key as "key[sep]ts"-value pair in snapshot path
 // and for acceleration store original key-value if ts is the latest
-func (ss *SuffixSnapshot) Save(key string, value string, ts typeutil.Timestamp) error {
+func (ss *SuffixSnapshot) Save(ctx context.Context, key string, value string, ts typeutil.Timestamp) error {
 	// if ts == 0, act like MetaKv
 	// will not update lastestTs since ts not not valid
 	if ts == 0 {
-		return ss.MetaKv.Save(key, value)
+		return ss.MetaKv.Save(ctx, key, value)
 	}
 
 	ss.Lock()
@@ -287,12 +288,12 @@ func (ss *SuffixSnapshot) Save(key string, value string, ts typeutil.Timestamp)
 
 	// provided key value is latest
 	// stores both tsKey and original key
-	after, err := ss.checkKeyTS(key, ts)
+	after, err := ss.checkKeyTS(ctx, key, ts)
 	if err != nil {
 		return err
 	}
 	if after {
-		err := ss.MetaKv.MultiSave(map[string]string{
+		err := ss.MetaKv.MultiSave(ctx, map[string]string{
 			key:   value,
 			tsKey: value,
 		})
@@ -304,14 +305,14 @@ func (ss *SuffixSnapshot) Save(key string, value string, ts typeutil.Timestamp)
 	}
 
 	// modifying history key, just save tskey-value
-	return ss.MetaKv.Save(tsKey, value)
+	return ss.MetaKv.Save(ctx, tsKey, value)
 }
 
-func (ss *SuffixSnapshot) Load(key string, ts typeutil.Timestamp) (string, error) {
+func (ss *SuffixSnapshot) Load(ctx context.Context, key string, ts typeutil.Timestamp) (string, error) {
 	// if ts == 0 or typeutil.MaxTimestamp, load latest by definition
 	// and with acceleration logic, just do load key will do
 	if ts == 0 || ts == typeutil.MaxTimestamp {
-		value, err := ss.MetaKv.Load(key)
+		value, err := ss.MetaKv.Load(ctx, key)
 		if ss.isTombstone(value) {
 			return "", errors.New("no value found")
 		}
@@ -319,7 +320,7 @@ func (ss *SuffixSnapshot) Load(key string, ts typeutil.Timestamp) (string, error
 	}
 
 	ss.Lock()
-	after, err := ss.checkKeyTS(key, ts)
+	after, err := ss.checkKeyTS(ctx, key, ts)
 	ss.Unlock()
 
 	ss.RLock()
@@ -329,7 +330,7 @@ func (ss *SuffixSnapshot) Load(key string, ts typeutil.Timestamp) (string, error
 		return "", err
 	}
 	if after {
-		value, err := ss.MetaKv.Load(key)
+		value, err := ss.MetaKv.Load(ctx, key)
 		if ss.isTombstone(value) {
 			return "", errors.New("no value found")
 		}
@@ -338,7 +339,7 @@ func (ss *SuffixSnapshot) Load(key string, ts typeutil.Timestamp) (string, error
 
 	// before ts, do time travel
 	// 1. load all tsKey with key/ prefix
-	keys, values, err := ss.MetaKv.LoadWithPrefix(ss.composeSnapshotPrefix(key))
+	keys, values, err := ss.MetaKv.LoadWithPrefix(ctx, ss.composeSnapshotPrefix(key))
 	if err != nil {
 		log.Warn("prefixSnapshot MetaKv LoadWithPrefix failed", zap.String("key", key), zap.Error(err))
 		return "", err
@@ -380,23 +381,23 @@ func (ss *SuffixSnapshot) Load(key string, ts typeutil.Timestamp) (string, error
 // MultiSave save multiple kvs
 // if ts == 0, act like MetaKv
 // each key-value will be treated using same logic like Save
-func (ss *SuffixSnapshot) MultiSave(kvs map[string]string, ts typeutil.Timestamp) error {
+func (ss *SuffixSnapshot) MultiSave(ctx context.Context, kvs map[string]string, ts typeutil.Timestamp) error {
 	// if ts == 0, act like MetaKv
 	if ts == 0 {
-		return ss.MetaKv.MultiSave(kvs)
+		return ss.MetaKv.MultiSave(ctx, kvs)
 	}
 	ss.Lock()
 	defer ss.Unlock()
 	var err error
 
 	// process each key, checks whether is the latest
-	execute, updateList, err := ss.generateSaveExecute(kvs, ts)
+	execute, updateList, err := ss.generateSaveExecute(ctx, kvs, ts)
 	if err != nil {
 		return err
 	}
 
 	// multi save execute map; if succeeds, update ts in the update list
-	err = ss.MetaKv.MultiSave(execute)
+	err = ss.MetaKv.MultiSave(ctx, execute)
 	if err == nil {
 		for _, key := range updateList {
 			ss.lastestTS[key] = ts
@@ -407,7 +408,7 @@ func (ss *SuffixSnapshot) MultiSave(kvs map[string]string, ts typeutil.Timestamp
 
 // generateSaveExecute examine each key is the after the corresponding latest
 // returns calculated execute map and update ts list
-func (ss *SuffixSnapshot) generateSaveExecute(kvs map[string]string, ts typeutil.Timestamp) (map[string]string, []string, error) {
+func (ss *SuffixSnapshot) generateSaveExecute(ctx context.Context, kvs map[string]string, ts typeutil.Timestamp) (map[string]string, []string, error) {
 	var after bool
 	var err error
 	execute := make(map[string]string)
@@ -416,7 +417,7 @@ func (ss *SuffixSnapshot) generateSaveExecute(kvs map[string]string, ts typeutil
 		tsKey := ss.composeTSKey(key, ts)
 		// provided key value is latest
 		// stores both tsKey and original key
-		after, err = ss.checkKeyTS(key, ts)
+		after, err = ss.checkKeyTS(ctx, key, ts)
 		if err != nil {
 			return nil, nil, err
 		}
@@ -432,10 +433,10 @@ func (ss *SuffixSnapshot) generateSaveExecute(kvs map[string]string, ts typeutil
 }
 
 // LoadWithPrefix load keys with provided prefix and returns value in the ts
-func (ss *SuffixSnapshot) LoadWithPrefix(key string, ts typeutil.Timestamp) ([]string, []string, error) {
+func (ss *SuffixSnapshot) LoadWithPrefix(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
 	// ts 0 case shall be treated as fetch latest/current value
 	if ts == 0 || ts == typeutil.MaxTimestamp {
-		keys, values, err := ss.MetaKv.LoadWithPrefix(key)
+		keys, values, err := ss.MetaKv.LoadWithPrefix(ctx, key)
 		fks := keys[:0]   // make([]string, 0, len(keys))
 		fvs := values[:0] // make([]string, 0, len(values))
 		// hide rootPrefix from return value
@@ -469,7 +470,7 @@ func (ss *SuffixSnapshot) LoadWithPrefix(key string, ts typeutil.Timestamp) ([]s
 		resultValues = append(resultValues, value)
 	}
 
-	err := ss.MetaKv.WalkWithPrefix(prefix, PaginationSize, func(k []byte, v []byte) error {
+	err := ss.MetaKv.WalkWithPrefix(ctx, prefix, PaginationSize, func(k []byte, v []byte) error {
 		sKey := string(k)
 		sValue := string(v)
 
@@ -506,24 +507,24 @@ func (ss *SuffixSnapshot) LoadWithPrefix(key string, ts typeutil.Timestamp) ([]s
 // MultiSaveAndRemove save muiltple kvs and remove as well
 // if ts == 0, act like MetaKv
 // each key-value will be treated in same logic like Save
-func (ss *SuffixSnapshot) MultiSaveAndRemove(saves map[string]string, removals []string, ts typeutil.Timestamp) error {
+func (ss *SuffixSnapshot) MultiSaveAndRemove(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error {
 	// if ts == 0, act like MetaKv
 	if ts == 0 {
-		return ss.MetaKv.MultiSaveAndRemove(saves, removals)
+		return ss.MetaKv.MultiSaveAndRemove(ctx, saves, removals)
 	}
 	ss.Lock()
 	defer ss.Unlock()
 	var err error
 
 	// process each key, checks whether is the latest
-	execute, updateList, err := ss.generateSaveExecute(saves, ts)
+	execute, updateList, err := ss.generateSaveExecute(ctx, saves, ts)
 	if err != nil {
 		return err
 	}
 
 	// load each removal, change execution to adding tombstones
 	for _, removal := range removals {
-		value, err := ss.MetaKv.Load(removal)
+		value, err := ss.MetaKv.Load(ctx, removal)
 		if err != nil {
 			log.Warn("SuffixSnapshot MetaKv Load failed", zap.String("key", removal), zap.Error(err))
 			if errors.Is(err, merr.ErrIoKeyNotFound) {
@@ -541,7 +542,7 @@ func (ss *SuffixSnapshot) MultiSaveAndRemove(saves map[string]string, removals [
 	}
 
 	// multi save execute map; if succeeds, update ts in the update list
-	err = ss.MetaKv.MultiSave(execute)
+	err = ss.MetaKv.MultiSave(ctx, execute)
 	if err == nil {
 		for _, key := range updateList {
 			ss.lastestTS[key] = ts
@@ -553,24 +554,24 @@ func (ss *SuffixSnapshot) MultiSaveAndRemove(saves map[string]string, removals [
 // MultiSaveAndRemoveWithPrefix save muiltple kvs and remove as well
 // if ts == 0, act like MetaKv
 // each key-value will be treated in same logic like Save
-func (ss *SuffixSnapshot) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string, ts typeutil.Timestamp) error {
+func (ss *SuffixSnapshot) MultiSaveAndRemoveWithPrefix(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error {
 	// if ts == 0, act like MetaKv
 	if ts == 0 {
-		return ss.MetaKv.MultiSaveAndRemoveWithPrefix(saves, removals)
+		return ss.MetaKv.MultiSaveAndRemoveWithPrefix(ctx, saves, removals)
 	}
 	ss.Lock()
 	defer ss.Unlock()
 	var err error
 
 	// process each key, checks whether is the latest
-	execute, updateList, err := ss.generateSaveExecute(saves, ts)
+	execute, updateList, err := ss.generateSaveExecute(ctx, saves, ts)
 	if err != nil {
 		return err
 	}
 
 	// load each removal, change execution to adding tombstones
 	for _, removal := range removals {
-		keys, values, err := ss.MetaKv.LoadWithPrefix(removal)
+		keys, values, err := ss.MetaKv.LoadWithPrefix(ctx, removal)
 		if err != nil {
 			log.Warn("SuffixSnapshot MetaKv LoadwithPrefix failed", zap.String("key", removal), zap.Error(err))
 			return err
@@ -589,7 +590,7 @@ func (ss *SuffixSnapshot) MultiSaveAndRemoveWithPrefix(saves map[string]string,
 	}
 
 	// multi save execute map; if succeeds, update ts in the update list
-	err = ss.MetaKv.MultiSave(execute)
+	err = ss.MetaKv.MultiSave(ctx, execute)
 	if err == nil {
 		for _, key := range updateList {
 			ss.lastestTS[key] = ts
@@ -603,7 +604,7 @@ func (ss *SuffixSnapshot) Close() {
 }
 
 // startBackgroundGC the data will clean up if key ts!=0 and expired
-func (ss *SuffixSnapshot) startBackgroundGC() {
+func (ss *SuffixSnapshot) startBackgroundGC(ctx context.Context) {
 	log.Debug("suffix snapshot GC goroutine start!")
 	ticker := time.NewTicker(60 * time.Minute)
 	defer ticker.Stop()
@@ -614,7 +615,7 @@ func (ss *SuffixSnapshot) startBackgroundGC() {
 			log.Warn("quit suffix snapshot GC goroutine!")
 			return
 		case now := <-ticker.C:
-			err := ss.removeExpiredKvs(now)
+			err := ss.removeExpiredKvs(ctx, now)
 			if err != nil {
 				log.Warn("remove expired data fail during GC", zap.Error(err))
 			}
@@ -635,7 +636,7 @@ func (ss *SuffixSnapshot) getOriginalKey(snapshotKey string) (string, error) {
 	return prefix[ss.snapshotLen:], nil
 }
 
-func (ss *SuffixSnapshot) batchRemoveExpiredKvs(keyGroup []string, originalKey string, includeOriginalKey bool) error {
+func (ss *SuffixSnapshot) batchRemoveExpiredKvs(ctx context.Context, keyGroup []string, originalKey string, includeOriginalKey bool) error {
 	if includeOriginalKey {
 		keyGroup = append(keyGroup, originalKey)
 	}
@@ -643,7 +644,7 @@ func (ss *SuffixSnapshot) batchRemoveExpiredKvs(keyGroup []string, originalKey s
 	// to protect txn finished with ascend order, reverse the latest kv with tombstone to tail of array
 	sort.Strings(keyGroup)
 	removeFn := func(partialKeys []string) error {
-		return ss.MetaKv.MultiRemove(partialKeys)
+		return ss.MetaKv.MultiRemove(ctx, partialKeys)
 	}
 	return etcd.RemoveByBatchWithLimit(keyGroup, util.MaxEtcdTxnNum, removeFn)
 }
@@ -651,7 +652,7 @@ func (ss *SuffixSnapshot) batchRemoveExpiredKvs(keyGroup []string, originalKey s
 // removeExpiredKvs removes expired key-value pairs from the snapshot
 // It walks through all keys with the snapshot prefix, groups them by original key,
 // and removes expired versions or all versions if the original key has been deleted
-func (ss *SuffixSnapshot) removeExpiredKvs(now time.Time) error {
+func (ss *SuffixSnapshot) removeExpiredKvs(ctx context.Context, now time.Time) error {
 	ttlTime := paramtable.Get().ServiceParam.MetaStoreCfg.SnapshotTTLSeconds.GetAsDuration(time.Second)
 	reserveTime := paramtable.Get().ServiceParam.MetaStoreCfg.SnapshotReserveTimeSeconds.GetAsDuration(time.Second)
 
@@ -667,7 +668,7 @@ func (ss *SuffixSnapshot) removeExpiredKvs(now time.Time) error {
 		}
 		if ss.isTombstone(latestOriginValue) {
 			// If deleted, remove all versions including the original key
-			return ss.batchRemoveExpiredKvs(candidateExpiredKeys, curOriginalKey, totalVersions == len(candidateExpiredKeys))
+			return ss.batchRemoveExpiredKvs(ctx, candidateExpiredKeys, curOriginalKey, totalVersions == len(candidateExpiredKeys))
 		}
 
 		// If not deleted, check for expired versions
@@ -680,13 +681,13 @@ func (ss *SuffixSnapshot) removeExpiredKvs(now time.Time) error {
 			}
 		}
 		if len(expiredKeys) > 0 {
-			return ss.batchRemoveExpiredKvs(expiredKeys, curOriginalKey, false)
+			return ss.batchRemoveExpiredKvs(ctx, expiredKeys, curOriginalKey, false)
 		}
 		return nil
 	}
 
 	// Walk through all keys with the snapshot prefix
-	err := ss.MetaKv.WalkWithPrefix(ss.snapshotPrefix, PaginationSize, func(k []byte, v []byte) error {
+	err := ss.MetaKv.WalkWithPrefix(ctx, ss.snapshotPrefix, PaginationSize, func(k []byte, v []byte) error {
 		key := ss.hideRootPrefix(string(k))
 		ts, ok := ss.isTSKey(key)
 		if !ok {
diff --git a/internal/metastore/kv/rootcoord/suffix_snapshot_test.go b/internal/metastore/kv/rootcoord/suffix_snapshot_test.go
index 326716919ff85..ec33cc364d508 100644
--- a/internal/metastore/kv/rootcoord/suffix_snapshot_test.go
+++ b/internal/metastore/kv/rootcoord/suffix_snapshot_test.go
@@ -17,6 +17,7 @@
 package rootcoord
 
 import (
+	"context"
 	"fmt"
 	"math/rand"
 	"os"
@@ -304,27 +305,27 @@ func Test_SuffixSnapshotLoad(t *testing.T) {
 	for i := 0; i < 20; i++ {
 		vtso = typeutil.Timestamp(100 + i*5)
 		ts := ftso()
-		err = ss.Save("key", fmt.Sprintf("value-%d", i), ts)
+		err = ss.Save(context.TODO(), "key", fmt.Sprintf("value-%d", i), ts)
 		assert.NoError(t, err)
 		assert.Equal(t, vtso, ts)
 	}
 	for i := 0; i < 20; i++ {
-		val, err := ss.Load("key", typeutil.Timestamp(100+i*5+2))
+		val, err := ss.Load(context.TODO(), "key", typeutil.Timestamp(100+i*5+2))
 		t.Log("ts:", typeutil.Timestamp(100+i*5+2), i, val)
 		assert.NoError(t, err)
 		assert.Equal(t, fmt.Sprintf("value-%d", i), val)
 	}
-	val, err := ss.Load("key", 0)
+	val, err := ss.Load(context.TODO(), "key", 0)
 	assert.NoError(t, err)
 	assert.Equal(t, "value-19", val)
 
 	for i := 0; i < 20; i++ {
-		val, err := ss.Load("key", typeutil.Timestamp(100+i*5+2))
+		val, err := ss.Load(context.TODO(), "key", typeutil.Timestamp(100+i*5+2))
 		assert.NoError(t, err)
 		assert.Equal(t, val, fmt.Sprintf("value-%d", i))
 	}
 
-	ss.RemoveWithPrefix("")
+	ss.RemoveWithPrefix(context.TODO(), "")
 }
 
 func Test_SuffixSnapshotMultiSave(t *testing.T) {
@@ -360,12 +361,12 @@ func Test_SuffixSnapshotMultiSave(t *testing.T) {
 		saves := map[string]string{"k1": fmt.Sprintf("v1-%d", i), "k2": fmt.Sprintf("v2-%d", i)}
 		vtso = typeutil.Timestamp(100 + i*5)
 		ts := ftso()
-		err = ss.MultiSave(saves, ts)
+		err = ss.MultiSave(context.TODO(), saves, ts)
 		assert.NoError(t, err)
 		assert.Equal(t, vtso, ts)
 	}
 	for i := 0; i < 20; i++ {
-		keys, vals, err := ss.LoadWithPrefix("k", typeutil.Timestamp(100+i*5+2))
+		keys, vals, err := ss.LoadWithPrefix(context.TODO(), "k", typeutil.Timestamp(100+i*5+2))
 		t.Log(i, keys, vals)
 		assert.NoError(t, err)
 		assert.Equal(t, len(keys), len(vals))
@@ -375,7 +376,7 @@ func Test_SuffixSnapshotMultiSave(t *testing.T) {
 		assert.Equal(t, vals[0], fmt.Sprintf("v1-%d", i))
 		assert.Equal(t, vals[1], fmt.Sprintf("v2-%d", i))
 	}
-	keys, vals, err := ss.LoadWithPrefix("k", 0)
+	keys, vals, err := ss.LoadWithPrefix(context.TODO(), "k", 0)
 	assert.NoError(t, err)
 	assert.Equal(t, len(keys), len(vals))
 	assert.Equal(t, len(keys), 2)
@@ -385,7 +386,7 @@ func Test_SuffixSnapshotMultiSave(t *testing.T) {
 	assert.Equal(t, vals[1], "v2-19")
 
 	for i := 0; i < 20; i++ {
-		keys, vals, err := ss.LoadWithPrefix("k", typeutil.Timestamp(100+i*5+2))
+		keys, vals, err := ss.LoadWithPrefix(context.TODO(), "k", typeutil.Timestamp(100+i*5+2))
 		assert.NoError(t, err)
 		assert.Equal(t, len(keys), len(vals))
 		assert.Equal(t, len(keys), 2)
@@ -393,9 +394,9 @@ func Test_SuffixSnapshotMultiSave(t *testing.T) {
 		assert.ElementsMatch(t, vals, []string{fmt.Sprintf("v1-%d", i), fmt.Sprintf("v2-%d", i)})
 	}
 	// mix non ts k-v
-	err = ss.Save("kextra", "extra-value", 0)
+	err = ss.Save(context.TODO(), "kextra", "extra-value", 0)
 	assert.NoError(t, err)
-	keys, vals, err = ss.LoadWithPrefix("k", typeutil.Timestamp(300))
+	keys, vals, err = ss.LoadWithPrefix(context.TODO(), "k", typeutil.Timestamp(300))
 	assert.NoError(t, err)
 	assert.Equal(t, len(keys), len(vals))
 	assert.Equal(t, len(keys), 2)
@@ -403,7 +404,7 @@ func Test_SuffixSnapshotMultiSave(t *testing.T) {
 	assert.ElementsMatch(t, vals, []string{"v1-19", "v2-19"})
 
 	// clean up
-	ss.RemoveWithPrefix("")
+	ss.RemoveWithPrefix(context.TODO(), "")
 }
 
 func Test_SuffixSnapshotRemoveExpiredKvs(t *testing.T) {
@@ -433,12 +434,12 @@ func Test_SuffixSnapshotRemoveExpiredKvs(t *testing.T) {
 	defer ss.Close()
 
 	saveFn := func(key, value string, ts typeutil.Timestamp) {
-		err = ss.Save(key, value, ts)
+		err = ss.Save(context.TODO(), key, value, ts)
 		assert.NoError(t, err)
 	}
 
 	multiSaveFn := func(kvs map[string]string, ts typeutil.Timestamp) {
-		err = ss.MultiSave(kvs, ts)
+		err = ss.MultiSave(context.TODO(), kvs, ts)
 		assert.NoError(t, err)
 	}
 
@@ -477,7 +478,7 @@ func Test_SuffixSnapshotRemoveExpiredKvs(t *testing.T) {
 
 	countPrefix := func(prefix string) int {
 		cnt := 0
-		err := etcdkv.WalkWithPrefix("", 10, func(key []byte, value []byte) error {
+		err := etcdkv.WalkWithPrefix(context.TODO(), "", 10, func(key []byte, value []byte) error {
 			cnt++
 			return nil
 		})
@@ -495,14 +496,14 @@ func Test_SuffixSnapshotRemoveExpiredKvs(t *testing.T) {
 		cnt := countPrefix(prefix)
 		assert.Equal(t, keyCnt*keyVersion+keyCnt, cnt)
 
-		err = ss.removeExpiredKvs(now)
+		err = ss.removeExpiredKvs(context.TODO(), now)
 		assert.NoError(t, err)
 
 		cnt = countPrefix(prefix)
 		assert.Equal(t, keyCnt*keyVersion+keyCnt-(expiredKCnt*keyVersion+expiredKCnt), cnt)
 
 		// clean all data
-		err := etcdkv.RemoveWithPrefix("")
+		err := etcdkv.RemoveWithPrefix(context.TODO(), "")
 		assert.NoError(t, err)
 	})
 
@@ -536,14 +537,14 @@ func Test_SuffixSnapshotRemoveExpiredKvs(t *testing.T) {
 		assert.Equal(t, 12, cnt)
 
 		// err = ss.removeExpiredKvs(now, time.Duration(50)*time.Millisecond)
-		err = ss.removeExpiredKvs(now)
+		err = ss.removeExpiredKvs(context.TODO(), now)
 		assert.NoError(t, err)
 
 		cnt = countPrefix(prefix)
 		assert.Equal(t, 4, cnt)
 
 		// clean all data
-		err := etcdkv.RemoveWithPrefix("")
+		err := etcdkv.RemoveWithPrefix(context.TODO(), "")
 		assert.NoError(t, err)
 	})
 
@@ -577,31 +578,31 @@ func Test_SuffixSnapshotRemoveExpiredKvs(t *testing.T) {
 		assert.Equal(t, 12, cnt)
 
 		// err = ss.removeExpiredKvs(now, time.Duration(50)*time.Millisecond)
-		err = ss.removeExpiredKvs(now)
+		err = ss.removeExpiredKvs(context.TODO(), now)
 		assert.NoError(t, err)
 
 		cnt = countPrefix(prefix)
 		assert.Equal(t, 1, cnt)
 
 		// clean all data
-		err := etcdkv.RemoveWithPrefix("")
+		err := etcdkv.RemoveWithPrefix(context.TODO(), "")
 		assert.NoError(t, err)
 	})
 
 	t.Run("parse ts fail", func(t *testing.T) {
 		prefix := fmt.Sprintf("prefix%d", rand.Int())
 		key := fmt.Sprintf("%s-%s", prefix, "ts_error-ts")
-		err = etcdkv.Save(ss.composeSnapshotPrefix(key), "")
+		err = etcdkv.Save(context.TODO(), ss.composeSnapshotPrefix(key), "")
 		assert.NoError(t, err)
 
-		err = ss.removeExpiredKvs(now)
+		err = ss.removeExpiredKvs(context.TODO(), now)
 		assert.NoError(t, err)
 
 		cnt := countPrefix(prefix)
 		assert.Equal(t, 1, cnt)
 
 		// clean all data
-		err := etcdkv.RemoveWithPrefix("")
+		err := etcdkv.RemoveWithPrefix(context.TODO(), "")
 		assert.NoError(t, err)
 	})
 
@@ -610,14 +611,14 @@ func Test_SuffixSnapshotRemoveExpiredKvs(t *testing.T) {
 		rootPath := "root/"
 		kv := mocks.NewMetaKv(t)
 		kv.EXPECT().
-			WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).
+			WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
 			Return(errors.New("error"))
 
 		ss, err := NewSuffixSnapshot(kv, sep, rootPath, snapshotPrefix)
 		assert.NotNil(t, ss)
 		assert.NoError(t, err)
 
-		err = ss.removeExpiredKvs(time.Now())
+		err = ss.removeExpiredKvs(context.TODO(), time.Now())
 		assert.Error(t, err)
 	})
 }
@@ -656,7 +657,7 @@ func Test_SuffixSnapshotMultiSaveAndRemoveWithPrefix(t *testing.T) {
 	for i := 0; i < 20; i++ {
 		vtso = typeutil.Timestamp(100 + i*5)
 		ts := ftso()
-		err = ss.Save(fmt.Sprintf("kd-%04d", i), fmt.Sprintf("value-%d", i), ts)
+		err = ss.Save(context.TODO(), fmt.Sprintf("kd-%04d", i), fmt.Sprintf("value-%d", i), ts)
 		assert.NoError(t, err)
 		assert.Equal(t, vtso, ts)
 	}
@@ -665,53 +666,53 @@ func Test_SuffixSnapshotMultiSaveAndRemoveWithPrefix(t *testing.T) {
 		dm := []string{fmt.Sprintf("kd-%04d", i-20)}
 		vtso = typeutil.Timestamp(100 + i*5)
 		ts := ftso()
-		err = ss.MultiSaveAndRemoveWithPrefix(sm, dm, ts)
+		err = ss.MultiSaveAndRemoveWithPrefix(context.TODO(), sm, dm, ts)
 		assert.NoError(t, err)
 		assert.Equal(t, vtso, ts)
 	}
 	for i := 0; i < 20; i++ {
-		val, err := ss.Load(fmt.Sprintf("kd-%04d", i), typeutil.Timestamp(100+i*5+2))
+		val, err := ss.Load(context.TODO(), fmt.Sprintf("kd-%04d", i), typeutil.Timestamp(100+i*5+2))
 		assert.NoError(t, err)
 		assert.Equal(t, fmt.Sprintf("value-%d", i), val)
-		_, vals, err := ss.LoadWithPrefix("kd-", typeutil.Timestamp(100+i*5+2))
+		_, vals, err := ss.LoadWithPrefix(context.TODO(), "kd-", typeutil.Timestamp(100+i*5+2))
 		assert.NoError(t, err)
 		assert.Equal(t, i+1, len(vals))
 	}
 	for i := 20; i < 40; i++ {
-		val, err := ss.Load("ks", typeutil.Timestamp(100+i*5+2))
+		val, err := ss.Load(context.TODO(), "ks", typeutil.Timestamp(100+i*5+2))
 		assert.NoError(t, err)
 		assert.Equal(t, fmt.Sprintf("value-%d", i), val)
-		_, vals, err := ss.LoadWithPrefix("kd-", typeutil.Timestamp(100+i*5+2))
+		_, vals, err := ss.LoadWithPrefix(context.TODO(), "kd-", typeutil.Timestamp(100+i*5+2))
 		assert.NoError(t, err)
 		assert.Equal(t, 39-i, len(vals))
 	}
 
 	for i := 0; i < 20; i++ {
-		val, err := ss.Load(fmt.Sprintf("kd-%04d", i), typeutil.Timestamp(100+i*5+2))
+		val, err := ss.Load(context.TODO(), fmt.Sprintf("kd-%04d", i), typeutil.Timestamp(100+i*5+2))
 		assert.NoError(t, err)
 		assert.Equal(t, fmt.Sprintf("value-%d", i), val)
-		_, vals, err := ss.LoadWithPrefix("kd-", typeutil.Timestamp(100+i*5+2))
+		_, vals, err := ss.LoadWithPrefix(context.TODO(), "kd-", typeutil.Timestamp(100+i*5+2))
 		assert.NoError(t, err)
 		assert.Equal(t, i+1, len(vals))
 	}
 	for i := 20; i < 40; i++ {
-		val, err := ss.Load("ks", typeutil.Timestamp(100+i*5+2))
+		val, err := ss.Load(context.TODO(), "ks", typeutil.Timestamp(100+i*5+2))
 		assert.NoError(t, err)
 		assert.Equal(t, fmt.Sprintf("value-%d", i), val)
-		_, vals, err := ss.LoadWithPrefix("kd-", typeutil.Timestamp(100+i*5+2))
+		_, vals, err := ss.LoadWithPrefix(context.TODO(), "kd-", typeutil.Timestamp(100+i*5+2))
 		assert.NoError(t, err)
 		assert.Equal(t, 39-i, len(vals))
 	}
 	// try to load
-	_, err = ss.Load("kd-0000", 500)
+	_, err = ss.Load(context.TODO(), "kd-0000", 500)
 	assert.Error(t, err)
-	_, err = ss.Load("kd-0000", 0)
+	_, err = ss.Load(context.TODO(), "kd-0000", 0)
 	assert.Error(t, err)
-	_, err = ss.Load("kd-0000", 1)
+	_, err = ss.Load(context.TODO(), "kd-0000", 1)
 	assert.Error(t, err)
 
 	// cleanup
-	ss.MultiSaveAndRemoveWithPrefix(map[string]string{}, []string{""}, 0)
+	ss.MultiSaveAndRemoveWithPrefix(context.TODO(), map[string]string{}, []string{""}, 0)
 }
 
 func Test_SuffixSnapshotMultiSaveAndRemove(t *testing.T) {
@@ -748,7 +749,7 @@ func Test_SuffixSnapshotMultiSaveAndRemove(t *testing.T) {
 	for i := 0; i < 20; i++ {
 		vtso = typeutil.Timestamp(100 + i*5)
 		ts := ftso()
-		err = ss.Save(fmt.Sprintf("kd-%04d", i), fmt.Sprintf("value-%d", i), ts)
+		err = ss.Save(context.TODO(), fmt.Sprintf("kd-%04d", i), fmt.Sprintf("value-%d", i), ts)
 		assert.NoError(t, err)
 		assert.Equal(t, vtso, ts)
 	}
@@ -757,37 +758,37 @@ func Test_SuffixSnapshotMultiSaveAndRemove(t *testing.T) {
 		dm := []string{fmt.Sprintf("kd-%04d", i-20)}
 		vtso = typeutil.Timestamp(100 + i*5)
 		ts := ftso()
-		err = ss.MultiSaveAndRemove(sm, dm, ts)
+		err = ss.MultiSaveAndRemove(context.TODO(), sm, dm, ts)
 		assert.NoError(t, err)
 		assert.Equal(t, vtso, ts)
 	}
 	for i := 0; i < 20; i++ {
-		val, err := ss.Load(fmt.Sprintf("kd-%04d", i), typeutil.Timestamp(100+i*5+2))
+		val, err := ss.Load(context.TODO(), fmt.Sprintf("kd-%04d", i), typeutil.Timestamp(100+i*5+2))
 		assert.NoError(t, err)
 		assert.Equal(t, fmt.Sprintf("value-%d", i), val)
-		_, vals, err := ss.LoadWithPrefix("kd-", typeutil.Timestamp(100+i*5+2))
+		_, vals, err := ss.LoadWithPrefix(context.TODO(), "kd-", typeutil.Timestamp(100+i*5+2))
 		assert.NoError(t, err)
 		assert.Equal(t, i+1, len(vals))
 	}
 	for i := 20; i < 40; i++ {
-		val, err := ss.Load("ks", typeutil.Timestamp(100+i*5+2))
+		val, err := ss.Load(context.TODO(), "ks", typeutil.Timestamp(100+i*5+2))
 		assert.NoError(t, err)
 		assert.Equal(t, fmt.Sprintf("value-%d", i), val)
-		_, vals, err := ss.LoadWithPrefix("kd-", typeutil.Timestamp(100+i*5+2))
+		_, vals, err := ss.LoadWithPrefix(context.TODO(), "kd-", typeutil.Timestamp(100+i*5+2))
 		assert.NoError(t, err)
 		assert.Equal(t, 39-i, len(vals))
 	}
 
 	// try to load
-	_, err = ss.Load("kd-0000", 500)
+	_, err = ss.Load(context.TODO(), "kd-0000", 500)
 	assert.Error(t, err)
-	_, err = ss.Load("kd-0000", 0)
+	_, err = ss.Load(context.TODO(), "kd-0000", 0)
 	assert.Error(t, err)
-	_, err = ss.Load("kd-0000", 1)
+	_, err = ss.Load(context.TODO(), "kd-0000", 1)
 	assert.Error(t, err)
 
 	// cleanup
-	ss.MultiSaveAndRemoveWithPrefix(map[string]string{}, []string{""}, 0)
+	ss.MultiSaveAndRemoveWithPrefix(context.TODO(), map[string]string{}, []string{""}, 0)
 }
 
 func TestSuffixSnapshot_LoadWithPrefix(t *testing.T) {
@@ -819,16 +820,16 @@ func TestSuffixSnapshot_LoadWithPrefix(t *testing.T) {
 	t.Run("parse ts fail", func(t *testing.T) {
 		prefix := fmt.Sprintf("prefix%d", rand.Int())
 		key := fmt.Sprintf("%s-%s", prefix, "ts_error-ts")
-		err = etcdkv.Save(ss.composeSnapshotPrefix(key), "")
+		err = etcdkv.Save(context.TODO(), ss.composeSnapshotPrefix(key), "")
 		assert.NoError(t, err)
 
-		keys, values, err := ss.LoadWithPrefix(prefix, 100)
+		keys, values, err := ss.LoadWithPrefix(context.TODO(), prefix, 100)
 		assert.NoError(t, err)
 		assert.Equal(t, 0, len(keys))
 		assert.Equal(t, 0, len(values))
 
 		// clean all data
-		err = etcdkv.RemoveWithPrefix("")
+		err = etcdkv.RemoveWithPrefix(context.TODO(), "")
 		assert.NoError(t, err)
 	})
 
@@ -837,14 +838,14 @@ func TestSuffixSnapshot_LoadWithPrefix(t *testing.T) {
 		rootPath := "root/"
 		kv := mocks.NewMetaKv(t)
 		kv.EXPECT().
-			WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).
+			WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
 			Return(errors.New("error"))
 
 		ss, err := NewSuffixSnapshot(kv, sep, rootPath, snapshotPrefix)
 		assert.NotNil(t, ss)
 		assert.NoError(t, err)
 
-		keys, values, err := ss.LoadWithPrefix("t", 100)
+		keys, values, err := ss.LoadWithPrefix(context.TODO(), "t", 100)
 		assert.Error(t, err)
 		assert.Nil(t, keys)
 		assert.Nil(t, values)
diff --git a/internal/metastore/kv/streamingcoord/kv_catalog.go b/internal/metastore/kv/streamingcoord/kv_catalog.go
index 390818a15c552..d3d804052e026 100644
--- a/internal/metastore/kv/streamingcoord/kv_catalog.go
+++ b/internal/metastore/kv/streamingcoord/kv_catalog.go
@@ -27,7 +27,7 @@ type catalog struct {
 
 // ListPChannels returns all pchannels
 func (c *catalog) ListPChannel(ctx context.Context) ([]*streamingpb.PChannelMeta, error) {
-	keys, values, err := c.metaKV.LoadWithPrefix(PChannelMeta)
+	keys, values, err := c.metaKV.LoadWithPrefix(ctx, PChannelMeta)
 	if err != nil {
 		return nil, err
 	}
@@ -56,7 +56,7 @@ func (c *catalog) SavePChannels(ctx context.Context, infos []*streamingpb.PChann
 		kvs[key] = string(v)
 	}
 	return etcd.SaveByBatchWithLimit(kvs, util.MaxEtcdTxnNum, func(partialKvs map[string]string) error {
-		return c.metaKV.MultiSave(partialKvs)
+		return c.metaKV.MultiSave(ctx, partialKvs)
 	})
 }
 
diff --git a/internal/metastore/kv/streamingcoord/kv_catalog_test.go b/internal/metastore/kv/streamingcoord/kv_catalog_test.go
index 2544b9c852965..227ad0469bca3 100644
--- a/internal/metastore/kv/streamingcoord/kv_catalog_test.go
+++ b/internal/metastore/kv/streamingcoord/kv_catalog_test.go
@@ -16,7 +16,7 @@ func TestCatalog(t *testing.T) {
 	kv := mock_kv.NewMockMetaKv(t)
 
 	kvStorage := make(map[string]string)
-	kv.EXPECT().LoadWithPrefix(mock.Anything).RunAndReturn(func(s string) ([]string, []string, error) {
+	kv.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, s string) ([]string, []string, error) {
 		keys := make([]string, 0, len(kvStorage))
 		vals := make([]string, 0, len(kvStorage))
 		for k, v := range kvStorage {
@@ -25,7 +25,7 @@ func TestCatalog(t *testing.T) {
 		}
 		return keys, vals, nil
 	})
-	kv.EXPECT().MultiSave(mock.Anything).RunAndReturn(func(kvs map[string]string) error {
+	kv.EXPECT().MultiSave(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, kvs map[string]string) error {
 		for k, v := range kvs {
 			kvStorage[k] = v
 		}
@@ -54,13 +54,13 @@ func TestCatalog(t *testing.T) {
 	assert.Len(t, metas, 2)
 
 	// error path.
-	kv.EXPECT().LoadWithPrefix(mock.Anything).Unset()
-	kv.EXPECT().LoadWithPrefix(mock.Anything).Return(nil, nil, errors.New("load error"))
+	kv.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Unset()
+	kv.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Return(nil, nil, errors.New("load error"))
 	metas, err = catalog.ListPChannel(context.Background())
 	assert.Error(t, err)
 	assert.Nil(t, metas)
 
-	kv.EXPECT().MultiSave(mock.Anything).Unset()
-	kv.EXPECT().MultiSave(mock.Anything).Return(errors.New("save error"))
+	kv.EXPECT().MultiSave(mock.Anything, mock.Anything).Unset()
+	kv.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(errors.New("save error"))
 	assert.Error(t, err)
 }
diff --git a/internal/metastore/kv/streamingnode/kv_catalog.go b/internal/metastore/kv/streamingnode/kv_catalog.go
index a7b4ac803f521..2d279ab6ea6ca 100644
--- a/internal/metastore/kv/streamingnode/kv_catalog.go
+++ b/internal/metastore/kv/streamingnode/kv_catalog.go
@@ -29,7 +29,7 @@ type catalog struct {
 
 func (c *catalog) ListSegmentAssignment(ctx context.Context, pChannelName string) ([]*streamingpb.SegmentAssignmentMeta, error) {
 	prefix := buildSegmentAssignmentMetaPath(pChannelName)
-	keys, values, err := c.metaKV.LoadWithPrefix(prefix)
+	keys, values, err := c.metaKV.LoadWithPrefix(ctx, prefix)
 	if err != nil {
 		return nil, err
 	}
@@ -66,7 +66,7 @@ func (c *catalog) SaveSegmentAssignments(ctx context.Context, pChannelName strin
 
 	if len(removes) > 0 {
 		if err := etcd.RemoveByBatchWithLimit(removes, util.MaxEtcdTxnNum, func(partialRemoves []string) error {
-			return c.metaKV.MultiRemove(partialRemoves)
+			return c.metaKV.MultiRemove(ctx, partialRemoves)
 		}); err != nil {
 			return err
 		}
@@ -74,7 +74,7 @@ func (c *catalog) SaveSegmentAssignments(ctx context.Context, pChannelName strin
 
 	if len(kvs) > 0 {
 		return etcd.SaveByBatchWithLimit(kvs, util.MaxEtcdTxnNum, func(partialKvs map[string]string) error {
-			return c.metaKV.MultiSave(partialKvs)
+			return c.metaKV.MultiSave(ctx, partialKvs)
 		})
 	}
 	return nil
diff --git a/internal/metastore/kv/streamingnode/kv_catalog_test.go b/internal/metastore/kv/streamingnode/kv_catalog_test.go
index 32aeadae0c788..35e8769664f57 100644
--- a/internal/metastore/kv/streamingnode/kv_catalog_test.go
+++ b/internal/metastore/kv/streamingnode/kv_catalog_test.go
@@ -19,15 +19,15 @@ func TestCatalog(t *testing.T) {
 	vs, err := proto.Marshal(&v)
 	assert.NoError(t, err)
 
-	kv.EXPECT().LoadWithPrefix(mock.Anything).Return([]string{k}, []string{string(vs)}, nil)
+	kv.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Return([]string{k}, []string{string(vs)}, nil)
 	catalog := NewCataLog(kv)
 	ctx := context.Background()
 	metas, err := catalog.ListSegmentAssignment(ctx, "p1")
 	assert.Len(t, metas, 1)
 	assert.NoError(t, err)
 
-	kv.EXPECT().MultiRemove(mock.Anything).Return(nil)
-	kv.EXPECT().MultiSave(mock.Anything).Return(nil)
+	kv.EXPECT().MultiRemove(mock.Anything, mock.Anything).Return(nil)
+	kv.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(nil)
 
 	err = catalog.SaveSegmentAssignments(ctx, "p1", []*streamingpb.SegmentAssignmentMeta{
 		{
diff --git a/pkg/kv/kv.go b/pkg/kv/kv.go
index c78aa0a960f7a..2386506b000f9 100644
--- a/pkg/kv/kv.go
+++ b/pkg/kv/kv.go
@@ -17,6 +17,8 @@
 package kv
 
 import (
+	"context"
+
 	clientv3 "go.etcd.io/etcd/client/v3"
 
 	"github.com/milvus-io/milvus/pkg/kv/predicates"
@@ -40,16 +42,16 @@ func NewCompareFailedError(err error) error {
 
 // BaseKV contains base operations of kv. Include save, load and remove.
 type BaseKV interface {
-	Load(key string) (string, error)
-	MultiLoad(keys []string) ([]string, error)
-	LoadWithPrefix(key string) ([]string, []string, error)
-	Save(key, value string) error
-	MultiSave(kvs map[string]string) error
-	Remove(key string) error
-	MultiRemove(keys []string) error
-	RemoveWithPrefix(key string) error
-	Has(key string) (bool, error)
-	HasPrefix(prefix string) (bool, error)
+	Load(ctx context.Context, key string) (string, error)
+	MultiLoad(ctx context.Context, keys []string) ([]string, error)
+	LoadWithPrefix(ctx context.Context, key string) ([]string, []string, error)
+	Save(ctx context.Context, key, value string) error
+	MultiSave(ctx context.Context, kvs map[string]string) error
+	Remove(ctx context.Context, key string) error
+	MultiRemove(ctx context.Context, keys []string) error
+	RemoveWithPrefix(ctx context.Context, key string) error
+	Has(ctx context.Context, key string) (bool, error)
+	HasPrefix(ctx context.Context, prefix string) (bool, error)
 	Close()
 }
 
@@ -58,8 +60,8 @@ type BaseKV interface {
 //go:generate mockery --name=TxnKV --with-expecter
 type TxnKV interface {
 	BaseKV
-	MultiSaveAndRemove(saves map[string]string, removals []string, preds ...predicates.Predicate) error
-	MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string, preds ...predicates.Predicate) error
+	MultiSaveAndRemove(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error
+	MultiSaveAndRemoveWithPrefix(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error
 }
 
 // MetaKv is TxnKV for metadata. It should save data with lease.
@@ -68,9 +70,9 @@ type TxnKV interface {
 type MetaKv interface {
 	TxnKV
 	GetPath(key string) string
-	LoadWithPrefix(key string) ([]string, []string, error)
-	CompareVersionAndSwap(key string, version int64, target string) (bool, error)
-	WalkWithPrefix(prefix string, paginationSize int, fn func([]byte, []byte) error) error
+	LoadWithPrefix(ctx context.Context, key string) ([]string, []string, error)
+	CompareVersionAndSwap(ctx context.Context, key string, version int64, target string) (bool, error)
+	WalkWithPrefix(ctx context.Context, prefix string, paginationSize int, fn func([]byte, []byte) error) error
 }
 
 // WatchKV is watchable MetaKv. As of today(2023/06/24), it's coupled with etcd.
@@ -78,19 +80,19 @@ type MetaKv interface {
 //go:generate mockery --name=WatchKV --with-expecter
 type WatchKV interface {
 	MetaKv
-	Watch(key string) clientv3.WatchChan
-	WatchWithPrefix(key string) clientv3.WatchChan
-	WatchWithRevision(key string, revision int64) clientv3.WatchChan
+	Watch(ctx context.Context, key string) clientv3.WatchChan
+	WatchWithPrefix(ctx context.Context, key string) clientv3.WatchChan
+	WatchWithRevision(ctx context.Context, key string, revision int64) clientv3.WatchChan
 }
 
 // SnapShotKV is TxnKV for snapshot data. It must save timestamp.
 //
 //go:generate mockery --name=SnapShotKV --with-expecter
 type SnapShotKV interface {
-	Save(key string, value string, ts typeutil.Timestamp) error
-	Load(key string, ts typeutil.Timestamp) (string, error)
-	MultiSave(kvs map[string]string, ts typeutil.Timestamp) error
-	LoadWithPrefix(key string, ts typeutil.Timestamp) ([]string, []string, error)
-	MultiSaveAndRemove(saves map[string]string, removals []string, ts typeutil.Timestamp) error
-	MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string, ts typeutil.Timestamp) error
+	Save(ctx context.Context, key string, value string, ts typeutil.Timestamp) error
+	Load(ctx context.Context, key string, ts typeutil.Timestamp) (string, error)
+	MultiSave(ctx context.Context, kvs map[string]string, ts typeutil.Timestamp) error
+	LoadWithPrefix(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error)
+	MultiSaveAndRemove(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error
+	MultiSaveAndRemoveWithPrefix(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error
 }
diff --git a/pkg/kv/rocksdb/rocksdb_kv.go b/pkg/kv/rocksdb/rocksdb_kv.go
index b81e7ead385f2..3b156c414f008 100644
--- a/pkg/kv/rocksdb/rocksdb_kv.go
+++ b/pkg/kv/rocksdb/rocksdb_kv.go
@@ -17,6 +17,7 @@
 package rocksdbkv
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/cockroachdb/errors"
@@ -97,7 +98,7 @@ func (kv *RocksdbKV) GetName() string {
 }
 
 // Load returns the value of specified key
-func (kv *RocksdbKV) Load(key string) (string, error) {
+func (kv *RocksdbKV) Load(ctx context.Context, key string) (string, error) {
 	if kv.DB == nil {
 		return "", fmt.Errorf("rocksdb instance is nil when load %s", key)
 	}
@@ -114,7 +115,7 @@ func (kv *RocksdbKV) Load(key string) (string, error) {
 	return string(value.Data()), nil
 }
 
-func (kv *RocksdbKV) LoadBytes(key string) ([]byte, error) {
+func (kv *RocksdbKV) LoadBytes(ctx context.Context, key string) ([]byte, error) {
 	if kv.DB == nil {
 		return nil, fmt.Errorf("rocksdb instance is nil when load %s", key)
 	}
@@ -139,7 +140,7 @@ func (kv *RocksdbKV) LoadBytes(key string) ([]byte, error) {
 
 // LoadWithPrefix returns a batch values of keys with a prefix
 // if prefix is "", then load every thing from the database
-func (kv *RocksdbKV) LoadWithPrefix(prefix string) ([]string, []string, error) {
+func (kv *RocksdbKV) LoadWithPrefix(ctx context.Context, prefix string) ([]string, []string, error) {
 	if kv.DB == nil {
 		return nil, nil, fmt.Errorf("rocksdb instance is nil when load %s", prefix)
 	}
@@ -164,7 +165,7 @@ func (kv *RocksdbKV) LoadWithPrefix(prefix string) ([]string, []string, error) {
 	return keys, values, nil
 }
 
-func (kv *RocksdbKV) Has(key string) (bool, error) {
+func (kv *RocksdbKV) Has(ctx context.Context, key string) (bool, error) {
 	if kv.DB == nil {
 		return false, fmt.Errorf("rocksdb instance is nil when check if has %s", key)
 	}
@@ -180,7 +181,7 @@ func (kv *RocksdbKV) Has(key string) (bool, error) {
 	return value.Size() != 0, nil
 }
 
-func (kv *RocksdbKV) HasPrefix(prefix string) (bool, error) {
+func (kv *RocksdbKV) HasPrefix(ctx context.Context, prefix string) (bool, error) {
 	if kv.DB == nil {
 		return false, fmt.Errorf("rocksdb instance is nil when check if has prefix %s", prefix)
 	}
@@ -194,7 +195,7 @@ func (kv *RocksdbKV) HasPrefix(prefix string) (bool, error) {
 	return iter.Valid(), nil
 }
 
-func (kv *RocksdbKV) LoadBytesWithPrefix(prefix string) ([]string, [][]byte, error) {
+func (kv *RocksdbKV) LoadBytesWithPrefix(ctx context.Context, prefix string) ([]string, [][]byte, error) {
 	if kv.DB == nil {
 		return nil, nil, fmt.Errorf("rocksdb instance is nil when load %s", prefix)
 	}
@@ -229,7 +230,7 @@ func (kv *RocksdbKV) LoadBytesWithPrefix(prefix string) ([]string, [][]byte, err
 }
 
 // MultiLoad load a batch of values by keys
-func (kv *RocksdbKV) MultiLoad(keys []string) ([]string, error) {
+func (kv *RocksdbKV) MultiLoad(ctx context.Context, keys []string) ([]string, error) {
 	if kv.DB == nil {
 		return nil, errors.New("rocksdb instance is nil when do MultiLoad")
 	}
@@ -253,7 +254,7 @@ func (kv *RocksdbKV) MultiLoad(keys []string) ([]string, error) {
 	return values, nil
 }
 
-func (kv *RocksdbKV) MultiLoadBytes(keys []string) ([][]byte, error) {
+func (kv *RocksdbKV) MultiLoadBytes(ctx context.Context, keys []string) ([][]byte, error) {
 	if kv.DB == nil {
 		return nil, errors.New("rocksdb instance is nil when do MultiLoad")
 	}
@@ -281,7 +282,7 @@ func (kv *RocksdbKV) MultiLoadBytes(keys []string) ([][]byte, error) {
 }
 
 // Save a pair of key-value
-func (kv *RocksdbKV) Save(key, value string) error {
+func (kv *RocksdbKV) Save(ctx context.Context, key, value string) error {
 	if kv.DB == nil {
 		return errors.New("rocksdb instance is nil when do save")
 	}
@@ -295,7 +296,7 @@ func (kv *RocksdbKV) Save(key, value string) error {
 	return kv.DB.Put(kv.WriteOptions, []byte(key), []byte(value))
 }
 
-func (kv *RocksdbKV) SaveBytes(key string, value []byte) error {
+func (kv *RocksdbKV) SaveBytes(ctx context.Context, key string, value []byte) error {
 	if kv.DB == nil {
 		return errors.New("rocksdb instance is nil when do save")
 	}
@@ -310,7 +311,7 @@ func (kv *RocksdbKV) SaveBytes(key string, value []byte) error {
 }
 
 // MultiSave a batch of key-values
-func (kv *RocksdbKV) MultiSave(kvs map[string]string) error {
+func (kv *RocksdbKV) MultiSave(ctx context.Context, kvs map[string]string) error {
 	if kv.DB == nil {
 		return errors.New("rocksdb instance is nil when do MultiSave")
 	}
@@ -325,7 +326,7 @@ func (kv *RocksdbKV) MultiSave(kvs map[string]string) error {
 	return kv.DB.Write(kv.WriteOptions, writeBatch)
 }
 
-func (kv *RocksdbKV) MultiSaveBytes(kvs map[string][]byte) error {
+func (kv *RocksdbKV) MultiSaveBytes(ctx context.Context, kvs map[string][]byte) error {
 	if kv.DB == nil {
 		return errors.New("rocksdb instance is nil when do MultiSave")
 	}
@@ -342,7 +343,7 @@ func (kv *RocksdbKV) MultiSaveBytes(kvs map[string][]byte) error {
 
 // RemoveWithPrefix removes a batch of key-values with specified prefix
 // If prefix is "", then all data in the rocksdb kv will be deleted
-func (kv *RocksdbKV) RemoveWithPrefix(prefix string) error {
+func (kv *RocksdbKV) RemoveWithPrefix(ctx context.Context, prefix string) error {
 	if kv.DB == nil {
 		return errors.New("rocksdb instance is nil when do RemoveWithPrefix")
 	}
@@ -355,17 +356,17 @@ func (kv *RocksdbKV) RemoveWithPrefix(prefix string) error {
 		// seek to the last key
 		iter.SeekToLast()
 		if iter.Valid() {
-			return kv.DeleteRange(prefix, typeutil.AddOne(string(iter.Key().Data())))
+			return kv.DeleteRange(ctx, prefix, typeutil.AddOne(string(iter.Key().Data())))
 		}
 		// nothing in the range, skip
 		return nil
 	}
 	prefixEnd := typeutil.AddOne(prefix)
-	return kv.DeleteRange(prefix, prefixEnd)
+	return kv.DeleteRange(ctx, prefix, prefixEnd)
 }
 
 // Remove is used to remove a pair of key-value
-func (kv *RocksdbKV) Remove(key string) error {
+func (kv *RocksdbKV) Remove(ctx context.Context, key string) error {
 	if kv.DB == nil {
 		return errors.New("rocksdb instance is nil when do Remove")
 	}
@@ -377,7 +378,7 @@ func (kv *RocksdbKV) Remove(key string) error {
 }
 
 // MultiRemove is used to remove a batch of key-values
-func (kv *RocksdbKV) MultiRemove(keys []string) error {
+func (kv *RocksdbKV) MultiRemove(ctx context.Context, keys []string) error {
 	if kv.DB == nil {
 		return errors.New("rocksdb instance is nil when do MultiRemove")
 	}
@@ -391,7 +392,7 @@ func (kv *RocksdbKV) MultiRemove(keys []string) error {
 }
 
 // MultiSaveAndRemove provides a transaction to execute a batch of operations
-func (kv *RocksdbKV) MultiSaveAndRemove(saves map[string]string, removals []string, preds ...predicates.Predicate) error {
+func (kv *RocksdbKV) MultiSaveAndRemove(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error {
 	if len(preds) > 0 {
 		return merr.WrapErrServiceUnavailable("predicates not supported")
 	}
@@ -411,7 +412,7 @@ func (kv *RocksdbKV) MultiSaveAndRemove(saves map[string]string, removals []stri
 }
 
 // DeleteRange remove a batch of key-values from startKey to endKey
-func (kv *RocksdbKV) DeleteRange(startKey, endKey string) error {
+func (kv *RocksdbKV) DeleteRange(ctx context.Context, startKey, endKey string) error {
 	if kv.DB == nil {
 		return errors.New("Rocksdb instance is nil when do DeleteRange")
 	}
@@ -426,7 +427,7 @@ func (kv *RocksdbKV) DeleteRange(startKey, endKey string) error {
 }
 
 // MultiSaveAndRemoveWithPrefix is used to execute a batch operators with the same prefix
-func (kv *RocksdbKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string, preds ...predicates.Predicate) error {
+func (kv *RocksdbKV) MultiSaveAndRemoveWithPrefix(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error {
 	if len(preds) > 0 {
 		return merr.WrapErrServiceUnavailable("predicates not supported")
 	}
diff --git a/pkg/kv/rocksdb/rocksdb_kv_test.go b/pkg/kv/rocksdb/rocksdb_kv_test.go
index 14c05a0309339..aeecb0a681d5f 100644
--- a/pkg/kv/rocksdb/rocksdb_kv_test.go
+++ b/pkg/kv/rocksdb/rocksdb_kv_test.go
@@ -17,6 +17,7 @@
 package rocksdbkv_test
 
 import (
+	"context"
 	"os"
 	"strconv"
 	"sync"
@@ -39,27 +40,27 @@ func TestRocksdbKV(t *testing.T) {
 	defer os.RemoveAll(name)
 	defer rocksdbKV.Close()
 	// Need to call RemoveWithPrefix
-	defer rocksdbKV.RemoveWithPrefix("")
+	defer rocksdbKV.RemoveWithPrefix(context.TODO(), "")
 
-	err = rocksdbKV.Save("abc", "123")
+	err = rocksdbKV.Save(context.TODO(), "abc", "123")
 	assert.NoError(t, err)
 
-	err = rocksdbKV.SaveBytes("abcd", []byte("1234"))
+	err = rocksdbKV.SaveBytes(context.TODO(), "abcd", []byte("1234"))
 	assert.NoError(t, err)
 
-	val, err := rocksdbKV.Load("abc")
+	val, err := rocksdbKV.Load(context.TODO(), "abc")
 	assert.NoError(t, err)
 	assert.Equal(t, val, "123")
-	value, err := rocksdbKV.LoadBytes("abc")
+	value, err := rocksdbKV.LoadBytes(context.TODO(), "abc")
 	assert.NoError(t, err)
 	assert.Equal(t, value, []byte("123"))
 
-	_, err = rocksdbKV.Load("")
+	_, err = rocksdbKV.Load(context.TODO(), "")
 	assert.Error(t, err)
-	_, err = rocksdbKV.LoadBytes("")
+	_, err = rocksdbKV.LoadBytes(context.TODO(), "")
 	assert.Error(t, err)
 
-	keys, vals, err := rocksdbKV.LoadWithPrefix("abc")
+	keys, vals, err := rocksdbKV.LoadWithPrefix(context.TODO(), "abc")
 	assert.NoError(t, err)
 	assert.Equal(t, len(keys), len(vals))
 	assert.Equal(t, len(keys), 2)
@@ -70,7 +71,7 @@ func TestRocksdbKV(t *testing.T) {
 	assert.Equal(t, vals[1], "1234")
 
 	var values [][]byte
-	keys, values, err = rocksdbKV.LoadBytesWithPrefix("abc")
+	keys, values, err = rocksdbKV.LoadBytesWithPrefix(context.TODO(), "abc")
 	assert.NoError(t, err)
 	assert.Equal(t, len(keys), len(vals))
 	assert.Equal(t, len(keys), 2)
@@ -80,27 +81,27 @@ func TestRocksdbKV(t *testing.T) {
 	assert.Equal(t, values[0], []byte("123"))
 	assert.Equal(t, values[1], []byte("1234"))
 
-	err = rocksdbKV.Save("key_1", "123")
+	err = rocksdbKV.Save(context.TODO(), "key_1", "123")
 	assert.NoError(t, err)
-	err = rocksdbKV.Save("key_2", "456")
+	err = rocksdbKV.Save(context.TODO(), "key_2", "456")
 	assert.NoError(t, err)
-	err = rocksdbKV.Save("key_3", "789")
+	err = rocksdbKV.Save(context.TODO(), "key_3", "789")
 	assert.NoError(t, err)
 
 	keys = []string{"key_1", "key_2"}
-	vals, err = rocksdbKV.MultiLoad(keys)
+	vals, err = rocksdbKV.MultiLoad(context.TODO(), keys)
 	assert.NoError(t, err)
 	assert.Equal(t, len(vals), len(keys))
 	assert.Equal(t, vals[0], "123")
 	assert.Equal(t, vals[1], "456")
 
-	values, err = rocksdbKV.MultiLoadBytes(keys)
+	values, err = rocksdbKV.MultiLoadBytes(context.TODO(), keys)
 	assert.NoError(t, err)
 	assert.Equal(t, len(values), len(keys))
 	assert.Equal(t, values[0], []byte("123"))
 	assert.Equal(t, values[1], []byte("456"))
 
-	err = rocksdbKV.MultiRemove(keys)
+	err = rocksdbKV.MultiRemove(context.TODO(), keys)
 	assert.NoError(t, err)
 
 	saves := map[string]string{
@@ -109,10 +110,10 @@ func TestRocksdbKV(t *testing.T) {
 		"s_3": "333",
 	}
 	removals := []string{"key_3"}
-	err = rocksdbKV.MultiSaveAndRemove(saves, removals)
+	err = rocksdbKV.MultiSaveAndRemove(context.TODO(), saves, removals)
 	assert.NoError(t, err)
 
-	err = rocksdbKV.DeleteRange("s_1", "s_3")
+	err = rocksdbKV.DeleteRange(context.TODO(), "s_1", "s_3")
 	assert.NoError(t, err)
 }
 
@@ -125,17 +126,17 @@ func TestRocksdbKV_Prefix(t *testing.T) {
 	defer os.RemoveAll(name)
 	defer rocksdbKV.Close()
 	// Need to call RemoveWithPrefix
-	defer rocksdbKV.RemoveWithPrefix("")
+	defer rocksdbKV.RemoveWithPrefix(context.TODO(), "")
 
 	kvs := map[string]string{
 		"abcd":    "123",
 		"abdd":    "1234",
 		"abddqqq": "1234555",
 	}
-	err = rocksdbKV.MultiSave(kvs)
+	err = rocksdbKV.MultiSave(context.TODO(), kvs)
 	assert.NoError(t, err)
 
-	keys, vals, err := rocksdbKV.LoadWithPrefix("abc")
+	keys, vals, err := rocksdbKV.LoadWithPrefix(context.TODO(), "abc")
 	assert.NoError(t, err)
 
 	assert.Equal(t, len(keys), 1)
@@ -145,50 +146,50 @@ func TestRocksdbKV_Prefix(t *testing.T) {
 
 	bytesKvs := map[string][]byte{}
 	for k, v := range kvs {
-		rocksdbKV.Remove(k)
+		rocksdbKV.Remove(context.TODO(), k)
 		bytesKvs[k] = []byte(v)
 	}
 
-	err = rocksdbKV.MultiSaveBytes(bytesKvs)
+	err = rocksdbKV.MultiSaveBytes(context.TODO(), bytesKvs)
 	assert.NoError(t, err)
 
 	var values [][]byte
-	keys, values, err = rocksdbKV.LoadBytesWithPrefix("abc")
+	keys, values, err = rocksdbKV.LoadBytesWithPrefix(context.TODO(), "abc")
 	assert.NoError(t, err)
 	assert.Equal(t, len(keys), 1)
 	assert.Equal(t, len(values), 1)
 	assert.Equal(t, keys[0], "abcd")
 	assert.Equal(t, values[0], []byte("123"))
 
-	err = rocksdbKV.RemoveWithPrefix("abc")
+	err = rocksdbKV.RemoveWithPrefix(context.TODO(), "abc")
 	assert.NoError(t, err)
-	val, err := rocksdbKV.Load("abc")
+	val, err := rocksdbKV.Load(context.TODO(), "abc")
 	assert.NoError(t, err)
 	assert.Equal(t, len(val), 0)
-	val, err = rocksdbKV.Load("abdd")
+	val, err = rocksdbKV.Load(context.TODO(), "abdd")
 	assert.NoError(t, err)
 	assert.Equal(t, val, "1234")
-	val, err = rocksdbKV.Load("abddqqq")
+	val, err = rocksdbKV.Load(context.TODO(), "abddqqq")
 	assert.NoError(t, err)
 	assert.Equal(t, val, "1234555")
 
 	// test remove ""
-	err = rocksdbKV.RemoveWithPrefix("")
+	err = rocksdbKV.RemoveWithPrefix(context.TODO(), "")
 	assert.NoError(t, err)
 
 	// test remove from an empty cf
-	err = rocksdbKV.RemoveWithPrefix("")
+	err = rocksdbKV.RemoveWithPrefix(context.TODO(), "")
 	assert.NoError(t, err)
 
-	val, err = rocksdbKV.Load("abddqqq")
+	val, err = rocksdbKV.Load(context.TODO(), "abddqqq")
 	assert.NoError(t, err)
 	assert.Equal(t, len(val), 0)
 
 	// test we can still save after drop
-	err = rocksdbKV.Save("abcd", "123")
+	err = rocksdbKV.Save(context.TODO(), "abcd", "123")
 	assert.NoError(t, err)
 
-	val, err = rocksdbKV.Load("abcd")
+	val, err = rocksdbKV.Load(context.TODO(), "abcd")
 	assert.NoError(t, err)
 	assert.Equal(t, val, "123")
 }
@@ -202,24 +203,24 @@ func TestRocksdbKV_Txn(t *testing.T) {
 	defer os.RemoveAll(name)
 	defer rocksdbKV.Close()
 	// Need to call RemoveWithPrefix
-	defer rocksdbKV.RemoveWithPrefix("")
+	defer rocksdbKV.RemoveWithPrefix(context.TODO(), "")
 
 	kvs := map[string]string{
 		"abcd":    "123",
 		"abdd":    "1234",
 		"abddqqq": "1234555",
 	}
-	err = rocksdbKV.MultiSave(kvs)
+	err = rocksdbKV.MultiSave(context.TODO(), kvs)
 	assert.NoError(t, err)
 
-	keys, vals, err := rocksdbKV.LoadWithPrefix("")
+	keys, vals, err := rocksdbKV.LoadWithPrefix(context.TODO(), "")
 	assert.NoError(t, err)
 	assert.Equal(t, len(keys), 3)
 	assert.Equal(t, len(vals), 3)
 
-	err = rocksdbKV.MultiSave(kvs)
+	err = rocksdbKV.MultiSave(context.TODO(), kvs)
 	assert.NoError(t, err)
-	keys, vals, err = rocksdbKV.LoadWithPrefix("")
+	keys, vals, err = rocksdbKV.LoadWithPrefix(context.TODO(), "")
 	assert.NoError(t, err)
 	assert.Equal(t, len(keys), 3)
 	assert.Equal(t, len(vals), 3)
@@ -229,8 +230,8 @@ func TestRocksdbKV_Txn(t *testing.T) {
 	kvs2 := map[string]string{
 		"abfad": "12345",
 	}
-	rocksdbKV.MultiSaveAndRemoveWithPrefix(kvs2, removePrefix)
-	keys, vals, err = rocksdbKV.LoadWithPrefix("")
+	rocksdbKV.MultiSaveAndRemoveWithPrefix(context.TODO(), kvs2, removePrefix)
+	keys, vals, err = rocksdbKV.LoadWithPrefix(context.TODO(), "")
 	assert.NoError(t, err)
 	assert.Equal(t, len(keys), 1)
 	assert.Equal(t, len(vals), 1)
@@ -242,7 +243,7 @@ func TestRocksdbKV_Goroutines(t *testing.T) {
 	assert.NoError(t, err)
 	defer os.RemoveAll(name)
 	defer rocksdbkv.Close()
-	defer rocksdbkv.RemoveWithPrefix("")
+	defer rocksdbkv.RemoveWithPrefix(context.TODO(), "")
 
 	var wg sync.WaitGroup
 	for i := 0; i < 10; i++ {
@@ -251,10 +252,10 @@ func TestRocksdbKV_Goroutines(t *testing.T) {
 			defer wg.Done()
 			key := "key_" + strconv.Itoa(i)
 			val := "val_" + strconv.Itoa(i)
-			err := rocksdbkv.Save(key, val)
+			err := rocksdbkv.Save(context.TODO(), key, val)
 			assert.NoError(t, err)
 
-			getVal, err := rocksdbkv.Load(key)
+			getVal, err := rocksdbkv.Load(context.TODO(), key)
 			assert.NoError(t, err)
 			assert.Equal(t, getVal, val)
 		}(i)
@@ -268,32 +269,32 @@ func TestRocksdbKV_DummyDB(t *testing.T) {
 	assert.NoError(t, err)
 	defer os.RemoveAll(name)
 	defer rocksdbkv.Close()
-	defer rocksdbkv.RemoveWithPrefix("")
+	defer rocksdbkv.RemoveWithPrefix(context.TODO(), "")
 
 	rocksdbkv.DB = nil
-	_, err = rocksdbkv.Load("")
+	_, err = rocksdbkv.Load(context.TODO(), "")
 	assert.Error(t, err)
-	_, _, err = rocksdbkv.LoadWithPrefix("")
+	_, _, err = rocksdbkv.LoadWithPrefix(context.TODO(), "")
 	assert.Error(t, err)
-	_, err = rocksdbkv.MultiLoad(nil)
+	_, err = rocksdbkv.MultiLoad(context.TODO(), nil)
 	assert.Error(t, err)
-	err = rocksdbkv.Save("", "")
+	err = rocksdbkv.Save(context.TODO(), "", "")
 	assert.Error(t, err)
-	err = rocksdbkv.MultiSave(nil)
+	err = rocksdbkv.MultiSave(context.TODO(), nil)
 	assert.Error(t, err)
-	err = rocksdbkv.RemoveWithPrefix("")
+	err = rocksdbkv.RemoveWithPrefix(context.TODO(), "")
 	assert.Error(t, err)
-	err = rocksdbkv.Remove("")
+	err = rocksdbkv.Remove(context.TODO(), "")
 	assert.Error(t, err)
-	err = rocksdbkv.MultiRemove(nil)
+	err = rocksdbkv.MultiRemove(context.TODO(), nil)
 	assert.Error(t, err)
-	err = rocksdbkv.MultiSaveAndRemove(nil, nil)
+	err = rocksdbkv.MultiSaveAndRemove(context.TODO(), nil, nil)
 	assert.Error(t, err)
-	err = rocksdbkv.DeleteRange("", "")
+	err = rocksdbkv.DeleteRange(context.TODO(), "", "")
 	assert.Error(t, err)
 
 	rocksdbkv.ReadOptions = nil
-	_, err = rocksdbkv.Load("dummy")
+	_, err = rocksdbkv.Load(context.TODO(), "dummy")
 	assert.Error(t, err)
 }
 
@@ -303,20 +304,20 @@ func TestRocksdbKV_CornerCase(t *testing.T) {
 	assert.NoError(t, err)
 	defer os.RemoveAll(name)
 	defer rocksdbkv.Close()
-	defer rocksdbkv.RemoveWithPrefix("")
-	_, err = rocksdbkv.Load("")
+	defer rocksdbkv.RemoveWithPrefix(context.TODO(), "")
+	_, err = rocksdbkv.Load(context.TODO(), "")
 	assert.Error(t, err)
-	keys, values, err := rocksdbkv.LoadWithPrefix("")
+	keys, values, err := rocksdbkv.LoadWithPrefix(context.TODO(), "")
 	assert.NoError(t, err)
 	assert.Equal(t, len(keys), 0)
 	assert.Equal(t, len(values), 0)
-	err = rocksdbkv.Save("", "")
+	err = rocksdbkv.Save(context.TODO(), "", "")
 	assert.Error(t, err)
-	err = rocksdbkv.Save("test", "")
+	err = rocksdbkv.Save(context.TODO(), "test", "")
 	assert.Error(t, err)
-	err = rocksdbkv.Remove("")
+	err = rocksdbkv.Remove(context.TODO(), "")
 	assert.Error(t, err)
-	err = rocksdbkv.DeleteRange("a", "a")
+	err = rocksdbkv.DeleteRange(context.TODO(), "a", "a")
 	assert.Error(t, err)
 }
 
@@ -325,19 +326,19 @@ func TestHas(t *testing.T) {
 	db, err := rocksdbkv.NewRocksdbKV(dir)
 	assert.NoError(t, err)
 	defer db.Close()
-	defer db.RemoveWithPrefix("")
+	defer db.RemoveWithPrefix(context.TODO(), "")
 
-	has, err := db.Has("key1")
+	has, err := db.Has(context.TODO(), "key1")
 	assert.NoError(t, err)
 	assert.False(t, has)
-	err = db.Save("key1", "value1")
+	err = db.Save(context.TODO(), "key1", "value1")
 	assert.NoError(t, err)
-	has, err = db.Has("key1")
+	has, err = db.Has(context.TODO(), "key1")
 	assert.NoError(t, err)
 	assert.True(t, has)
-	err = db.Remove("key1")
+	err = db.Remove(context.TODO(), "key1")
 	assert.NoError(t, err)
-	has, err = db.Has("key1")
+	has, err = db.Has(context.TODO(), "key1")
 	assert.NoError(t, err)
 	assert.False(t, has)
 }
@@ -347,23 +348,23 @@ func TestHasPrefix(t *testing.T) {
 	db, err := rocksdbkv.NewRocksdbKV(dir)
 	assert.NoError(t, err)
 	defer db.Close()
-	defer db.RemoveWithPrefix("")
+	defer db.RemoveWithPrefix(context.TODO(), "")
 
-	has, err := db.HasPrefix("key")
+	has, err := db.HasPrefix(context.TODO(), "key")
 	assert.NoError(t, err)
 	assert.False(t, has)
 
-	err = db.Save("key1", "value1")
+	err = db.Save(context.TODO(), "key1", "value1")
 	assert.NoError(t, err)
 
-	has, err = db.HasPrefix("key")
+	has, err = db.HasPrefix(context.TODO(), "key")
 	assert.NoError(t, err)
 	assert.True(t, has)
 
-	err = db.Remove("key1")
+	err = db.Remove(context.TODO(), "key1")
 	assert.NoError(t, err)
 
-	has, err = db.HasPrefix("key")
+	has, err = db.HasPrefix(context.TODO(), "key")
 	assert.NoError(t, err)
 	assert.False(t, has)
 }
@@ -374,13 +375,13 @@ func TestPredicates(t *testing.T) {
 
 	require.NoError(t, err)
 	defer db.Close()
-	defer db.RemoveWithPrefix("")
+	defer db.RemoveWithPrefix(context.TODO(), "")
 
-	err = db.MultiSaveAndRemove(map[string]string{}, []string{}, predicates.ValueEqual("a", "b"))
+	err = db.MultiSaveAndRemove(context.TODO(), map[string]string{}, []string{}, predicates.ValueEqual("a", "b"))
 	assert.Error(t, err)
 	assert.ErrorIs(t, err, merr.ErrServiceUnavailable)
 
-	err = db.MultiSaveAndRemoveWithPrefix(map[string]string{}, []string{}, predicates.ValueEqual("a", "b"))
+	err = db.MultiSaveAndRemoveWithPrefix(context.TODO(), map[string]string{}, []string{}, predicates.ValueEqual("a", "b"))
 	assert.Error(t, err)
 	assert.ErrorIs(t, err, merr.ErrServiceUnavailable)
 }
diff --git a/pkg/mocks/mock_kv/mock_MetaKv.go b/pkg/mocks/mock_kv/mock_MetaKv.go
index 52702fd8cfb27..683ef85f18496 100644
--- a/pkg/mocks/mock_kv/mock_MetaKv.go
+++ b/pkg/mocks/mock_kv/mock_MetaKv.go
@@ -3,8 +3,11 @@
 package mock_kv
 
 import (
-	predicates "github.com/milvus-io/milvus/pkg/kv/predicates"
+	context "context"
+
 	mock "github.com/stretchr/testify/mock"
+
+	predicates "github.com/milvus-io/milvus/pkg/kv/predicates"
 )
 
 // MockMetaKv is an autogenerated mock type for the MetaKv type
@@ -52,9 +55,9 @@ func (_c *MockMetaKv_Close_Call) RunAndReturn(run func()) *MockMetaKv_Close_Call
 	return _c
 }
 
-// CompareVersionAndSwap provides a mock function with given fields: key, version, target
-func (_m *MockMetaKv) CompareVersionAndSwap(key string, version int64, target string) (bool, error) {
-	ret := _m.Called(key, version, target)
+// CompareVersionAndSwap provides a mock function with given fields: ctx, key, version, target
+func (_m *MockMetaKv) CompareVersionAndSwap(ctx context.Context, key string, version int64, target string) (bool, error) {
+	ret := _m.Called(ctx, key, version, target)
 
 	if len(ret) == 0 {
 		panic("no return value specified for CompareVersionAndSwap")
@@ -62,17 +65,17 @@ func (_m *MockMetaKv) CompareVersionAndSwap(key string, version int64, target st
 
 	var r0 bool
 	var r1 error
-	if rf, ok := ret.Get(0).(func(string, int64, string) (bool, error)); ok {
-		return rf(key, version, target)
+	if rf, ok := ret.Get(0).(func(context.Context, string, int64, string) (bool, error)); ok {
+		return rf(ctx, key, version, target)
 	}
-	if rf, ok := ret.Get(0).(func(string, int64, string) bool); ok {
-		r0 = rf(key, version, target)
+	if rf, ok := ret.Get(0).(func(context.Context, string, int64, string) bool); ok {
+		r0 = rf(ctx, key, version, target)
 	} else {
 		r0 = ret.Get(0).(bool)
 	}
 
-	if rf, ok := ret.Get(1).(func(string, int64, string) error); ok {
-		r1 = rf(key, version, target)
+	if rf, ok := ret.Get(1).(func(context.Context, string, int64, string) error); ok {
+		r1 = rf(ctx, key, version, target)
 	} else {
 		r1 = ret.Error(1)
 	}
@@ -86,16 +89,17 @@ type MockMetaKv_CompareVersionAndSwap_Call struct {
 }
 
 // CompareVersionAndSwap is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
 //   - version int64
 //   - target string
-func (_e *MockMetaKv_Expecter) CompareVersionAndSwap(key interface{}, version interface{}, target interface{}) *MockMetaKv_CompareVersionAndSwap_Call {
-	return &MockMetaKv_CompareVersionAndSwap_Call{Call: _e.mock.On("CompareVersionAndSwap", key, version, target)}
+func (_e *MockMetaKv_Expecter) CompareVersionAndSwap(ctx interface{}, key interface{}, version interface{}, target interface{}) *MockMetaKv_CompareVersionAndSwap_Call {
+	return &MockMetaKv_CompareVersionAndSwap_Call{Call: _e.mock.On("CompareVersionAndSwap", ctx, key, version, target)}
 }
 
-func (_c *MockMetaKv_CompareVersionAndSwap_Call) Run(run func(key string, version int64, target string)) *MockMetaKv_CompareVersionAndSwap_Call {
+func (_c *MockMetaKv_CompareVersionAndSwap_Call) Run(run func(ctx context.Context, key string, version int64, target string)) *MockMetaKv_CompareVersionAndSwap_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string), args[1].(int64), args[2].(string))
+		run(args[0].(context.Context), args[1].(string), args[2].(int64), args[3].(string))
 	})
 	return _c
 }
@@ -105,7 +109,7 @@ func (_c *MockMetaKv_CompareVersionAndSwap_Call) Return(_a0 bool, _a1 error) *Mo
 	return _c
 }
 
-func (_c *MockMetaKv_CompareVersionAndSwap_Call) RunAndReturn(run func(string, int64, string) (bool, error)) *MockMetaKv_CompareVersionAndSwap_Call {
+func (_c *MockMetaKv_CompareVersionAndSwap_Call) RunAndReturn(run func(context.Context, string, int64, string) (bool, error)) *MockMetaKv_CompareVersionAndSwap_Call {
 	_c.Call.Return(run)
 	return _c
 }
@@ -156,9 +160,9 @@ func (_c *MockMetaKv_GetPath_Call) RunAndReturn(run func(string) string) *MockMe
 	return _c
 }
 
-// Has provides a mock function with given fields: key
-func (_m *MockMetaKv) Has(key string) (bool, error) {
-	ret := _m.Called(key)
+// Has provides a mock function with given fields: ctx, key
+func (_m *MockMetaKv) Has(ctx context.Context, key string) (bool, error) {
+	ret := _m.Called(ctx, key)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Has")
@@ -166,17 +170,17 @@ func (_m *MockMetaKv) Has(key string) (bool, error) {
 
 	var r0 bool
 	var r1 error
-	if rf, ok := ret.Get(0).(func(string) (bool, error)); ok {
-		return rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) (bool, error)); ok {
+		return rf(ctx, key)
 	}
-	if rf, ok := ret.Get(0).(func(string) bool); ok {
-		r0 = rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) bool); ok {
+		r0 = rf(ctx, key)
 	} else {
 		r0 = ret.Get(0).(bool)
 	}
 
-	if rf, ok := ret.Get(1).(func(string) error); ok {
-		r1 = rf(key)
+	if rf, ok := ret.Get(1).(func(context.Context, string) error); ok {
+		r1 = rf(ctx, key)
 	} else {
 		r1 = ret.Error(1)
 	}
@@ -190,14 +194,15 @@ type MockMetaKv_Has_Call struct {
 }
 
 // Has is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
-func (_e *MockMetaKv_Expecter) Has(key interface{}) *MockMetaKv_Has_Call {
-	return &MockMetaKv_Has_Call{Call: _e.mock.On("Has", key)}
+func (_e *MockMetaKv_Expecter) Has(ctx interface{}, key interface{}) *MockMetaKv_Has_Call {
+	return &MockMetaKv_Has_Call{Call: _e.mock.On("Has", ctx, key)}
 }
 
-func (_c *MockMetaKv_Has_Call) Run(run func(key string)) *MockMetaKv_Has_Call {
+func (_c *MockMetaKv_Has_Call) Run(run func(ctx context.Context, key string)) *MockMetaKv_Has_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string))
+		run(args[0].(context.Context), args[1].(string))
 	})
 	return _c
 }
@@ -207,14 +212,14 @@ func (_c *MockMetaKv_Has_Call) Return(_a0 bool, _a1 error) *MockMetaKv_Has_Call
 	return _c
 }
 
-func (_c *MockMetaKv_Has_Call) RunAndReturn(run func(string) (bool, error)) *MockMetaKv_Has_Call {
+func (_c *MockMetaKv_Has_Call) RunAndReturn(run func(context.Context, string) (bool, error)) *MockMetaKv_Has_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// HasPrefix provides a mock function with given fields: prefix
-func (_m *MockMetaKv) HasPrefix(prefix string) (bool, error) {
-	ret := _m.Called(prefix)
+// HasPrefix provides a mock function with given fields: ctx, prefix
+func (_m *MockMetaKv) HasPrefix(ctx context.Context, prefix string) (bool, error) {
+	ret := _m.Called(ctx, prefix)
 
 	if len(ret) == 0 {
 		panic("no return value specified for HasPrefix")
@@ -222,17 +227,17 @@ func (_m *MockMetaKv) HasPrefix(prefix string) (bool, error) {
 
 	var r0 bool
 	var r1 error
-	if rf, ok := ret.Get(0).(func(string) (bool, error)); ok {
-		return rf(prefix)
+	if rf, ok := ret.Get(0).(func(context.Context, string) (bool, error)); ok {
+		return rf(ctx, prefix)
 	}
-	if rf, ok := ret.Get(0).(func(string) bool); ok {
-		r0 = rf(prefix)
+	if rf, ok := ret.Get(0).(func(context.Context, string) bool); ok {
+		r0 = rf(ctx, prefix)
 	} else {
 		r0 = ret.Get(0).(bool)
 	}
 
-	if rf, ok := ret.Get(1).(func(string) error); ok {
-		r1 = rf(prefix)
+	if rf, ok := ret.Get(1).(func(context.Context, string) error); ok {
+		r1 = rf(ctx, prefix)
 	} else {
 		r1 = ret.Error(1)
 	}
@@ -246,14 +251,15 @@ type MockMetaKv_HasPrefix_Call struct {
 }
 
 // HasPrefix is a helper method to define mock.On call
+//   - ctx context.Context
 //   - prefix string
-func (_e *MockMetaKv_Expecter) HasPrefix(prefix interface{}) *MockMetaKv_HasPrefix_Call {
-	return &MockMetaKv_HasPrefix_Call{Call: _e.mock.On("HasPrefix", prefix)}
+func (_e *MockMetaKv_Expecter) HasPrefix(ctx interface{}, prefix interface{}) *MockMetaKv_HasPrefix_Call {
+	return &MockMetaKv_HasPrefix_Call{Call: _e.mock.On("HasPrefix", ctx, prefix)}
 }
 
-func (_c *MockMetaKv_HasPrefix_Call) Run(run func(prefix string)) *MockMetaKv_HasPrefix_Call {
+func (_c *MockMetaKv_HasPrefix_Call) Run(run func(ctx context.Context, prefix string)) *MockMetaKv_HasPrefix_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string))
+		run(args[0].(context.Context), args[1].(string))
 	})
 	return _c
 }
@@ -263,14 +269,14 @@ func (_c *MockMetaKv_HasPrefix_Call) Return(_a0 bool, _a1 error) *MockMetaKv_Has
 	return _c
 }
 
-func (_c *MockMetaKv_HasPrefix_Call) RunAndReturn(run func(string) (bool, error)) *MockMetaKv_HasPrefix_Call {
+func (_c *MockMetaKv_HasPrefix_Call) RunAndReturn(run func(context.Context, string) (bool, error)) *MockMetaKv_HasPrefix_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// Load provides a mock function with given fields: key
-func (_m *MockMetaKv) Load(key string) (string, error) {
-	ret := _m.Called(key)
+// Load provides a mock function with given fields: ctx, key
+func (_m *MockMetaKv) Load(ctx context.Context, key string) (string, error) {
+	ret := _m.Called(ctx, key)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Load")
@@ -278,17 +284,17 @@ func (_m *MockMetaKv) Load(key string) (string, error) {
 
 	var r0 string
 	var r1 error
-	if rf, ok := ret.Get(0).(func(string) (string, error)); ok {
-		return rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) (string, error)); ok {
+		return rf(ctx, key)
 	}
-	if rf, ok := ret.Get(0).(func(string) string); ok {
-		r0 = rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) string); ok {
+		r0 = rf(ctx, key)
 	} else {
 		r0 = ret.Get(0).(string)
 	}
 
-	if rf, ok := ret.Get(1).(func(string) error); ok {
-		r1 = rf(key)
+	if rf, ok := ret.Get(1).(func(context.Context, string) error); ok {
+		r1 = rf(ctx, key)
 	} else {
 		r1 = ret.Error(1)
 	}
@@ -302,14 +308,15 @@ type MockMetaKv_Load_Call struct {
 }
 
 // Load is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
-func (_e *MockMetaKv_Expecter) Load(key interface{}) *MockMetaKv_Load_Call {
-	return &MockMetaKv_Load_Call{Call: _e.mock.On("Load", key)}
+func (_e *MockMetaKv_Expecter) Load(ctx interface{}, key interface{}) *MockMetaKv_Load_Call {
+	return &MockMetaKv_Load_Call{Call: _e.mock.On("Load", ctx, key)}
 }
 
-func (_c *MockMetaKv_Load_Call) Run(run func(key string)) *MockMetaKv_Load_Call {
+func (_c *MockMetaKv_Load_Call) Run(run func(ctx context.Context, key string)) *MockMetaKv_Load_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string))
+		run(args[0].(context.Context), args[1].(string))
 	})
 	return _c
 }
@@ -319,14 +326,14 @@ func (_c *MockMetaKv_Load_Call) Return(_a0 string, _a1 error) *MockMetaKv_Load_C
 	return _c
 }
 
-func (_c *MockMetaKv_Load_Call) RunAndReturn(run func(string) (string, error)) *MockMetaKv_Load_Call {
+func (_c *MockMetaKv_Load_Call) RunAndReturn(run func(context.Context, string) (string, error)) *MockMetaKv_Load_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// LoadWithPrefix provides a mock function with given fields: key
-func (_m *MockMetaKv) LoadWithPrefix(key string) ([]string, []string, error) {
-	ret := _m.Called(key)
+// LoadWithPrefix provides a mock function with given fields: ctx, key
+func (_m *MockMetaKv) LoadWithPrefix(ctx context.Context, key string) ([]string, []string, error) {
+	ret := _m.Called(ctx, key)
 
 	if len(ret) == 0 {
 		panic("no return value specified for LoadWithPrefix")
@@ -335,27 +342,27 @@ func (_m *MockMetaKv) LoadWithPrefix(key string) ([]string, []string, error) {
 	var r0 []string
 	var r1 []string
 	var r2 error
-	if rf, ok := ret.Get(0).(func(string) ([]string, []string, error)); ok {
-		return rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) ([]string, []string, error)); ok {
+		return rf(ctx, key)
 	}
-	if rf, ok := ret.Get(0).(func(string) []string); ok {
-		r0 = rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) []string); ok {
+		r0 = rf(ctx, key)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).([]string)
 		}
 	}
 
-	if rf, ok := ret.Get(1).(func(string) []string); ok {
-		r1 = rf(key)
+	if rf, ok := ret.Get(1).(func(context.Context, string) []string); ok {
+		r1 = rf(ctx, key)
 	} else {
 		if ret.Get(1) != nil {
 			r1 = ret.Get(1).([]string)
 		}
 	}
 
-	if rf, ok := ret.Get(2).(func(string) error); ok {
-		r2 = rf(key)
+	if rf, ok := ret.Get(2).(func(context.Context, string) error); ok {
+		r2 = rf(ctx, key)
 	} else {
 		r2 = ret.Error(2)
 	}
@@ -369,14 +376,15 @@ type MockMetaKv_LoadWithPrefix_Call struct {
 }
 
 // LoadWithPrefix is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
-func (_e *MockMetaKv_Expecter) LoadWithPrefix(key interface{}) *MockMetaKv_LoadWithPrefix_Call {
-	return &MockMetaKv_LoadWithPrefix_Call{Call: _e.mock.On("LoadWithPrefix", key)}
+func (_e *MockMetaKv_Expecter) LoadWithPrefix(ctx interface{}, key interface{}) *MockMetaKv_LoadWithPrefix_Call {
+	return &MockMetaKv_LoadWithPrefix_Call{Call: _e.mock.On("LoadWithPrefix", ctx, key)}
 }
 
-func (_c *MockMetaKv_LoadWithPrefix_Call) Run(run func(key string)) *MockMetaKv_LoadWithPrefix_Call {
+func (_c *MockMetaKv_LoadWithPrefix_Call) Run(run func(ctx context.Context, key string)) *MockMetaKv_LoadWithPrefix_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string))
+		run(args[0].(context.Context), args[1].(string))
 	})
 	return _c
 }
@@ -386,14 +394,14 @@ func (_c *MockMetaKv_LoadWithPrefix_Call) Return(_a0 []string, _a1 []string, _a2
 	return _c
 }
 
-func (_c *MockMetaKv_LoadWithPrefix_Call) RunAndReturn(run func(string) ([]string, []string, error)) *MockMetaKv_LoadWithPrefix_Call {
+func (_c *MockMetaKv_LoadWithPrefix_Call) RunAndReturn(run func(context.Context, string) ([]string, []string, error)) *MockMetaKv_LoadWithPrefix_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// MultiLoad provides a mock function with given fields: keys
-func (_m *MockMetaKv) MultiLoad(keys []string) ([]string, error) {
-	ret := _m.Called(keys)
+// MultiLoad provides a mock function with given fields: ctx, keys
+func (_m *MockMetaKv) MultiLoad(ctx context.Context, keys []string) ([]string, error) {
+	ret := _m.Called(ctx, keys)
 
 	if len(ret) == 0 {
 		panic("no return value specified for MultiLoad")
@@ -401,19 +409,19 @@ func (_m *MockMetaKv) MultiLoad(keys []string) ([]string, error) {
 
 	var r0 []string
 	var r1 error
-	if rf, ok := ret.Get(0).(func([]string) ([]string, error)); ok {
-		return rf(keys)
+	if rf, ok := ret.Get(0).(func(context.Context, []string) ([]string, error)); ok {
+		return rf(ctx, keys)
 	}
-	if rf, ok := ret.Get(0).(func([]string) []string); ok {
-		r0 = rf(keys)
+	if rf, ok := ret.Get(0).(func(context.Context, []string) []string); ok {
+		r0 = rf(ctx, keys)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).([]string)
 		}
 	}
 
-	if rf, ok := ret.Get(1).(func([]string) error); ok {
-		r1 = rf(keys)
+	if rf, ok := ret.Get(1).(func(context.Context, []string) error); ok {
+		r1 = rf(ctx, keys)
 	} else {
 		r1 = ret.Error(1)
 	}
@@ -427,14 +435,15 @@ type MockMetaKv_MultiLoad_Call struct {
 }
 
 // MultiLoad is a helper method to define mock.On call
+//   - ctx context.Context
 //   - keys []string
-func (_e *MockMetaKv_Expecter) MultiLoad(keys interface{}) *MockMetaKv_MultiLoad_Call {
-	return &MockMetaKv_MultiLoad_Call{Call: _e.mock.On("MultiLoad", keys)}
+func (_e *MockMetaKv_Expecter) MultiLoad(ctx interface{}, keys interface{}) *MockMetaKv_MultiLoad_Call {
+	return &MockMetaKv_MultiLoad_Call{Call: _e.mock.On("MultiLoad", ctx, keys)}
 }
 
-func (_c *MockMetaKv_MultiLoad_Call) Run(run func(keys []string)) *MockMetaKv_MultiLoad_Call {
+func (_c *MockMetaKv_MultiLoad_Call) Run(run func(ctx context.Context, keys []string)) *MockMetaKv_MultiLoad_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].([]string))
+		run(args[0].(context.Context), args[1].([]string))
 	})
 	return _c
 }
@@ -444,22 +453,22 @@ func (_c *MockMetaKv_MultiLoad_Call) Return(_a0 []string, _a1 error) *MockMetaKv
 	return _c
 }
 
-func (_c *MockMetaKv_MultiLoad_Call) RunAndReturn(run func([]string) ([]string, error)) *MockMetaKv_MultiLoad_Call {
+func (_c *MockMetaKv_MultiLoad_Call) RunAndReturn(run func(context.Context, []string) ([]string, error)) *MockMetaKv_MultiLoad_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// MultiRemove provides a mock function with given fields: keys
-func (_m *MockMetaKv) MultiRemove(keys []string) error {
-	ret := _m.Called(keys)
+// MultiRemove provides a mock function with given fields: ctx, keys
+func (_m *MockMetaKv) MultiRemove(ctx context.Context, keys []string) error {
+	ret := _m.Called(ctx, keys)
 
 	if len(ret) == 0 {
 		panic("no return value specified for MultiRemove")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func([]string) error); ok {
-		r0 = rf(keys)
+	if rf, ok := ret.Get(0).(func(context.Context, []string) error); ok {
+		r0 = rf(ctx, keys)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -473,14 +482,15 @@ type MockMetaKv_MultiRemove_Call struct {
 }
 
 // MultiRemove is a helper method to define mock.On call
+//   - ctx context.Context
 //   - keys []string
-func (_e *MockMetaKv_Expecter) MultiRemove(keys interface{}) *MockMetaKv_MultiRemove_Call {
-	return &MockMetaKv_MultiRemove_Call{Call: _e.mock.On("MultiRemove", keys)}
+func (_e *MockMetaKv_Expecter) MultiRemove(ctx interface{}, keys interface{}) *MockMetaKv_MultiRemove_Call {
+	return &MockMetaKv_MultiRemove_Call{Call: _e.mock.On("MultiRemove", ctx, keys)}
 }
 
-func (_c *MockMetaKv_MultiRemove_Call) Run(run func(keys []string)) *MockMetaKv_MultiRemove_Call {
+func (_c *MockMetaKv_MultiRemove_Call) Run(run func(ctx context.Context, keys []string)) *MockMetaKv_MultiRemove_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].([]string))
+		run(args[0].(context.Context), args[1].([]string))
 	})
 	return _c
 }
@@ -490,22 +500,22 @@ func (_c *MockMetaKv_MultiRemove_Call) Return(_a0 error) *MockMetaKv_MultiRemove
 	return _c
 }
 
-func (_c *MockMetaKv_MultiRemove_Call) RunAndReturn(run func([]string) error) *MockMetaKv_MultiRemove_Call {
+func (_c *MockMetaKv_MultiRemove_Call) RunAndReturn(run func(context.Context, []string) error) *MockMetaKv_MultiRemove_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// MultiSave provides a mock function with given fields: kvs
-func (_m *MockMetaKv) MultiSave(kvs map[string]string) error {
-	ret := _m.Called(kvs)
+// MultiSave provides a mock function with given fields: ctx, kvs
+func (_m *MockMetaKv) MultiSave(ctx context.Context, kvs map[string]string) error {
+	ret := _m.Called(ctx, kvs)
 
 	if len(ret) == 0 {
 		panic("no return value specified for MultiSave")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(map[string]string) error); ok {
-		r0 = rf(kvs)
+	if rf, ok := ret.Get(0).(func(context.Context, map[string]string) error); ok {
+		r0 = rf(ctx, kvs)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -519,14 +529,15 @@ type MockMetaKv_MultiSave_Call struct {
 }
 
 // MultiSave is a helper method to define mock.On call
+//   - ctx context.Context
 //   - kvs map[string]string
-func (_e *MockMetaKv_Expecter) MultiSave(kvs interface{}) *MockMetaKv_MultiSave_Call {
-	return &MockMetaKv_MultiSave_Call{Call: _e.mock.On("MultiSave", kvs)}
+func (_e *MockMetaKv_Expecter) MultiSave(ctx interface{}, kvs interface{}) *MockMetaKv_MultiSave_Call {
+	return &MockMetaKv_MultiSave_Call{Call: _e.mock.On("MultiSave", ctx, kvs)}
 }
 
-func (_c *MockMetaKv_MultiSave_Call) Run(run func(kvs map[string]string)) *MockMetaKv_MultiSave_Call {
+func (_c *MockMetaKv_MultiSave_Call) Run(run func(ctx context.Context, kvs map[string]string)) *MockMetaKv_MultiSave_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(map[string]string))
+		run(args[0].(context.Context), args[1].(map[string]string))
 	})
 	return _c
 }
@@ -536,19 +547,19 @@ func (_c *MockMetaKv_MultiSave_Call) Return(_a0 error) *MockMetaKv_MultiSave_Cal
 	return _c
 }
 
-func (_c *MockMetaKv_MultiSave_Call) RunAndReturn(run func(map[string]string) error) *MockMetaKv_MultiSave_Call {
+func (_c *MockMetaKv_MultiSave_Call) RunAndReturn(run func(context.Context, map[string]string) error) *MockMetaKv_MultiSave_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// MultiSaveAndRemove provides a mock function with given fields: saves, removals, preds
-func (_m *MockMetaKv) MultiSaveAndRemove(saves map[string]string, removals []string, preds ...predicates.Predicate) error {
+// MultiSaveAndRemove provides a mock function with given fields: ctx, saves, removals, preds
+func (_m *MockMetaKv) MultiSaveAndRemove(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error {
 	_va := make([]interface{}, len(preds))
 	for _i := range preds {
 		_va[_i] = preds[_i]
 	}
 	var _ca []interface{}
-	_ca = append(_ca, saves, removals)
+	_ca = append(_ca, ctx, saves, removals)
 	_ca = append(_ca, _va...)
 	ret := _m.Called(_ca...)
 
@@ -557,8 +568,8 @@ func (_m *MockMetaKv) MultiSaveAndRemove(saves map[string]string, removals []str
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(map[string]string, []string, ...predicates.Predicate) error); ok {
-		r0 = rf(saves, removals, preds...)
+	if rf, ok := ret.Get(0).(func(context.Context, map[string]string, []string, ...predicates.Predicate) error); ok {
+		r0 = rf(ctx, saves, removals, preds...)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -572,23 +583,24 @@ type MockMetaKv_MultiSaveAndRemove_Call struct {
 }
 
 // MultiSaveAndRemove is a helper method to define mock.On call
+//   - ctx context.Context
 //   - saves map[string]string
 //   - removals []string
 //   - preds ...predicates.Predicate
-func (_e *MockMetaKv_Expecter) MultiSaveAndRemove(saves interface{}, removals interface{}, preds ...interface{}) *MockMetaKv_MultiSaveAndRemove_Call {
+func (_e *MockMetaKv_Expecter) MultiSaveAndRemove(ctx interface{}, saves interface{}, removals interface{}, preds ...interface{}) *MockMetaKv_MultiSaveAndRemove_Call {
 	return &MockMetaKv_MultiSaveAndRemove_Call{Call: _e.mock.On("MultiSaveAndRemove",
-		append([]interface{}{saves, removals}, preds...)...)}
+		append([]interface{}{ctx, saves, removals}, preds...)...)}
 }
 
-func (_c *MockMetaKv_MultiSaveAndRemove_Call) Run(run func(saves map[string]string, removals []string, preds ...predicates.Predicate)) *MockMetaKv_MultiSaveAndRemove_Call {
+func (_c *MockMetaKv_MultiSaveAndRemove_Call) Run(run func(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate)) *MockMetaKv_MultiSaveAndRemove_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		variadicArgs := make([]predicates.Predicate, len(args)-2)
-		for i, a := range args[2:] {
+		variadicArgs := make([]predicates.Predicate, len(args)-3)
+		for i, a := range args[3:] {
 			if a != nil {
 				variadicArgs[i] = a.(predicates.Predicate)
 			}
 		}
-		run(args[0].(map[string]string), args[1].([]string), variadicArgs...)
+		run(args[0].(context.Context), args[1].(map[string]string), args[2].([]string), variadicArgs...)
 	})
 	return _c
 }
@@ -598,19 +610,19 @@ func (_c *MockMetaKv_MultiSaveAndRemove_Call) Return(_a0 error) *MockMetaKv_Mult
 	return _c
 }
 
-func (_c *MockMetaKv_MultiSaveAndRemove_Call) RunAndReturn(run func(map[string]string, []string, ...predicates.Predicate) error) *MockMetaKv_MultiSaveAndRemove_Call {
+func (_c *MockMetaKv_MultiSaveAndRemove_Call) RunAndReturn(run func(context.Context, map[string]string, []string, ...predicates.Predicate) error) *MockMetaKv_MultiSaveAndRemove_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// MultiSaveAndRemoveWithPrefix provides a mock function with given fields: saves, removals, preds
-func (_m *MockMetaKv) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string, preds ...predicates.Predicate) error {
+// MultiSaveAndRemoveWithPrefix provides a mock function with given fields: ctx, saves, removals, preds
+func (_m *MockMetaKv) MultiSaveAndRemoveWithPrefix(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate) error {
 	_va := make([]interface{}, len(preds))
 	for _i := range preds {
 		_va[_i] = preds[_i]
 	}
 	var _ca []interface{}
-	_ca = append(_ca, saves, removals)
+	_ca = append(_ca, ctx, saves, removals)
 	_ca = append(_ca, _va...)
 	ret := _m.Called(_ca...)
 
@@ -619,8 +631,8 @@ func (_m *MockMetaKv) MultiSaveAndRemoveWithPrefix(saves map[string]string, remo
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(map[string]string, []string, ...predicates.Predicate) error); ok {
-		r0 = rf(saves, removals, preds...)
+	if rf, ok := ret.Get(0).(func(context.Context, map[string]string, []string, ...predicates.Predicate) error); ok {
+		r0 = rf(ctx, saves, removals, preds...)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -634,23 +646,24 @@ type MockMetaKv_MultiSaveAndRemoveWithPrefix_Call struct {
 }
 
 // MultiSaveAndRemoveWithPrefix is a helper method to define mock.On call
+//   - ctx context.Context
 //   - saves map[string]string
 //   - removals []string
 //   - preds ...predicates.Predicate
-func (_e *MockMetaKv_Expecter) MultiSaveAndRemoveWithPrefix(saves interface{}, removals interface{}, preds ...interface{}) *MockMetaKv_MultiSaveAndRemoveWithPrefix_Call {
+func (_e *MockMetaKv_Expecter) MultiSaveAndRemoveWithPrefix(ctx interface{}, saves interface{}, removals interface{}, preds ...interface{}) *MockMetaKv_MultiSaveAndRemoveWithPrefix_Call {
 	return &MockMetaKv_MultiSaveAndRemoveWithPrefix_Call{Call: _e.mock.On("MultiSaveAndRemoveWithPrefix",
-		append([]interface{}{saves, removals}, preds...)...)}
+		append([]interface{}{ctx, saves, removals}, preds...)...)}
 }
 
-func (_c *MockMetaKv_MultiSaveAndRemoveWithPrefix_Call) Run(run func(saves map[string]string, removals []string, preds ...predicates.Predicate)) *MockMetaKv_MultiSaveAndRemoveWithPrefix_Call {
+func (_c *MockMetaKv_MultiSaveAndRemoveWithPrefix_Call) Run(run func(ctx context.Context, saves map[string]string, removals []string, preds ...predicates.Predicate)) *MockMetaKv_MultiSaveAndRemoveWithPrefix_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		variadicArgs := make([]predicates.Predicate, len(args)-2)
-		for i, a := range args[2:] {
+		variadicArgs := make([]predicates.Predicate, len(args)-3)
+		for i, a := range args[3:] {
 			if a != nil {
 				variadicArgs[i] = a.(predicates.Predicate)
 			}
 		}
-		run(args[0].(map[string]string), args[1].([]string), variadicArgs...)
+		run(args[0].(context.Context), args[1].(map[string]string), args[2].([]string), variadicArgs...)
 	})
 	return _c
 }
@@ -660,22 +673,22 @@ func (_c *MockMetaKv_MultiSaveAndRemoveWithPrefix_Call) Return(_a0 error) *MockM
 	return _c
 }
 
-func (_c *MockMetaKv_MultiSaveAndRemoveWithPrefix_Call) RunAndReturn(run func(map[string]string, []string, ...predicates.Predicate) error) *MockMetaKv_MultiSaveAndRemoveWithPrefix_Call {
+func (_c *MockMetaKv_MultiSaveAndRemoveWithPrefix_Call) RunAndReturn(run func(context.Context, map[string]string, []string, ...predicates.Predicate) error) *MockMetaKv_MultiSaveAndRemoveWithPrefix_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// Remove provides a mock function with given fields: key
-func (_m *MockMetaKv) Remove(key string) error {
-	ret := _m.Called(key)
+// Remove provides a mock function with given fields: ctx, key
+func (_m *MockMetaKv) Remove(ctx context.Context, key string) error {
+	ret := _m.Called(ctx, key)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Remove")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(string) error); ok {
-		r0 = rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) error); ok {
+		r0 = rf(ctx, key)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -689,14 +702,15 @@ type MockMetaKv_Remove_Call struct {
 }
 
 // Remove is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
-func (_e *MockMetaKv_Expecter) Remove(key interface{}) *MockMetaKv_Remove_Call {
-	return &MockMetaKv_Remove_Call{Call: _e.mock.On("Remove", key)}
+func (_e *MockMetaKv_Expecter) Remove(ctx interface{}, key interface{}) *MockMetaKv_Remove_Call {
+	return &MockMetaKv_Remove_Call{Call: _e.mock.On("Remove", ctx, key)}
 }
 
-func (_c *MockMetaKv_Remove_Call) Run(run func(key string)) *MockMetaKv_Remove_Call {
+func (_c *MockMetaKv_Remove_Call) Run(run func(ctx context.Context, key string)) *MockMetaKv_Remove_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string))
+		run(args[0].(context.Context), args[1].(string))
 	})
 	return _c
 }
@@ -706,22 +720,22 @@ func (_c *MockMetaKv_Remove_Call) Return(_a0 error) *MockMetaKv_Remove_Call {
 	return _c
 }
 
-func (_c *MockMetaKv_Remove_Call) RunAndReturn(run func(string) error) *MockMetaKv_Remove_Call {
+func (_c *MockMetaKv_Remove_Call) RunAndReturn(run func(context.Context, string) error) *MockMetaKv_Remove_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// RemoveWithPrefix provides a mock function with given fields: key
-func (_m *MockMetaKv) RemoveWithPrefix(key string) error {
-	ret := _m.Called(key)
+// RemoveWithPrefix provides a mock function with given fields: ctx, key
+func (_m *MockMetaKv) RemoveWithPrefix(ctx context.Context, key string) error {
+	ret := _m.Called(ctx, key)
 
 	if len(ret) == 0 {
 		panic("no return value specified for RemoveWithPrefix")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(string) error); ok {
-		r0 = rf(key)
+	if rf, ok := ret.Get(0).(func(context.Context, string) error); ok {
+		r0 = rf(ctx, key)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -735,14 +749,15 @@ type MockMetaKv_RemoveWithPrefix_Call struct {
 }
 
 // RemoveWithPrefix is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
-func (_e *MockMetaKv_Expecter) RemoveWithPrefix(key interface{}) *MockMetaKv_RemoveWithPrefix_Call {
-	return &MockMetaKv_RemoveWithPrefix_Call{Call: _e.mock.On("RemoveWithPrefix", key)}
+func (_e *MockMetaKv_Expecter) RemoveWithPrefix(ctx interface{}, key interface{}) *MockMetaKv_RemoveWithPrefix_Call {
+	return &MockMetaKv_RemoveWithPrefix_Call{Call: _e.mock.On("RemoveWithPrefix", ctx, key)}
 }
 
-func (_c *MockMetaKv_RemoveWithPrefix_Call) Run(run func(key string)) *MockMetaKv_RemoveWithPrefix_Call {
+func (_c *MockMetaKv_RemoveWithPrefix_Call) Run(run func(ctx context.Context, key string)) *MockMetaKv_RemoveWithPrefix_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string))
+		run(args[0].(context.Context), args[1].(string))
 	})
 	return _c
 }
@@ -752,22 +767,22 @@ func (_c *MockMetaKv_RemoveWithPrefix_Call) Return(_a0 error) *MockMetaKv_Remove
 	return _c
 }
 
-func (_c *MockMetaKv_RemoveWithPrefix_Call) RunAndReturn(run func(string) error) *MockMetaKv_RemoveWithPrefix_Call {
+func (_c *MockMetaKv_RemoveWithPrefix_Call) RunAndReturn(run func(context.Context, string) error) *MockMetaKv_RemoveWithPrefix_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// Save provides a mock function with given fields: key, value
-func (_m *MockMetaKv) Save(key string, value string) error {
-	ret := _m.Called(key, value)
+// Save provides a mock function with given fields: ctx, key, value
+func (_m *MockMetaKv) Save(ctx context.Context, key string, value string) error {
+	ret := _m.Called(ctx, key, value)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Save")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(string, string) error); ok {
-		r0 = rf(key, value)
+	if rf, ok := ret.Get(0).(func(context.Context, string, string) error); ok {
+		r0 = rf(ctx, key, value)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -781,15 +796,16 @@ type MockMetaKv_Save_Call struct {
 }
 
 // Save is a helper method to define mock.On call
+//   - ctx context.Context
 //   - key string
 //   - value string
-func (_e *MockMetaKv_Expecter) Save(key interface{}, value interface{}) *MockMetaKv_Save_Call {
-	return &MockMetaKv_Save_Call{Call: _e.mock.On("Save", key, value)}
+func (_e *MockMetaKv_Expecter) Save(ctx interface{}, key interface{}, value interface{}) *MockMetaKv_Save_Call {
+	return &MockMetaKv_Save_Call{Call: _e.mock.On("Save", ctx, key, value)}
 }
 
-func (_c *MockMetaKv_Save_Call) Run(run func(key string, value string)) *MockMetaKv_Save_Call {
+func (_c *MockMetaKv_Save_Call) Run(run func(ctx context.Context, key string, value string)) *MockMetaKv_Save_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string), args[1].(string))
+		run(args[0].(context.Context), args[1].(string), args[2].(string))
 	})
 	return _c
 }
@@ -799,22 +815,22 @@ func (_c *MockMetaKv_Save_Call) Return(_a0 error) *MockMetaKv_Save_Call {
 	return _c
 }
 
-func (_c *MockMetaKv_Save_Call) RunAndReturn(run func(string, string) error) *MockMetaKv_Save_Call {
+func (_c *MockMetaKv_Save_Call) RunAndReturn(run func(context.Context, string, string) error) *MockMetaKv_Save_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// WalkWithPrefix provides a mock function with given fields: prefix, paginationSize, fn
-func (_m *MockMetaKv) WalkWithPrefix(prefix string, paginationSize int, fn func([]byte, []byte) error) error {
-	ret := _m.Called(prefix, paginationSize, fn)
+// WalkWithPrefix provides a mock function with given fields: ctx, prefix, paginationSize, fn
+func (_m *MockMetaKv) WalkWithPrefix(ctx context.Context, prefix string, paginationSize int, fn func([]byte, []byte) error) error {
+	ret := _m.Called(ctx, prefix, paginationSize, fn)
 
 	if len(ret) == 0 {
 		panic("no return value specified for WalkWithPrefix")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(string, int, func([]byte, []byte) error) error); ok {
-		r0 = rf(prefix, paginationSize, fn)
+	if rf, ok := ret.Get(0).(func(context.Context, string, int, func([]byte, []byte) error) error); ok {
+		r0 = rf(ctx, prefix, paginationSize, fn)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -828,16 +844,17 @@ type MockMetaKv_WalkWithPrefix_Call struct {
 }
 
 // WalkWithPrefix is a helper method to define mock.On call
+//   - ctx context.Context
 //   - prefix string
 //   - paginationSize int
 //   - fn func([]byte , []byte) error
-func (_e *MockMetaKv_Expecter) WalkWithPrefix(prefix interface{}, paginationSize interface{}, fn interface{}) *MockMetaKv_WalkWithPrefix_Call {
-	return &MockMetaKv_WalkWithPrefix_Call{Call: _e.mock.On("WalkWithPrefix", prefix, paginationSize, fn)}
+func (_e *MockMetaKv_Expecter) WalkWithPrefix(ctx interface{}, prefix interface{}, paginationSize interface{}, fn interface{}) *MockMetaKv_WalkWithPrefix_Call {
+	return &MockMetaKv_WalkWithPrefix_Call{Call: _e.mock.On("WalkWithPrefix", ctx, prefix, paginationSize, fn)}
 }
 
-func (_c *MockMetaKv_WalkWithPrefix_Call) Run(run func(prefix string, paginationSize int, fn func([]byte, []byte) error)) *MockMetaKv_WalkWithPrefix_Call {
+func (_c *MockMetaKv_WalkWithPrefix_Call) Run(run func(ctx context.Context, prefix string, paginationSize int, fn func([]byte, []byte) error)) *MockMetaKv_WalkWithPrefix_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string), args[1].(int), args[2].(func([]byte, []byte) error))
+		run(args[0].(context.Context), args[1].(string), args[2].(int), args[3].(func([]byte, []byte) error))
 	})
 	return _c
 }
@@ -847,7 +864,7 @@ func (_c *MockMetaKv_WalkWithPrefix_Call) Return(_a0 error) *MockMetaKv_WalkWith
 	return _c
 }
 
-func (_c *MockMetaKv_WalkWithPrefix_Call) RunAndReturn(run func(string, int, func([]byte, []byte) error) error) *MockMetaKv_WalkWithPrefix_Call {
+func (_c *MockMetaKv_WalkWithPrefix_Call) RunAndReturn(run func(context.Context, string, int, func([]byte, []byte) error) error) *MockMetaKv_WalkWithPrefix_Call {
 	_c.Call.Return(run)
 	return _c
 }
diff --git a/pkg/mq/mqimpl/rocksmq/server/rocksmq_impl.go b/pkg/mq/mqimpl/rocksmq/server/rocksmq_impl.go
index 195604b5eeb7f..43554e1620d2a 100644
--- a/pkg/mq/mqimpl/rocksmq/server/rocksmq_impl.go
+++ b/pkg/mq/mqimpl/rocksmq/server/rocksmq_impl.go
@@ -12,6 +12,7 @@
 package server
 
 import (
+	"context"
 	"fmt"
 	"path"
 	"strconv"
@@ -356,7 +357,7 @@ func (rmq *rocksmq) Info() bool {
 		}
 
 		pageTsSizeKey := constructKey(PageTsTitle, topic)
-		pages, _, err := rmq.kv.LoadWithPrefix(pageTsSizeKey)
+		pages, _, err := rmq.kv.LoadWithPrefix(context.TODO(), pageTsSizeKey)
 		if err != nil {
 			log.Error("Rocksmq get page num failed", zap.String("topic", topic))
 			rtn = false
@@ -364,7 +365,7 @@ func (rmq *rocksmq) Info() bool {
 		}
 
 		msgSizeKey := MessageSizeTitle + topic
-		msgSizeVal, err := rmq.kv.Load(msgSizeKey)
+		msgSizeVal, err := rmq.kv.Load(context.TODO(), msgSizeKey)
 		if err != nil {
 			log.Error("Rocksmq get last page size failed", zap.String("topic", topic))
 			rtn = false
@@ -405,7 +406,7 @@ func (rmq *rocksmq) CreateTopic(topicName string) error {
 
 	// topicIDKey is the only identifier of a topic
 	topicIDKey := TopicIDTitle + topicName
-	val, err := rmq.kv.Load(topicIDKey)
+	val, err := rmq.kv.Load(context.TODO(), topicIDKey)
 	if err != nil {
 		return err
 	}
@@ -429,7 +430,7 @@ func (rmq *rocksmq) CreateTopic(topicName string) error {
 	// Initialize topic id to its creating time, we don't really use it for now
 	nowTs := strconv.FormatInt(time.Now().Unix(), 10)
 	kvs[topicIDKey] = nowTs
-	if err = rmq.kv.MultiSave(kvs); err != nil {
+	if err = rmq.kv.MultiSave(context.TODO(), kvs); err != nil {
 		return retry.Unrecoverable(err)
 	}
 
@@ -459,28 +460,28 @@ func (rmq *rocksmq) DestroyTopic(topicName string) error {
 
 	// clean the topic data it self
 	fixTopicName := topicName + "/"
-	err := rmq.kv.RemoveWithPrefix(fixTopicName)
+	err := rmq.kv.RemoveWithPrefix(context.TODO(), fixTopicName)
 	if err != nil {
 		return err
 	}
 
 	// clean page size info
 	pageMsgSizeKey := constructKey(PageMsgSizeTitle, topicName)
-	err = rmq.kv.RemoveWithPrefix(pageMsgSizeKey)
+	err = rmq.kv.RemoveWithPrefix(context.TODO(), pageMsgSizeKey)
 	if err != nil {
 		return err
 	}
 
 	// clean page ts info
 	pageMsgTsKey := constructKey(PageTsTitle, topicName)
-	err = rmq.kv.RemoveWithPrefix(pageMsgTsKey)
+	err = rmq.kv.RemoveWithPrefix(context.TODO(), pageMsgTsKey)
 	if err != nil {
 		return err
 	}
 
 	// cleaned acked ts info
 	ackedTsKey := constructKey(AckedTsTitle, topicName)
-	err = rmq.kv.RemoveWithPrefix(ackedTsKey)
+	err = rmq.kv.RemoveWithPrefix(context.TODO(), ackedTsKey)
 	if err != nil {
 		return err
 	}
@@ -492,7 +493,7 @@ func (rmq *rocksmq) DestroyTopic(topicName string) error {
 	var removedKeys []string
 	removedKeys = append(removedKeys, topicIDKey, msgSizeKey)
 	// Batch remove, atomic operation
-	err = rmq.kv.MultiRemove(removedKeys)
+	err = rmq.kv.MultiRemove(context.TODO(), removedKeys)
 	if err != nil {
 		return err
 	}
@@ -702,7 +703,7 @@ func (rmq *rocksmq) Produce(topicName string, messages []ProducerMessage) ([]Uni
 func (rmq *rocksmq) updatePageInfo(topicName string, msgIDs []UniqueID, msgSizes map[UniqueID]int64) error {
 	params := paramtable.Get()
 	msgSizeKey := MessageSizeTitle + topicName
-	msgSizeVal, err := rmq.kv.Load(msgSizeKey)
+	msgSizeVal, err := rmq.kv.Load(context.TODO(), msgSizeKey)
 	if err != nil {
 		return err
 	}
@@ -731,7 +732,7 @@ func (rmq *rocksmq) updatePageInfo(topicName string, msgIDs []UniqueID, msgSizes
 		}
 	}
 	mutateBuffer[msgSizeKey] = strconv.FormatInt(curMsgSize, 10)
-	err = rmq.kv.MultiSave(mutateBuffer)
+	err = rmq.kv.MultiSave(context.TODO(), mutateBuffer)
 	return err
 }
 
@@ -1125,7 +1126,7 @@ func (rmq *rocksmq) updateAckedInfo(topicName, groupName string, firstID UniqueI
 				ackedTsKvs[pageAckedTsKey] = nowTs
 			}
 		}
-		err := rmq.kv.MultiSave(ackedTsKvs)
+		err := rmq.kv.MultiSave(context.TODO(), ackedTsKvs)
 		if err != nil {
 			return err
 		}
diff --git a/pkg/mq/mqimpl/rocksmq/server/rocksmq_retention.go b/pkg/mq/mqimpl/rocksmq/server/rocksmq_retention.go
index 98ce3d03bcc22..fe58c8cc6433d 100644
--- a/pkg/mq/mqimpl/rocksmq/server/rocksmq_retention.go
+++ b/pkg/mq/mqimpl/rocksmq/server/rocksmq_retention.go
@@ -12,6 +12,7 @@
 package server
 
 import (
+	"context"
 	"fmt"
 	"path"
 	"strconv"
@@ -55,7 +56,7 @@ func initRetentionInfo(kv *rocksdbkv.RocksdbKV, db *gorocksdb.DB) (*retentionInf
 		closeWg:           sync.WaitGroup{},
 	}
 	// Get topic from topic begin id
-	topicKeys, _, err := ri.kv.LoadWithPrefix(TopicIDTitle)
+	topicKeys, _, err := ri.kv.LoadWithPrefix(context.TODO(), TopicIDTitle)
 	if err != nil {
 		return nil, err
 	}
@@ -164,7 +165,7 @@ func (ri *retentionInfo) expiredCleanUp(topic string) error {
 			return err
 		}
 		ackedTsKey := fixedAckedTsKey + "/" + strconv.FormatInt(pageID, 10)
-		ackedTsVal, err := ri.kv.Load(ackedTsKey)
+		ackedTsVal, err := ri.kv.Load(context.TODO(), ackedTsKey)
 		if err != nil {
 			return err
 		}
@@ -265,7 +266,7 @@ func (ri *retentionInfo) calculateTopicAckedSize(topic string) (int64, error) {
 
 		// check if page is acked
 		ackedTsKey := fixedAckedTsKey + "/" + strconv.FormatInt(pageID, 10)
-		ackedTsVal, err := ri.kv.Load(ackedTsKey)
+		ackedTsVal, err := ri.kv.Load(context.TODO(), ackedTsKey)
 		if err != nil {
 			return -1, err
 		}
diff --git a/pkg/mq/mqimpl/rocksmq/server/rocksmq_retention_test.go b/pkg/mq/mqimpl/rocksmq/server/rocksmq_retention_test.go
index fcb7b143f7585..1266d4068c550 100644
--- a/pkg/mq/mqimpl/rocksmq/server/rocksmq_retention_test.go
+++ b/pkg/mq/mqimpl/rocksmq/server/rocksmq_retention_test.go
@@ -12,6 +12,7 @@
 package server
 
 import (
+	"context"
 	"os"
 	"strconv"
 	"testing"
@@ -97,24 +98,24 @@ func TestRmqRetention_Basic(t *testing.T) {
 
 	// test acked size acked ts and other meta are updated as expect
 	msgSizeKey := MessageSizeTitle + topicName
-	msgSizeVal, err := rmq.kv.Load(msgSizeKey)
+	msgSizeVal, err := rmq.kv.Load(context.TODO(), msgSizeKey)
 	assert.NoError(t, err)
 	assert.Equal(t, msgSizeVal, "0")
 
 	pageMsgSizeKey := constructKey(PageMsgSizeTitle, topicName)
-	keys, values, err := rmq.kv.LoadWithPrefix(pageMsgSizeKey)
+	keys, values, err := rmq.kv.LoadWithPrefix(context.TODO(), pageMsgSizeKey)
 	assert.NoError(t, err)
 	assert.Equal(t, len(keys), 0)
 	assert.Equal(t, len(values), 0)
 
 	pageTsSizeKey := constructKey(PageTsTitle, topicName)
-	keys, values, err = rmq.kv.LoadWithPrefix(pageTsSizeKey)
+	keys, values, err = rmq.kv.LoadWithPrefix(context.TODO(), pageTsSizeKey)
 	assert.NoError(t, err)
 	assert.Equal(t, len(keys), 0)
 	assert.Equal(t, len(values), 0)
 
 	aclTsSizeKey := constructKey(AckedTsTitle, topicName)
-	keys, values, err = rmq.kv.LoadWithPrefix(aclTsSizeKey)
+	keys, values, err = rmq.kv.LoadWithPrefix(context.TODO(), aclTsSizeKey)
 	assert.NoError(t, err)
 	assert.Equal(t, len(keys), 0)
 	assert.Equal(t, len(values), 0)
@@ -183,7 +184,7 @@ func TestRmqRetention_NotConsumed(t *testing.T) {
 	id := cMsgs[0].MsgID
 
 	aclTsSizeKey := constructKey(AckedTsTitle, topicName)
-	keys, values, err := rmq.kv.LoadWithPrefix(aclTsSizeKey)
+	keys, values, err := rmq.kv.LoadWithPrefix(context.TODO(), aclTsSizeKey)
 	assert.NoError(t, err)
 	assert.Equal(t, len(keys), 2)
 	assert.Equal(t, len(values), 2)
@@ -202,26 +203,26 @@ func TestRmqRetention_NotConsumed(t *testing.T) {
 
 	// test acked size acked ts and other meta are updated as expect
 	msgSizeKey := MessageSizeTitle + topicName
-	msgSizeVal, err := rmq.kv.Load(msgSizeKey)
+	msgSizeVal, err := rmq.kv.Load(context.TODO(), msgSizeKey)
 	assert.NoError(t, err)
 	assert.Equal(t, msgSizeVal, "0")
 
 	// should only clean 2 pages
 	pageMsgSizeKey := constructKey(PageMsgSizeTitle, topicName)
-	keys, values, err = rmq.kv.LoadWithPrefix(pageMsgSizeKey)
+	keys, values, err = rmq.kv.LoadWithPrefix(context.TODO(), pageMsgSizeKey)
 	assert.NoError(t, err)
 	assert.Equal(t, len(keys), 48)
 	assert.Equal(t, len(values), 48)
 
 	pageTsSizeKey := constructKey(PageTsTitle, topicName)
-	keys, values, err = rmq.kv.LoadWithPrefix(pageTsSizeKey)
+	keys, values, err = rmq.kv.LoadWithPrefix(context.TODO(), pageTsSizeKey)
 	assert.NoError(t, err)
 
 	assert.Equal(t, len(keys), 48)
 	assert.Equal(t, len(values), 48)
 
 	aclTsSizeKey = constructKey(AckedTsTitle, topicName)
-	keys, values, err = rmq.kv.LoadWithPrefix(aclTsSizeKey)
+	keys, values, err = rmq.kv.LoadWithPrefix(context.TODO(), aclTsSizeKey)
 	assert.NoError(t, err)
 	assert.Equal(t, len(keys), 0)
 	assert.Equal(t, len(values), 0)
@@ -317,46 +318,46 @@ func TestRmqRetention_MultipleTopic(t *testing.T) {
 
 	// test acked size acked ts and other meta are updated as expect
 	msgSizeKey := MessageSizeTitle + topicName
-	msgSizeVal, err := rmq.kv.Load(msgSizeKey)
+	msgSizeVal, err := rmq.kv.Load(context.TODO(), msgSizeKey)
 	assert.NoError(t, err)
 	assert.Equal(t, msgSizeVal, "0")
 
 	// 100 page left, each entity is a page
 	pageMsgSizeKey := constructKey(PageMsgSizeTitle, "topic_a")
-	keys, values, err := rmq.kv.LoadWithPrefix(pageMsgSizeKey)
+	keys, values, err := rmq.kv.LoadWithPrefix(context.TODO(), pageMsgSizeKey)
 	assert.NoError(t, err)
 	assert.Equal(t, len(keys), 0)
 	assert.Equal(t, len(values), 0)
 
 	pageTsSizeKey := constructKey(PageTsTitle, "topic_a")
-	keys, values, err = rmq.kv.LoadWithPrefix(pageTsSizeKey)
+	keys, values, err = rmq.kv.LoadWithPrefix(context.TODO(), pageTsSizeKey)
 	assert.NoError(t, err)
 
 	assert.Equal(t, len(keys), 0)
 	assert.Equal(t, len(values), 0)
 
 	aclTsSizeKey := constructKey(AckedTsTitle, "topic_a")
-	keys, values, err = rmq.kv.LoadWithPrefix(aclTsSizeKey)
+	keys, values, err = rmq.kv.LoadWithPrefix(context.TODO(), aclTsSizeKey)
 	assert.NoError(t, err)
 	assert.Equal(t, len(keys), 0)
 	assert.Equal(t, len(values), 0)
 
 	// for topic B, nothing has been cleadn
 	pageMsgSizeKey = constructKey(PageMsgSizeTitle, "topic_b")
-	keys, values, err = rmq.kv.LoadWithPrefix(pageMsgSizeKey)
+	keys, values, err = rmq.kv.LoadWithPrefix(context.TODO(), pageMsgSizeKey)
 	assert.NoError(t, err)
 	assert.Equal(t, len(keys), 50)
 	assert.Equal(t, len(values), 50)
 
 	pageTsSizeKey = constructKey(PageTsTitle, "topic_b")
-	keys, values, err = rmq.kv.LoadWithPrefix(pageTsSizeKey)
+	keys, values, err = rmq.kv.LoadWithPrefix(context.TODO(), pageTsSizeKey)
 	assert.NoError(t, err)
 
 	assert.Equal(t, len(keys), 50)
 	assert.Equal(t, len(values), 50)
 
 	aclTsSizeKey = constructKey(AckedTsTitle, "topic_b")
-	keys, values, err = rmq.kv.LoadWithPrefix(aclTsSizeKey)
+	keys, values, err = rmq.kv.LoadWithPrefix(context.TODO(), aclTsSizeKey)
 	assert.NoError(t, err)
 	assert.Equal(t, len(keys), 0)
 	assert.Equal(t, len(values), 0)
@@ -539,26 +540,26 @@ func TestRmqRetention_PageTimeExpire(t *testing.T) {
 
 	// test acked size acked ts and other meta are updated as expect
 	msgSizeKey := MessageSizeTitle + topicName
-	msgSizeVal, err := rmq.kv.Load(msgSizeKey)
+	msgSizeVal, err := rmq.kv.Load(context.TODO(), msgSizeKey)
 	assert.NoError(t, err)
 	assert.Equal(t, msgSizeVal, "0")
 
 	// 100 page left, each entity is a page
 	pageMsgSizeKey := constructKey(PageMsgSizeTitle, topicName)
-	keys, values, err := rmq.kv.LoadWithPrefix(pageMsgSizeKey)
+	keys, values, err := rmq.kv.LoadWithPrefix(context.TODO(), pageMsgSizeKey)
 	assert.NoError(t, err)
 	assert.Equal(t, len(keys), 100)
 	assert.Equal(t, len(values), 100)
 
 	pageTsSizeKey := constructKey(PageTsTitle, topicName)
-	keys, values, err = rmq.kv.LoadWithPrefix(pageTsSizeKey)
+	keys, values, err = rmq.kv.LoadWithPrefix(context.TODO(), pageTsSizeKey)
 	assert.NoError(t, err)
 
 	assert.Equal(t, len(keys), 100)
 	assert.Equal(t, len(values), 100)
 
 	aclTsSizeKey := constructKey(AckedTsTitle, topicName)
-	keys, values, err = rmq.kv.LoadWithPrefix(aclTsSizeKey)
+	keys, values, err = rmq.kv.LoadWithPrefix(context.TODO(), aclTsSizeKey)
 	assert.NoError(t, err)
 	assert.Equal(t, len(keys), 100)
 	assert.Equal(t, len(values), 100)