perf: optimize listpack removeNext cost

xgzlucario · Jul 16, 2024 · fec9d2f · fec9d2f
1 parent 5db120f
commit fec9d2f
Show file tree

Hide file tree

Showing 4 changed files with 73 additions and 11 deletions.
diff --git a/internal/hash/benchmark/main.go b/internal/hash/benchmark/main.go
@@ -33,10 +33,10 @@ func main() {
 	fmt.Println(c)
 
 	start := time.Now()
+	m := map[int]any{}
 
 	switch c {
 	case "map":
-		m := map[int]any{}
 		for i := 0; i < 10000; i++ {
 			hm := hash.NewMap()
 			for i := 0; i < 512; i++ {
@@ -47,7 +47,6 @@ func main() {
 		}
 
 	case "zipmap":
-		m := map[int]any{}
 		for i := 0; i < 10000; i++ {
 			hm := hash.NewZipMap()
 			for i := 0; i < 512; i++ {
@@ -58,7 +57,6 @@ func main() {
 		}
 
 	case "zipmap-compressed":
-		m := map[int]any{}
 		for i := 0; i < 10000; i++ {
 			hm := hash.NewZipMap()
 			for i := 0; i < 512; i++ {

diff --git a/internal/hash/map_test.go b/internal/hash/map_test.go
@@ -35,6 +35,11 @@ func testMapI(m MapI, t *testing.T) {
 	assert.True(ok)
 	assert.Equal(string(val), "val3")
 
+	val, ok = m.Get("key999")
+	assert.False(ok)
+	var nilBytes []byte
+	assert.Equal(val, nilBytes)
+
 	// set(update)
 	assert.False(m.Set("key1", []byte("newval1")))
 	assert.False(m.Set("key2", []byte("newval2")))
@@ -53,6 +58,10 @@ func testMapI(m MapI, t *testing.T) {
 	assert.True(ok)
 	assert.Equal(string(val), "newval3")
 
+	val, ok = m.Get("key999")
+	assert.False(ok)
+	assert.Equal(val, nilBytes)
+
 	// scan
 	count := 0
 	m.Scan(func(key string, val []byte) {
@@ -77,3 +86,34 @@ func testMapI(m MapI, t *testing.T) {
 	// len
 	assert.Equal(m.Len(), 0)
 }
+
+func TestToMap(t *testing.T) {
+	assert := assert.New(t)
+
+	m := NewZipMap()
+	m.Set("key1", []byte("value1"))
+	m.Set("key2", []byte("value2"))
+	m.Set("key3", []byte("value3"))
+
+	m.Compress()
+	m.Decompress()
+
+	nm := m.ToMap()
+
+	// scan
+	count := 0
+	nm.Scan(func(key string, val []byte) {
+		switch key {
+		case "key1":
+			assert.Equal(val, []byte("value1"))
+		case "key2":
+			assert.Equal(val, []byte("value2"))
+		case "key3":
+			assert.Equal(val, []byte("value3"))
+		default:
+			panic("error")
+		}
+		count++
+	})
+	assert.Equal(count, 3)
+}
diff --git a/internal/hash/set_test.go b/internal/hash/set_test.go
@@ -63,3 +63,29 @@ func testSetI(m SetI, t *testing.T) {
 	// len
 	assert.Equal(m.Len(), 0)
 }
+
+func TestToSet(t *testing.T) {
+	assert := assert.New(t)
+
+	m := NewZipSet()
+	m.Add("key1")
+	m.Add("key2")
+	m.Add("key3")
+
+	m.Compress()
+	m.Decompress()
+
+	nm := m.ToSet()
+
+	// scan
+	count := 0
+	nm.Scan(func(key string) {
+		switch key {
+		case "key1", "key2", "key3":
+		default:
+			panic("error")
+		}
+		count++
+	})
+	assert.Equal(count, 3)
+}
diff --git a/internal/list/listpack.go b/internal/list/listpack.go
@@ -77,6 +77,7 @@ func (lp *ListPack) Compress() {
 		return
 	}
 	lp.data = encoder.EncodeAll(lp.data, make([]byte, 0, len(lp.data)/3))
+	lp.data = slices.Clip(lp.data)
 	lp.compress = true
 }
 
@@ -156,14 +157,7 @@ func (it *lpIterator) Prev() []byte {
 }
 
 func (it *lpIterator) Insert(datas ...string) {
-	if it.IsLast() {
-		for _, data := range datas {
-			it.data = appendEntry(it.data, data)
-			it.size++
-		}
-		return
-	}
-	alloc := bpool.Get(maxListPackSize)[:0]
+	var alloc []byte
 	for _, data := range datas {
 		alloc = appendEntry(alloc, data)
 		it.size++
@@ -212,9 +206,13 @@ func (it *lpIterator) ReplaceNext(key string) {
 	alloc := appendEntry(nil, key)
 	it.data = slices.Replace(it.data, before, after, alloc...)
 	it.index = before
+	bpool.Put(alloc)
 }
 
 func appendEntry(dst []byte, data string) []byte {
+	if dst == nil {
+		dst = bpool.Get(len(data) + 10)[:0]
+	}
 	before := len(dst)
 	dst = appendUvarint(dst, len(data), false)
 	dst = append(dst, data...)