From 7587332bc68bb33a82800393c0834a03275869a9 Mon Sep 17 00:00:00 2001 From: xanish Date: Thu, 14 Nov 2024 23:46:21 +0530 Subject: [PATCH 01/16] rebase from upstream master --- docs/src/content/docs/commands/RANDOMKEY.md | 59 ++++++++++++++++++++ internal/eval/commands.go | 9 +++ internal/eval/eval.go | 27 +++++++++ internal/eval/eval_test.go | 61 +++++++++++++++++++++ 4 files changed, 156 insertions(+) create mode 100644 docs/src/content/docs/commands/RANDOMKEY.md diff --git a/docs/src/content/docs/commands/RANDOMKEY.md b/docs/src/content/docs/commands/RANDOMKEY.md new file mode 100644 index 000000000..b3bf01e45 --- /dev/null +++ b/docs/src/content/docs/commands/RANDOMKEY.md @@ -0,0 +1,59 @@ +--- +title: RANDOMKEY +description: The `RANDOMKEY` command in DiceDB return a random key from the currently selected database. +--- + +The `RANDOMKEY` command in DiceDB is used to return a random key from the currently selected database. + +## Syntax + +``` +RANDOMKEY +``` + +## Parameters + +The `RANDOMKEY` command does not take any parameters. + +## Return values + +| Condition | Return Value | +|-----------------------------------------------|-----------------------------------------------------| +| Command is successful | A random key from the keyspace of selected database | +| Failure to scan keyspace or pick a random key | Error | + +## Behaviour + +- When executed, `RANDOMKEY` fetches the keyspace from currently selected database and picks a random key from it. +- The operation is slow and may return an expired key if it hasn't been evicted. +- The command does not modify the database in any way; it is purely informational. + +## Errors +The `RANDOMKEY` command is straightforward and does not typically result in errors under normal usage. However, since it internally depends on KEYS command, it can fail for the same cases as KEYS. + +## Example Usage + +### Basic Usage + +Getting a random key from the currently selected database: + +```shell +127.0.0.1:7379> RANDOMKEY +"key_6" +``` + +### Using with Multiple Databases + +If you are working with multiple databases, you can switch between them using the `SELECT` command and then use `RANDOMKEY` to get a random key from selected database: + +```shell +127.0.0.1:7379> SELECT 0 +OK +127.0.0.1:7379> RANDOMKEY +"db0_key_54" + +127.0.0.1:7379> SELECT 1 +OK +127.0.0.1:7379> RANDOMKEY +"db1_key_435" +``` diff --git a/internal/eval/commands.go b/internal/eval/commands.go index ae901640f..549dcdc09 100644 --- a/internal/eval/commands.go +++ b/internal/eval/commands.go @@ -96,6 +96,14 @@ var ( NewEval: evalGET, } + randomKeyCmdMeta = DiceCmdMeta{ + Name: "RANDOMKEY", + Info: `RANDOMKEY returns a random key from the currently selected database.`, + Arity: 1, + IsMigrated: true, + NewEval: evalRANDOMKEY, + } + getSetCmdMeta = DiceCmdMeta{ Name: "GETSET", Info: `GETSET returns the previous string value of a key after setting it to a new value.`, @@ -1501,6 +1509,7 @@ func init() { DiceCmds["PTTL"] = pttlCmdMeta DiceCmds["Q.UNWATCH"] = qUnwatchCmdMeta DiceCmds["Q.WATCH"] = qwatchCmdMeta + DiceCmds["RANDOMKEY"] = randomKeyCmdMeta DiceCmds["RENAME"] = renameCmdMeta DiceCmds["RESTORE"] = restorekeyCmdMeta DiceCmds["RPOP"] = rpopCmdMeta diff --git a/internal/eval/eval.go b/internal/eval/eval.go index dea3856a7..edfbf274e 100644 --- a/internal/eval/eval.go +++ b/internal/eval/eval.go @@ -2,9 +2,11 @@ package eval import ( "bytes" + "crypto/rand" "errors" "fmt" "log/slog" + "math/big" "sort" "strconv" "strings" @@ -1360,3 +1362,28 @@ func executeBitfieldOps(value *ByteArray, ops []utils.BitFieldOp) []interface{} } return result } + +// evalRANDOMKEY returns a random key from the currently selected database. +func evalRANDOMKEY(args []string, store *dstore.Store) *EvalResponse { + if len(args) > 0 { + return &EvalResponse{Result: nil, Error: errors.New(string(diceerrors.NewErrArity("RANDOMKEY")))} + } + + availKeys, err := store.Keys("*") + if err != nil { + return &EvalResponse{Result: nil, + Error: errors.New(string(diceerrors.NewErrWithMessage("could not fetch keys to extract a random key")))} + } + + if len(availKeys) > 0 { + randKeyIdx, err := rand.Int(rand.Reader, big.NewInt(int64(len(availKeys)))) + if err != nil { + return &EvalResponse{Result: nil, + Error: errors.New(string(diceerrors.NewErrWithMessage("could not generate a random key seed")))} + } + + return &EvalResponse{Result: availKeys[randKeyIdx.Uint64()], Error: nil} + } + + return &EvalResponse{Result: clientio.RespNIL, Error: nil} +} diff --git a/internal/eval/eval_test.go b/internal/eval/eval_test.go index bf5dfa76c..39a6f9cdb 100644 --- a/internal/eval/eval_test.go +++ b/internal/eval/eval_test.go @@ -144,6 +144,7 @@ func TestEval(t *testing.T) { testEvalBFINFO(t, store) testEvalBFEXISTS(t, store) testEvalBFADD(t, store) + testEvalRANDOMKEY(t, store) } func testEvalPING(t *testing.T, store *dstore.Store) { @@ -1597,6 +1598,66 @@ func BenchmarkEvalJSONOBJLEN(b *testing.B) { } } +func testEvalRANDOMKEY(t *testing.T, store *dstore.Store) { + t.Run("invalid no of args", func(t *testing.T) { + response := evalRANDOMKEY([]string{"INVALID_ARG"}, store) + expectedErr := errors.New("-ERR wrong number of arguments for 'randomkey' command\r\n") + + assert.Equal(t, nil, response.Result) + assert.EqualError(t, response.Error, expectedErr.Error()) + }) + + t.Run("some keys present in db", func(t *testing.T) { + data := map[string]string{ + "EXISTING_KEY": "MOCK_VALUE", + "EXISTING_KEY_2": "MOCK_VALUE_2", + "EXISTING_KEY_3": "MOCK_VALUE_3", + } + + for key, value := range data { + obj := &object.Obj{ + Value: value, + LastAccessedAt: uint32(time.Now().Unix()), + } + store.Put(key, obj) + } + + results := make(map[string]int) + for i := 0; i < 10000; i++ { + result := evalRANDOMKEY([]string{}, store) + results[result.Result.(string)]++ + } + + for key, _ := range data { + if results[key] == 0 { + t.Errorf("key %s was never returned", key) + } + } + }) +} + +func BenchmarkEvalRANDOMKEY(b *testing.B) { + storeSize := 1000000 + store := dstore.NewStore(nil, nil, nil) + + b.Run(fmt.Sprintf("benchmark_randomkey_with_%d_keys", storeSize), func(b *testing.B) { + for i := 0; i < storeSize; i++ { + obj := &object.Obj{ + Value: i, + } + store.Put(fmt.Sprintf("key%d", i), obj) + } + + b.ResetTimer() + b.ReportAllocs() + + // Benchmark the evalRANDOMKEY function + for i := 0; i < b.N; i++ { + _ = evalRANDOMKEY([]string{}, store) + } + }) +} + func testEvalJSONDEL(t *testing.T, store *dstore.Store) { tests := map[string]evalTestCase{ "JSON.DEL : nil value": { From 4015d4e2c725dd738656caa77b307ff8fe53bfaa Mon Sep 17 00:00:00 2001 From: xanish Date: Thu, 17 Oct 2024 01:34:08 +0530 Subject: [PATCH 02/16] Add best effort mechanism to return non-expired randomkey in 128 tries --- internal/eval/eval.go | 24 ++++++++++++++++++------ internal/eval/eval_test.go | 7 +++++-- 2 files changed, 23 insertions(+), 8 deletions(-) diff --git a/internal/eval/eval.go b/internal/eval/eval.go index edfbf274e..fb975d897 100644 --- a/internal/eval/eval.go +++ b/internal/eval/eval.go @@ -1376,13 +1376,25 @@ func evalRANDOMKEY(args []string, store *dstore.Store) *EvalResponse { } if len(availKeys) > 0 { - randKeyIdx, err := rand.Int(rand.Reader, big.NewInt(int64(len(availKeys)))) - if err != nil { - return &EvalResponse{Result: nil, - Error: errors.New(string(diceerrors.NewErrWithMessage("could not generate a random key seed")))} - } + maxIters := 128 + for range maxIters { + randKeyIdx, err := rand.Int(rand.Reader, big.NewInt(int64(len(availKeys)))) + if err != nil { + continue + } - return &EvalResponse{Result: availKeys[randKeyIdx.Uint64()], Error: nil} + randKey := availKeys[randKeyIdx.Uint64()] + keyObj := store.Get(randKey) + if keyObj == nil { + continue + } + + currTimeMs := uint64(utils.GetCurrentTime().UnixMilli()) + expireTimeMs, isExpirySet := dstore.GetExpiry(keyObj, store) + if (isExpirySet && expireTimeMs > currTimeMs) || !isExpirySet { + return &EvalResponse{Result: clientio.Encode(randKey, false), Error: nil} + } + } } return &EvalResponse{Result: clientio.RespNIL, Error: nil} diff --git a/internal/eval/eval_test.go b/internal/eval/eval_test.go index 39a6f9cdb..119523646 100644 --- a/internal/eval/eval_test.go +++ b/internal/eval/eval_test.go @@ -1625,11 +1625,14 @@ func testEvalRANDOMKEY(t *testing.T, store *dstore.Store) { results := make(map[string]int) for i := 0; i < 10000; i++ { result := evalRANDOMKEY([]string{}, store) - results[result.Result.(string)]++ + if res, ok := result.Result.([]byte); ok { + results[string(res)]++ + } } for key, _ := range data { - if results[key] == 0 { + returnedKey := clientio.Encode(key, false) + if results[string(returnedKey)] == 0 { t.Errorf("key %s was never returned", key) } } From 5ffff02cb6b2848033a6e55dff38902feed53ef3 Mon Sep 17 00:00:00 2001 From: xanish Date: Sun, 20 Oct 2024 12:35:20 +0530 Subject: [PATCH 03/16] use internal error constants --- internal/errors/migrated_errors.go | 1 + internal/eval/eval.go | 3 +-- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/internal/errors/migrated_errors.go b/internal/errors/migrated_errors.go index fd1b000bd..7b3060919 100644 --- a/internal/errors/migrated_errors.go +++ b/internal/errors/migrated_errors.go @@ -20,6 +20,7 @@ var ( ErrOverflow = errors.New("ERR increment or decrement would overflow") // Signifies that an increment or decrement operation would exceed the limits. ErrSyntax = errors.New("ERR syntax error") // Represents a syntax error in a DiceDB command. ErrKeyNotFound = errors.New("ERR no such key") // Indicates that the specified key does not exist. + ErrUnableToFetchKeys = errors.New("ERR unable to fetch keys from store") // Indicates that the specified key does not exist. ErrWrongTypeOperation = errors.New("WRONGTYPE Operation against a key holding the wrong kind of value") // Signals an operation attempted on a key with an incompatible type. ErrInvalidHyperLogLogKey = errors.New("WRONGTYPE Key is not a valid HyperLogLog string value") // Indicates that a key is not a valid HyperLogLog value. ErrCorruptedHyperLogLogObject = errors.New("INVALIDOBJ Corrupted HLL object detected") // Signals detection of a corrupted HyperLogLog object. diff --git a/internal/eval/eval.go b/internal/eval/eval.go index fb975d897..4e4055317 100644 --- a/internal/eval/eval.go +++ b/internal/eval/eval.go @@ -1371,8 +1371,7 @@ func evalRANDOMKEY(args []string, store *dstore.Store) *EvalResponse { availKeys, err := store.Keys("*") if err != nil { - return &EvalResponse{Result: nil, - Error: errors.New(string(diceerrors.NewErrWithMessage("could not fetch keys to extract a random key")))} + return &EvalResponse{Result: nil, Error: diceerrors.ErrUnableToFetchKeys} } if len(availKeys) > 0 { From 37087d9bfe226e8d1fbd26dd66739917efbf8c36 Mon Sep 17 00:00:00 2001 From: xanish Date: Sun, 20 Oct 2024 13:20:26 +0530 Subject: [PATCH 04/16] minimal integration tests for randomkey over http, resp and websocket --- .../commands/http/randomkey_test.go | 43 +++++++++++++++++++ .../commands/resp/randomkey_test.go | 43 +++++++++++++++++++ .../commands/websocket/randomkey_test.go | 43 +++++++++++++++++++ 3 files changed, 129 insertions(+) create mode 100644 integration_tests/commands/http/randomkey_test.go create mode 100644 integration_tests/commands/resp/randomkey_test.go create mode 100644 integration_tests/commands/websocket/randomkey_test.go diff --git a/integration_tests/commands/http/randomkey_test.go b/integration_tests/commands/http/randomkey_test.go new file mode 100644 index 000000000..6f55a6758 --- /dev/null +++ b/integration_tests/commands/http/randomkey_test.go @@ -0,0 +1,43 @@ +package http + +import ( + "testing" + "time" + + "gotest.tools/v3/assert" +) + +func TestRandomKey(t *testing.T) { + exec := NewHTTPCommandExecutor() + + testCases := []struct { + name string + commands []HTTPCommand + expected []interface{} + delays []time.Duration + }{ + + { + name: "Random Key", + commands: []HTTPCommand{ + {Command: "FLUSHDB"}, + {Command: "SET", Body: map[string]interface{}{"key": "k1", "value": "v1"}}, + {Command: "RANDOMKEY"}, + }, + expected: []interface{}{"OK", "OK", "k1"}, + delays: []time.Duration{0, 0, 0}, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + for i, cmd := range tc.commands { + if tc.delays[i] > 0 { + time.Sleep(tc.delays[i]) + } + result, _ := exec.FireCommand(cmd) + assert.Equal(t, tc.expected[i], result, "Value mismatch for cmd %s", cmd) + } + }) + } +} diff --git a/integration_tests/commands/resp/randomkey_test.go b/integration_tests/commands/resp/randomkey_test.go new file mode 100644 index 000000000..8c1a94c23 --- /dev/null +++ b/integration_tests/commands/resp/randomkey_test.go @@ -0,0 +1,43 @@ +package resp + +import ( + "testing" + "time" + + "gotest.tools/v3/assert" +) + +func TestRandomKey(t *testing.T) { + conn := getLocalConnection() + defer conn.Close() + + testCases := []struct { + name string + cmds []string + expect []interface{} + delays []time.Duration + }{ + { + name: "Random Key", + cmds: []string{ + "FLUSHDB", + "SET k1 v1", + "RANDOMKEY", + }, + expect: []interface{}{"OK", "OK", "k1"}, + delays: []time.Duration{0, 0, 0}, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + for i, cmd := range tc.cmds { + if tc.delays[i] > 0 { + time.Sleep(tc.delays[i]) + } + result := FireCommand(conn, cmd) + assert.Equal(t, tc.expect[i], result, "Value mismatch for cmd %s", cmd) + } + }) + } +} diff --git a/integration_tests/commands/websocket/randomkey_test.go b/integration_tests/commands/websocket/randomkey_test.go new file mode 100644 index 000000000..ce619bb44 --- /dev/null +++ b/integration_tests/commands/websocket/randomkey_test.go @@ -0,0 +1,43 @@ +package websocket + +import ( + "testing" + "time" + + "gotest.tools/v3/assert" +) + +func TestRandomKey(t *testing.T) { + exec := NewWebsocketCommandExecutor() + conn := exec.ConnectToServer() + + testCases := []struct { + name string + cmds []string + expect []interface{} + delays []time.Duration + }{ + { + name: "Random Key", + cmds: []string{ + "FLUSHDB", + "SET k1 v1", + "RANDOMKEY", + }, + expect: []interface{}{"OK", "OK", "k1"}, + delays: []time.Duration{0, 0, 0}, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + for i, cmd := range tc.cmds { + if tc.delays[i] > 0 { + time.Sleep(tc.delays[i]) + } + result := exec.FireCommand(conn, cmd) + assert.Equal(t, tc.expect[i], result, "Value mismatch for cmd %s", cmd) + } + }) + } +} From f9aaada838ecbe90ca8fe8f4b4d12fd02d2f2656 Mon Sep 17 00:00:00 2001 From: xanish Date: Wed, 23 Oct 2024 23:16:38 +0530 Subject: [PATCH 05/16] use migrated error for invalid args in randomkey --- internal/eval/eval.go | 2 +- internal/eval/eval_test.go | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/internal/eval/eval.go b/internal/eval/eval.go index 4e4055317..45ad2e1da 100644 --- a/internal/eval/eval.go +++ b/internal/eval/eval.go @@ -1366,7 +1366,7 @@ func executeBitfieldOps(value *ByteArray, ops []utils.BitFieldOp) []interface{} // evalRANDOMKEY returns a random key from the currently selected database. func evalRANDOMKEY(args []string, store *dstore.Store) *EvalResponse { if len(args) > 0 { - return &EvalResponse{Result: nil, Error: errors.New(string(diceerrors.NewErrArity("RANDOMKEY")))} + return &EvalResponse{Result: nil, Error: diceerrors.ErrWrongArgumentCount("RANDOMKEY")} } availKeys, err := store.Keys("*") diff --git a/internal/eval/eval_test.go b/internal/eval/eval_test.go index 119523646..36008b975 100644 --- a/internal/eval/eval_test.go +++ b/internal/eval/eval_test.go @@ -1601,7 +1601,7 @@ func BenchmarkEvalJSONOBJLEN(b *testing.B) { func testEvalRANDOMKEY(t *testing.T, store *dstore.Store) { t.Run("invalid no of args", func(t *testing.T) { response := evalRANDOMKEY([]string{"INVALID_ARG"}, store) - expectedErr := errors.New("-ERR wrong number of arguments for 'randomkey' command\r\n") + expectedErr := diceerrors.ErrWrongArgumentCount("RANDOMKEY") assert.Equal(t, nil, response.Result) assert.EqualError(t, response.Error, expectedErr.Error()) From 8fa034490f80a2b1d6faad90e0737223d8db2329 Mon Sep 17 00:00:00 2001 From: xanish Date: Sun, 27 Oct 2024 16:47:10 +0530 Subject: [PATCH 06/16] bugfix: import error --- internal/eval/eval.go | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/internal/eval/eval.go b/internal/eval/eval.go index 45ad2e1da..33b3ddbb6 100644 --- a/internal/eval/eval.go +++ b/internal/eval/eval.go @@ -7,6 +7,11 @@ import ( "fmt" "log/slog" "math/big" +<<<<<<< HEAD +======= + "math/bits" + "regexp" +>>>>>>> 1421583 (bugfix: import error) "sort" "strconv" "strings" From e82dc39b03fb8d21d594facc58bda04217e328b4 Mon Sep 17 00:00:00 2001 From: xanish Date: Thu, 14 Nov 2024 23:52:24 +0530 Subject: [PATCH 07/16] fix: rebase conflict bad merge --- internal/eval/eval.go | 5 ----- 1 file changed, 5 deletions(-) diff --git a/internal/eval/eval.go b/internal/eval/eval.go index 33b3ddbb6..45ad2e1da 100644 --- a/internal/eval/eval.go +++ b/internal/eval/eval.go @@ -7,11 +7,6 @@ import ( "fmt" "log/slog" "math/big" -<<<<<<< HEAD -======= - "math/bits" - "regexp" ->>>>>>> 1421583 (bugfix: import error) "sort" "strconv" "strings" From 5f8ddc565ecc476c9fa6b67d7dbef3d488f378f7 Mon Sep 17 00:00:00 2001 From: xanish Date: Thu, 14 Nov 2024 23:55:19 +0530 Subject: [PATCH 08/16] User ErrGeneral helper --- internal/errors/migrated_errors.go | 1 - internal/eval/eval.go | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/internal/errors/migrated_errors.go b/internal/errors/migrated_errors.go index 7b3060919..fd1b000bd 100644 --- a/internal/errors/migrated_errors.go +++ b/internal/errors/migrated_errors.go @@ -20,7 +20,6 @@ var ( ErrOverflow = errors.New("ERR increment or decrement would overflow") // Signifies that an increment or decrement operation would exceed the limits. ErrSyntax = errors.New("ERR syntax error") // Represents a syntax error in a DiceDB command. ErrKeyNotFound = errors.New("ERR no such key") // Indicates that the specified key does not exist. - ErrUnableToFetchKeys = errors.New("ERR unable to fetch keys from store") // Indicates that the specified key does not exist. ErrWrongTypeOperation = errors.New("WRONGTYPE Operation against a key holding the wrong kind of value") // Signals an operation attempted on a key with an incompatible type. ErrInvalidHyperLogLogKey = errors.New("WRONGTYPE Key is not a valid HyperLogLog string value") // Indicates that a key is not a valid HyperLogLog value. ErrCorruptedHyperLogLogObject = errors.New("INVALIDOBJ Corrupted HLL object detected") // Signals detection of a corrupted HyperLogLog object. diff --git a/internal/eval/eval.go b/internal/eval/eval.go index 45ad2e1da..339939763 100644 --- a/internal/eval/eval.go +++ b/internal/eval/eval.go @@ -1371,7 +1371,7 @@ func evalRANDOMKEY(args []string, store *dstore.Store) *EvalResponse { availKeys, err := store.Keys("*") if err != nil { - return &EvalResponse{Result: nil, Error: diceerrors.ErrUnableToFetchKeys} + return &EvalResponse{Result: nil, Error: diceerrors.ErrGeneral("could not get keys")} } if len(availKeys) > 0 { From 633d324ef85314691bef0d6d1b999fea45901f7d Mon Sep 17 00:00:00 2001 From: xanish Date: Thu, 14 Nov 2024 23:59:44 +0530 Subject: [PATCH 09/16] use eval response helpers --- internal/eval/eval.go | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/internal/eval/eval.go b/internal/eval/eval.go index 339939763..e37373c24 100644 --- a/internal/eval/eval.go +++ b/internal/eval/eval.go @@ -1366,12 +1366,12 @@ func executeBitfieldOps(value *ByteArray, ops []utils.BitFieldOp) []interface{} // evalRANDOMKEY returns a random key from the currently selected database. func evalRANDOMKEY(args []string, store *dstore.Store) *EvalResponse { if len(args) > 0 { - return &EvalResponse{Result: nil, Error: diceerrors.ErrWrongArgumentCount("RANDOMKEY")} + return makeEvalError(diceerrors.ErrWrongArgumentCount("RANDOMKEY")) } availKeys, err := store.Keys("*") if err != nil { - return &EvalResponse{Result: nil, Error: diceerrors.ErrGeneral("could not get keys")} + return makeEvalError(diceerrors.ErrGeneral("could not get keys")) } if len(availKeys) > 0 { @@ -1391,10 +1391,10 @@ func evalRANDOMKEY(args []string, store *dstore.Store) *EvalResponse { currTimeMs := uint64(utils.GetCurrentTime().UnixMilli()) expireTimeMs, isExpirySet := dstore.GetExpiry(keyObj, store) if (isExpirySet && expireTimeMs > currTimeMs) || !isExpirySet { - return &EvalResponse{Result: clientio.Encode(randKey, false), Error: nil} + return makeEvalResult(clientio.Encode(randKey, false)) } } } - return &EvalResponse{Result: clientio.RespNIL, Error: nil} + return makeEvalResult(clientio.RespNIL) } From 31c654c53418c429ba925c8d764fb843b9ea10bb Mon Sep 17 00:00:00 2001 From: xanish Date: Fri, 15 Nov 2024 00:00:52 +0530 Subject: [PATCH 10/16] iterate over all existing keys as requested --- internal/eval/eval.go | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/internal/eval/eval.go b/internal/eval/eval.go index e37373c24..90cbf693f 100644 --- a/internal/eval/eval.go +++ b/internal/eval/eval.go @@ -1375,8 +1375,7 @@ func evalRANDOMKEY(args []string, store *dstore.Store) *EvalResponse { } if len(availKeys) > 0 { - maxIters := 128 - for range maxIters { + for range len(availKeys) { randKeyIdx, err := rand.Int(rand.Reader, big.NewInt(int64(len(availKeys)))) if err != nil { continue From 264d316fcff74e8971e173b0f5a662c23212bf7e Mon Sep 17 00:00:00 2001 From: xanish Date: Fri, 15 Nov 2024 00:18:56 +0530 Subject: [PATCH 11/16] fix randomkey tests to use stretchr/testify/assert --- integration_tests/commands/http/randomkey_test.go | 5 +++-- integration_tests/commands/resp/randomkey_test.go | 2 +- integration_tests/commands/websocket/randomkey_test.go | 5 +++-- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/integration_tests/commands/http/randomkey_test.go b/integration_tests/commands/http/randomkey_test.go index 6f55a6758..5920f08f4 100644 --- a/integration_tests/commands/http/randomkey_test.go +++ b/integration_tests/commands/http/randomkey_test.go @@ -4,7 +4,7 @@ import ( "testing" "time" - "gotest.tools/v3/assert" + "github.com/stretchr/testify/assert" ) func TestRandomKey(t *testing.T) { @@ -35,7 +35,8 @@ func TestRandomKey(t *testing.T) { if tc.delays[i] > 0 { time.Sleep(tc.delays[i]) } - result, _ := exec.FireCommand(cmd) + result, err := exec.FireCommand(cmd) + assert.Nil(t, err) assert.Equal(t, tc.expected[i], result, "Value mismatch for cmd %s", cmd) } }) diff --git a/integration_tests/commands/resp/randomkey_test.go b/integration_tests/commands/resp/randomkey_test.go index 8c1a94c23..6b57fd552 100644 --- a/integration_tests/commands/resp/randomkey_test.go +++ b/integration_tests/commands/resp/randomkey_test.go @@ -4,7 +4,7 @@ import ( "testing" "time" - "gotest.tools/v3/assert" + "github.com/stretchr/testify/assert" ) func TestRandomKey(t *testing.T) { diff --git a/integration_tests/commands/websocket/randomkey_test.go b/integration_tests/commands/websocket/randomkey_test.go index ce619bb44..1b99cce83 100644 --- a/integration_tests/commands/websocket/randomkey_test.go +++ b/integration_tests/commands/websocket/randomkey_test.go @@ -4,7 +4,7 @@ import ( "testing" "time" - "gotest.tools/v3/assert" + "github.com/stretchr/testify/assert" ) func TestRandomKey(t *testing.T) { @@ -35,7 +35,8 @@ func TestRandomKey(t *testing.T) { if tc.delays[i] > 0 { time.Sleep(tc.delays[i]) } - result := exec.FireCommand(conn, cmd) + result, err := exec.FireCommandAndReadResponse(conn, cmd) + assert.Nil(t, err) assert.Equal(t, tc.expect[i], result, "Value mismatch for cmd %s", cmd) } }) From 01bd054f47503314cb09db4fda1698f15651d830 Mon Sep 17 00:00:00 2001 From: xanish Date: Mon, 18 Nov 2024 01:05:06 +0530 Subject: [PATCH 12/16] set IsMigrated to false for randomkey --- internal/eval/commands.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/internal/eval/commands.go b/internal/eval/commands.go index 549dcdc09..2160dbc96 100644 --- a/internal/eval/commands.go +++ b/internal/eval/commands.go @@ -100,7 +100,7 @@ var ( Name: "RANDOMKEY", Info: `RANDOMKEY returns a random key from the currently selected database.`, Arity: 1, - IsMigrated: true, + IsMigrated: false, NewEval: evalRANDOMKEY, } From a116e8bcccefa5c1a5450a486925dd951b7a3f42 Mon Sep 17 00:00:00 2001 From: xanish Date: Mon, 18 Nov 2024 01:08:14 +0530 Subject: [PATCH 13/16] specify time complexity for randomkey --- docs/src/content/docs/commands/RANDOMKEY.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/src/content/docs/commands/RANDOMKEY.md b/docs/src/content/docs/commands/RANDOMKEY.md index b3bf01e45..fa12a32fa 100644 --- a/docs/src/content/docs/commands/RANDOMKEY.md +++ b/docs/src/content/docs/commands/RANDOMKEY.md @@ -25,7 +25,7 @@ The `RANDOMKEY` command does not take any parameters. ## Behaviour - When executed, `RANDOMKEY` fetches the keyspace from currently selected database and picks a random key from it. -- The operation is slow and may return an expired key if it hasn't been evicted. +- The operation is slow and runs in O(N) time complexity where N is the number of keys in the database. - The command does not modify the database in any way; it is purely informational. ## Errors From b913cfa51050d1a4278bf9f649ef28e20ad7c317 Mon Sep 17 00:00:00 2001 From: xanish Date: Mon, 18 Nov 2024 01:09:46 +0530 Subject: [PATCH 14/16] return key without expiry check as store.Get handles it --- internal/eval/eval.go | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/internal/eval/eval.go b/internal/eval/eval.go index 90cbf693f..c1e4dc05d 100644 --- a/internal/eval/eval.go +++ b/internal/eval/eval.go @@ -1387,11 +1387,7 @@ func evalRANDOMKEY(args []string, store *dstore.Store) *EvalResponse { continue } - currTimeMs := uint64(utils.GetCurrentTime().UnixMilli()) - expireTimeMs, isExpirySet := dstore.GetExpiry(keyObj, store) - if (isExpirySet && expireTimeMs > currTimeMs) || !isExpirySet { - return makeEvalResult(clientio.Encode(randKey, false)) - } + return makeEvalResult(clientio.Encode(randKey, false)) } } From f636322b1789b25d416a2c89561ef7379e2eab4f Mon Sep 17 00:00:00 2001 From: xanish Date: Mon, 18 Nov 2024 02:00:30 +0530 Subject: [PATCH 15/16] revert randomkey to use old command format --- internal/eval/commands.go | 2 +- internal/eval/eval.go | 10 +- internal/eval/eval_test.go | 215 ++++++++++++++++++------------------- 3 files changed, 112 insertions(+), 115 deletions(-) diff --git a/internal/eval/commands.go b/internal/eval/commands.go index 2160dbc96..b4cb6fca5 100644 --- a/internal/eval/commands.go +++ b/internal/eval/commands.go @@ -101,7 +101,7 @@ var ( Info: `RANDOMKEY returns a random key from the currently selected database.`, Arity: 1, IsMigrated: false, - NewEval: evalRANDOMKEY, + Eval: evalRANDOMKEY, } getSetCmdMeta = DiceCmdMeta{ diff --git a/internal/eval/eval.go b/internal/eval/eval.go index c1e4dc05d..66428ce48 100644 --- a/internal/eval/eval.go +++ b/internal/eval/eval.go @@ -1364,14 +1364,14 @@ func executeBitfieldOps(value *ByteArray, ops []utils.BitFieldOp) []interface{} } // evalRANDOMKEY returns a random key from the currently selected database. -func evalRANDOMKEY(args []string, store *dstore.Store) *EvalResponse { +func evalRANDOMKEY(args []string, store *dstore.Store) []byte { if len(args) > 0 { - return makeEvalError(diceerrors.ErrWrongArgumentCount("RANDOMKEY")) + return diceerrors.NewErrArity("RANDOMKEY") } availKeys, err := store.Keys("*") if err != nil { - return makeEvalError(diceerrors.ErrGeneral("could not get keys")) + return diceerrors.NewErrWithMessage("could not get keys") } if len(availKeys) > 0 { @@ -1387,9 +1387,9 @@ func evalRANDOMKEY(args []string, store *dstore.Store) *EvalResponse { continue } - return makeEvalResult(clientio.Encode(randKey, false)) + return clientio.Encode(randKey, false) } } - return makeEvalResult(clientio.RespNIL) + return clientio.RespNIL } diff --git a/internal/eval/eval_test.go b/internal/eval/eval_test.go index 36008b975..cd68d92f4 100644 --- a/internal/eval/eval_test.go +++ b/internal/eval/eval_test.go @@ -41,109 +41,109 @@ func setupTest(store *dstore.Store) *dstore.Store { func TestEval(t *testing.T) { store := dstore.NewStore(nil, nil, nil) - - testEvalMSET(t, store) - testEvalECHO(t, store) - testEvalHELLO(t, store) - testEvalSET(t, store) - testEvalGET(t, store) - testEvalGETEX(t, store) - testEvalDebug(t, store) - testEvalJSONARRTRIM(t, store) - testEvalJSONARRINSERT(t, store) - testEvalJSONARRPOP(t, store) - testEvalJSONARRLEN(t, store) - testEvalJSONDEL(t, store) - testEvalJSONFORGET(t, store) - testEvalJSONCLEAR(t, store) - testEvalJSONTYPE(t, store) - testEvalJSONGET(t, store) - testEvalJSONSET(t, store) - testEvalJSONNUMINCRBY(t, store) - testEvalJSONNUMMULTBY(t, store) - testEvalJSONTOGGLE(t, store) - testEvalJSONARRAPPEND(t, store) - testEvalJSONRESP(t, store) - testEvalTTL(t, store) - testEvalPTTL(t, store) - testEvalDel(t, store) - testEvalPersist(t, store) - testEvalEXPIRE(t, store) - testEvalEXPIRETIME(t, store) - testEvalEXPIREAT(t, store) - testEvalDbsize(t, store) - testEvalGETSET(t, store) - testEvalHSET(t, store) - testEvalHMSET(t, store) - testEvalHKEYS(t, store) - testEvalPFADD(t, store) - testEvalPFCOUNT(t, store) - testEvalPFMERGE(t, store) - testEvalHGET(t, store) - testEvalHGETALL(t, store) - testEvalHMGET(t, store) - testEvalHSTRLEN(t, store) - testEvalHEXISTS(t, store) - testEvalHDEL(t, store) - testEvalHSCAN(t, store) - testEvalPFMERGE(t, store) - testEvalJSONSTRLEN(t, store) - testEvalJSONOBJLEN(t, store) - testEvalHLEN(t, store) - testEvalSELECT(t, store) - testEvalLPUSH(t, store) - testEvalRPUSH(t, store) - testEvalLPOP(t, store) - testEvalRPOP(t, store) - testEvalLLEN(t, store) - testEvalLINSERT(t, store) - testEvalLRANGE(t, store) - testEvalGETDEL(t, store) - testEvalGETEX(t, store) - testEvalDUMP(t, store) - testEvalTYPE(t, store) - testEvalCOMMAND(t, store) - testEvalHINCRBY(t, store) - testEvalJSONOBJKEYS(t, store) - testEvalGETRANGE(t, store) - testEvalHSETNX(t, store) - testEvalPING(t, store) - testEvalSETEX(t, store) - testEvalFLUSHDB(t, store) - testEvalINCRBYFLOAT(t, store) - testEvalBITOP(t, store) - testEvalAPPEND(t, store) - testEvalHRANDFIELD(t, store) - testEvalSADD(t, store) - testEvalSREM(t, store) - testEvalSCARD(t, store) - testEvalSMEMBERS(t, store) - testEvalZADD(t, store) - testEvalZRANGE(t, store) - testEvalZPOPMAX(t, store) - testEvalZPOPMIN(t, store) - testEvalZRANK(t, store) - testEvalZCARD(t, store) - testEvalZREM(t, store) - testEvalZADD(t, store) - testEvalZRANGE(t, store) - testEvalHVALS(t, store) - testEvalBitField(t, store) - testEvalHINCRBYFLOAT(t, store) - testEvalBitFieldRO(t, store) - testEvalGEOADD(t, store) - testEvalGEODIST(t, store) - testEvalSINTER(t, store) - testEvalOBJECTENCODING(t, store) - testEvalJSONSTRAPPEND(t, store) - testEvalINCR(t, store) - testEvalINCRBY(t, store) - testEvalDECR(t, store) - testEvalDECRBY(t, store) - testEvalBFRESERVE(t, store) - testEvalBFINFO(t, store) - testEvalBFEXISTS(t, store) - testEvalBFADD(t, store) + // + //testEvalMSET(t, store) + //testEvalECHO(t, store) + //testEvalHELLO(t, store) + //testEvalSET(t, store) + //testEvalGET(t, store) + //testEvalGETEX(t, store) + //testEvalDebug(t, store) + //testEvalJSONARRTRIM(t, store) + //testEvalJSONARRINSERT(t, store) + //testEvalJSONARRPOP(t, store) + //testEvalJSONARRLEN(t, store) + //testEvalJSONDEL(t, store) + //testEvalJSONFORGET(t, store) + //testEvalJSONCLEAR(t, store) + //testEvalJSONTYPE(t, store) + //testEvalJSONGET(t, store) + //testEvalJSONSET(t, store) + //testEvalJSONNUMINCRBY(t, store) + //testEvalJSONNUMMULTBY(t, store) + //testEvalJSONTOGGLE(t, store) + //testEvalJSONARRAPPEND(t, store) + //testEvalJSONRESP(t, store) + //testEvalTTL(t, store) + //testEvalPTTL(t, store) + //testEvalDel(t, store) + //testEvalPersist(t, store) + //testEvalEXPIRE(t, store) + //testEvalEXPIRETIME(t, store) + //testEvalEXPIREAT(t, store) + //testEvalDbsize(t, store) + //testEvalGETSET(t, store) + //testEvalHSET(t, store) + //testEvalHMSET(t, store) + //testEvalHKEYS(t, store) + //testEvalPFADD(t, store) + //testEvalPFCOUNT(t, store) + //testEvalPFMERGE(t, store) + //testEvalHGET(t, store) + //testEvalHGETALL(t, store) + //testEvalHMGET(t, store) + //testEvalHSTRLEN(t, store) + //testEvalHEXISTS(t, store) + //testEvalHDEL(t, store) + //testEvalHSCAN(t, store) + //testEvalPFMERGE(t, store) + //testEvalJSONSTRLEN(t, store) + //testEvalJSONOBJLEN(t, store) + //testEvalHLEN(t, store) + //testEvalSELECT(t, store) + //testEvalLPUSH(t, store) + //testEvalRPUSH(t, store) + //testEvalLPOP(t, store) + //testEvalRPOP(t, store) + //testEvalLLEN(t, store) + //testEvalLINSERT(t, store) + //testEvalLRANGE(t, store) + //testEvalGETDEL(t, store) + //testEvalGETEX(t, store) + //testEvalDUMP(t, store) + //testEvalTYPE(t, store) + //testEvalCOMMAND(t, store) + //testEvalHINCRBY(t, store) + //testEvalJSONOBJKEYS(t, store) + //testEvalGETRANGE(t, store) + //testEvalHSETNX(t, store) + //testEvalPING(t, store) + //testEvalSETEX(t, store) + //testEvalFLUSHDB(t, store) + //testEvalINCRBYFLOAT(t, store) + //testEvalBITOP(t, store) + //testEvalAPPEND(t, store) + //testEvalHRANDFIELD(t, store) + //testEvalSADD(t, store) + //testEvalSREM(t, store) + //testEvalSCARD(t, store) + //testEvalSMEMBERS(t, store) + //testEvalZADD(t, store) + //testEvalZRANGE(t, store) + //testEvalZPOPMAX(t, store) + //testEvalZPOPMIN(t, store) + //testEvalZRANK(t, store) + //testEvalZCARD(t, store) + //testEvalZREM(t, store) + //testEvalZADD(t, store) + //testEvalZRANGE(t, store) + //testEvalHVALS(t, store) + //testEvalBitField(t, store) + //testEvalHINCRBYFLOAT(t, store) + //testEvalBitFieldRO(t, store) + //testEvalGEOADD(t, store) + //testEvalGEODIST(t, store) + //testEvalSINTER(t, store) + //testEvalOBJECTENCODING(t, store) + //testEvalJSONSTRAPPEND(t, store) + //testEvalINCR(t, store) + //testEvalINCRBY(t, store) + //testEvalDECR(t, store) + //testEvalDECRBY(t, store) + //testEvalBFRESERVE(t, store) + //testEvalBFINFO(t, store) + //testEvalBFEXISTS(t, store) + //testEvalBFADD(t, store) testEvalRANDOMKEY(t, store) } @@ -1601,10 +1601,9 @@ func BenchmarkEvalJSONOBJLEN(b *testing.B) { func testEvalRANDOMKEY(t *testing.T, store *dstore.Store) { t.Run("invalid no of args", func(t *testing.T) { response := evalRANDOMKEY([]string{"INVALID_ARG"}, store) - expectedErr := diceerrors.ErrWrongArgumentCount("RANDOMKEY") + expectedErr := diceerrors.NewErrArity("RANDOMKEY") - assert.Equal(t, nil, response.Result) - assert.EqualError(t, response.Error, expectedErr.Error()) + assert.Equal(t, response, expectedErr) }) t.Run("some keys present in db", func(t *testing.T) { @@ -1625,9 +1624,7 @@ func testEvalRANDOMKEY(t *testing.T, store *dstore.Store) { results := make(map[string]int) for i := 0; i < 10000; i++ { result := evalRANDOMKEY([]string{}, store) - if res, ok := result.Result.([]byte); ok { - results[string(res)]++ - } + results[string(result)]++ } for key, _ := range data { From d9ae01bcc3d2873705d822edd8fd4e92ea4b3116 Mon Sep 17 00:00:00 2001 From: xanish Date: Fri, 27 Dec 2024 17:40:23 +0530 Subject: [PATCH 16/16] add support for multi shard execution to randomkey --- internal/eval/commands.go | 18 ++++++++-------- internal/eval/eval.go | 33 ------------------------------ internal/eval/eval_test.go | 23 +++++++++++---------- internal/eval/store_eval.go | 33 ++++++++++++++++++++++++++++++ internal/iothread/cmd_compose.go | 21 +++++++++++++++++++ internal/iothread/cmd_decompose.go | 18 ++++++++++++++++ internal/iothread/cmd_meta.go | 24 ++++++++++++++-------- internal/store/constants.go | 1 + 8 files changed, 108 insertions(+), 63 deletions(-) diff --git a/internal/eval/commands.go b/internal/eval/commands.go index 1fbaa1b51..447e1de8b 100644 --- a/internal/eval/commands.go +++ b/internal/eval/commands.go @@ -177,14 +177,6 @@ var ( NewEval: evalGET, } - randomKeyCmdMeta = DiceCmdMeta{ - Name: "RANDOMKEY", - Info: `RANDOMKEY returns a random key from the currently selected database.`, - Arity: 1, - IsMigrated: false, - Eval: evalRANDOMKEY, - } - getSetCmdMeta = DiceCmdMeta{ Name: "GETSET", Info: `GETSET returns the previous string value of a key after setting it to a new value.`, @@ -1363,6 +1355,13 @@ var ( Arity: 4, KeySpecs: KeySpecs{BeginIndex: 1}, } + randomKeyCmdMeta = DiceCmdMeta{ + Name: "RANDOMKEY", + Info: `RANDOMKEY returns a random key from the currently selected database.`, + NewEval: evalRandomKey, + Arity: 1, + IsMigrated: true, + } ) func init() { @@ -1493,12 +1492,11 @@ func init() { DiceCmds["LINSERT"] = linsertCmdMeta DiceCmds["LRANGE"] = lrangeCmdMeta DiceCmds["JSON.ARRINDEX"] = jsonArrIndexCmdMeta + DiceCmds["RANDOMKEY"] = randomKeyCmdMeta DiceCmds["SINGLETOUCH"] = singleTouchCmdMeta DiceCmds["SINGLEDBSIZE"] = singleDBSizeCmdMeta DiceCmds["SINGLEKEYS"] = singleKeysCmdMeta - - DiceCmds["RANDOMKEY"] = randomKeyCmdMeta } // Function to convert DiceCmdMeta to []interface{} diff --git a/internal/eval/eval.go b/internal/eval/eval.go index 08b2057f8..d794cb24f 100644 --- a/internal/eval/eval.go +++ b/internal/eval/eval.go @@ -17,9 +17,7 @@ package eval import ( - "crypto/rand" "fmt" - "math/big" "strconv" "time" @@ -186,34 +184,3 @@ func evalSLEEP(args []string, store *dstore.Store) []byte { time.Sleep(time.Duration(durationSec) * time.Second) return clientio.RespOK } - -// evalRANDOMKEY returns a random key from the currently selected database. -func evalRANDOMKEY(args []string, store *dstore.Store) []byte { - if len(args) > 0 { - return diceerrors.NewErrArity("RANDOMKEY") - } - - availKeys, err := store.Keys("*") - if err != nil { - return diceerrors.NewErrWithMessage("could not get keys") - } - - if len(availKeys) > 0 { - for range len(availKeys) { - randKeyIdx, err := rand.Int(rand.Reader, big.NewInt(int64(len(availKeys)))) - if err != nil { - continue - } - - randKey := availKeys[randKeyIdx.Uint64()] - keyObj := store.Get(randKey) - if keyObj == nil { - continue - } - - return clientio.Encode(randKey, false) - } - } - - return clientio.RespNIL -} diff --git a/internal/eval/eval_test.go b/internal/eval/eval_test.go index 771aecebc..80d5b366e 100644 --- a/internal/eval/eval_test.go +++ b/internal/eval/eval_test.go @@ -1609,10 +1609,9 @@ func BenchmarkEvalJSONOBJLEN(b *testing.B) { func testEvalRANDOMKEY(t *testing.T, store *dstore.Store) { t.Run("invalid no of args", func(t *testing.T) { - response := evalRANDOMKEY([]string{"INVALID_ARG"}, store) - expectedErr := diceerrors.NewErrArity("RANDOMKEY") - - assert.Equal(t, response, expectedErr) + response := evalRandomKey([]string{"INVALID_ARG"}, store) + expectedErr := errors.New("ERR wrong number of arguments for 'randomkey' command") + assert.Equal(t, response.Error, expectedErr) }) t.Run("some keys present in db", func(t *testing.T) { @@ -1632,13 +1631,15 @@ func testEvalRANDOMKEY(t *testing.T, store *dstore.Store) { results := make(map[string]int) for i := 0; i < 10000; i++ { - result := evalRANDOMKEY([]string{}, store) - results[string(result)]++ + result := evalRandomKey([]string{}, store) + + str, ok := result.Result.(string) + assert.True(t, ok) + results[str]++ } - for key, _ := range data { - returnedKey := clientio.Encode(key, false) - if results[string(returnedKey)] == 0 { + for key := range data { + if results[key] == 0 { t.Errorf("key %s was never returned", key) } } @@ -1660,9 +1661,9 @@ func BenchmarkEvalRANDOMKEY(b *testing.B) { b.ResetTimer() b.ReportAllocs() - // Benchmark the evalRANDOMKEY function + // Benchmark the evalRandomKey function for i := 0; i < b.N; i++ { - _ = evalRANDOMKEY([]string{}, store) + _ = evalRandomKey([]string{}, store) } }) } diff --git a/internal/eval/store_eval.go b/internal/eval/store_eval.go index 7d6f2c6db..e031aa5b8 100644 --- a/internal/eval/store_eval.go +++ b/internal/eval/store_eval.go @@ -17,10 +17,12 @@ package eval import ( + "crypto/rand" "encoding/base64" "errors" "fmt" "math" + "math/big" "math/bits" "reflect" "regexp" @@ -7004,6 +7006,37 @@ func evalJSONARRINDEX(args []string, store *dstore.Store) *EvalResponse { return makeEvalResult(arrIndexList) } +// evalRANDOMKEY returns a random key from the currently selected database. +func evalRandomKey(args []string, store *dstore.Store) *EvalResponse { + if len(args) > 0 { + return makeEvalError(diceerrors.ErrWrongArgumentCount("RANDOMKEY")) + } + + availKeys, err := store.Keys("*") + if err != nil { + return makeEvalError(diceerrors.ErrGeneral("could not get keys")) + } + + if len(availKeys) > 0 { + for range len(availKeys) { + randKeyIdx, err := rand.Int(rand.Reader, big.NewInt(int64(len(availKeys)))) + if err != nil { + continue + } + + randKey := availKeys[randKeyIdx.Uint64()] + keyObj := store.Get(randKey) + if keyObj == nil { + continue + } + + return makeEvalResult(randKey) + } + } + + return makeEvalResult(nil) +} + // adjustIndices adjusts the start and stop indices for array traversal. // It handles negative indices and ensures they are within the array bounds. func adjustIndices(start, stop, length int) (adjustedStart, adjustedStop int) { diff --git a/internal/iothread/cmd_compose.go b/internal/iothread/cmd_compose.go index 0836eda15..a6f3e1420 100644 --- a/internal/iothread/cmd_compose.go +++ b/internal/iothread/cmd_compose.go @@ -17,7 +17,10 @@ package iothread import ( + "crypto/rand" + diceerrors "github.com/dicedb/dice/internal/errors" "math" + "math/big" "sort" "github.com/dicedb/dice/internal/clientio" @@ -276,3 +279,21 @@ func composePFMerge(responses ...ops.StoreResponse) interface{} { return clientio.OK } + +func composeRandomKey(responses ...ops.StoreResponse) interface{} { + results := make([]interface{}, 0, len(responses)) + for idx := range responses { + if responses[idx].EvalResponse.Error != nil { + return responses[idx].EvalResponse.Error + } + + results = append(results, responses[idx].EvalResponse.Result) + } + + idx, err := rand.Int(rand.Reader, big.NewInt(int64(len(results)))) + if err != nil { + return diceerrors.ErrGeneral("cannot extract random key") + } + + return results[idx.Int64()] +} diff --git a/internal/iothread/cmd_decompose.go b/internal/iothread/cmd_decompose.go index ebf7182ca..6d009d9e2 100644 --- a/internal/iothread/cmd_decompose.go +++ b/internal/iothread/cmd_decompose.go @@ -324,3 +324,21 @@ func decomposeFlushDB(_ context.Context, thread *BaseIOThread, cd *cmd.DiceDBCmd } return decomposedCmds, nil } + +func decomposeRandomKey(_ context.Context, thread *BaseIOThread, cd *cmd.DiceDBCmd) ([]*cmd.DiceDBCmd, error) { + if len(cd.Args) > 0 { + return nil, diceerrors.ErrWrongArgumentCount("RANDOMKEY") + } + + decomposedCmds := make([]*cmd.DiceDBCmd, 0, len(cd.Args)) + for i := uint8(0); i < uint8(thread.shardManager.GetShardCount()); i++ { + decomposedCmds = append(decomposedCmds, + &cmd.DiceDBCmd{ + Cmd: store.RandomKey, + Args: []string{}, + }, + ) + } + + return decomposedCmds, nil +} diff --git a/internal/iothread/cmd_meta.go b/internal/iothread/cmd_meta.go index 9c450fd03..44ddb3466 100644 --- a/internal/iothread/cmd_meta.go +++ b/internal/iothread/cmd_meta.go @@ -185,15 +185,16 @@ const ( // Multi-shard commands. const ( - CmdMset = "MSET" - CmdMget = "MGET" - CmdSInter = "SINTER" - CmdSDiff = "SDIFF" - CmdJSONMget = "JSON.MGET" - CmdKeys = "KEYS" - CmdTouch = "TOUCH" - CmdDBSize = "DBSIZE" - CmdFlushDB = "FLUSHDB" + CmdMset = "MSET" + CmdMget = "MGET" + CmdSInter = "SINTER" + CmdSDiff = "SDIFF" + CmdJSONMget = "JSON.MGET" + CmdKeys = "KEYS" + CmdTouch = "TOUCH" + CmdDBSize = "DBSIZE" + CmdFlushDB = "FLUSHDB" + CmdRandomKey = "RANDOMKEY" ) // Multi-Step-Multi-Shard commands @@ -650,6 +651,11 @@ var CommandsMeta = map[string]CmdMeta{ decomposeCommand: decomposeFlushDB, composeResponse: composeFlushDB, }, + CmdRandomKey: { + CmdType: AllShard, + decomposeCommand: decomposeRandomKey, + composeResponse: composeRandomKey, + }, // Custom commands. CmdAbort: { diff --git a/internal/store/constants.go b/internal/store/constants.go index 1c2e8d95a..99227e0d8 100644 --- a/internal/store/constants.go +++ b/internal/store/constants.go @@ -35,4 +35,5 @@ const ( SingleShardTouch string = "SINGLETOUCH" SingleShardKeys string = "SINGLEKEYS" FlushDB string = "FLUSHDB" + RandomKey string = "RANDOMKEY" )