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" )