multi exec and discardmulti/execute

1 files changed, 48 insertions, 5 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