add support for multi shard execution to randomkey

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{}

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
-}

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

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

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()]
+}

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
+}

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: {

internal/store/constants.go

@@ -35,4 +35,5 @@ const (
 	SingleShardTouch string = "SINGLETOUCH"
 	SingleShardKeys  string = "SINGLEKEYS"
 	FlushDB          string = "FLUSHDB"
+	RandomKey        string = "RANDOMKEY"
 )