diff --git a/README.md b/README.md index c08dd13..fa61867 100644 --- a/README.md +++ b/README.md @@ -66,6 +66,18 @@ Note: the key shouldn't be repeated for the executed command. **Reply:** Null if not equal or for non existing key when the `XX` flag is used. On success, the reply depends on the actual command executed. +## `CHOP key count` + +Removes characters from a value of a String key. + +* If 'count' is positive, it removes from the right of the string; +* if 'count' is negative, it removes from the left of the string. +* If |count| is greater than the string length, the value is set as an empty string. +* If 'count' is zero, then the value remains unchanged. +* It is an error if the value is not a String. + +**Reply:** Integer, the length of the string after the chop operation. + ## `PREPEND key value` > Time complexity: O(1). The amortized time complexity is O(1) assuming the prepended value is small and the already present value is of any size, since the dynamic string library used by Redis will double the free space available on every reallocation. @@ -102,6 +114,15 @@ Credit: Meni Katz **Reply:** Integer, the length of the String after it was modified. +## `STRPOP key` + +Gets the value at `key`, deletes `key`, then returns the value. +It is equivalent to an atomic [GET](https://redis.io/commands/get) and [DEL](https://redis.io/commands/del). + + * An error is returned if the key exists and does not hold a string. + +**Reply:** String, the value at `key` or NULL if `key` didn't exist. + # rxhashes This module provides extended Redis Hashes commands. diff --git a/src/rxstrings.c b/src/rxstrings.c index 1820a9d..2951f9e 100644 --- a/src/rxstrings.c +++ b/src/rxstrings.c @@ -135,6 +135,76 @@ int CheckAndCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { return REDISMODULE_OK; } +/* +* CHOP key count +* Removes characters from a value of a String key. +* If 'count' is positive, it removes from the right of the string; +* if 'count' is negative, it removes from the left of the string. +* If |count| is greater than the string length, the value is set as an empty string. +* If 'count' is zero, then the value remains unchanged. +* It is an error if the value is not a String. +* Integer Reply: the length of the string after the chop operation. +*/ +int ChopCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { + if (argc != 3) { + return RedisModule_WrongArity(ctx); + } + RedisModule_AutoMemory(ctx); + + /* Obtain key and extract offset argument. */ + RedisModuleKey *key = + RedisModule_OpenKey(ctx, argv[1], REDISMODULE_READ | REDISMODULE_WRITE); + long long count; + if (RedisModule_StringToLongLong(argv[2], &count) != REDISMODULE_OK) { + RedisModule_ReplyWithError(ctx, "ERR invalid count"); + return REDISMODULE_ERR; + } + + /* Key must be a string. */ + if (RedisModule_KeyType(key) != REDISMODULE_KEYTYPE_STRING) { + RedisModule_ReplyWithError(ctx, REDISMODULE_ERRORMSG_WRONGTYPE); + return REDISMODULE_ERR; + } + + /* Calculate new length, protecting against overchop */ + size_t valLen = RedisModule_ValueLength(key); + if (count == 0) { + RedisModule_ReplyWithLongLong(ctx, valLen); + return REDISMODULE_OK; + } + size_t absCount = (count < 0) ? -count : count; + size_t newLen = (valLen < absCount) ? 0 : (valLen - absCount); + + /* Work around https://github.com/antirez/redis/issues/3717 + * We use RM_StringSet to an empty string instead */ + if (newLen == 0) { + RedisModuleString* emptyStr = RedisModule_CreateString(ctx, "", 0); + if (RedisModule_StringSet(key, emptyStr) != REDISMODULE_OK) { + RedisModule_ReplyWithError(ctx, "ERR RM_SetString empty failed"); + return REDISMODULE_ERR; + } + RedisModule_ReplyWithLongLong(ctx, 0); + return REDISMODULE_OK; + } + + /* If we chop from the left, we need to shift the characters, + * but only bother if we aren't zeroing the value */ + if (count < 0 && newLen != 0) { + size_t dmaLen; + char* valPtr = RedisModule_StringDMA(key, &dmaLen, REDISMODULE_READ|REDISMODULE_WRITE); + memmove(valPtr, valPtr + absCount, newLen); + } + + /* Truncate to new length */ + if (RedisModule_StringTruncate(key, newLen) != REDISMODULE_OK) { + RedisModule_ReplyWithError(ctx, "ERR RM_StringTruncate failed"); + return REDISMODULE_ERR; + } + + RedisModule_ReplyWithLongLong(ctx, newLen); + return REDISMODULE_OK; +} + /* * PREPEND key value * Prepends a value to a String key. @@ -167,7 +237,7 @@ int PrependCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { /* Otherwise key must be a string. */ if (RedisModule_KeyType(key) != REDISMODULE_KEYTYPE_STRING) { RedisModule_ReplyWithError(ctx, REDISMODULE_ERRORMSG_WRONGTYPE); - return REDISMODULE_OK; + return REDISMODULE_ERR; } /* Prepend the string: 1) expand string, 2) shift oldVal via memmove, 3) prepend arg */ @@ -393,6 +463,49 @@ int SetRangeRandCommand(RedisModuleCtx *ctx, RedisModuleString **argv, return REDISMODULE_OK; } +/* +* STRPOP key +* Gets the value at key, deletes the key, then return the value. +* Equivalent to an atomic GET and DEL. +* An error is returned if there is a value stored that is not a string. +* String Reply: the value of key, or nil if key does not exist. +*/ +int StrPopCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { + if (argc != 2) { + return RedisModule_WrongArity(ctx); + } + RedisModule_AutoMemory(ctx); + + /* Obtain key */ + RedisModuleKey *key = + RedisModule_OpenKey(ctx, argv[1], REDISMODULE_READ | REDISMODULE_WRITE); + + if (RedisModule_KeyType(key) == REDISMODULE_KEYTYPE_EMPTY) { + RedisModule_ReplyWithNull(ctx); + return REDISMODULE_OK; + } + + /* Key must be a string */ + if (RedisModule_KeyType(key) != REDISMODULE_KEYTYPE_STRING) { + RedisModule_ReplyWithError(ctx, REDISMODULE_ERRORMSG_WRONGTYPE); + return REDISMODULE_ERR; + } + + /* GET the string */ + size_t dmaLen; + char* valPtr = RedisModule_StringDMA(key, &dmaLen, REDISMODULE_READ); + RedisModuleString* valStr = RedisModule_CreateString(ctx, valPtr, dmaLen); + + /* DELete the key */ + if (RedisModule_DeleteKey(key) != REDISMODULE_OK) { + RedisModule_ReplyWithError(ctx, "ERR DeleteKey failed"); + return REDISMODULE_ERR; + } + + RedisModule_ReplyWithString(ctx, valStr); + return REDISMODULE_OK; +} + int testCheckAnd(RedisModuleCtx *ctx) { RedisModuleCallReply *r; @@ -409,6 +522,52 @@ int testCheckAnd(RedisModuleCtx *ctx) { return 0; } +int testChop(RedisModuleCtx *ctx) { + RedisModuleCallReply *r; + + r = RedisModule_Call(ctx, "set", "cc", "foo", "abcde"); + RMUtil_AssertReplyEquals(r,"OK"); + r = RedisModule_Call(ctx, "chop", "cc", "foo", "-2"); + RMUtil_Assert(RedisModule_CallReplyInteger(r) == 3); + r = RedisModule_Call(ctx, "get", "c", "foo"); + RMUtil_AssertReplyEquals(r,"cde"); + r = RedisModule_Call(ctx, "FLUSHALL", ""); + + r = RedisModule_Call(ctx, "set", "cc", "foo", "abcde"); + RMUtil_AssertReplyEquals(r,"OK"); + r = RedisModule_Call(ctx, "chop", "cc", "foo", "2"); + RMUtil_Assert(RedisModule_CallReplyInteger(r) == 3); + r = RedisModule_Call(ctx, "get", "c", "foo"); + RMUtil_AssertReplyEquals(r,"abc"); + r = RedisModule_Call(ctx, "FLUSHALL", ""); + + r = RedisModule_Call(ctx, "set", "cc", "foo", "abcde"); + RMUtil_AssertReplyEquals(r,"OK"); + r = RedisModule_Call(ctx, "chop", "cc", "foo", "-10"); + RMUtil_Assert(RedisModule_CallReplyInteger(r) == 0); + r = RedisModule_Call(ctx, "get", "c", "foo"); + RMUtil_AssertReplyEquals(r,""); + r = RedisModule_Call(ctx, "FLUSHALL", ""); + + r = RedisModule_Call(ctx, "set", "cc", "foo", "abcde"); + RMUtil_AssertReplyEquals(r,"OK"); + r = RedisModule_Call(ctx, "chop", "cc", "foo", "10"); + RMUtil_Assert(RedisModule_CallReplyInteger(r) == 0); + r = RedisModule_Call(ctx, "get", "c", "foo"); + RMUtil_AssertReplyEquals(r,""); + r = RedisModule_Call(ctx, "FLUSHALL", ""); + + r = RedisModule_Call(ctx, "set", "cc", "foo", "abcde"); + RMUtil_AssertReplyEquals(r,"OK"); + r = RedisModule_Call(ctx, "chop", "cc", "foo", "0"); + RMUtil_Assert(RedisModule_CallReplyInteger(r) == 5); + r = RedisModule_Call(ctx, "get", "c", "foo"); + RMUtil_AssertReplyEquals(r,"abcde"); + r = RedisModule_Call(ctx, "FLUSHALL", ""); + + return 0; +} + int testPrepend(RedisModuleCtx *ctx) { RedisModuleCallReply *r; @@ -433,6 +592,22 @@ int testSetRangeRand(RedisModuleCtx *ctx) { return 0; } +int testStrPop(RedisModuleCtx *ctx) { + RedisModuleCallReply *r; + + r = RedisModule_Call(ctx, "set", "cc", "foo", "abcdef"); + RMUtil_AssertReplyEquals(r,"OK"); + r = RedisModule_Call(ctx, "strpop", "c", "foo"); + RMUtil_AssertReplyEquals(r,"abcdef"); + r = RedisModule_Call(ctx, "exists", "c", "foo"); + RMUtil_Assert(RedisModule_CallReplyInteger(r) == 0); + r = RedisModule_Call(ctx, "strpop", "c", "bar"); + RMUtil_Assert(RedisModule_CallReplyType(r) == REDISMODULE_REPLY_NULL); + r = RedisModule_Call(ctx, "FLUSHALL", ""); + + return 0; +} + int TestModule(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { RedisModule_AutoMemory(ctx); @@ -445,8 +620,10 @@ int TestModule(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { } RMUtil_Test(testCheckAnd); + RMUtil_Test(testChop); RMUtil_Test(testPrepend); RMUtil_Test(testSetRangeRand); + RMUtil_Test(testStrPop); RedisModule_ReplyWithSimpleString(ctx, "PASS"); return REDISMODULE_OK; @@ -460,6 +637,10 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx) { if (RedisModule_CreateCommand(ctx, "checkand", CheckAndCommand, "write deny-oom", 1, 1, 1) == REDISMODULE_ERR) return REDISMODULE_ERR; + if (RedisModule_CreateCommand(ctx, "chop", ChopCommand, + "write fast deny-oom", 1, 1, + 1) == REDISMODULE_ERR) + return REDISMODULE_ERR; if (RedisModule_CreateCommand(ctx, "prepend", PrependCommand, "write fast deny-oom", 1, 1, 1) == REDISMODULE_ERR) @@ -468,9 +649,13 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx) { "write fast deny-oom", 1, 1, 1) == REDISMODULE_ERR) return REDISMODULE_ERR; + if (RedisModule_CreateCommand(ctx, "strpop", StrPopCommand, + "fast deny-oom", 1, 1, + 1) == REDISMODULE_ERR) + return REDISMODULE_ERR; if (RedisModule_CreateCommand(ctx, "rxstrings.test", TestModule, "write", 0, 0, 0) == REDISMODULE_ERR) return REDISMODULE_ERR; return REDISMODULE_OK; -} \ No newline at end of file +}