basic_multithread_kvs.c

/* Basic multithread key value store in c 
 * with test cases.
 * */


#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <pthread.h>

#define NUM_BUCKETS 1024
#define VALUE_SIZE 32
#define NUM_THREADS 4
#define OPS_PER_THREAD 5

typedef struct Entry {
    size_t key;
    char value[VALUE_SIZE];
    struct Entry* next;
} Entry;

typedef struct {
    Entry* buckets[NUM_BUCKETS];
    pthread_mutex_t locks[NUM_BUCKETS];
} KVStore;

// ---------------- KV Store Functions ---------------- //

size_t hash(size_t key) {
    return key % NUM_BUCKETS;
}

KVStore* kvstore_create() {
    KVStore* store = malloc(sizeof(KVStore));
    for (size_t i = 0; i < NUM_BUCKETS; i++) {
        store->buckets[i] = NULL;
        pthread_mutex_init(&store->locks[i], NULL);
    }
    return store;
}

void kv_set(KVStore* store, size_t key, const char* value) {
    size_t idx = hash(key);
    pthread_mutex_lock(&store->locks[idx]);

    Entry* curr = store->buckets[idx];
    while (curr) {
        if (curr->key == key) {
            strncpy(curr->value, value, VALUE_SIZE);
            printf("[SET] Updated key %zu -> %s\n", key, value);
            pthread_mutex_unlock(&store->locks[idx]);
            return;
        }
        curr = curr->next;
    }

    Entry* new_entry = malloc(sizeof(Entry));
    new_entry->key = key;
    strncpy(new_entry->value, value, VALUE_SIZE);
    new_entry->next = store->buckets[idx];
    store->buckets[idx] = new_entry;

    printf("[SET] Inserted key %zu -> %s\n", key, value);
    pthread_mutex_unlock(&store->locks[idx]);
}

int kv_get(KVStore* store, size_t key, char* out_value) {
    if (!store) {
        printf("[GET] ERROR: KV store is NULL\n");
        return 0;
    }

    size_t idx = hash(key);
    pthread_mutex_lock(&store->locks[idx]);

    Entry* curr = store->buckets[idx];
    while (curr) {
        if (curr->key == key) {
            strncpy(out_value, curr->value, VALUE_SIZE);
            pthread_mutex_unlock(&store->locks[idx]);
            printf("[GET] Key %zu -> %s\n", key, out_value);
            return 1;
        }
        curr = curr->next;
    }

    pthread_mutex_unlock(&store->locks[idx]);
    printf("[GET] Key %zu not found\n", key);
    return 0;
}

void kvstore_destroy(KVStore** store_ptr) {
    if (!store_ptr || !*store_ptr) return;

    KVStore* store = *store_ptr;
    printf("[DESTROY] Destroying KV store...\n");
    for (size_t i = 0; i < NUM_BUCKETS; i++) {
        Entry* curr = store->buckets[i];
        while (curr) {
            Entry* tmp = curr;
            curr = curr->next;
            free(tmp);
        }
        pthread_mutex_destroy(&store->locks[i]);
    }
    free(store);
    *store_ptr = NULL;
    printf("[DESTROY] KV store destroyed and pointer set to NULL\n");
}

// ---------------- Multi-thread Test ---------------- //

typedef struct {
    KVStore* store;
    int thread_id;
} ThreadArg;

void* thread_task(void* arg) {
    ThreadArg* t = (ThreadArg*)arg;
    char value[VALUE_SIZE];

    for (size_t i = 0; i < OPS_PER_THREAD; i++) {
        size_t key = t->thread_id * OPS_PER_THREAD + i;
        char val[VALUE_SIZE];
        snprintf(val, VALUE_SIZE, "T%d_val%zu", t->thread_id, i);
        kv_set(t->store, key, val);

        // Randomly get some keys (could hit same keys across threads)
        size_t get_key = rand() % (NUM_THREADS * OPS_PER_THREAD);
        kv_get(t->store, get_key, value);
    }

    return NULL;
}

void test_multithreaded() {
    printf("\n=== Test: Multi-threaded KV Store ===\n");
    KVStore* store = kvstore_create();
    pthread_t threads[NUM_THREADS];
    ThreadArg args[NUM_THREADS];

    for (int i = 0; i < NUM_THREADS; i++) {
        args[i].store = store;
        args[i].thread_id = i;
        pthread_create(&threads[i], NULL, thread_task, &args[i]);
    }

    for (int i = 0; i < NUM_THREADS; i++) {
        pthread_join(threads[i], NULL);
    }

    kvstore_destroy(&store);
}

// ---------------- Main ---------------- //

int main() {
    test_multithreaded();
    printf("\nAll tests done!\n");
    return 0;
}