From b7511c63debff02f55241dc583ca6b66d8d34ec5 Mon Sep 17 00:00:00 2001
From: guangzhixu <guangzhixu@deepglint.com>
Date: Fri, 30 Aug 2024 12:17:22 +0800
Subject: [PATCH] feat: set command support NX

---
 command.go      | 18 +++++++++++++++---
 command_test.go |  9 +++++++++
 rotom.go        |  1 -
 3 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/command.go b/command.go
index 6b4f7b7..3044635 100644
--- a/command.go
+++ b/command.go
@@ -15,6 +15,7 @@ import (
 var (
 	WITH_SCORES = "WITHSCORES"
 	KEEP_TTL    = "KEEPTTL"
+	NX          = "NX"
 	EX          = "EX"
 	PX          = "PX"
 )
@@ -97,8 +98,10 @@ func setCommand(writer *RESPWriter, args []RESP) {
 	var ttl int64
 
 	for len(extra) > 0 {
+		arg := extra[0].ToStringUnsafe()
+
 		// EX
-		if equalFold(extra[0].ToStringUnsafe(), EX) && len(extra) >= 2 {
+		if equalFold(arg, EX) && len(extra) >= 2 {
 			n, err := extra[1].ToInt()
 			if err != nil {
 				writer.WriteError(errParseInteger)
@@ -108,7 +111,7 @@ func setCommand(writer *RESPWriter, args []RESP) {
 			extra = extra[2:]
 
 			// PX
-		} else if equalFold(extra[0].ToStringUnsafe(), PX) && len(extra) >= 2 {
+		} else if equalFold(arg, PX) && len(extra) >= 2 {
 			n, err := extra[1].ToInt()
 			if err != nil {
 				writer.WriteError(errParseInteger)
@@ -118,9 +121,18 @@ func setCommand(writer *RESPWriter, args []RESP) {
 			extra = extra[2:]
 
 			// KEEPTTL
-		} else if equalFold(extra[0].ToStringUnsafe(), KEEP_TTL) {
+		} else if equalFold(arg, KEEP_TTL) {
+			extra = extra[1:]
 			ttl = -1
 
+			// NX
+		} else if equalFold(arg, NX) {
+			if _, ttl := db.dict.Get(key); ttl != dict.KEY_NOT_EXIST {
+				writer.WriteNull()
+				return
+			}
+			extra = extra[1:]
+
 		} else {
 			writer.WriteError(errSyntax)
 			return
diff --git a/command_test.go b/command_test.go
index 02a5c89..c0974e4 100644
--- a/command_test.go
+++ b/command_test.go
@@ -58,6 +58,15 @@ func TestCommand(t *testing.T) {
 
 		n, _ := rdb.Del(ctx, "foo", "none").Result()
 		assert.Equal(n, int64(1))
+
+		// setnx
+		ok, err := rdb.SetNX(ctx, "key-nx", "123", redis.KeepTTL).Result()
+		assert.Nil(err)
+		assert.True(ok)
+
+		ok, err = rdb.SetNX(ctx, "key-nx", "123", redis.KeepTTL).Result()
+		assert.Nil(err)
+		assert.False(ok)
 	})
 
 	t.Run("error-get", func(t *testing.T) {
diff --git a/rotom.go b/rotom.go
index cf1d00e..26d77c2 100644
--- a/rotom.go
+++ b/rotom.go
@@ -119,7 +119,6 @@ READ:
 	client.recvx += n
 
 	if readSize == 0 {
-		log.Info().Msgf("client %d read query empty, now free", fd)
 		freeClient(client)
 		return
 	}