src/memory_prefetch.c - unstable - githubmirror/redis

src/memory_prefetch.c
526 строк23 KB
Filipe Oliveira (Redis)
Batched MGET/MSET dict prefetch with dictType-driven payload hints (#15133)
11 май 2026, 05:22
11 май 2026, 05:2262551a7
/*
 * This file utilizes prefetching keys and data for multiple commands in a batch,
 * to improve performance by amortizing memory access costs across multiple operations.
 *
 * Copyright (c) 2025-Present, Redis Ltd. and contributors.
 * All rights reserved.
 *
 * Copyright (c) 2024-present, Valkey contributors.
 * All rights reserved.
 *
 * Licensed under your choice of (a) the Redis Source Available License 2.0
 * (RSALv2); or (b) the Server Side Public License v1 (SSPLv1); or (c) the
 * GNU Affero General Public License v3 (AGPLv3).
 *
 * Portions of this file are available under BSD3 terms; see REDISCONTRIBUTIONS for more information.
 */

#include "memory_prefetch.h"
#include "server.h"
#include "dict.h"

/* --------------------------------------------------------------------------
 * Dict prefetching state machine
 * -------------------------------------------------------------------------- */

typedef enum { HT_IDX_FIRST = 0, HT_IDX_SECOND = 1, HT_IDX_INVALID = -1 } dictHtIdx;

typedef enum {
    PREFETCH_BUCKET,        /* Initial state, determines which hash table to use and prefetch the table's bucket */
    PREFETCH_ENTRY,         /* prefetch entries associated with the given key's hash */
    PREFETCH_ENTRY_KEY,     /* dictType-driven prefetch of the entry's key payload (for keyCompare) */
    PREFETCH_ENTRY_VALUE,   /* compare keys; on match, dictType-driven prefetch of the value payload */
    PREFETCH_DONE           /* Indicates that prefetching for this key is complete */
} dictPrefetchState;

/* Per-key state of an in-flight, software-pipelined dictFind, advanced one
 * stage at a time by dictPrefetcher (see below). The non-state fields mirror
 * the locals that a synchronous dictFind would otherwise carry across one
 * bucket walk. */
typedef struct dictPrefetchLookup {
    dictPrefetchState state;  /* Current FSM stage of this lookup */
    dictHtIdx ht_idx;         /* Index of the current hash table (0 or 1 for rehashing) */
    uint64_t bucket_idx;      /* Index of the bucket in the current hash table */
    uint64_t key_hash;        /* Hash value of the key being looked up */
    dictEntry *current_entry; /* Pointer to the current entry being processed */
} dictPrefetchLookup;

/* dictPrefetcher drives a batch of dictPrefetchLookup objects through the
 * prefetch FSM, yielding to the next in-flight lookup each time a prefetch
 * is issued — so one lookup's memory stall overlaps another's work. The
 * state machine itself is fully dict-pure: any key/value payload prefetching
 * is delegated to the dictType->prefetchEntryKey / prefetchEntryValue
 * callbacks of each key's dict. The same prefetcher is used by both the
 * cross-command batch path and the intra-command dictPrefetchKeys() API. */
typedef struct dictPrefetcher {
    size_t cur_idx;              /* Cursor; advances on each prefetch issue */
    size_t nkeys;                /* Total key lookups in this batch */
    size_t remaining;            /* Number of in-flight key lookups (not yet PREFETCH_DONE) */
    void **keys;                 /* Array of key pointers (sds) */
    dict **dicts;                /* Per-key dictionary pointers */
    dictPrefetchLookup *lookups; /* Per-key lookup state, capacity == max_keys */
    size_t max_keys;             /* Capacity of lookups[] */
} dictPrefetcher;

/******************************** State machine diagram for the dict prefetch operation. ******************************
                                                           │
                                                         start
                                                           │
                                                  ┌────────▼─────────┐
                                       ┌─────────►│  PREFETCH_BUCKET ├────►────────┐
                                       │          └────────┬─────────┘            no more tables -> done
                                       |             bucket|found                  |
                                       │                   |                       │
        entry not found - goto next table         ┌────────▼────────┐              │
                                       └────◄─────┤ PREFETCH_ENTRY  |              ▼
                                    ┌────────────►└────────┬────────┘              │
                                    |                 Entry│found                  │
                                    │                      |                       │
                                    |          ┌───────────▼─────────────┐         │
                                    │          |   PREFETCH_ENTRY_KEY    |         ▼
                                    │          └───────────┬─────────────┘         │
        key mismatch - goto next entry                     |                       |
                                    │          ┌───────────▼─────────────┐         │
                                    └──────◄───│   PREFETCH_ENTRY_VALUE  │         ▼
                                               └───────────┬─────────────┘         │
                                                           |                       │
                                                 ┌───────-─▼─────────────┐         │
                                                 │     PREFETCH_DONE     │◄────────┘
                                                 └───────────────────────┘

**********************************************************************************************************************/

/* Issue a software prefetch for `addr`, then yield to the next lookup by
 * advancing the cursor. */
static inline void dictPrefetchAdvance(dictPrefetcher *p, void *addr) {
    redis_prefetch_read(addr);
    if (++p->cur_idx >= p->nkeys) p->cur_idx = 0;
}

static inline void dictPrefetchMarkDone(dictPrefetcher *p, dictPrefetchLookup *lk) {
    lk->state = PREFETCH_DONE;
    p->remaining--;
    server.stat_total_prefetch_entries++;
}

/* Return the next in-flight lookup that still needs work, or NULL if all done. */
static inline dictPrefetchLookup *dictPrefetchNextInFlight(dictPrefetcher *p) {
    if (p->remaining == 0) return NULL;
    while (p->lookups[p->cur_idx].state == PREFETCH_DONE) {
        if (++p->cur_idx >= p->nkeys) p->cur_idx = 0;
    }
    return &p->lookups[p->cur_idx];
}

/* Prefetch the bucket of the next hash table index.
 * If no tables are left, move to the PREFETCH_DONE state. */
static void dictPrefetchBucket(dictPrefetcher *p, dictPrefetchLookup *lk) {
    size_t i = p->cur_idx;
    dict *d = p->dicts[i];

    /* Determine which hash table to use */
    if (lk->ht_idx == HT_IDX_INVALID) {
        lk->ht_idx = HT_IDX_FIRST;
    } else if (lk->ht_idx == HT_IDX_FIRST && dictIsRehashing(d)) {
        lk->ht_idx = HT_IDX_SECOND;
    } else {
        /* No more tables left - mark as done. */
        dictPrefetchMarkDone(p, lk);
        return;
    }

    /* Prefetch the bucket */
    lk->bucket_idx = lk->key_hash & DICTHT_SIZE_MASK(d->ht_size_exp[lk->ht_idx]);
    dictPrefetchAdvance(p, &d->ht_table[lk->ht_idx][lk->bucket_idx]);
    lk->current_entry = NULL;
    lk->state = PREFETCH_ENTRY;
}

/* Prefetch the entry in the bucket and move to the PREFETCH_ENTRY_KEY state.
 * If no more entries in the bucket, move to the PREFETCH_BUCKET state to look at the next table. */
static void dictPrefetchEntry(dictPrefetcher *p, dictPrefetchLookup *lk) {
    size_t i = p->cur_idx;

    if (lk->current_entry) {
        /* We already found an entry in the bucket - move to the next entry */
        lk->current_entry = dictGetNext(lk->current_entry);
    } else {
        /* Go to the first entry in the bucket */
        lk->current_entry = p->dicts[i]->ht_table[lk->ht_idx][lk->bucket_idx];
    }

    if (lk->current_entry) {
        dictPrefetchAdvance(p, lk->current_entry);
        lk->state = PREFETCH_ENTRY_KEY;
    } else {
        /* No entry found in the bucket - try the bucket in the next table */
        lk->state = PREFETCH_BUCKET;
    }
}

/* Bring the entry's key payload into cache via the dictType callback,
 * then move to PREFETCH_ENTRY_VALUE where the keyCompare runs. If the
 * dict provides no callback, the entry alone already carries everything
 * keyCompare needs. */
static void dictPrefetchEntryKey(dictPrefetcher *p, dictPrefetchLookup *lk) {
    dictType *type = p->dicts[p->cur_idx]->type;
    lk->state = PREFETCH_ENTRY_VALUE;
    if (type->prefetchEntryKey) {
        void *addr = type->prefetchEntryKey(lk->current_entry);
        if (addr) dictPrefetchAdvance(p, addr);
    }
}

/* Compare the entry's stored key against the lookup key. On match, ask
 * the dictType to prefetch the value-side payload (if any) and mark the
 * lookup done. On mismatch, walk to the next entry in the chain.
 *
 * The entry's stored key may be in a different shape than the lookup key
 * (e.g. dbDictType stores a kvobj but keyCompare wants the sds). When that
 * is the case the dict provides keyFromStoredKey to convert; otherwise the
 * stored key is already in comparable form. This mirrors what
 * dictFindLinkInternal does. */
static void dictPrefetchEntryValue(dictPrefetcher *p, dictPrefetchLookup *lk) {
    size_t i = p->cur_idx;
    dict *d = p->dicts[i];
    dictType *type = d->type;
    const void *stored_key = dictGetKey(lk->current_entry);
    const void *cmp_key = type->keyFromStoredKey ? type->keyFromStoredKey(stored_key) : stored_key;

    /* 1. If this is the last element, we assume a hit and don't compare the keys
     * 2. The stored entry matches the lookup key. */
    if ((!dictGetNext(lk->current_entry) && !dictIsRehashing(d)) ||
        dictCompareKeys(d, p->keys[i], cmp_key))
    {
        if (type->prefetchEntryValue) {
            void *addr = type->prefetchEntryValue(lk->current_entry);
            if (addr) dictPrefetchAdvance(p, addr);
        }
        dictPrefetchMarkDone(p, lk);
    } else {
        /* Not found in the current entry, move to the next entry */
        lk->state = PREFETCH_ENTRY;
    }
}

/* Allocate the per-key lookup array. The prefetcher can then be reused across
 * many batches by repeated dictPrefetcherReset / dictPrefetcherRun calls. */
static void dictPrefetcherInit(dictPrefetcher *p, size_t max_keys) {
    p->lookups = zcalloc(max_keys * sizeof(dictPrefetchLookup));
    p->max_keys = max_keys;
}

static void dictPrefetcherFree(dictPrefetcher *p) {
    zfree(p->lookups);
    p->lookups = NULL;
    p->max_keys = 0;
}

/* Configure the prefetcher for a single batch and seed every lookup's
 * starting state. dicts/keys must remain valid until dictPrefetcherRun
 * returns; only the pointers are stored. */
static void dictPrefetcherReset(dictPrefetcher *p, dict **dicts, void **keys, size_t nkeys) {
    serverAssert(nkeys <= p->max_keys);
    p->dicts = dicts;
    p->keys = keys;
    p->nkeys = nkeys;
    p->cur_idx = 0;

    size_t remaining = 0;
    for (size_t i = 0; i < nkeys; i++) {
        dictPrefetchLookup *lk = &p->lookups[i];
        if (!dicts[i] || dictSize(dicts[i]) == 0) {
            lk->state = PREFETCH_DONE;
            continue;
        }

        /* We skip prefetch during loading, so ht_table[0] should never be NULL
         * when dictSize() > 0 (which only happens mid-dictEmpty via _dictReset). */
        serverAssert(dicts[i]->ht_table[0]);

        lk->ht_idx = HT_IDX_INVALID;
        lk->current_entry = NULL;
        lk->state = PREFETCH_BUCKET;
        lk->key_hash = dictGetHash(dicts[i], keys[i]);
        remaining++;
    }
    p->remaining = remaining;
}

/* Drive the prefetch state machine across all dict lookups until every lookup
 * reaches PREFETCH_DONE.
 *
 * Conceptually each dict lookup is a dictFind broken into four stages:
 *   bucket → entry → entry key payload → entry value payload
 * If the key is not found in ht[0] and the dict is mid-rehash, the lookup
 * loops back to the bucket stage to retry against ht[1].
 *
 * Instead of waiting for each stage's memory access to complete, the FSM
 * issues a prefetch and yields to another in-flight lookup, hiding the
 * memory access latency.
 *
 * Any prefetching of the entry's key payload (e.g. an out-of-line kvobj head)
 * and the entry's value payload (e.g. kv->ptr for a RAW string) is delegated
 * to dictType->prefetchEntryKey and prefetchEntryValue respectively. */
static void dictPrefetcherRun(dictPrefetcher *p) {
    dictPrefetchLookup *lk;
    while ((lk = dictPrefetchNextInFlight(p))) {
        switch (lk->state) {
            case PREFETCH_BUCKET:      dictPrefetchBucket(p, lk); break;
            case PREFETCH_ENTRY:       dictPrefetchEntry(p, lk); break;
            case PREFETCH_ENTRY_KEY:   dictPrefetchEntryKey(p, lk); break;
            case PREFETCH_ENTRY_VALUE: dictPrefetchEntryValue(p, lk); break;
            default: serverPanic("Unknown prefetch state %d", lk->state);
        }
    }
}

/* --------------------------------------------------------------------------
 * Intra-command prefetch API
 * --------------------------------------------------------------------------
 * dictPrefetchKeys() allows a single multi-key command (e.g. MGET) to
 * prefetch dict data for a batch of its own keys, reusing the same state
 * machine that the cross-command path uses.
 *
 * Typical usage from a command implementation:
 *
 *   #define BATCH 16
 *   void myMultiKeyCommand(client *c) {
 *       dict *d = kvstoreGetDict(c->db->keys, slot);
 *       for (int j = 0; j < numkeys; j += BATCH) {
 *           int n = MIN(BATCH, numkeys - j);
 *           void *keys[BATCH]; dict *dicts[BATCH];
 *           for (int k = 0; k < n; k++) {
 *               keys[k] = c->argv[j+k+1]->ptr;
 *               dicts[k] = d;
 *           }
 *           dictPrefetchKeys(dicts, keys, n);
 *           // Now process these n keys — dict bucket / entry / key payload
 *           // (and value payload, if dictType->prefetchEntryValue is set)
 *           // are warm in cache.
 *       }
 *   }
 * ----------------------------------------------------------------------- */
void dictPrefetchKeys(dict **dicts, void **keys, size_t nkeys) {
    /* Single-key prefetch has no benefit — nothing to interleave with.
     * Callers passing nkeys==1 (e.g. tail of a multi-key batch) should
     * fall through to a direct lookup. */
    if (nkeys <= 1) return;

    /* Guard the fixed-size stack array below; callers must batch larger
     * inputs into chunks of DICT_PREFETCH_MAX_SIZE or smaller. */
    serverAssert(nkeys <= DICT_PREFETCH_MAX_SIZE);
    server.stat_total_prefetch_batches++;

    dictPrefetchLookup lookups[DICT_PREFETCH_MAX_SIZE];
    dictPrefetcher p = { .lookups = lookups, .max_keys = nkeys };
    dictPrefetcherReset(&p, dicts, keys, nkeys);
    dictPrefetcherRun(&p);
}

/* --------------------------------------------------------------------------
 * Cross-command batch prefetching
 * -------------------------------------------------------------------------- */

/* PrefetchCommandsBatch structure holds the state of the current batch of client commands being processed. */
typedef struct PrefetchCommandsBatch {
    size_t key_count;               /* Number of keys in the current batch */
    size_t client_count;            /* Number of clients in the current batch */
    size_t pcmd_count;              /* Number of pending commands in the current batch */
    size_t max_prefetch_size;       /* Maximum number of keys to prefetch in a batch */
    void **keys;                    /* Array of keys to prefetch in the current batch */
    client **clients;               /* Array of clients in the current batch */
    pendingCommand **pending_cmds;  /* Array of pending commands in the current batch */
    dict **keys_dicts;              /* Main dict for each key */
    dictPrefetcher prefetcher;      /* Initialized once; reset and reused per batch. */
} PrefetchCommandsBatch;

static PrefetchCommandsBatch *batch = NULL;

void freePrefetchCommandsBatch(void) {
    if (batch == NULL) {
        return;
    }

    zfree(batch->clients);
    zfree(batch->pending_cmds);
    zfree(batch->keys);
    zfree(batch->keys_dicts);
    dictPrefetcherFree(&batch->prefetcher);
    zfree(batch);
    batch = NULL;
}

void prefetchCommandsBatchInit(void) {
    serverAssert(!batch);

    /* To avoid prefetching small batches, we set the max size to twice
     * the configured size, so if not exceeding twice the limit, we can
     * prefetch all of it. See also `determinePrefetchCount` */
    size_t max_prefetch_size = server.prefetch_batch_max_size * 2;

    if (max_prefetch_size == 0) {
        return;
    }

    batch = zcalloc(sizeof(PrefetchCommandsBatch));
    batch->max_prefetch_size = max_prefetch_size;
    batch->clients = zcalloc(max_prefetch_size * sizeof(client *));
    batch->pending_cmds = zcalloc(max_prefetch_size * sizeof(pendingCommand *));
    batch->keys = zcalloc(max_prefetch_size * sizeof(void *));
    batch->keys_dicts = zcalloc(max_prefetch_size * sizeof(dict *));
    dictPrefetcherInit(&batch->prefetcher, max_prefetch_size);
}

void onMaxBatchSizeChange(void) {
    if (batch && batch->client_count > 0) {
        /* We need to process the current batch before updating the size */
        return;
    }

    freePrefetchCommandsBatch();
    prefetchCommandsBatchInit();
}

void resetCommandsBatch(void) {
    if (batch == NULL) {
        /* Handle the case where prefetching becomes enabled from disabled. */
        if (server.prefetch_batch_max_size) prefetchCommandsBatchInit();
        return;
    }

    batch->key_count = 0;
    batch->client_count = 0;
    batch->pcmd_count = 0;

    /* Handle the case where the max prefetch size has been changed. */
    if (batch->max_prefetch_size != (size_t)server.prefetch_batch_max_size * 2) {
        onMaxBatchSizeChange();
    }
}

/* Prefetching in very small batches tends to be ineffective because the technique
 * relies on a small gap—typically a few CPU cycles—between issuing the prefetch
 * and performing the actual memory access. If the batch is too small, this delay
 * cannot be effectively inserted, and the prefetching yields little to no benefit.
 *
 * To avoid wasting effort, when the remaining data is small (less than twice the
 * maximum batch size), we simply prefetch all of it at once. Otherwise, we only
 * prefetch a limited portion, capped at the configured maximum. */
int determinePrefetchCount(int len) {
    if (!batch) return 0;

    /* The batch max size is double of the configured size. */
    int config_size = batch->max_prefetch_size / 2;
    return len < (int)batch->max_prefetch_size ? len : config_size;
}

/* Prefetch command-related data:
 * 1. Prefetch the command arguments allocated by the I/O thread to bring them
 *    closer to the L1 cache.
 * 2. Prefetch the io_deferred_objects for all clients.
 * 3. Prefetch the keys and values for all commands in the current batch from
 *    the main dictionaries. */
void prefetchCommands(void) {
    if (!batch || server.loading) return;

    /* Prefetch argv's for all pending commands */
    for (size_t i = 0; i < batch->pcmd_count; i++) {
        pendingCommand *pcmd = batch->pending_cmds[i];
        if (unlikely(pcmd->argc <= 0)) continue;
        for (int j = 0; j < pcmd->argc; j++) {
            redis_prefetch_read(pcmd->argv[j]);
        }
    }

    /* Prefetch the argv->ptr if required */
    for (size_t i = 0; i < batch->pcmd_count; i++) {
        pendingCommand *pcmd = batch->pending_cmds[i];
        if (unlikely(pcmd->argc <= 1)) continue;
        /* Skip the first argument (command name), as it's typically short */
        for (int j = 1; j < pcmd->argc; j++) {
            if (pcmd->argv[j]->encoding == OBJ_ENCODING_RAW) {
                redis_prefetch_read(pcmd->argv[j]->ptr);
            }
        }
    }

    /* Prefetch io_deferred_objects for all clients */
    for (size_t i = 0; i < batch->client_count; i++) {
        client *c = batch->clients[i];
        if (!c->io_deferred_objects || c->io_deferred_objects_num == 0) continue;
        for (int j = 0; j < c->io_deferred_objects_num; j++)
            redis_prefetch_read(c->io_deferred_objects[j]);
    }

    /* Get the keys ptrs - we do it here after the key obj was prefetched. */
    for (size_t i = 0; i < batch->key_count; i++) {
        batch->keys[i] = ((robj *)batch->keys[i])->ptr;
    }

    /* Prefetch dict keys for all commands.
     * Prefetching is beneficial only if there are more than one key. */
    if (batch->key_count > 1) {
        server.stat_total_prefetch_batches++;
        /* Prefetch keys from the main dict — value-side prefetch (if any)
         * is driven by dbDictType->prefetchEntryValue. */
        dictPrefetcherReset(&batch->prefetcher, batch->keys_dicts, batch->keys, batch->key_count);
        dictPrefetcherRun(&batch->prefetcher);
    }
}

/* Adds the client's command to the current batch.
 *
 * Returns C_OK if the command was added successfully, C_ERR otherwise. */
int addCommandToBatch(client *c) {
    if (unlikely(!batch)) return C_ERR;

    /* If the batch is full, process it.
     * We also check the client count to handle cases where
     * no keys exist for the clients' commands. */
    if (batch->client_count == batch->max_prefetch_size ||
        batch->key_count == batch->max_prefetch_size ||
        batch->pcmd_count == batch->max_prefetch_size)
    {
        return C_ERR;
    }

    /* Avoid partial prefetching: if the batch already has commands and adding this
     * client's ready commands would likely exceed the batch size limit, reject
     * the entire client. This is a conservative estimate using command count as
     * a proxy for key count to ensure all keys from a client are either fully
     * prefetched together or not prefetched at all. */
    if (batch->pcmd_count > 0 &&
        (c->pending_cmds.ready_len + batch->key_count > batch->max_prefetch_size ||
         c->pending_cmds.ready_len + batch->pcmd_count > batch->max_prefetch_size))
    {
        return C_ERR;
    }

    batch->clients[batch->client_count++] = c;

    pendingCommand *pcmd = c->pending_cmds.head;
    while (pcmd != NULL && batch->pcmd_count < batch->max_prefetch_size) {
        if (pcmd->next) redis_prefetch_read(pcmd->next);

        /* Skip commands that have not been preprocessed, or have errors. */
        if ((pcmd->flags & PENDING_CMD_FLAG_INCOMPLETE) || !pcmd->cmd || pcmd->read_error) break;

        batch->pending_cmds[batch->pcmd_count++] = pcmd;

        serverAssert(pcmd->flags & PENDING_CMD_KEYS_RESULT_VALID);
        dict *cmd_dict = kvstoreGetDict(c->db->keys, pcmd->slot > 0 ? pcmd->slot : 0);
        for (int i = 0; i < pcmd->keys_result.numkeys && batch->key_count < batch->max_prefetch_size; i++) {
            batch->keys[batch->key_count] = pcmd->argv[pcmd->keys_result.keys[i].pos];
            batch->keys_dicts[batch->key_count] = cmd_dict;
            batch->key_count++;
            /* Mark the command as prefetched so the intra-command prefetch
             * path skips it. Even on a partial batch, running both paths
             * would just contend for cache bandwidth. */
            pcmd->flags |= PENDING_CMD_KEYS_PREFETCHED;
        }
        pcmd = pcmd->next;
    }

    return C_OK;
}