From 2764dc3768675912ca21af111fc6b795270707d2 Mon Sep 17 00:00:00 2001 From: Oran Agra Date: Sun, 28 May 2023 10:58:29 +0300 Subject: [PATCH] Optimize MSETNX to avoid double lookup (#11944) This is a redo of #11594 which got reverted in #11940 It improves performance by avoiding double lookup of the the key. --- src/db.c | 34 +++++++++++++++++++++++++--------- src/server.h | 1 + src/t_string.c | 6 +++++- tests/unit/type/string.tcl | 5 +++++ 4 files changed, 36 insertions(+), 10 deletions(-) diff --git a/src/db.c b/src/db.c index 04afc3173..cf5788774 100644 --- a/src/db.c +++ b/src/db.c @@ -47,6 +47,7 @@ int expireIfNeeded(redisDb *db, robj *key, int flags); int keyIsExpired(redisDb *db, robj *key); +static void dbSetValue(redisDb *db, robj *key, robj *val, int overwrite, dictEntry *de); /* Update LFU when an object is accessed. * Firstly, decrement the counter if the decrement time is reached. @@ -187,11 +188,17 @@ robj *lookupKeyWriteOrReply(client *c, robj *key, robj *reply) { /* Add the key to the DB. It's up to the caller to increment the reference * counter of the value if needed. * - * The program is aborted if the key already exists. */ -void dbAdd(redisDb *db, robj *key, robj *val) { - sds copy = sdsdup(key->ptr); - dictEntry *de = dictAddRaw(db->dict, copy, NULL); + * If the update_if_existing argument is false, the the program is aborted + * if the key already exists, otherwise, it can fall back to dbOverwite. */ +static void dbAddInternal(redisDb *db, robj *key, robj *val, int update_if_existing) { + dictEntry *existing; + dictEntry *de = dictAddRaw(db->dict, key->ptr, &existing); + if (update_if_existing && existing) { + dbSetValue(db, key, val, 1, existing); + return; + } serverAssertWithInfo(NULL, key, de != NULL); + dictSetKey(db->dict, de, sdsdup(key->ptr)); initObjectLRUOrLFU(val); dictSetVal(db->dict, de, val); signalKeyAsReady(db, key, val->type); @@ -199,6 +206,10 @@ void dbAdd(redisDb *db, robj *key, robj *val) { notifyKeyspaceEvent(NOTIFY_NEW,"new",key,db->id); } +void dbAdd(redisDb *db, robj *key, robj *val) { + dbAddInternal(db, key, val, 0); +} + /* This is a special version of dbAdd() that is used only when loading * keys from the RDB file: the key is passed as an SDS string that is * retained by the function (and not freed by the caller). @@ -228,10 +239,11 @@ int dbAddRDBLoad(redisDb *db, sds key, robj *val) { * replacement (in which case we need to emit deletion signals), or just an * update of a value of an existing key (when false). * + * The dictEntry input is optional, can be used if we already have one. + * * The program is aborted if the key was not already present. */ -static void dbSetValue(redisDb *db, robj *key, robj *val, int overwrite) { - dictEntry *de = dictFind(db->dict,key->ptr); - +static void dbSetValue(redisDb *db, robj *key, robj *val, int overwrite, dictEntry *de) { + if (!de) de = dictFind(db->dict,key->ptr); serverAssertWithInfo(NULL,key,de != NULL); robj *old = dictGetVal(de); @@ -264,7 +276,7 @@ static void dbSetValue(redisDb *db, robj *key, robj *val, int overwrite) { /* Replace an existing key with a new value, we just replace value and don't * emit any events */ void dbReplaceValue(redisDb *db, robj *key, robj *val) { - dbSetValue(db, key, val, 0); + dbSetValue(db, key, val, 0, NULL); } /* High level Set operation. This function can be used in order to set @@ -285,13 +297,17 @@ void setKey(client *c, redisDb *db, robj *key, robj *val, int flags) { if (flags & SETKEY_ALREADY_EXIST) keyfound = 1; + else if (flags & SETKEY_ADD_OR_UPDATE) + keyfound = -1; else if (!(flags & SETKEY_DOESNT_EXIST)) keyfound = (lookupKeyWrite(db,key) != NULL); if (!keyfound) { dbAdd(db,key,val); + } else if (keyfound<0) { + dbAddInternal(db,key,val,1); } else { - dbSetValue(db,key,val,1); + dbSetValue(db,key,val,1,NULL); } incrRefCount(val); if (!(flags & SETKEY_KEEPTTL)) removeExpire(db,key); diff --git a/src/server.h b/src/server.h index 7bf099499..b0cf27927 100644 --- a/src/server.h +++ b/src/server.h @@ -3247,6 +3247,7 @@ void dbReplaceValue(redisDb *db, robj *key, robj *val); #define SETKEY_NO_SIGNAL 2 #define SETKEY_ALREADY_EXIST 4 #define SETKEY_DOESNT_EXIST 8 +#define SETKEY_ADD_OR_UPDATE 16 /* Key most likely doesn't exists */ void setKey(client *c, redisDb *db, robj *key, robj *val, int flags); robj *dbRandomKey(redisDb *db); int dbGenericDelete(redisDb *db, robj *key, int async, int flags); diff --git a/src/t_string.c b/src/t_string.c index 96bbf606d..02785fabb 100644 --- a/src/t_string.c +++ b/src/t_string.c @@ -576,10 +576,14 @@ void msetGenericCommand(client *c, int nx) { } } + int setkey_flags = nx ? SETKEY_DOESNT_EXIST : 0; for (j = 1; j < c->argc; j += 2) { c->argv[j+1] = tryObjectEncoding(c->argv[j+1]); - setKey(c, c->db, c->argv[j], c->argv[j + 1], 0); + setKey(c, c->db, c->argv[j], c->argv[j + 1], setkey_flags); notifyKeyspaceEvent(NOTIFY_STRING,"set",c->argv[j],c->db->id); + /* In MSETNX, It could be that we're overriding the same key, we can't be sure it doesn't exist. */ + if (nx) + setkey_flags = SETKEY_ADD_OR_UPDATE; } server.dirty += (c->argc-1)/2; addReply(c, nx ? shared.cone : shared.ok); diff --git a/tests/unit/type/string.tcl b/tests/unit/type/string.tcl index 68c360b97..5a8af7b53 100644 --- a/tests/unit/type/string.tcl +++ b/tests/unit/type/string.tcl @@ -234,6 +234,11 @@ start_server {tags {"string"}} { assert_error {*wrong number of arguments for 'msetnx' command} {r msetnx x{t} 20 y{t} "foo bar" z{t}} } + test {MSET with already existing - same key twice} { + r set x{t} x + list [r mset x{t} xxx x{t} yyy] [r get x{t}] + } {OK yyy} + test {MSETNX with already existent key} { list [r msetnx x1{t} xxx y2{t} yyy x{t} 20] [r exists x1{t}] [r exists y2{t}] } {0 0 0}