feat: set support ex & px

command.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/xgzlucario/rotom/internal/dict"
 	"github.com/xgzlucario/rotom/internal/hash"
@@ -13,6 +14,9 @@ import (
 
 var (
 	WITH_SCORES = "WITHSCORES"
+	KEEP_TTL    = "KEEPTTL"
+	EX          = "EX"
+	PX          = "PX"
 )
 
 type Command struct {
@@ -89,7 +93,41 @@ func pingCommand(writer *RESPWriter, _ []RESP) {
 func setCommand(writer *RESPWriter, args []RESP) {
 	key := args[0].ToString()
 	value := args[1].Clone()
-	db.dict.Set(key, value)
+	extra := args[2:]
+	var ttl int64
+
+	for len(extra) > 0 {
+		// EX
+		if equalFold(extra[0].ToStringUnsafe(), EX) && len(extra) >= 2 {
+			n, err := extra[1].ToInt()
+			if err != nil {
+				writer.WriteError(errParseInteger)
+				return
+			}
+			ttl = time.Now().Add(time.Second * time.Duration(n)).UnixNano()
+			extra = extra[2:]
+
+			// PX
+		} else if equalFold(extra[0].ToStringUnsafe(), PX) && len(extra) >= 2 {
+			n, err := extra[1].ToInt()
+			if err != nil {
+				writer.WriteError(errParseInteger)
+				return
+			}
+			ttl = time.Now().Add(time.Millisecond * time.Duration(n)).UnixNano()
+			extra = extra[2:]
+
+			// KEEPTTL
+		} else if equalFold(extra[0].ToStringUnsafe(), KEEP_TTL) {
+			ttl = -1
+
+		} else {
+			writer.WriteError(errSyntax)
+			return
+		}
+	}
+
+	db.dict.SetWithTTL(key, value, ttl)
 	writer.WriteString("OK")
 }
 
@@ -549,7 +587,6 @@ func fetch[T any](key string, new func() T, setnx ...bool) (T, error) {
 				object.SetData(zm.ToSet())
 			}
 		}
-
 		return v, nil
 	}
 

command_test.go

@@ -60,6 +60,32 @@ func TestCommand(t *testing.T) {
 		assert.Equal(n, int64(1))
 	})
 
+	t.Run("setex", func(t *testing.T) {
+		res, _ := rdb.Set(ctx, "foo", "bar", time.Second).Result()
+		assert.Equal(res, "OK")
+
+		res, _ = rdb.Get(ctx, "foo").Result()
+		assert.Equal(res, "bar")
+
+		time.Sleep(time.Second + time.Millisecond)
+
+		_, err := rdb.Get(ctx, "foo").Result()
+		assert.Equal(err, redis.Nil)
+	})
+
+	t.Run("setpx", func(t *testing.T) {
+		res, _ := rdb.Set(ctx, "foo", "bar", time.Millisecond*100).Result()
+		assert.Equal(res, "OK")
+
+		res, _ = rdb.Get(ctx, "foo").Result()
+		assert.Equal(res, "bar")
+
+		time.Sleep(time.Millisecond * 101)
+
+		_, err := rdb.Get(ctx, "foo").Result()
+		assert.Equal(err, redis.Nil)
+	})
+
 	t.Run("pipline", func(t *testing.T) {
 		pip := rdb.Pipeline()
 		pip.RPush(ctx, "ls-pip", "A", "B", "C")

errors.go

@@ -11,5 +11,5 @@ var (
 	errInvalidArguments = errors.New("ERR invalid number of arguments")
 	errUnknownCommand   = errors.New("ERR unknown command")
 	errOOM              = errors.New("ERR command not allowed when out of memory")
-	// errSyntax           = errors.New("ERR syntax error")
+	errSyntax           = errors.New("ERR syntax error")
 )

internal/dict/dict.go

@@ -77,13 +77,16 @@ func (dict *Dict) Set(key string, data any) {
 }
 
 func (dict *Dict) SetWithTTL(key string, data any, ttl int64) {
-	dict.data.Put(key, &Object{
+	object := &Object{
 		typ:         typeOfData(data),
 		lastAccessd: _sec.Load(),
 		data:        data,
-		hasTTL:      true,
-	})
-	dict.expire.Put(key, ttl)
+	}
+	if ttl > 0 {
+		dict.expire.Put(key, ttl)
+		object.hasTTL = true
+	}
+	dict.data.Put(key, object)
 }
 
 func (dict *Dict) Delete(key string) bool {