milvus-io
/
milvus
mirror of https://github.com/milvus-io/milvus.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add ByteSlice Method for memkv (#15774) 

Signed-off-by: Letian Jiang <letian.jiang@zilliz.com>
pull/15814/head
Letian Jiang 2022-03-02 17:29:56 +08:00 committed by GitHub
parent 63eec7ffc8
commit 442e791128
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 350 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

192
internal/kv/mem/mem_kv.go
View File

@ -38,10 +38,38 @@ func NewMemoryKV() *MemoryKV {
}
}
type memoryKVItem struct {
key, value string
type Value interface {
String() string
ByteSlice() []byte
}
type StringValue string
func (v StringValue) String() string {
return string(v)
}
func (v StringValue) ByteSlice() []byte {
return []byte(v)
}
type ByteSliceValue []byte
func (v ByteSliceValue) String() string {
return string(v)
}
func (v ByteSliceValue) ByteSlice() []byte {
return v
}
type memoryKVItem struct {
key string
value Value
}
var _ btree.Item = (*memoryKVItem)(nil)
// Less returns true if the item is less than the given one.
func (s memoryKVItem) Less(than btree.Item) bool {
return s.key < than.(memoryKVItem).key
@ -51,24 +79,48 @@ func (s memoryKVItem) Less(than btree.Item) bool {
func (kv *MemoryKV) Load(key string) (string, error) {
kv.RLock()
defer kv.RUnlock()
item := kv.tree.Get(memoryKVItem{key, ""})
item := kv.tree.Get(memoryKVItem{key: key})
// TODO，load unexisted key behavior is weird
if item == nil {
return "", nil
}
return item.(memoryKVItem).value, nil
return item.(memoryKVItem).value.String(), nil
}
// LoadBytes loads an object with @key.
func (kv *MemoryKV) LoadBytes(key string) ([]byte, error) {
kv.RLock()
defer kv.RUnlock()
item := kv.tree.Get(memoryKVItem{key: key})
// TODO，load unexisted key behavior is weird
if item == nil {
return []byte{}, nil
}
return item.(memoryKVItem).value.ByteSlice(), nil
}
// LoadWithDefault loads an object with @key. If the object does not exist, @defaultValue will be returned.
func (kv *MemoryKV) LoadWithDefault(key, defaultValue string) string {
kv.RLock()
defer kv.RUnlock()
item := kv.tree.Get(memoryKVItem{key, ""})
item := kv.tree.Get(memoryKVItem{key: key})
if item == nil {
return defaultValue
}
return item.(memoryKVItem).value
return item.(memoryKVItem).value.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 {
kv.RLock()
defer kv.RUnlock()
item := kv.tree.Get(memoryKVItem{key: key})
if item == nil {
return defaultValue
}
return item.(memoryKVItem).value.ByteSlice()
}
// LoadRange loads objects with range @startKey to @endKey with @limit number of objects.
@ -77,9 +129,26 @@ func (kv *MemoryKV) LoadRange(key, endKey string, limit int) ([]string, []string
defer kv.RUnlock()
keys := make([]string, 0, limit)
values := make([]string, 0, limit)
kv.tree.AscendRange(memoryKVItem{key, ""}, memoryKVItem{endKey, ""}, func(item btree.Item) bool {
kv.tree.AscendRange(memoryKVItem{key: key}, memoryKVItem{key: endKey}, func(item btree.Item) bool {
keys = append(keys, item.(memoryKVItem).key)
values = append(values, item.(memoryKVItem).value)
values = append(values, item.(memoryKVItem).value.String())
if limit > 0 {
return len(keys) < limit
}
return true
})
return keys, values, nil
}
// 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) {
kv.RLock()
defer kv.RUnlock()
keys := make([]string, 0, limit)
values := make([][]byte, 0, limit)
kv.tree.AscendRange(memoryKVItem{key: key}, memoryKVItem{key: endKey}, func(item btree.Item) bool {
keys = append(keys, item.(memoryKVItem).key)
values = append(values, item.(memoryKVItem).value.ByteSlice())
if limit > 0 {
return len(keys) < limit
}
@ -92,7 +161,15 @@ func (kv *MemoryKV) LoadRange(key, endKey string, limit int) ([]string, []string
func (kv *MemoryKV) Save(key, value string) error {
kv.Lock()
defer kv.Unlock()
kv.tree.ReplaceOrInsert(memoryKVItem{key, value})
kv.tree.ReplaceOrInsert(memoryKVItem{key, StringValue(value)})
return nil
}
// SaveBytes object with @key to btree. Object value is @value.
func (kv *MemoryKV) SaveBytes(key string, value []byte) error {
kv.Lock()
defer kv.Unlock()
kv.tree.ReplaceOrInsert(memoryKVItem{key, ByteSliceValue(value)})
return nil
}
@ -101,7 +178,7 @@ func (kv *MemoryKV) Remove(key string) error {
kv.Lock()
defer kv.Unlock()
kv.tree.Delete(memoryKVItem{key, ""})
kv.tree.Delete(memoryKVItem{key: key})
return nil
}
@ -111,8 +188,20 @@ func (kv *MemoryKV) MultiLoad(keys []string) ([]string, error) {
defer kv.RUnlock()
result := make([]string, 0, len(keys))
for _, key := range keys {
item := kv.tree.Get(memoryKVItem{key, ""})
result = append(result, item.(memoryKVItem).value)
item := kv.tree.Get(memoryKVItem{key: key})
result = append(result, item.(memoryKVItem).value.String())
}
return result, nil
}
// MultiLoadBytes loads objects with multi @keys.
func (kv *MemoryKV) MultiLoadBytes(keys []string) ([][]byte, error) {
kv.RLock()
defer kv.RUnlock()
result := make([][]byte, 0, len(keys))
for _, key := range keys {
item := kv.tree.Get(memoryKVItem{key: key})
result = append(result, item.(memoryKVItem).value.ByteSlice())
}
return result, nil
}
@ -122,7 +211,17 @@ func (kv *MemoryKV) MultiSave(kvs map[string]string) error {
kv.Lock()
defer kv.Unlock()
for key, value := range kvs {
kv.tree.ReplaceOrInsert(memoryKVItem{key, value})
kv.tree.ReplaceOrInsert(memoryKVItem{key, StringValue(value)})
}
return nil
}
// MultiSaveBytes saves given key-value pairs in MemoryKV atomicly.
func (kv *MemoryKV) MultiSaveBytes(kvs map[string][]byte) error {
kv.Lock()
defer kv.Unlock()
for key, value := range kvs {
kv.tree.ReplaceOrInsert(memoryKVItem{key, ByteSliceValue(value)})
}
return nil
}
@ -132,7 +231,7 @@ func (kv *MemoryKV) MultiRemove(keys []string) error {
kv.Lock()
defer kv.Unlock()
for _, key := range keys {
kv.tree.Delete(memoryKVItem{key, ""})
kv.tree.Delete(memoryKVItem{key: key})
}
return nil
}
@ -142,10 +241,23 @@ func (kv *MemoryKV) MultiSaveAndRemove(saves map[string]string, removals []strin
kv.Lock()
defer kv.Unlock()
for key, value := range saves {
kv.tree.ReplaceOrInsert(memoryKVItem{key, value})
kv.tree.ReplaceOrInsert(memoryKVItem{key, StringValue(value)})
}
for _, key := range removals {
kv.tree.Delete(memoryKVItem{key, ""})
kv.tree.Delete(memoryKVItem{key: key})
}
return nil
}
// MultiSaveBytesAndRemove saves and removes given key-value pairs in MemoryKV atomicly.
func (kv *MemoryKV) MultiSaveBytesAndRemove(saves map[string][]byte, removals []string) error {
kv.Lock()
defer kv.Unlock()
for key, value := range saves {
kv.tree.ReplaceOrInsert(memoryKVItem{key, ByteSliceValue(value)})
}
for _, key := range removals {
kv.tree.Delete(memoryKVItem{key: key})
}
return nil
}
@ -161,7 +273,25 @@ func (kv *MemoryKV) LoadWithPrefix(key string) ([]string, []string, error) {
kv.tree.Ascend(func(i btree.Item) bool {
if strings.HasPrefix(i.(memoryKVItem).key, key) {
keys = append(keys, i.(memoryKVItem).key)
values = append(values, i.(memoryKVItem).value)
values = append(values, i.(memoryKVItem).value.String())
}
return true
})
return keys, values, nil
}
// LoadBytesWithPrefix returns all keys & values with given prefix.
func (kv *MemoryKV) LoadBytesWithPrefix(key string) ([]string, [][]byte, error) {
kv.Lock()
defer kv.Unlock()
var keys []string
var values [][]byte
kv.tree.Ascend(func(i btree.Item) bool {
if strings.HasPrefix(i.(memoryKVItem).key, key) {
keys = append(keys, i.(memoryKVItem).key)
values = append(values, i.(memoryKVItem).value.ByteSlice())
}
return true
})
@ -177,7 +307,7 @@ func (kv *MemoryKV) MultiRemoveWithPrefix(keys []string) error {
panic("not implement")
}
// MultiSaveAndRemoveWithPrefix saves key-value pairs in @saves, & remove key with prefix in @removals in MemoryKV atomicly.
// 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) error {
kv.Lock()
defer kv.Unlock()
@ -196,7 +326,31 @@ func (kv *MemoryKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, remova
}
for key, value := range saves {
kv.tree.ReplaceOrInsert(memoryKVItem{key, value})
kv.tree.ReplaceOrInsert(memoryKVItem{key, StringValue(value)})
}
return nil
}
// 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 {
kv.Lock()
defer kv.Unlock()
var keys []memoryKVItem
for _, key := range removals {
kv.tree.Ascend(func(i btree.Item) bool {
if strings.HasPrefix(i.(memoryKVItem).key, key) {
keys = append(keys, i.(memoryKVItem))
}
return true
})
}
for _, item := range keys {
kv.tree.Delete(item)
}
for key, value := range saves {
kv.tree.ReplaceOrInsert(memoryKVItem{key, ByteSliceValue(value)})
}
return nil
}

177
internal/kv/mem/mem_kv_test.go
View File

@ -22,6 +22,183 @@ import (
"github.com/stretchr/testify/assert"
)
func TestMemoryKV_SaveAndLoadBytes(t *testing.T) {
mem := NewMemoryKV()
key := "key"
value := []byte("value")
err := mem.SaveBytes(key, value)
assert.NoError(t, err)
_value, err := mem.LoadBytes(key)
assert.NoError(t, err)
assert.Equal(t, value, _value)
noKey := "no_key"
_value, err = mem.LoadBytes(noKey)
assert.NoError(t, err)
assert.Empty(t, _value)
}
func TestMemoryKV_LoadBytesRange(t *testing.T) {
saveAndLoadBytesTests := []struct {
key string
value []byte
}{
{"test1", []byte("value1")},
{"test2", []byte("value2")},
{"test1/a", []byte("value_a")},
{"test1/b", []byte("value_b")},
}
mem := NewMemoryKV()
for _, kv := range saveAndLoadBytesTests {
err := mem.SaveBytes(kv.key, kv.value)
assert.NoError(t, err)
}
keys, values, err := mem.LoadBytesRange("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)
assert.Equal(t, len(keys), 1)
assert.Equal(t, len(values), 1)
assert.NoError(t, err)
}
func TestMemoryKV_LoadBytesWithDefault(t *testing.T) {
mem := NewMemoryKV()
key := "key"
value := []byte("value")
err := mem.SaveBytes(key, value)
assert.NoError(t, err)
_default := []byte("default")
_value := mem.LoadBytesWithDefault(key, _default)
assert.Equal(t, value, _value)
noKey := "no_key"
_value = mem.LoadBytesWithDefault(noKey, _default)
assert.Equal(t, _value, _default)
}
func TestMemoryKV_LoadBytesWithPrefix(t *testing.T) {
saveAndLoadBytesTests := []struct {
key string
value []byte
}{
{"test1", []byte("value1")},
{"test2", []byte("value2")},
{"test1/a", []byte("value_a")},
{"test1/b", []byte("value_b")},
}
mem := NewMemoryKV()
for _, kv := range saveAndLoadBytesTests {
err := mem.SaveBytes(kv.key, kv.value)
assert.NoError(t, err)
}
keys, values, err := mem.LoadBytesWithPrefix("test1")
assert.Equal(t, len(keys), 3)
assert.Equal(t, len(values), 3)
assert.NoError(t, err)
keys, values, err = mem.LoadBytesWithPrefix("test")
assert.Equal(t, len(keys), 4)
assert.Equal(t, len(values), 4)
assert.NoError(t, err)
keys, values, err = mem.LoadBytesWithPrefix("a")
assert.Equal(t, len(keys), 0)
assert.Equal(t, len(values), 0)
assert.NoError(t, err)
}
func TestMemoryKV_MultiSaveBytes(t *testing.T) {
saveAndLoadBytesTests := map[string][]byte{
"test1": []byte("value1"),
"test2": []byte("value2"),
"test1/a": []byte("value_a"),
"test1/b": []byte("value_b"),
}
var keys []string
var values [][]byte
for k, v := range saveAndLoadBytesTests {
keys = append(keys, k)
values = append(values, v)
}
mem := NewMemoryKV()
err := mem.MultiSaveBytes(saveAndLoadBytesTests)
assert.NoError(t, err)
_values, err := mem.MultiLoadBytes(keys)
assert.Equal(t, values, _values)
assert.NoError(t, err)
_values, err = mem.MultiLoadBytes([]string{})
assert.Empty(t, _values)
assert.NoError(t, err)
}
func TestMemoryKV_MultiSaveBytesAndRemove(t *testing.T) {
saveAndLoadBytesTests := map[string][]byte{
"test1": []byte("value1"),
"test2": []byte("value2"),
"test1/a": []byte("value_a"),
"test1/b": []byte("value_b"),
}
var keys []string
var values [][]byte
for k, v := range saveAndLoadBytesTests {
keys = append(keys, k)
values = append(values, v)
}
mem := NewMemoryKV()
err := mem.MultiSaveBytesAndRemove(saveAndLoadBytesTests, []string{keys[0]})
assert.NoError(t, err)
_value, err := mem.LoadBytes(keys[0])
assert.Empty(t, _value)
assert.NoError(t, err)
_values, err := mem.MultiLoadBytes(keys[1:])
assert.Equal(t, values[1:], _values)
assert.NoError(t, err)
}
func TestMemoryKV_MultiSaveBytesAndRemoveWithPrefix(t *testing.T) {
saveAndLoadBytesTests := map[string][]byte{
"test1": []byte("value1"),
"test2": []byte("value2"),
"test1/a": []byte("value_a"),
"test1/b": []byte("value_b"),
}
var keys []string
var values [][]byte
for k, v := range saveAndLoadBytesTests {
keys = append(keys, k)
values = append(values, v)
}
mem := NewMemoryKV()
err := mem.MultiSaveBytesAndRemoveWithPrefix(saveAndLoadBytesTests, []string{"test1"})
assert.NoError(t, err)
_keys, _values, err := mem.LoadBytesWithPrefix("test")
assert.ElementsMatch(t, keys, _keys)
assert.ElementsMatch(t, values, _values)
assert.NoError(t, err)
}
func TestMemoryKV_LoadPartial(t *testing.T) {
memKV := NewMemoryKV()