multi exec and discardmulti/execute

9 files changed, 205 insertions, 35 deletions

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