From 6cb3313793bf45e197b3048f0abd2f600bf97bb5 Mon Sep 17 00:00:00 2001 From: alex Date: Sun, 12 Jul 2026 15:00:16 +0200 Subject: multi exec and discard --- pkg/commands/command-registry.go | 53 ++++++++++++++++++++++++++++--- pkg/commands/dbsize.go | 12 +++---- pkg/commands/discard.go | 18 +++++++++++ pkg/commands/exec.go | 68 ++++++++++++++++++++++++++++++++++++++++ pkg/commands/flushall.go | 10 +++--- pkg/commands/keys.go | 32 +++++++++---------- pkg/commands/multi.go | 18 +++++++++++ pkg/commands/session.go | 18 +++++++++++ pkg/core/shards.go | 11 +++++++ 9 files changed, 205 insertions(+), 35 deletions(-) create mode 100644 pkg/commands/discard.go create mode 100644 pkg/commands/exec.go create mode 100644 pkg/commands/multi.go create mode 100644 pkg/commands/session.go (limited to 'pkg') diff --git a/pkg/commands/command-registry.go b/pkg/commands/command-registry.go index 19c8acc..961f5ff 100644 --- a/pkg/commands/command-registry.go +++ b/pkg/commands/command-registry.go @@ -24,12 +24,12 @@ func AllSubsequentKeys(args []string) []string { func NoKeys(args []string) []string { return nil } - func RenameKeys(args []string) []string { return []string{args[1], args[2]} } type BaseCommand struct{ + LocksAll bool MinArgs int ExtractKeys func(args []string)[]string Execute func(args []string, getShard func(k string) *core.Shard) interface{} @@ -41,10 +41,21 @@ func DispatchBaseCommand(db core.RedisDB ,cmd BaseCommand, args []string) interf return errors.New("wrong number of arguments") } - keys := cmd.ExtractKeys(args) - - db.Lock(keys) - defer db.Unlock(keys) + if cmd.LocksAll { + for _, shard := range db.Shards { + shard.Mu.Lock() + } + defer func() { + for i := len(db.Shards) - 1; i >= 0; i-- { + db.Shards[i].Mu.Unlock() + } + }() + } else { + keys := cmd.ExtractKeys(args) + db.Lock(keys) + defer db.Unlock(keys) + } + return cmd.Execute(args, db.GetShard) } @@ -63,6 +74,29 @@ func DispatchPubSub(conn net.Conn, h *pubsub.Hub, cmd PubSubCommand, args []stri } } +type SessionCommand struct { + MinArgs int + Execute func(db core.RedisDB,session *ClientSession, conn net.Conn, args []string) interface{} +} + +func DispatchSession(db core.RedisDB,conn net.Conn, s *ClientSession, cmd SessionCommand, args []string) { + if len(args) < cmd.MinArgs { + conn.Write(core.SerializeRESP(errors.New("wrong number of arguments"))) + return + } + if res := cmd.Execute(db,s, conn, args); res != nil { + conn.Write(core.SerializeRESP(res)) + } +} + +func QueueBaseCommand(conn net.Conn, s *ClientSession, args []string, minArgs int) { + if len(args) < minArgs { + conn.Write(core.SerializeRESP(errors.New("wrong number of arguments"))) + return + } + s.Queue = append(s.Queue, args) + conn.Write(core.SerializeRESP(core.SimpleString("QUEUED"))) +} var BaseRegistry = map[string]BaseCommand{ "GET": { @@ -126,6 +160,7 @@ var BaseRegistry = map[string]BaseCommand{ Execute: Lrange, }, "FLUSHALL": { + LocksAll: true, MinArgs: 1, ExtractKeys: NoKeys, Execute: Flushall, @@ -166,11 +201,13 @@ var BaseRegistry = map[string]BaseCommand{ Execute: Type, }, "DBSIZE": { + LocksAll: true, MinArgs: 1, ExtractKeys: NoKeys, Execute: Dbsize, }, "KEYS": { + LocksAll: true, MinArgs: 2, ExtractKeys: NoKeys, Execute: Keys, @@ -191,4 +228,10 @@ var PubSubRegistry = map[string]PubSubCommand{ "SUBSCRIBE": {MinArgs: 2, Execute: Subscribe}, "UNSUBSCRIBE": {MinArgs: 1, Execute: Unsubscribe}, "PUBLISH": {MinArgs: 3, Execute: Publish}, +} + +var SessionRegistry = map[string]SessionCommand{ + "MULTI": {MinArgs: 1, Execute: Multi}, + "EXEC": {MinArgs: 1, Execute: Exec}, + "DISCARD": {MinArgs: 1, Execute: Discard}, } \ No newline at end of file diff --git a/pkg/commands/dbsize.go b/pkg/commands/dbsize.go index dbe39a4..b7086a5 100644 --- a/pkg/commands/dbsize.go +++ b/pkg/commands/dbsize.go @@ -5,11 +5,9 @@ import ( ) func Dbsize(args []string, getShard func(k string) *core.Shard) interface{} { - return dbInstance.ExecuteReadAll(func(shards []*core.Shard) interface{} { - count := 0 - for _, shard := range shards { - count += len(shard.Data) - } - return count - }) + count := 0 + for _, shard := range dbInstance.Shards { + count += len(shard.Data) + } + return count } diff --git a/pkg/commands/discard.go b/pkg/commands/discard.go new file mode 100644 index 0000000..e96d2c0 --- /dev/null +++ b/pkg/commands/discard.go @@ -0,0 +1,18 @@ +package commands + +import ( + "errors" + "net" + "redisClone/pkg/core" +) + +func Discard(db core.RedisDB, session *ClientSession, conn net.Conn, execArgs []string) interface{} { + if !session.InTransaction { + return errors.New("DISCARD without MULTI") + } + + session.InTransaction = false + session.Queue = nil + + return core.SimpleString("OK") +} diff --git a/pkg/commands/exec.go b/pkg/commands/exec.go new file mode 100644 index 0000000..a260f69 --- /dev/null +++ b/pkg/commands/exec.go @@ -0,0 +1,68 @@ +package commands + +import ( + "errors" + "net" + "redisClone/pkg/core" + "sort" +) + +func Exec(db core.RedisDB, session *ClientSession, conn net.Conn, execArgs []string) interface{} { + if session.Queue == nil { + return errors.New("EXEC without MULTI") + } + + shardMap := make(map[int]*core.Shard) + lockedAll := false + var sortedIDs []int + + for _, cmdArgs := range session.Queue { + cmdName := cmdArgs[0] + if cmd, ok := BaseRegistry[cmdName]; ok { + if cmd.LocksAll { + db.LockAll() + lockedAll = true + break + } + for _, key := range cmd.ExtractKeys(cmdArgs) { + shard := db.GetShard(key) + shardMap[shard.Id] = shard + } + } + } + + if !lockedAll { + for id := range shardMap { + sortedIDs = append(sortedIDs, id) + } + sort.Ints(sortedIDs) + for _, id := range sortedIDs { + shardMap[id].Mu.Lock() + } + + defer func() { + for i := len(sortedIDs) - 1; i >= 0; i-- { + shardMap[sortedIDs[i]].Mu.Unlock() + } + }() + } else { + defer db.UnlockAll() + } + + var results []interface{} + + for _, cmdArgs := range session.Queue { + cmdName := cmdArgs[0] + if cmd, ok := BaseRegistry[cmdName]; ok { + // Capture the result of the command + res := cmd.Execute(cmdArgs, db.GetShard) + results = append(results, res) + } else { + results = append(results, "ERR unknown command") + } + } + + session.Queue = nil + + return results +} \ No newline at end of file diff --git a/pkg/commands/flushall.go b/pkg/commands/flushall.go index 9f1316a..e5625c4 100644 --- a/pkg/commands/flushall.go +++ b/pkg/commands/flushall.go @@ -5,10 +5,8 @@ import ( ) func Flushall(args []string, getShard func(k string) *core.Shard) interface{} { - return dbInstance.ExecuteAll(func(shards []*core.Shard) interface{} { - for _, shard := range shards { - shard.Data = make(map[string]core.Item) - } - return core.SimpleString("OK") - }) + for _, shard := range dbInstance.Shards { + shard.Data = make(map[string]core.Item) + } + return core.SimpleString("OK") } diff --git a/pkg/commands/keys.go b/pkg/commands/keys.go index b233964..40d9277 100644 --- a/pkg/commands/keys.go +++ b/pkg/commands/keys.go @@ -19,23 +19,21 @@ func Keys(args []string, getShard func(k string) *core.Shard) interface{} { limit = val } - return dbInstance.ExecuteReadAll(func(shards []*core.Shard) interface{} { - var matches []string - for _, shard := range shards { - for key := range shard.Data { - matched, err := filepath.Match(pattern, key) - if err != nil { - return errors.New("illegal glob pattern") - } - if matched { - matches = append(matches, key) - } + var matches []string + for _, shard := range dbInstance.Shards { + for key := range shard.Data { + matched, err := filepath.Match(pattern, key) + if err != nil { + return errors.New("illegal glob pattern") + } + if matched { + matches = append(matches, key) } } - sort.Strings(matches) - if limit != -1 && len(matches) > limit { - matches = matches[:limit] - } - return matches - }) + } + sort.Strings(matches) + if limit != -1 && len(matches) > limit { + matches = matches[:limit] + } + return matches } diff --git a/pkg/commands/multi.go b/pkg/commands/multi.go new file mode 100644 index 0000000..2b0fcb7 --- /dev/null +++ b/pkg/commands/multi.go @@ -0,0 +1,18 @@ +package commands + +import ( + "errors" + "net" + "redisClone/pkg/core" +) + +func Multi(db core.RedisDB, session *ClientSession, conn net.Conn, execArgs []string) interface{} { + if session.InTransaction { + return errors.New("MULTI calls can not be nested") + } + session.InTransaction = true + + session.Queue = make([][]string, 0) + + return core.SimpleString("OK") +} \ No newline at end of file diff --git a/pkg/commands/session.go b/pkg/commands/session.go new file mode 100644 index 0000000..e7a39fe --- /dev/null +++ b/pkg/commands/session.go @@ -0,0 +1,18 @@ +package commands + +type ClientSession struct{ + Queue [][]string + InTransaction bool +} + +func NewClientSession() *ClientSession { + return &ClientSession{ + InTransaction: false, + Queue: make([][]string, 0), + } +} + +func (s *ClientSession) Reset() { + s.InTransaction = false + s.Queue = nil +} \ No newline at end of file diff --git a/pkg/core/shards.go b/pkg/core/shards.go index 5ede61c..1732be5 100644 --- a/pkg/core/shards.go +++ b/pkg/core/shards.go @@ -131,4 +131,15 @@ func ( db *RedisDB) Unlock(keys []string){ for i := len(sortedIDs)-1; i >= 0; i --{ shardMap[sortedIDs[i]].Mu.Unlock() } +} + +func (db *RedisDB)LockAll(){ + for _, shard := range db.Shards { + shard.Mu.Lock() + } +} +func (db *RedisDB)UnlockAll(){ + for i := len(db.Shards)-1;i >= 0; i --{ + db.Shards[i].Mu.Unlock() + } } \ No newline at end of file -- cgit v1.2.3