Repository: valkey-io/valkey
File: src/valkey-benchmark.c (C)

Code before:

    /* Build the request buffer:
     * Queue N requests accordingly to the pipeline size, or simply clone
     * the example client buffer. */
    c->obuf = sdsempty();
    /* Prefix the request buffer with AUTH and/or SELECT commands, if applicable.
     * These commands are discarded after the first response, so if the client is
     * reused the commands will not be used again. */
    c->prefix_pending = 0;
    if (config.conn_info.auth) {
        char *buf = NULL;
        int len;
        if (config.conn_info.user == NULL)
            len = valkeyFormatCommand(&buf, "AUTH %s", config.conn_info.auth);
        else
            len = valkeyFormatCommand(&buf, "AUTH %s %s", config.conn_info.user, config.conn_info.auth);
        c->obuf = sdscatlen(c->obuf, buf, len);
        free(buf);
        c->prefix_pending++;
    }

    if (config.enable_tracking) {
        char *buf = NULL;
        int len = valkeyFormatCommand(&buf, "CLIENT TRACKING on");
        c->obuf = sdscatlen(c->obuf, buf, len);
        free(buf);
        c->prefix_pending++;
    }

    /* If a DB number different than zero is selected, prefix our request
     * buffer with the SELECT command, that will be discarded the first
     * time the replies are received, so if the client is reused the
     * SELECT command will not be used again. */
    if (config.conn_info.input_dbnum) {
        c->obuf = sdscatprintf(c->obuf, "*2\r\n$6\r\nSELECT\r\n$%d\r\n%s\r\n", (int)sdslen(config.input_dbnumstr),
                               config.input_dbnumstr);
        c->prefix_pending++;
    }

    if (config.resp3) {
        char *buf = NULL;
        int len = valkeyFormatCommand(&buf, "HELLO 3");
        c->obuf = sdscatlen(c->obuf, buf, len);
        free(buf);
        c->prefix_pending++;
    }

    if (config.cluster_mode && (config.read_from_replica == FROM_REPLICA_ONLY || config.read_from_replica == FROM_ALL)) {
        char *buf = NULL;
        int len;
        len = valkeyFormatCommand(&buf, "READONLY");

Review comment (suggestion):
I just realized there can be a problem if this feature is combined with

* #2057

We need one token for each command, right? If there are multiple commands in a sequence given on the command line, then the total number of commands in the query buffer sent at once is `config.pipeline * c->seqlen`. Maybe we need to do some follow-up about this...

Code after:
            break;
        }
        // When client acquires a token, try to write with `reuse`.
        c->paused = 0;
        c->reuse = 1;
        writeHandler(eventLoop, c->context->fd, c, AE_WRITABLE);
        listDelNode(paused_clients, ln);
    }

    // If there are no more paused clients, remove the event.
    if (delay == 0) {
        return AE_NOMORE;
    }
    return delay;
}

static void writeHandler(aeEventLoop *el, int fd, void *privdata, int mask) {
    client c = privdata;
    UNUSED(el);
    UNUSED(fd);
    UNUSED(mask);

    // When benchmark with rps control, and client is not reuse, try to acquire a token.
    if (config.rps > 0 && c->reuse == 0) {
        /* Acquire a token from the token bucket. */
        long long delay = acquireTokenOrWait(config.pipeline);

        if (delay) {
            int thread_id = c->thread_id;
            int paused_clients_count = 0;

            c->paused = 1;
            aeDeleteFileEvent(el, c->context->fd, AE_WRITABLE);

            benchmarkThread *thread = NULL;
            if (thread_id < 0) {
                paused_clients_count = listLength(config.paused_clients);
                listAddNodeTail(config.paused_clients, c);
            } else {
                thread = config.threads[thread_id];
                paused_clients_count = listLength(thread->paused_clients);
                listAddNodeTail(thread->paused_clients, c);
            }
            if (paused_clients_count == 0) {
                /* Create a time event to awaken the client. */
                aeCreateTimeEvent(el, delay, awakenPausedClient, (void *)thread, NULL);
            }
            return;
        }
    }
    c->reuse = 0;