Add RENAME, RENAMENX and RANDOMKEY keyspace commands

docs/commands.md

@@ -35,6 +35,9 @@ conventions; arguments in `<>` are required, `[]` optional. RAMen is single
 | `PERSIST <key>` | Remove the TTL. |
 | `KEYS <pattern>` | Glob-match keys (`*`, `?`, `[...]`). |
 | `TYPE <key>` | `string`/`hash`/`list`/`set`/`zset`/`vector`/`none`. |
+| `RENAME <key> <newkey>` | Rename a key, overwriting `newkey`; the TTL moves with it. Errors if the key is missing. |
+| `RENAMENX <key> <newkey>` | Rename only if `newkey` does not exist; `1` if renamed, `0` if it already exists. |
+| `RANDOMKEY` | Return a random key, or nil when the keyspace is empty. |
 
 ## Strings
 

internal/server/commands.go

@@ -41,6 +41,9 @@ func (s *Server) registerCommands() {
 		"PERSIST":    (*conn).cmdPersist,
 		"KEYS":       (*conn).cmdKeys,
 		"TYPE":       (*conn).cmdType,
+		"RENAME":     (*conn).cmdRename,
+		"RENAMENX":   (*conn).cmdRenameNX,
+		"RANDOMKEY":  (*conn).cmdRandomKey,
 
 		// strings
 		"GET":         (*conn).cmdGet,
@@ -347,6 +350,37 @@ func (c *conn) cmdType(args []string) error {
 	return c.writeSimple(c.s.store.Type(args[1]))
 }
 
+func (c *conn) cmdRename(args []string) error {
+	if len(args) != 3 {
+		return c.wrongArgs("rename")
+	}
+	if err := c.s.store.Rename(args[1], args[2]); err != nil {
+		return c.storeErr(err)
+	}
+	return c.writeSimple("OK")
+}
+
+func (c *conn) cmdRenameNX(args []string) error {
+	if len(args) != 3 {
+		return c.wrongArgs("renamenx")
+	}
+	renamed, err := c.s.store.RenameNX(args[1], args[2])
+	if err != nil {
+		return c.storeErr(err)
+	}
+	return c.writeInt(boolToInt(renamed))
+}
+
+func (c *conn) cmdRandomKey(args []string) error {
+	if len(args) != 1 {
+		return c.wrongArgs("randomkey")
+	}
+	if k, ok := c.s.store.RandomKey(); ok {
+		return c.writeBulk(k)
+	}
+	return c.writeNull()
+}
+
 func boolToInt(b bool) int64 {
 	if b {
 		return 1

internal/server/server_test.go

@@ -847,3 +847,105 @@ func TestMSetNX(t *testing.T) {
 	mustError(t, cli, "MSETNX", "k", "v", "k2") // dangling key
 	mustError(t, cli, "MSETNX")                 // nothing
 }
+
+func TestRename(t *testing.T) {
+	cli, cleanup := startTestServer(t)
+	defer cleanup()
+
+	// a missing source is an error, even renaming it onto itself
+	mustError(t, cli, "RENAME", "nope", "dst")
+	mustError(t, cli, "RENAME", "ghost", "ghost")
+
+	// basic rename returns OK and moves the value
+	mustDo(t, cli, "SET", "a", "1")
+	if r := mustDo(t, cli, "RENAME", "a", "b"); r != "OK" {
+		t.Fatalf("RENAME = %v", r)
+	}
+	if r := mustDo(t, cli, "GET", "b"); r != "1" {
+		t.Fatalf("GET b after RENAME = %v", r)
+	}
+	if r := mustDo(t, cli, "EXISTS", "a"); r != int64(0) {
+		t.Fatalf("source still exists after RENAME = %v", r)
+	}
+
+	// the TTL travels with the key
+	mustDo(t, cli, "SET", "t", "v")
+	mustDo(t, cli, "EXPIRE", "t", "100")
+	mustDo(t, cli, "RENAME", "t", "t2")
+	if r := mustDo(t, cli, "TTL", "t2").(int64); r <= 0 {
+		t.Fatalf("TTL after RENAME = %v, want a positive TTL", r)
+	}
+
+	// renaming over an existing key overwrites it
+	mustDo(t, cli, "SET", "s1", "new")
+	mustDo(t, cli, "SET", "d1", "old")
+	mustDo(t, cli, "RENAME", "s1", "d1")
+	if r := mustDo(t, cli, "GET", "d1"); r != "new" {
+		t.Fatalf("RENAME overwrite = %v", r)
+	}
+
+	// rename to itself succeeds
+	mustDo(t, cli, "SET", "self", "v")
+	if r := mustDo(t, cli, "RENAME", "self", "self"); r != "OK" {
+		t.Fatalf("RENAME self = %v", r)
+	}
+
+	mustError(t, cli, "RENAME", "a")           // arity: missing newkey
+	mustError(t, cli, "RENAME")                // arity: nothing
+	mustError(t, cli, "RENAME", "a", "b", "c") // arity: too many
+}
+
+func TestRenameNX(t *testing.T) {
+	cli, cleanup := startTestServer(t)
+	defer cleanup()
+
+	mustError(t, cli, "RENAMENX", "nope", "dst")    // missing source
+	mustError(t, cli, "RENAMENX", "ghost", "ghost") // missing source, even self->self
+
+	// destination free: renames and reports 1
+	mustDo(t, cli, "SET", "a", "1")
+	if r := mustDo(t, cli, "RENAMENX", "a", "b"); r != int64(1) {
+		t.Fatalf("RENAMENX to free dst = %v, want 1", r)
+	}
+	if r := mustDo(t, cli, "GET", "b"); r != "1" {
+		t.Fatalf("GET b after RENAMENX = %v", r)
+	}
+
+	// destination exists: reports 0 and nothing moves
+	mustDo(t, cli, "SET", "x", "xv")
+	mustDo(t, cli, "SET", "y", "yv")
+	if r := mustDo(t, cli, "RENAMENX", "x", "y"); r != int64(0) {
+		t.Fatalf("RENAMENX to existing dst = %v, want 0", r)
+	}
+	if r := mustDo(t, cli, "GET", "x"); r != "xv" {
+		t.Fatalf("src changed after a failed RENAMENX = %v", r)
+	}
+
+	mustError(t, cli, "RENAMENX", "x")           // arity
+	mustError(t, cli, "RENAMENX", "x", "y", "z") // arity
+}
+
+func TestRandomKey(t *testing.T) {
+	cli, cleanup := startTestServer(t)
+	defer cleanup()
+
+	// an empty keyspace returns nil
+	if r, err := cli.Do("RANDOMKEY"); err != nil || r != nil {
+		t.Fatalf("RANDOMKEY on empty keyspace = %v (err %v), want nil", r, err)
+	}
+
+	// with keys present it returns one of them
+	keys := map[string]bool{"a": true, "b": true, "c": true}
+	for k := range keys {
+		mustDo(t, cli, "SET", k, "v")
+	}
+	for i := 0; i < 50; i++ {
+		r := mustDo(t, cli, "RANDOMKEY")
+		k, ok := r.(string)
+		if !ok || !keys[k] {
+			t.Fatalf("RANDOMKEY = %v, want one of a/b/c", r)
+		}
+	}
+
+	mustError(t, cli, "RANDOMKEY", "extra") // arity: takes no arguments
+}

internal/store/snapshot.go

@@ -30,12 +30,29 @@ type VecRecord struct {
 }
 
 // Export returns a snapshot of every live key. It is safe to call while the
-// store is serving traffic; each shard is read-locked in turn.
+// store is serving traffic. Every shard is read-locked for the whole scan so
+// the snapshot is a globally consistent view: RENAME moves a key from one shard
+// to another, so releasing each shard before reading the next could otherwise
+// record the key twice or miss it entirely. Locks are taken in shard-index
+// order, matching every other multi-shard operation, so this cannot deadlock.
+//
+// Values that writers mutate in place (hash maps, list slices) are copied while
+// the locks are held, so the returned records share no memory with live data
+// and can be serialised after Export returns without racing concurrent writes.
+// The other types are already detached: sets and zsets are built fresh here,
+// and vector items are replaced wholesale on write, never mutated in place.
 func (s *Store) Export() []Record {
 	now := s.now()
-	var out []Record
 	for _, sh := range s.shards {
 		sh.mu.RLock()
+	}
+	defer func() {
+		for _, sh := range s.shards {
+			sh.mu.RUnlock()
+		}
+	}()
+	var out []Record
+	for _, sh := range s.shards {
 		for k, e := range sh.m {
 			if e.expired(now) {
 				continue
@@ -48,9 +65,15 @@ func (s *Store) Export() []Record {
 			case string:
 				rec.Type, rec.Str = "string", v
 			case map[string]string:
-				rec.Type, rec.Hash = "hash", v
+				rec.Type = "hash"
+				h := make(map[string]string, len(v))
+				for f, fv := range v {
+					h[f] = fv
+				}
+				rec.Hash = h
 			case *list:
-				rec.Type, rec.List = "list", v.items
+				rec.Type = "list"
+				rec.List = append([]string(nil), v.items...)
 			case map[string]struct{}:
 				rec.Type = "set"
 				for m := range v {
@@ -68,7 +91,6 @@ func (s *Store) Export() []Record {
 			}
 			out = append(out, rec)
 		}
-		sh.mu.RUnlock()
 	}
 	return out
 }

internal/store/snapshot_test.go

@@ -0,0 +1,114 @@
+package store
+
+import (
+	"encoding/gob"
+	"io"
+	"strconv"
+	"sync"
+	"testing"
+)
+
+// twoKeysOnDifferentShards returns two key names that hash to different shards,
+// so a RENAME between them actually moves the entry across shards.
+func twoKeysOnDifferentShards(s *Store) (string, string) {
+	base := "k0"
+	bi := s.shardIndex(base)
+	for i := 1; ; i++ {
+		other := "k" + strconv.Itoa(i)
+		if s.shardIndex(other) != bi {
+			return base, other
+		}
+	}
+}
+
+// TestExportConsistentUnderRename guards against the snapshot losing or
+// duplicating a key that RENAME moves between shards. Only one of {a, b} exists
+// at any instant (they are ping-ponged), so every Export must observe exactly
+// one of them. Before Export locked all shards for the whole scan, a snapshot
+// overlapping a cross-shard rename could see the key under both names or under
+// neither. Run with -race to also catch torn map access.
+func TestExportConsistentUnderRename(t *testing.T) {
+	s := New()
+	a, b := twoKeysOnDifferentShards(s)
+	s.Set(a, "v", SetOptions{})
+
+	stop := make(chan struct{})
+	var pinger sync.WaitGroup
+	// Deferred so an early t.Fatalf still stops the pinger; otherwise it could
+	// keep calling t.Errorf after the test returns and mask the real failure.
+	defer pinger.Wait()
+	defer close(stop)
+
+	// Ping-pong the single key between the two shards until stopped.
+	pinger.Add(1)
+	go func() {
+		defer pinger.Done()
+		src, dst := a, b
+		for {
+			select {
+			case <-stop:
+				return
+			default:
+			}
+			if err := s.Rename(src, dst); err != nil {
+				t.Errorf("Rename(%s,%s): %v", src, dst, err)
+				return
+			}
+			src, dst = dst, src
+		}
+	}()
+
+	// Snapshot repeatedly; each snapshot must contain exactly one of a/b.
+	for i := 0; i < 200000; i++ {
+		n := 0
+		for _, rec := range s.Export() {
+			if rec.Key == a || rec.Key == b {
+				n++
+			}
+		}
+		if n != 1 {
+			t.Fatalf("snapshot saw the renamed key %d times, want exactly 1", n)
+		}
+	}
+}
+
+// TestExportCopiesMutableValues guards against snapshot records aliasing live
+// hash maps and list slices. Serialisation runs after Export has released
+// every lock, so an aliased map iterated by the encoder while a writer mutates
+// it kills the whole process with "concurrent map iteration and map write".
+// Encode snapshots while hammering the exported keys; run with -race to also
+// catch aliasing that does not crash outright.
+func TestExportCopiesMutableValues(t *testing.T) {
+	s := New()
+	s.HSet("h", map[string]string{"f0": "v"})
+	s.push("l", true, []string{"x"})
+
+	stop := make(chan struct{})
+	var writer sync.WaitGroup
+	defer writer.Wait()
+	defer close(stop)
+
+	writer.Add(1)
+	go func() {
+		defer writer.Done()
+		for i := 0; ; i++ {
+			select {
+			case <-stop:
+				return
+			default:
+			}
+			s.HSet("h", map[string]string{"f" + strconv.Itoa(i%64): "v"})
+			if err := s.LSet("l", 0, strconv.Itoa(i)); err != nil {
+				t.Errorf("LSet: %v", err)
+				return
+			}
+		}
+	}()
+
+	enc := gob.NewEncoder(io.Discard)
+	for i := 0; i < 2000; i++ {
+		if err := enc.Encode(s.Export()); err != nil {
+			t.Fatalf("encode snapshot: %v", err)
+		}
+	}
+}