feat: add zipmap, zipset based on listpack

README.md

@@ -4,17 +4,15 @@
 
 ## 介绍
 
-这里是 rotom，一个使用 Go 编写的 tiny Redis Server。基于 IO 多路复用技术，还原了 Redis 中的 AeLoop 核心事件循环机制。
-
-rotom 基于 [godis](https://github.com/archeryue/godis) 项目
+这里是 rotom，一个使用 Go 编写的 tiny Redis Server。基于 IO 多路复用还原了 Redis 中的 AeLoop 核心事件循环机制。
 
 ### 实现特性
 
 1. 基于 epoll 网络模型，还原了 Redis 中的 AeLoop 单线程事件循环
 2. 兼容 Redis RESP 协议，你可以使用任何 redis 客户端连接 rotom
-3. 实现了 dict, quicklist, hash, set, zset 数据结构
+3. 实现了 dict, quicklist(listpack), hash(map, zipmap), set(mapset, zipset), zset 数据结构
 4. AOF 支持
-5. 支持 17 种常用命令
+5. 支持 18 种常用命令
 
 ### 原理介绍
 
@@ -64,7 +62,7 @@ $ go run .
 
 ```
 REPOSITORY       TAG           IMAGE ID       CREATED         SIZE
-rotom            latest        22f42ce9ae0e   8 seconds ago   18.6MB
+rotom            latest        22f42ce9ae0e   8 seconds ago   18.8MB
 ```
 
 然后启动容器：

command.go

@@ -37,6 +37,7 @@ var cmdTable []*Command = []*Command{
 	{"rpop", rpopCommand, 1, true},
 	{"lpop", lpopCommand, 1, true},
 	{"sadd", saddCommand, 2, true},
+	{"srem", sremCommand, 2, true},
 	{"spop", spopCommand, 1, true},
 	{"zadd", zaddCommand, 3, true},
 	{"ping", pingCommand, 0, false},
@@ -342,6 +343,24 @@ func saddCommand(writer *RESPWriter, args []RESP) {
 	writer.WriteInteger(newItems)
 }
 
+func sremCommand(writer *RESPWriter, args []RESP) {
+	key := args[0].ToString()
+
+	set, err := fetchSet(key)
+	if err != nil {
+		writer.WriteError(err)
+		return
+	}
+
+	var count int
+	for _, arg := range args[1:] {
+		if set.Remove(arg.ToStringUnsafe()) {
+			count++
+		}
+	}
+	writer.WriteInteger(count)
+}
+
 func spopCommand(writer *RESPWriter, args []RESP) {
 	key := args[0].ToString()
 

command_test.go

@@ -137,13 +137,19 @@ func TestCommand(t *testing.T) {
 		n, _ := rdb.SAdd(ctx, "set", "k1", "k2", "k3").Result()
 		assert.Equal(n, int64(3))
 
+		// spop
 		for i := 0; i < 3; i++ {
 			val, _ := rdb.SPop(ctx, "set").Result()
 			assert.NotEqual(val, "")
 		}
 
 		_, err := rdb.SPop(ctx, "set").Result()
 		assert.Equal(err, redis.Nil)
+
+		// srem
+		rdb.SAdd(ctx, "set", "k1", "k2", "k3").Result()
+		res, _ := rdb.SRem(ctx, "set", "k1", "k2", "k999").Result()
+		assert.Equal(res, int64(2))
 	})
 
 	t.Run("zset", func(t *testing.T) {

go.mod

@@ -4,6 +4,7 @@ go 1.22
 
 require (
 	github.com/cockroachdb/swiss v0.0.0-20240612210725-f4de07ae6964
+	github.com/deckarep/golang-set/v2 v2.6.0
 	github.com/klauspost/compress v1.17.9
 	github.com/redis/go-redis/v9 v9.5.2
 	github.com/rs/zerolog v1.33.0

go.sum

@@ -12,6 +12,8 @@ github.com/coreos/go-systemd/v22 v22.5.0/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSV
 github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/deckarep/golang-set/v2 v2.6.0 h1:XfcQbWM1LlMB8BsJ8N9vW5ehnnPVIw0je80NsVHagjM=
+github.com/deckarep/golang-set/v2 v2.6.0/go.mod h1:VAky9rY/yGXJOLEDv3OMci+7wtDpOF4IN+y82NBOac4=
 github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78=
 github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc=
 github.com/edsrzf/mmap-go v1.1.0 h1:6EUwBLQ/Mcr1EYLE4Tn1VdW1A4ckqCQWZBw8Hr0kjpQ=

internal/hash/bench_test.go

@@ -0,0 +1,52 @@
+package hash
+
+import (
+	"fmt"
+	"testing"
+)
+
+func genKey(i int) string {
+	return fmt.Sprintf("%08x", i)
+}
+
+func BenchmarkMap(b *testing.B) {
+	benchMapI("map", func() MapI { return NewMap() }, b)
+	benchMapI("zipmap", func() MapI { return NewZipMap() }, b)
+}
+
+func benchMapI(name string, newf func() MapI, b *testing.B) {
+	b.Run(name+"-set", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			m := newf()
+			for i := 0; i < 512; i++ {
+				k := genKey(i)
+				m.Set(k, []byte(k))
+			}
+		}
+	})
+	b.Run(name+"-update", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			m := newf()
+			for i := 0; i < 512; i++ {
+				k := genKey(0)
+				m.Set(k, []byte(k))
+			}
+		}
+	})
+}
+
+func BenchmarkSet(b *testing.B) {
+	benchSetI("set", func() SetI { return NewSet() }, b)
+	benchSetI("zipset", func() SetI { return NewZipSet() }, b)
+}
+
+func benchSetI(name string, newf func() SetI, b *testing.B) {
+	b.Run(name+"-add", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			m := newf()
+			for i := 0; i < 512; i++ {
+				m.Add(genKey(i))
+			}
+		}
+	})
+}

internal/hash/map.go

@@ -2,19 +2,24 @@ package hash
 
 import (
 	"github.com/cockroachdb/swiss"
-	"github.com/xgzlucario/rotom/internal/pkg"
 )
 
-var (
-	mapAllocator = pkg.NewAllocator[string, []byte]()
-)
+type MapI interface {
+	Set(key string, val []byte) (ok bool)
+	Get(key string) (val []byte, ok bool)
+	Remove(key string) (ok bool)
+	Len() int
+	Scan(iterator func(key string, val []byte))
+}
+
+var _ MapI = (*Map)(nil)
 
 type Map struct {
 	m *swiss.Map[string, []byte]
 }
 
 func NewMap() *Map {
-	return &Map{m: swiss.New(8, swiss.WithAllocator(mapAllocator))}
+	return &Map{m: swiss.New[string, []byte](8)}
 }
 
 func (m *Map) Get(key string) ([]byte, bool) {
@@ -43,7 +48,3 @@ func (m *Map) Scan(fn func(key string, value []byte)) {
 		return true
 	})
 }
-
-func (m *Map) Free() {
-	m.m.Close()
-}

internal/hash/map_test.go

@@ -0,0 +1,64 @@
+package hash
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestMap(t *testing.T) {
+	testMapI(NewMap(), t)
+	testMapI(NewZipMap(), t)
+}
+
+func testMapI(m MapI, t *testing.T) {
+	assert := assert.New(t)
+
+	// set
+	assert.True(m.Set("key1", []byte("val1")))
+	assert.True(m.Set("key2", []byte("val2")))
+	assert.True(m.Set("key3", []byte("val3")))
+
+	// len
+	assert.Equal(m.Len(), 3)
+
+	// get
+	val, ok := m.Get("key1")
+	assert.True(ok)
+	assert.Equal(string(val), "val1")
+
+	val, ok = m.Get("key2")
+	assert.True(ok)
+	assert.Equal(string(val), "val2")
+
+	val, ok = m.Get("key3")
+	assert.True(ok)
+	assert.Equal(string(val), "val3")
+
+	// set(update)
+	assert.False(m.Set("key1", []byte("newval1")))
+	assert.False(m.Set("key2", []byte("newval2")))
+	assert.False(m.Set("key3", []byte("newval3")))
+
+	// get(update)
+	val, ok = m.Get("key1")
+	assert.True(ok)
+	assert.Equal(string(val), "newval1")
+
+	val, ok = m.Get("key2")
+	assert.True(ok)
+	assert.Equal(string(val), "newval2")
+
+	val, ok = m.Get("key3")
+	assert.True(ok)
+	assert.Equal(string(val), "newval3")
+
+	// remove
+	assert.True(m.Remove("key1"))
+	assert.True(m.Remove("key2"))
+	assert.True(m.Remove("key3"))
+	assert.False(m.Remove("notexist"))
+
+	// len
+	assert.Equal(m.Len(), 0)
+}

internal/hash/set.go

@@ -1,39 +1,34 @@
 package hash
 
 import (
-	"github.com/cockroachdb/swiss"
-	"github.com/xgzlucario/rotom/internal/pkg"
+	mapset "github.com/deckarep/golang-set/v2"
 )
 
-var (
-	setAllocator = pkg.NewAllocator[string, struct{}]()
-)
+type SetI interface {
+	Add(key string) (ok bool)
+	Remove(key string) (ok bool)
+	Pop() (key string, ok bool)
+	Len() int
+}
+
+var _ SetI = (*Set)(nil)
 
 type Set struct {
-	m *swiss.Map[string, struct{}]
+	mapset.Set[string]
 }
 
 func NewSet() *Set {
-	return &Set{m: swiss.New(8, swiss.WithAllocator(setAllocator))}
+	return &Set{mapset.NewThreadUnsafeSet[string]()}
 }
 
-func (s *Set) Add(key string) bool {
-	if _, ok := s.m.Get(key); ok {
+func (s Set) Remove(key string) bool {
+	if !s.ContainsOne(key) {
 		return false
 	}
-	s.m.Put(key, struct{}{})
+	s.Set.Remove(key)
 	return true
 }
 
-func (s *Set) Pop() (item string, ok bool) {
-	s.m.All(func(key string, _ struct{}) bool {
-		s.m.Delete(key)
-		item, ok = key, true
-		return false
-	})
-	return
-}
-
-func (s *Set) Free() {
-	s.m.Close()
+func (s Set) Len() int {
+	return s.Cardinality()
 }

internal/hash/zipmap.go

@@ -0,0 +1,71 @@
+package hash
+
+import (
+	"unsafe"
+
+	"github.com/xgzlucario/rotom/internal/list"
+)
+
+var _ MapI = (*ZipMap)(nil)
+
+// ZipMap store datas as [key1, val1, key2, val2...] in listpack.
+type ZipMap struct {
+	m *list.ListPack
+}
+
+func NewZipMap() *ZipMap {
+	return &ZipMap{list.NewListPack()}
+}
+
+func (zm *ZipMap) Set(key string, val []byte) (newField bool) {
+	it := zm.m.Iterator().SeekLast()
+	for !it.IsFirst() {
+		it.Prev()
+		keyBytes := it.Prev()
+		if key == b2s(keyBytes) {
+			// update val
+			it.Next()
+			it.ReplaceNext(b2s(val))
+			return false
+		}
+	}
+	zm.m.RPush(key, b2s(val))
+	return true
+}
+
+func (zm *ZipMap) Get(key string) ([]byte, bool) {
+	it := zm.m.Iterator().SeekLast()
+	for !it.IsFirst() {
+		valBytes := it.Prev()
+		keyBytes := it.Prev()
+		if key == b2s(keyBytes) {
+			return valBytes, true
+		}
+	}
+	return nil, false
+}
+
+func (zm *ZipMap) Remove(key string) bool {
+	it := zm.m.Iterator().SeekLast()
+	for !it.IsFirst() {
+		it.Prev()
+		keyBytes := it.Prev()
+		if key == b2s(keyBytes) {
+			it.RemoveNext()
+			it.RemoveNext()
+			return true
+		}
+	}
+	return false
+}
+
+func (zm *ZipMap) Len() int {
+	return zm.m.Size() / 2
+}
+
+func (zm *ZipMap) Scan(fn func(string, []byte)) {
+}
+
+func b2s(b []byte) string {
+	return *(*string)(unsafe.Pointer(&b))
+}