Skip to content

Commit

Permalink
Move prepareClientToWrite out of loop for HGETALL command (valkey-io#…
Browse files Browse the repository at this point in the history 
…1119)

Similar to valkey-io#860 but this is for HGETALL families (HGETALL/HKEYS/HVALS).
This patch moves `prepareClientToWrite` out of the loop to reduce the
function overhead.

Signed-off-by: Masahiro Ide <[email protected]>
Co-authored-by: Madelyn Olson <[email protected]>
  • Loading branch information
@imasahiro @madolson @eifrah-aws
2 people authored and eifrah-aws committed Oct 20, 2024
1 parent c86cdd6 commit b37f749
Show file tree
Hide file tree
Showing 3 changed files with 59 additions and 35 deletions.
25 changes: 21 additions & 4 deletions src/networking.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1043,7 +1043,8 @@ void addReplyArrayLen(client *c, long length) {
addReplyAggregateLen(c, length, '*');
}

void addWritePreparedReplyArrayLen(writePreparedClient *c, long length) {
void addWritePreparedReplyArrayLen(writePreparedClient *wpc, long length) {
client *c = (client *)wpc;
serverAssert(length >= 0);
_addReplyLongLongWithPrefix(c, length, '*');
}
Expand All @@ -1054,6 +1055,13 @@ void addReplyMapLen(client *c, long length) {
addReplyAggregateLen(c, length, prefix);
}

void addWritePreparedReplyMapLen(writePreparedClient *wpc, long length) {
client *c = (client *)wpc;
int prefix = c->resp == 2 ? '*' : '%';
if (c->resp == 2) length *= 2;
_addReplyLongLongWithPrefix(c, length, prefix);
}

void addReplySetLen(client *c, long length) {
int prefix = c->resp == 2 ? '*' : '~';
addReplyAggregateLen(c, length, prefix);
Expand Down Expand Up @@ -1120,7 +1128,8 @@ void addReplyBulkCBuffer(client *c, const void *p, size_t len) {
_addReplyToBufferOrList(c, "\r\n", 2);
}

void addWritePreparedReplyBulkCBuffer(writePreparedClient *c, const void *p, size_t len) {
void addWritePreparedReplyBulkCBuffer(writePreparedClient *wpc, const void *p, size_t len) {
client *c = (client *)wpc;
_addReplyLongLongWithPrefix(c, len, '$');
_addReplyToBufferOrList(c, p, len);
_addReplyToBufferOrList(c, "\r\n", 2);
Expand All @@ -1138,6 +1147,14 @@ void addReplyBulkSds(client *c, sds s) {
_addReplyToBufferOrList(c, "\r\n", 2);
}

void addWritePreparedReplyBulkSds(writePreparedClient *wpc, sds s) {
client *c = (client *)wpc;
_addReplyLongLongWithPrefix(c, sdslen(s), '$');
_addReplyToBufferOrList(c, s, sdslen(s));
sdsfree(s);
_addReplyToBufferOrList(c, "\r\n", 2);
}

/* Set sds to a deferred reply (for symmetry with addReplyBulkSds it also frees the sds) */
void setDeferredReplyBulkSds(client *c, void *node, sds s) {
sds reply = sdscatprintf(sdsempty(), "$%d\r\n%s\r\n", (unsigned)sdslen(s), s);
Expand All @@ -1164,12 +1181,12 @@ void addReplyBulkLongLong(client *c, long long ll) {
addReplyBulkCBuffer(c, buf, len);
}

void addWritePreparedReplyBulkLongLong(writePreparedClient *c, long long ll) {
void addWritePreparedReplyBulkLongLong(writePreparedClient *wpc, long long ll) {
char buf[64];
int len;

len = ll2string(buf, 64, ll);
addWritePreparedReplyBulkCBuffer(c, buf, len);
addWritePreparedReplyBulkCBuffer(wpc, buf, len);
}

/* Reply with a verbatim type having the specified extension.
Expand Down
4 changes: 3 additions & 1 deletion src/server.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1379,7 +1379,7 @@ typedef struct client {
* skip certain types of initialization. This type is used to indicate a client that has been initialized
* and can be used with addWritePreparedReply* functions. A client can be cast into this type with
* prepareClientForFutureWrites(client *c). */
typedef client writePreparedClient;
typedef struct writePreparedClient writePreparedClient;

/* ACL information */
typedef struct aclInfo {
Expand Down Expand Up @@ -2789,6 +2789,7 @@ void addReply(client *c, robj *obj);
void addReplyStatusLength(client *c, const char *s, size_t len);
void addReplySds(client *c, sds s);
void addReplyBulkSds(client *c, sds s);
void addWritePreparedReplyBulkSds(writePreparedClient *c, sds s);
void setDeferredReplyBulkSds(client *c, void *node, sds s);
void addReplyErrorObject(client *c, robj *err);
void addReplyOrErrorObject(client *c, robj *reply);
Expand All @@ -2808,6 +2809,7 @@ void addReplyLongLong(client *c, long long ll);
void addReplyArrayLen(client *c, long length);
void addWritePreparedReplyArrayLen(writePreparedClient *c, long length);
void addReplyMapLen(client *c, long length);
void addWritePreparedReplyMapLen(writePreparedClient *c, long length);
void addReplySetLen(client *c, long length);
void addReplyAttributeLen(client *c, long length);
void addReplyPushLen(client *c, long length);
Expand Down
65 changes: 35 additions & 30 deletions src/t_hash.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -788,20 +788,20 @@ void hstrlenCommand(client *c) {
addReplyLongLong(c, hashTypeGetValueLength(o, c->argv[2]->ptr));
}

static void addHashIteratorCursorToReply(client *c, hashTypeIterator *hi, int what) {
static void addHashIteratorCursorToReply(writePreparedClient *wpc, hashTypeIterator *hi, int what) {
if (hi->encoding == OBJ_ENCODING_LISTPACK) {
unsigned char *vstr = NULL;
unsigned int vlen = UINT_MAX;
long long vll = LLONG_MAX;

hashTypeCurrentFromListpack(hi, what, &vstr, &vlen, &vll);
if (vstr)
addReplyBulkCBuffer(c, vstr, vlen);
addWritePreparedReplyBulkCBuffer(wpc, vstr, vlen);
else
addReplyBulkLongLong(c, vll);
addWritePreparedReplyBulkLongLong(wpc, vll);
} else if (hi->encoding == OBJ_ENCODING_HT) {
sds value = hashTypeCurrentFromHashTable(hi, what);
addReplyBulkCBuffer(c, value, sdslen(value));
addWritePreparedReplyBulkCBuffer(wpc, value, sdslen(value));
} else {
serverPanic("Unknown hash encoding");
}
Expand All @@ -815,23 +815,25 @@ void genericHgetallCommand(client *c, int flags) {
robj *emptyResp = (flags & OBJ_HASH_KEY && flags & OBJ_HASH_VALUE) ? shared.emptymap[c->resp] : shared.emptyarray;
if ((o = lookupKeyReadOrReply(c, c->argv[1], emptyResp)) == NULL || checkType(c, o, OBJ_HASH)) return;

writePreparedClient *wpc = prepareClientForFutureWrites(c);
if (!wpc) return;
/* We return a map if the user requested keys and values, like in the
* HGETALL case. Otherwise to use a flat array makes more sense. */
length = hashTypeLength(o);
if (flags & OBJ_HASH_KEY && flags & OBJ_HASH_VALUE) {
addReplyMapLen(c, length);
addWritePreparedReplyMapLen(wpc, length);
} else {
addReplyArrayLen(c, length);
addWritePreparedReplyArrayLen(wpc, length);
}

hashTypeInitIterator(o, &hi);
while (hashTypeNext(&hi) != C_ERR) {
if (flags & OBJ_HASH_KEY) {
addHashIteratorCursorToReply(c, &hi, OBJ_HASH_KEY);
addHashIteratorCursorToReply(wpc, &hi, OBJ_HASH_KEY);
count++;
}
if (flags & OBJ_HASH_VALUE) {
addHashIteratorCursorToReply(c, &hi, OBJ_HASH_VALUE);
addHashIteratorCursorToReply(wpc, &hi, OBJ_HASH_VALUE);
count++;
}
}
Expand Down Expand Up @@ -871,18 +873,19 @@ void hscanCommand(client *c) {
scanGenericCommand(c, o, cursor);
}

static void hrandfieldReplyWithListpack(client *c, unsigned int count, listpackEntry *keys, listpackEntry *vals) {
static void hrandfieldReplyWithListpack(writePreparedClient *wpc, unsigned int count, listpackEntry *keys, listpackEntry *vals) {
client *c = (client *)wpc;
for (unsigned long i = 0; i < count; i++) {
if (vals && c->resp > 2) addReplyArrayLen(c, 2);
if (vals && c->resp > 2) addWritePreparedReplyArrayLen(wpc, 2);
if (keys[i].sval)
addReplyBulkCBuffer(c, keys[i].sval, keys[i].slen);
addWritePreparedReplyBulkCBuffer(wpc, keys[i].sval, keys[i].slen);
else
addReplyBulkLongLong(c, keys[i].lval);
addWritePreparedReplyBulkLongLong(wpc, keys[i].lval);
if (vals) {
if (vals[i].sval)
addReplyBulkCBuffer(c, vals[i].sval, vals[i].slen);
addWritePreparedReplyBulkCBuffer(wpc, vals[i].sval, vals[i].slen);
else
addReplyBulkLongLong(c, vals[i].lval);
addWritePreparedReplyBulkLongLong(wpc, vals[i].lval);
}
}
}
Expand Down Expand Up @@ -918,25 +921,27 @@ void hrandfieldWithCountCommand(client *c, long l, int withvalues) {
return;
}

writePreparedClient *wpc = prepareClientForFutureWrites(c);
if (!wpc) return;
/* CASE 1: The count was negative, so the extraction method is just:
* "return N random elements" sampling the whole set every time.
* This case is trivial and can be served without auxiliary data
* structures. This case is the only one that also needs to return the
* elements in random order. */
if (!uniq || count == 1) {
if (withvalues && c->resp == 2)
addReplyArrayLen(c, count * 2);
addWritePreparedReplyArrayLen(wpc, count * 2);
else
addReplyArrayLen(c, count);
addWritePreparedReplyArrayLen(wpc, count);
if (hash->encoding == OBJ_ENCODING_HT) {
sds key, value;
while (count--) {
dictEntry *de = dictGetFairRandomKey(hash->ptr);
key = dictGetKey(de);
value = dictGetVal(de);
if (withvalues && c->resp > 2) addReplyArrayLen(c, 2);
addReplyBulkCBuffer(c, key, sdslen(key));
if (withvalues) addReplyBulkCBuffer(c, value, sdslen(value));
if (withvalues && c->resp > 2) addWritePreparedReplyArrayLen(wpc, 2);
addWritePreparedReplyBulkCBuffer(wpc, key, sdslen(key));
if (withvalues) addWritePreparedReplyBulkCBuffer(wpc, value, sdslen(value));
if (c->flag.close_asap) break;
}
} else if (hash->encoding == OBJ_ENCODING_LISTPACK) {
Expand All @@ -950,7 +955,7 @@ void hrandfieldWithCountCommand(client *c, long l, int withvalues) {
sample_count = count > limit ? limit : count;
count -= sample_count;
lpRandomPairs(hash->ptr, sample_count, keys, vals);
hrandfieldReplyWithListpack(c, sample_count, keys, vals);
hrandfieldReplyWithListpack(wpc, sample_count, keys, vals);
if (c->flag.close_asap) break;
}
zfree(keys);
Expand All @@ -962,9 +967,9 @@ void hrandfieldWithCountCommand(client *c, long l, int withvalues) {
/* Initiate reply count, RESP3 responds with nested array, RESP2 with flat one. */
long reply_size = count < size ? count : size;
if (withvalues && c->resp == 2)
addReplyArrayLen(c, reply_size * 2);
addWritePreparedReplyArrayLen(wpc, reply_size * 2);
else
addReplyArrayLen(c, reply_size);
addWritePreparedReplyArrayLen(wpc, reply_size);

/* CASE 2:
* The number of requested elements is greater than the number of
Expand All @@ -973,9 +978,9 @@ void hrandfieldWithCountCommand(client *c, long l, int withvalues) {
hashTypeIterator hi;
hashTypeInitIterator(hash, &hi);
while (hashTypeNext(&hi) != C_ERR) {
if (withvalues && c->resp > 2) addReplyArrayLen(c, 2);
addHashIteratorCursorToReply(c, &hi, OBJ_HASH_KEY);
if (withvalues) addHashIteratorCursorToReply(c, &hi, OBJ_HASH_VALUE);
if (withvalues && c->resp > 2) addWritePreparedReplyArrayLen(wpc, 2);
addHashIteratorCursorToReply(wpc, &hi, OBJ_HASH_KEY);
if (withvalues) addHashIteratorCursorToReply(wpc, &hi, OBJ_HASH_VALUE);
}
hashTypeResetIterator(&hi);
return;
Expand All @@ -994,7 +999,7 @@ void hrandfieldWithCountCommand(client *c, long l, int withvalues) {
keys = zmalloc(sizeof(listpackEntry) * count);
if (withvalues) vals = zmalloc(sizeof(listpackEntry) * count);
serverAssert(lpRandomPairsUnique(hash->ptr, count, keys, vals) == count);
hrandfieldReplyWithListpack(c, count, keys, vals);
hrandfieldReplyWithListpack(wpc, count, keys, vals);
zfree(keys);
zfree(vals);
return;
Expand Down Expand Up @@ -1048,9 +1053,9 @@ void hrandfieldWithCountCommand(client *c, long l, int withvalues) {
while ((de = dictNext(di)) != NULL) {
sds key = dictGetKey(de);
sds value = dictGetVal(de);
if (withvalues && c->resp > 2) addReplyArrayLen(c, 2);
addReplyBulkSds(c, key);
if (withvalues) addReplyBulkSds(c, value);
if (withvalues && c->resp > 2) addWritePreparedReplyArrayLen(wpc, 2);
addWritePreparedReplyBulkSds(wpc, key);
if (withvalues) addWritePreparedReplyBulkSds(wpc, value);
}

dictReleaseIterator(di);
Expand Down Expand Up @@ -1081,7 +1086,7 @@ void hrandfieldWithCountCommand(client *c, long l, int withvalues) {
added++;

/* We can reply right away, so that we don't need to store the value in the dict. */
if (withvalues && c->resp > 2) addReplyArrayLen(c, 2);
if (withvalues && c->resp > 2) addWritePreparedReplyArrayLen(wpc, 2);
hashReplyFromListpackEntry(c, &key);
if (withvalues) hashReplyFromListpackEntry(c, &value);
}
Expand Down

0 comments on commit b37f749

Please sign in to comment.