Implement options for SET command, fix key resolution bug

OmkarPh · OmkarPh · commit ceee39bd83f2 · 2024-03-30T02:18:39.000+05:30
Signed-off-by: Omkar Phansopkar &lt;omkarphansopkar@gmail.com&gt;
diff --git a/README.md b/README.md
@@ -54,18 +54,18 @@ Detailed documentation - https://redis.io/commands/
 
 | Command  | Syntax                                   | Example                                                   | Description                                     |
 |----------|------------------------------------------|-----------------------------------------------------------|-------------------------------------------------|
-| SET      | SET <key> <value>                        | redis-cli SET name omkar                                  | Set the string value of a key                   |
-| GET      | GET <key>                                | redis-cli GET name                                        | Get the value of a key                          |
-| DEL      | DEL key [key ...]                        | redis-cli DEL name<br/>redis-cli DEL name age             | Delete one or more keys                         |
-| INCR     | INCR key                                 | redis-cli INCR age                                        | Increment the integer value of a key            |
-| DECR     | DECR key                                 | redis-cli DECR age                                        | Decrement the integer value of a key            |
-| EXISTS   | EXISTS key [key ...]                     | redis-cli EXISTS name<br/>redis-cli EXISTS name age       | Check if a key exists                           |
-| EXPIRE   | EXPIRE key seconds [NX / XX / GT / LT]   | redis-cli EXPIRE name 20<br/>redis-cli EXPIRE name 20 NX  | Set a key's time to live in seconds             |
-| PERSIST  | PERSIST key                              | redis-cli PERSIST name                                    | Remove the expiration from a key                |
-| TTL      | TTL key                                  | redis-cli TTL key                                         | Get the time to live for a key (in seconds)     |
-| TYPE     | TYPE key                                 | redis-cli TYPE name                                       | Determine the type stored at a key              |
-| PING     | PING                                     | redis-cli PING                                            | Ping the server                                 |
-| ECHO     | ECHO <message>                           | redis-cli ECHO "Hello world"                              | Echo the given string                           |
+| SET      | **SET key value** [NX / XX] [GET]<br/>[EX seconds / PX milliseconds<br/> / EXAT unix-time-seconds / PXAT unix-time-milliseconds / KEEPTTL]                        | redis-cli SET name omkar<br/>redis-cli SET name omkar GET KEEPTTL              | Set the string value of a key                   |
+| GET      | **GET key**                                  | redis-cli GET name                                        | Get the value of a key                          |
+| DEL      | **DEL key** [key ...]                        | redis-cli DEL name<br/>redis-cli DEL name age             | Delete one or more keys                         |
+| INCR     | **INCR key**                                 | redis-cli INCR age                                        | Increment the integer value of a key            |
+| DECR     | **DECR key**                                 | redis-cli DECR age                                        | Decrement the integer value of a key            |
+| EXISTS   | **EXISTS key** [key ...]                     | redis-cli EXISTS name<br/>redis-cli EXISTS name age       | Check if a key exists                           |
+| EXPIRE   | **EXPIRE key seconds** [NX / XX / GT / LT]   | redis-cli EXPIRE name 20<br/>redis-cli EXPIRE name 20 NX  | Set a key's time to live in seconds             |
+| PERSIST  | **PERSIST key **                             | redis-cli PERSIST name                                    | Remove the expiration from a key                |
+| TTL      | **TTL key**                                  | redis-cli TTL key                                         | Get the time to live for a key (in seconds)     |
+| TYPE     | **TYPE key**                                 | redis-cli TYPE name                                       | Determine the type stored at a key              |
+| PING     | **PING**                                     | redis-cli PING                                            | Ping the server                                 |
+| ECHO     | **ECHO message**                           | redis-cli ECHO "Hello world"                              | Echo the given string                           |
 
 
 
@@ -91,6 +91,9 @@ go build -o build/redis-lite-server -v
 
 ## Benchmarks
 
+> [!NOTE]  
+> These are results from Macbook air M1 8gb
+
 ```bash
 redis-benchmark -t SET,GET,INCR -q
 ```
diff --git a/config/redisConfig.go b/config/redisConfig.go
@@ -88,7 +88,7 @@ func (rc *RedisConfig) SetDefaultConfig() {
 	rc.Params["save"] = "3600 1 300 100 60 10000"
 	rc.Params["appendonly"] = "no"
 	rc.Params["kv_store"] = "sharded" // Options - "simple", "sharded"
-	rc.Params["shardfactor"] = "10"
+	rc.Params["shardfactor"] = "32"
 }
 
 func (rc *RedisConfig) GetParam(key string) (string, bool) {
diff --git a/core/actions/decr.go b/core/actions/decr.go
@@ -28,6 +28,6 @@ func (action *DecrAction) Execute(kvStore *store.KvStore, redisConfig *config.Re
 		return [][]byte{resp.ResolveResponse(errString, resp.Response_ERRORS)}, errors.New(errString)
 	}
 
-	(*kvStore).Set(key, newValueString)
+	(*kvStore).Set(key, newValueString, store.SetOptions{})
 	return [][]byte{resp.ResolveResponse(newValueString, resp.Response_INTEGERS)}, nil
 }
diff --git a/core/actions/incr.go b/core/actions/incr.go
@@ -28,6 +28,6 @@ func (action *IncrAction) Execute(kvStore *store.KvStore, redisConfig *config.Re
 		return [][]byte{resp.ResolveResponse(errString, resp.Response_ERRORS)}, errors.New(errString)
 	}
 
-	(*kvStore).Set(key, newValueString)
+	(*kvStore).Set(key, newValueString, store.SetOptions{})
 	return [][]byte{resp.ResolveResponse(newValueString, resp.Response_INTEGERS)}, nil
 }
diff --git a/core/actions/set.go b/core/actions/set.go
@@ -14,7 +14,7 @@ import (
 type SetAction struct{}
 
 func (action *SetAction) Execute(kvStore *store.KvStore, redisConfig *config.RedisConfig, args ...string) ([][]byte, error) {
-	if len(args) != 2 {
+	if len(args) < 2 {
 		errString := "ERR wrong number of arguments for 'SET' command"
 		return [][]byte{resp.ResolveResponse(errString, resp.Response_ERRORS)}, errors.New(errString)
 	}
@@ -23,6 +23,33 @@ func (action *SetAction) Execute(kvStore *store.KvStore, redisConfig *config.Red
 	value := args[1]
 	slog.Debug(fmt.Sprintf("Set action (%s => %s)\n", key, value))
 
-	(*kvStore).Set(key, value)
-	return [][]byte{resp.ResolveResponse("OK", resp.Response_SIMPLE_STRING)}, nil
+	optionsResolved, err := utils.ResolveSetOptions(args[2:]...)
+
+	if err != nil {
+		return [][]byte{resp.ResolveResponse(err.Error(), resp.Response_ERRORS)}, err
+	}
+
+	options := store.SetOptions(optionsResolved)
+
+	if (optionsResolved.NX && optionsResolved.XX) || (optionsResolved.ExpireDuration && optionsResolved.ExpireTimestamp) || (optionsResolved.KEEPTTL && (optionsResolved.ExpireDuration || optionsResolved.ExpireTimestamp)) {
+		errString := "ERR syntax error"
+		return [][]byte{resp.ResolveResponse(errString, resp.Response_ERRORS)}, errors.New(errString)
+	}
+
+	set, prevValue, err := (*kvStore).Set(key, value, options)
+
+	if err != nil {
+		return [][]byte{resp.ResolveResponse(err.Error(), resp.Response_ERRORS)}, err
+	}
+
+	if set {
+		if optionsResolved.GET {
+			if prevValue == "" {
+				return [][]byte{resp.ResolveResponse(nil, resp.Response_NULL)}, nil
+			}
+			return [][]byte{resp.ResolveResponse(prevValue, resp.Response_BULK_STRINGS)}, nil
+		}
+		return [][]byte{resp.ResolveResponse("OK", resp.Response_SIMPLE_STRING)}, nil
+	}
+	return [][]byte{resp.ResolveResponse(nil, resp.Response_NULL)}, nil
 }
diff --git a/main.go b/main.go
@@ -4,13 +4,20 @@ import (
 	"fmt"
 	"log/slog"
 	"net"
+	"os"
 
 	"github.com/OmkarPh/redis-lite/config"
 	"github.com/OmkarPh/redis-lite/core"
 	"github.com/OmkarPh/redis-lite/store"
 )
 
 func main() {
+
+	fmt.Println("Redis-lite server v0.2")
+	fmt.Println("Port:", config.PORT, ", PID:", os.Getpid())
+	fmt.Println("Visit - https://github.com/OmkarPh/redis-server-lite")
+	fmt.Println()
+
 	slog.SetLogLoggerLevel(config.DefaultLoggerLevel)
 
 	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", config.PORT))
@@ -23,11 +30,10 @@ func main() {
 	redisConfig := config.NewRedisConfig()
 	kvStore := store.NewKvStore(redisConfig)
 
-	fmt.Println("Redis-lite server is up & running on port", config.PORT)
-	fmt.Println()
-
 	go core.CleanExpiredKeys(kvStore)
 
+	fmt.Println()
+
 	for {
 		conn, err := listener.Accept()
 		if err != nil {
diff --git a/store/kvStore.go b/store/kvStore.go
@@ -23,6 +23,18 @@ type StoredValue struct {
 	// Value interface{}
 	Expiry time.Time
 }
+type SetOptions struct {
+	NX                        bool
+	XX                        bool
+	GET                       bool
+	ExpireDuration            bool // EX or PX specified
+	ExpiryTimeSeconds         int64
+	ExpiryTimeMiliSeconds     int64
+	ExpireTimestamp           bool // EXAT or PXAT specified
+	ExpiryUnixTimeSeconds     int64
+	ExpiryUnixTimeMiliSeconds int64
+	KEEPTTL                   bool
+}
 type ExpireOptions struct {
 	NX bool
 	XX bool
@@ -31,7 +43,7 @@ type ExpireOptions struct {
 }
 type KvStore interface {
 	Get(key string) (string, bool)
-	Set(key string, value string)
+	Set(key string, value string, options SetOptions) (bool, string, error)
 	Del(key string) bool
 	DeleteIfExpired(keysCount int) int
 	Incr(key string) (string, error)
diff --git a/store/shardedKvStore.go b/store/shardedKvStore.go
@@ -43,7 +43,13 @@ func (kvStore *ShardedKvStore) resolveShardIdx(key string) uint32 {
 
 	h := murmur3.New32()
 	h.Write([]byte(key))
-	shardIdx := h.Sum32() % kvStore.shardFactor
+
+	var shardIdx uint32
+	if kvStore.shardFactor&(kvStore.shardFactor-1) == 0 {
+		shardIdx = h.Sum32() & (kvStore.shardFactor - 1) // Faster than modulus (when shardfactor is power of 2)
+	} else {
+		shardIdx = h.Sum32() % kvStore.shardFactor
+	}
 
 	slog.Debug(fmt.Sprintf("Shard idx for %s => %d", key, shardIdx))
 	// shardIdxCache[key] = shardIdx
@@ -64,14 +70,41 @@ func (kvStore *ShardedKvStore) Get(key string) (string, bool) {
 	return "", false
 }
 
-func (kvStore *ShardedKvStore) Set(key string, value string) {
+func (kvStore *ShardedKvStore) Set(key string, value string, options SetOptions) (bool, string, error) {
 	shardIdx := kvStore.resolveShardIdx(key)
 	kvStore.mutex[shardIdx].Lock()
 	defer kvStore.mutex[shardIdx].Unlock()
+
+	existingData, exists := kvStore.kv_stores[shardIdx][key]
+
+	// Convert options to ExpirationTimeOptions
+	expirationOptions := utils.ExpirationTimeOptions{
+		NX:                        options.NX,
+		XX:                        options.XX,
+		ExpireDuration:            options.ExpireDuration,
+		ExpiryTimeSeconds:         options.ExpiryTimeSeconds,
+		ExpiryTimeMiliSeconds:     options.ExpiryTimeMiliSeconds,
+		ExpireTimestamp:           options.ExpireTimestamp,
+		ExpiryUnixTimeSeconds:     options.ExpiryUnixTimeSeconds,
+		ExpiryUnixTimeMiliSeconds: options.ExpiryUnixTimeMiliSeconds,
+		KEEPTTL:                   options.KEEPTTL,
+	}
+
+	if (options.NX && exists) || (options.XX && !exists) {
+		return false, "", nil
+	}
+
+	expiryTime, canSet, err := utils.ResolveExpirationTime(expirationOptions, exists, existingData.Expiry)
+
+	if !canSet {
+		return false, "", err
+	}
+
 	kvStore.kv_stores[shardIdx][key] = StoredValue{
 		Value:  value,
-		Expiry: time.Time{},
+		Expiry: expiryTime,
 	}
+	return true, existingData.Value, nil
 }
 
 func (kvStore *ShardedKvStore) Del(key string) bool {
diff --git a/store/simpleKvStore.go b/store/simpleKvStore.go
@@ -38,13 +38,40 @@ func (kvStore *SimpleKvStore) Get(key string) (string, bool) {
 	return "", false
 }
 
-func (kvStore *SimpleKvStore) Set(key string, value string) {
+func (kvStore *SimpleKvStore) Set(key string, value string, options SetOptions) (bool, string, error) {
 	kvStore.mutex.Lock()
 	defer kvStore.mutex.Unlock()
+
+	existingData, exists := kvStore.kv_store[key]
+
+	// Convert options to ExpirationTimeOptions
+	expirationOptions := utils.ExpirationTimeOptions{
+		NX:                        options.NX,
+		XX:                        options.XX,
+		ExpireDuration:            options.ExpireDuration,
+		ExpiryTimeSeconds:         options.ExpiryTimeSeconds,
+		ExpiryTimeMiliSeconds:     options.ExpiryTimeMiliSeconds,
+		ExpireTimestamp:           options.ExpireTimestamp,
+		ExpiryUnixTimeSeconds:     options.ExpiryUnixTimeSeconds,
+		ExpiryUnixTimeMiliSeconds: options.ExpiryUnixTimeMiliSeconds,
+		KEEPTTL:                   options.KEEPTTL,
+	}
+
+	if (options.NX && exists) || (options.XX && !exists) {
+		return false, "", nil
+	}
+
+	expiryTime, canSet, err := utils.ResolveExpirationTime(expirationOptions, exists, existingData.Expiry)
+
+	if !canSet {
+		return false, "", err
+	}
+
 	kvStore.kv_store[key] = StoredValue{
 		Value:  value,
-		Expiry: time.Time{},
+		Expiry: expiryTime,
 	}
+	return true, existingData.Value, nil
 }
 
 func (kvStore *SimpleKvStore) Del(key string) bool {
diff --git a/utils/expiration.go b/utils/expiration.go
@@ -1,6 +1,7 @@
 package utils
 
 import (
+	"errors"
 	"time"
 )
 
@@ -11,3 +12,48 @@ type ValueWithExpiration struct {
 func IsExpired(expiration time.Time) bool {
 	return !expiration.IsZero() && expiration.Before(time.Now())
 }
+
+type ExpirationTimeOptions struct {
+	NX                        bool
+	XX                        bool
+	ExpireDuration            bool // EX or PX specified
+	ExpiryTimeSeconds         int64
+	ExpiryTimeMiliSeconds     int64
+	ExpireTimestamp           bool // EXAT or PXAT specified
+	ExpiryUnixTimeSeconds     int64
+	ExpiryUnixTimeMiliSeconds int64
+	KEEPTTL                   bool
+}
+
+func ResolveExpirationTime(options ExpirationTimeOptions, exists bool, existingExpiry time.Time) (time.Time, bool, error) {
+	expiryTime := time.Time{}
+
+	if options.KEEPTTL && exists {
+		expiryTime = existingExpiry
+	}
+
+	if options.ExpireDuration {
+		if options.ExpiryTimeSeconds != -1 {
+			expiryTime = time.Now().Add(time.Second * time.Duration(options.ExpiryTimeSeconds))
+		} else if options.ExpiryTimeMiliSeconds != -1 {
+			expiryTime = time.Now().Add(time.Millisecond * time.Duration(options.ExpiryTimeMiliSeconds))
+		} else {
+			errString := "ERR invalid expire duration in SET"
+			return expiryTime, false, errors.New(errString)
+		}
+	}
+
+	if options.ExpireTimestamp {
+		if options.ExpiryUnixTimeSeconds != -1 {
+			expiryTime = time.Unix(options.ExpiryUnixTimeSeconds, 0)
+		} else if options.ExpiryUnixTimeMiliSeconds != -1 {
+			nanoseconds := options.ExpiryUnixTimeMiliSeconds * 1000000 * int64(time.Nanosecond)
+			expiryTime = time.Unix(0, nanoseconds)
+		} else {
+			errString := "ERR invalid expire timestamp in SET"
+			return expiryTime, false, errors.New(errString)
+		}
+	}
+
+	return expiryTime, true, nil
+}
diff --git a/utils/keys.go b/utils/keys.go
@@ -11,8 +11,8 @@ func GenerateRandomKey() string {
 }
 
 func ResolvePossibleKeyDirectives(key string) string {
-	key = strings.ToLower(key)
-	if key == "key:__rand_int__" || key == "__rand_int__" {
+	normalisedKey := strings.ToLower(key)
+	if normalisedKey == "key:__rand_int__" || normalisedKey == "__rand_int__" {
 		return GenerateRandomKey()
 	}
 	return key
diff --git a/utils/optionsResolvers.go b/utils/optionsResolvers.go

Original file line number	Diff line number	Diff line change
`@@ -88,7 +88,7 @@ func (rc *RedisConfig) SetDefaultConfig() {`
`88`	`88`	`rc.Params["save"] = "3600 1 300 100 60 10000"`
`89`	`89`	`rc.Params["appendonly"] = "no"`
`90`	`90`	`rc.Params["kv_store"] = "sharded" // Options - "simple", "sharded"`
`91`		`- rc.Params["shardfactor"] = "10"`
	`91`	`+ rc.Params["shardfactor"] = "32"`
`92`	`92`	`}`
`93`	`93`
`94`	`94`	`func (rc *RedisConfig) GetParam(key string) (string, bool) {`
Original file line number	Diff line number	Diff line change
`@@ -28,6 +28,6 @@ func (action DecrAction) Execute(kvStore store.KvStore, redisConfig *config.Re`
`28`	`28`	`return [][]byte{resp.ResolveResponse(errString, resp.Response_ERRORS)}, errors.New(errString)`
`29`	`29`	`}`
`30`	`30`
`31`		`- (*kvStore).Set(key, newValueString)`
	`31`	`+ (*kvStore).Set(key, newValueString, store.SetOptions{})`
`32`	`32`	`return [][]byte{resp.ResolveResponse(newValueString, resp.Response_INTEGERS)}, nil`
`33`	`33`	`}`
Original file line number	Diff line number	Diff line change
`@@ -28,6 +28,6 @@ func (action IncrAction) Execute(kvStore store.KvStore, redisConfig *config.Re`
`28`	`28`	`return [][]byte{resp.ResolveResponse(errString, resp.Response_ERRORS)}, errors.New(errString)`
`29`	`29`	`}`
`30`	`30`
`31`		`- (*kvStore).Set(key, newValueString)`
	`31`	`+ (*kvStore).Set(key, newValueString, store.SetOptions{})`
`32`	`32`	`return [][]byte{resp.ResolveResponse(newValueString, resp.Response_INTEGERS)}, nil`
`33`	`33`	`}`
Original file line number	Diff line number	Diff line change
`@@ -11,8 +11,8 @@ func GenerateRandomKey() string {`
`11`	`11`	`}`
`12`	`12`
`13`	`13`	`func ResolvePossibleKeyDirectives(key string) string {`
`14`		`- key = strings.ToLower(key)`
`15`		`- if key == "key:__rand_int__" \|\| key == "__rand_int__" {`
	`14`	`+ normalisedKey := strings.ToLower(key)`
	`15`	`+ if normalisedKey == "key:__rand_int__" \|\| normalisedKey == "__rand_int__" {`
`16`	`16`	`return GenerateRandomKey()`
`17`	`17`	`}`
`18`	`18`	`return key`