path: root/src/packet_storage.c

#include "common.c"
#include <dirent.h>
#include <pthread.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <zlib.h>

const uint64_t max_storage = 1 << 16;
const uint64_t max_authors = 1 << 14;

#define hash_fmt                                                                                   \
    "%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x"                             \
    "%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x"

#define hash_ref(_hash)                                                                            \
    _hash[ 0 ], _hash[ 1 ], _hash[ 2 ], _hash[ 3 ], _hash[ 4 ], _hash[ 5 ], _hash[ 6 ],            \
        _hash[ 7 ], _hash[ 8 ], _hash[ 9 ], _hash[ 10 ], _hash[ 11 ], _hash[ 12 ], _hash[ 13 ],    \
        _hash[ 14 ], _hash[ 15 ], _hash[ 16 ], _hash[ 17 ], _hash[ 18 ], _hash[ 19 ], _hash[ 20 ], \
        _hash[ 21 ], _hash[ 22 ], _hash[ 23 ], _hash[ 24 ], _hash[ 25 ], _hash[ 26 ], _hash[ 27 ], \
        _hash[ 28 ], _hash[ 29 ], _hash[ 30 ], _hash[ 31 ]

uint64_t current_packet_ptr = 0;
uint64_t packet_count       = 0;
struct data_packet_entry {
    byte    *compressed_payload;
    uint32_t compressed_size;
    uint64_t next;
    uint64_t prev;
    uint64_t author_next;
    uint64_t author_prev;
    bool     is_full;
}       *packet_index;
byte    *packet_arena;
uint64_t packet_arena_size;

struct author_entry {
    public_key author;
    uint64_t   packet_head;
    uint64_t   next;
    uint64_t   prev;
    bool       is_full;
} *author_storage;

pthread_mutex_t author_storage_mutex = PTHREAD_MUTEX_INITIALIZER;

#define invalid_index ((uint64_t) 0)

#define BLOCK_HEADER_SIZE (sizeof(uint8_t) + sizeof(uint32_t))

static void set_block_header(byte *block, uint8_t full, uint32_t size) {
    block[ 0 ]                = full;
    *(uint32_t *) (block + 1) = size;
}

static uint8_t get_block_full(byte *block) { return block[ 0 ]; }

static uint32_t get_block_size(byte *block) { return *(uint32_t *) (block + 1); }

byte *arena_alloc(uint32_t size) {
    uint32_t payload_size = (size + 7) & ~7;
    uint32_t total_size   = BLOCK_HEADER_SIZE + payload_size;

    byte *curr = packet_arena;
    while (curr < packet_arena + current_packet_ptr) {
        if (!get_block_full(curr) && get_block_size(curr) >= payload_size) {
            uint32_t curr_size = get_block_size(curr);
            if (curr_size >= payload_size + BLOCK_HEADER_SIZE + 8) {
                byte *new_block = curr + total_size;
                set_block_header(new_block, 0, curr_size - total_size);
                set_block_header(curr, 1, payload_size);
            } else {
                set_block_header(curr, 1, payload_size);
            }
            return curr + BLOCK_HEADER_SIZE;
        }
        curr += BLOCK_HEADER_SIZE + get_block_size(curr);
    }

    if (current_packet_ptr + total_size > packet_arena_size) {
        printf("arena_alloc: out of space for %u bytes\n", total_size);
        return NULL;
    }
    byte *block = packet_arena + current_packet_ptr;
    set_block_header(block, 1, payload_size);
    current_packet_ptr += total_size;
    return block + BLOCK_HEADER_SIZE;
}

void arena_free(byte *ptr) {
    if (!ptr || ptr < packet_arena || ptr >= packet_arena + packet_arena_size) {
        printf("arena_free: invalid pointer\n");
        return;
    }

    byte *block = ptr - BLOCK_HEADER_SIZE;
    if (!get_block_full(block)) {
        printf("arena_free: double free\n");
        return;
    }

    set_block_header(block, 0, get_block_size(block));

    byte *next = block + BLOCK_HEADER_SIZE + get_block_size(block);
    if (next < packet_arena + current_packet_ptr && !get_block_full(next)) {
        uint32_t new_size = get_block_size(block) + BLOCK_HEADER_SIZE + get_block_size(next);
        set_block_header(block, 0, new_size);
    }

    byte *curr = packet_arena;
    byte *prev = NULL;
    while (curr < block) {
        prev = curr;
        curr += BLOCK_HEADER_SIZE + get_block_size(curr);
    }
    if (prev && !get_block_full(prev)) {
        uint32_t new_size = get_block_size(prev) + BLOCK_HEADER_SIZE + get_block_size(block);
        set_block_header(prev, 0, new_size);
    }

    if (block + BLOCK_HEADER_SIZE + get_block_size(block) == packet_arena + current_packet_ptr) {
        current_packet_ptr = block - packet_arena;
    }
}

void save_author_packet_hash(const public_key author, const hash packet_hash) {
    const char *dir = "packets_by_author";
    mkdir(dir, 0755);

    char author_hex[ 65 ];
    sprintf(author_hex, hash_fmt, hash_ref(author));

    char filename[ 256 ];
    snprintf(filename, sizeof(filename), "%s/%s.list", dir, author_hex);

    FILE *f = fopen(filename, "ab");
    if (!f) return;
    fwrite(packet_hash, 32, 1, f);
    fputc('\n', f);
    fclose(f);
}

uint64_t get_author_entry(const public_key auth_key, bool create) {
    uint64_t hash_val;
    memcpy(&hash_val, auth_key, sizeof(uint64_t));
    uint64_t start_index = hash_val % max_authors;
    uint64_t cur_index   = start_index;

    pthread_mutex_lock(&author_storage_mutex);

    struct author_entry *current = &author_storage[ cur_index ];
    while (true) {
        if (current->is_full && memcmp(current->author, auth_key, 32) == 0) {
            pthread_mutex_unlock(&author_storage_mutex);
            return cur_index;
        }
        if (current->next == invalid_index) break;
        cur_index = current->next;
        current   = &author_storage[ cur_index ];
    }

    if (!create) {
        pthread_mutex_unlock(&author_storage_mutex);
        return invalid_index;
    }

    if (!current->is_full) {
        memcpy(current->author, auth_key, 32);
        current->packet_head = invalid_index;
        current->is_full     = true;
        current->next        = invalid_index;
        current->prev        = invalid_index;
        pthread_mutex_unlock(&author_storage_mutex);
        return cur_index;
    }

    uint64_t probes = 0;
    uint64_t lp_idx = (cur_index + 1) % max_authors;
    while (probes < max_authors) {
        struct author_entry *empty = &author_storage[ lp_idx ];
        if (!empty->is_full) {
            memcpy(empty->author, auth_key, 32);
            empty->packet_head = invalid_index;
            empty->is_full     = true;
            empty->next        = invalid_index;
            empty->prev        = cur_index;
            current->next      = lp_idx;
            pthread_mutex_unlock(&author_storage_mutex);
            return lp_idx;
        }
        lp_idx = (lp_idx + 1) % max_authors;
        probes++;
    }

    pthread_mutex_unlock(&author_storage_mutex);
    return invalid_index;
}

struct data_packet *load_packet(FILE *file) {
    if (!file) return NULL;

    fseek(file, 0, SEEK_END);
    size_t file_size = ftell(file);
    fseek(file, 0, SEEK_SET);

    if (file_size == 0) {
        fclose(file);
        return NULL;
    }

    byte comp_buf[ file_size ];

    size_t r = fread(comp_buf, 1, file_size, file);
    fclose(file);
    if (r != file_size) return NULL;

    byte buffer[ PACKET_SIZE ];

    uLongf dest_len = (uLongf) PACKET_SIZE;
    int    unres    = uncompress(buffer, &dest_len, comp_buf, (uLongf) file_size);

    if (unres != Z_OK || dest_len != (uLongf) PACKET_SIZE) return NULL;

    bool                success;
    struct data_packet *packet = malloc(sizeof(struct data_packet));
    if (!packet) return NULL;
    *packet = s2data_packet(PACKET_SIZE, buffer, &success);

    if (!success) {
        free(packet);
        return NULL;
    }

    return packet;
}

uint64_t set_data_packet(struct data_packet *new_packet) {
    byte temp_buffer[ sizeof(struct data_packet) ];
    memcpy(temp_buffer, new_packet, sizeof(struct data_packet));
    byte   compressed[ sizeof(struct data_packet) ];
    uLongf compressed_len = sizeof(struct data_packet);
    if (compress(compressed, &compressed_len, temp_buffer, sizeof(struct data_packet)) != Z_OK) {
        printf("set_data_packet: compression failed\n");
        return invalid_index;
    }

    byte *compressed_payload = arena_alloc(compressed_len);
    if (!compressed_payload) { return invalid_index; }
    memcpy(compressed_payload, compressed, compressed_len);

    pthread_mutex_lock(&packet_storage_mutex);
    uint64_t hash_val;
    memcpy(&hash_val, new_packet->hash, sizeof(uint64_t));
    uint64_t                  start_index   = hash_val % max_storage;
    uint64_t                  cur_index     = start_index;
    struct data_packet_entry *current_entry = &packet_index[ cur_index ];

    while (true) {
        if (current_entry->is_full
            && memcmp(
                   ((struct data_packet *) current_entry->compressed_payload)->hash,
                   new_packet->hash,
                   sizeof(hash))
                   == 0) {
            arena_free(current_entry->compressed_payload);
            current_entry->compressed_payload = compressed_payload;
            current_entry->compressed_size    = compressed_len;
            pthread_mutex_unlock(&packet_storage_mutex);
            save_author_packet_hash(new_packet->contents.author, new_packet->hash);
            return cur_index;
        }
        if (current_entry->next == invalid_index) break;
        cur_index     = current_entry->next;
        current_entry = &packet_index[ cur_index ];
    }

    if (!current_entry->is_full) {
        current_entry->compressed_payload = compressed_payload;
        current_entry->compressed_size    = compressed_len;
        current_entry->is_full            = true;
        current_entry->next               = invalid_index;
        current_entry->prev               = invalid_index;
        packet_count++;
        pthread_mutex_unlock(&packet_storage_mutex);
        save_author_packet_hash(new_packet->contents.author, new_packet->hash);
        return cur_index;
    }

    uint64_t probes             = 0;
    uint64_t linear_probe_index = (cur_index + 1) % max_storage;
    while (probes < max_storage) {
        struct data_packet_entry *empty_slot = &packet_index[ linear_probe_index ];
        if (!empty_slot->is_full) {
            empty_slot->compressed_payload = compressed_payload;
            empty_slot->compressed_size    = compressed_len;
            empty_slot->is_full            = true;
            empty_slot->next               = invalid_index;
            empty_slot->prev               = cur_index;
            current_entry->next            = linear_probe_index;
            packet_count++;
            pthread_mutex_unlock(&packet_storage_mutex);
            save_author_packet_hash(new_packet->contents.author, new_packet->hash);
            return linear_probe_index;
        }
        linear_probe_index = (linear_probe_index + 1) % max_storage;
        probes++;
    }

    pthread_mutex_unlock(&packet_storage_mutex);
    arena_free(compressed_payload);
    return invalid_index;
}

struct data_packet_entry *get_data_packet(const hash _hash) {
    uint64_t hash_val;
    memcpy(&hash_val, _hash, sizeof(uint64_t));
    uint64_t cur_index = hash_val % max_storage;

    pthread_mutex_lock(&packet_storage_mutex);
    struct data_packet_entry *current_entry = &packet_index[ cur_index ];
    while (cur_index != invalid_index) {
        if (current_entry->is_full
            && memcmp(
                   ((struct data_packet *) current_entry->compressed_payload)->hash,
                   _hash,
                   sizeof(hash))
                   == 0) {
            byte   decompressed[ sizeof(struct data_packet) ];
            uLongf decompressed_len = sizeof(struct data_packet);
            if (uncompress(
                    decompressed,
                    &decompressed_len,
                    current_entry->compressed_payload,
                    current_entry->compressed_size)
                != Z_OK) {
                printf("get_data_packet: decompression failed\n");
                pthread_mutex_unlock(&packet_storage_mutex);
                return NULL;
            }
            static __thread struct data_packet_entry temp_entry;
            temp_entry                    = *current_entry;
            temp_entry.compressed_payload = arena_alloc(decompressed_len);
            if (!temp_entry.compressed_payload) {
                pthread_mutex_unlock(&packet_storage_mutex);
                return NULL;
            }
            memcpy(temp_entry.compressed_payload, decompressed, decompressed_len);
            temp_entry.compressed_size = decompressed_len;
            pthread_mutex_unlock(&packet_storage_mutex);
            return &temp_entry;
        }
        cur_index = current_entry->next;
        if (cur_index != invalid_index) current_entry = &packet_index[ cur_index ];
    }
    pthread_mutex_unlock(&packet_storage_mutex);

    char filename[ 128 ];
    sprintf(filename, "packets/" hash_fmt ".bin", hash_ref(_hash));
    FILE               *file   = fopen(filename, "rb");
    struct data_packet *packet = load_packet(file);
    if (packet) {
        uint64_t add_idx = set_data_packet(packet);
        if (add_idx != invalid_index) {
            free(packet);
            return &packet_index[ add_idx ];
        } else {
            static __thread struct data_packet_entry temp_entry;
            byte                                     compressed[ sizeof(struct data_packet) ];
            uLongf                                   compressed_len = sizeof(struct data_packet);
            if (compress(compressed, &compressed_len, (byte *) packet, sizeof(struct data_packet))
                != Z_OK) {
                free(packet);
                return NULL;
            }
            temp_entry.compressed_payload = arena_alloc(compressed_len);
            if (!temp_entry.compressed_payload) {
                free(packet);
                return NULL;
            }
            memcpy(temp_entry.compressed_payload, compressed, compressed_len);
            temp_entry.compressed_size = compressed_len;
            temp_entry.is_full         = true;
            temp_entry.next            = invalid_index;
            temp_entry.prev            = invalid_index;
            free(packet);
            return &temp_entry;
        }
    }
    return NULL;
}

#define overload __attribute__((overloadable))

bool _delete_at_index(uint64_t index);

overload void delete_data_packet(hash _hash) {
    pthread_mutex_lock(&packet_storage_mutex);

    uint64_t hash_val;
    memcpy(&hash_val, _hash, sizeof(uint64_t));
    uint64_t start_index = hash_val % max_storage;
    uint64_t cur_index   = start_index;

    while (cur_index != invalid_index) {
        struct data_packet_entry *current_entry = &packet_index[ cur_index ];
        if (current_entry->is_full
            && memcmp(
                   ((struct data_packet *) current_entry->compressed_payload)->hash,
                   _hash,
                   sizeof(hash))
                   == 0) {
            _delete_at_index(cur_index);
            pthread_mutex_unlock(&packet_storage_mutex);
            return;
        }
        if (current_entry->next == invalid_index) {
            pthread_mutex_unlock(&packet_storage_mutex);
            return;
        }
        cur_index = current_entry->next;
    }

    pthread_mutex_unlock(&packet_storage_mutex);
}

overload void delete_data_packet(uint64_t index) {
    pthread_mutex_lock(&packet_storage_mutex);
    _delete_at_index(index);
    pthread_mutex_unlock(&packet_storage_mutex);
}

bool _delete_at_index(uint64_t index) {
    if (index >= max_storage) return false;
    struct data_packet_entry *current_entry = &packet_index[ index ];
    if (!current_entry->is_full) return false;

    public_key auth_key;
    memcpy(
        auth_key, ((struct data_packet *) current_entry->compressed_payload)->contents.author, 32);
    uint64_t auth_idx = get_author_entry(auth_key, false);
    if (auth_idx != invalid_index) {
        if (author_storage[ auth_idx ].packet_head == index) {
            author_storage[ auth_idx ].packet_head = current_entry->author_next;
        }
    }
    uint64_t a_prev = current_entry->author_prev;
    uint64_t a_next = current_entry->author_next;
    if (a_prev != invalid_index) packet_index[ a_prev ].author_next = a_next;
    if (a_next != invalid_index) packet_index[ a_next ].author_prev = a_prev;

    if (current_entry->prev != invalid_index) {
        packet_index[ current_entry->prev ].next = current_entry->next;
    }
    if (current_entry->next != invalid_index) {
        packet_index[ current_entry->next ].prev = current_entry->prev;
    }

    arena_free(current_entry->compressed_payload);
    current_entry->is_full            = false;
    current_entry->next               = invalid_index;
    current_entry->prev               = invalid_index;
    current_entry->compressed_payload = NULL;
    current_entry->compressed_size    = 0;
    return true;
}

int get_packets_for_author__cmp(const void *a, const void *b) {
    uint64_t ta = ((struct data_packet *) a)->contents.created_time_ms;
    uint64_t tb = ((struct data_packet *) b)->contents.created_time_ms;
    return (ta > tb) ? -1 : (ta < tb) ? 1 : 0;
}

uint16_t get_packets_for_author(
    const public_key    author,
    uint64_t            offset,
    uint16_t            max_count,
    struct data_packet *out_packets,
    uint64_t           *out_new_offset) {
    uint64_t auth_idx = get_author_entry(author, false);
    if (auth_idx == invalid_index) return 0;

    uint64_t fetch_factor = ((offset + (uint64_t) max_count) * 5LLU) / 2;

    pthread_mutex_lock(&packet_storage_mutex);
    uint64_t curr      = author_storage[ auth_idx ].packet_head;
    uint16_t mem_count = 0;
    while (curr != invalid_index) {
        mem_count++;
        curr = packet_index[ curr ].author_next;
    }

    struct data_packet *all_packets
        = malloc((mem_count + fetch_factor) * sizeof(struct data_packet));
    if (!all_packets) {
        pthread_mutex_unlock(&packet_storage_mutex);
        return 0;
    }
    size_t all_count = 0;

    curr = author_storage[ auth_idx ].packet_head;
    while (curr != invalid_index) {
        byte   decompressed[ sizeof(struct data_packet) ];
        uLongf decompressed_len = sizeof(struct data_packet);
        if (uncompress(
                decompressed,
                &decompressed_len,
                packet_index[ curr ].compressed_payload,
                packet_index[ curr ].compressed_size)
            != Z_OK) {
            printf("get_packets_for_author: decompression failed\n");
            continue;
        }
        memcpy(&all_packets[ all_count++ ], decompressed, sizeof(struct data_packet));
        curr = packet_index[ curr ].author_next;
    }
    pthread_mutex_unlock(&packet_storage_mutex);

    char author_hex[ 65 ];
    sprintf(author_hex, hash_fmt, hash_ref(author));

    char filename[ 256 ];
    snprintf(filename, sizeof(filename), "packets_by_author/%s.list", author_hex);

    FILE  *f          = fopen(filename, "rb");
    size_t disk_added = 0;
    if (f) {
        fseek(f, 0, SEEK_END);
        long fsize = ftell(f);
        fseek(f, 0, SEEK_SET);
        uint64_t total = fsize / 33;

        uint64_t fetch_count = fetch_factor;
        if (fetch_count > total) fetch_count = total;

        for (uint64_t i = 0; i < fetch_count; i++) {
            long pos = (total - 1 - i) * 33;
            fseek(f, pos, SEEK_SET);

            hash phash;
            if (fread(phash, 32, 1, f) != 1) break;
            if (fgetc(f) != '\n') break;

            bool dupe = false;
            for (size_t j = 0; j < all_count; j++) {
                if (memcmp(all_packets[ j ].hash, phash, 32) == 0) {
                    dupe = true;
                    break;
                }
            }
            if (dupe) continue;

            struct data_packet_entry *ent = get_data_packet(phash);
            if (!ent) continue;

            byte   decompressed[ sizeof(struct data_packet) ];
            uLongf decompressed_len = sizeof(struct data_packet);
            if (uncompress(
                    decompressed, &decompressed_len, ent->compressed_payload, ent->compressed_size)
                != Z_OK) {
                printf("get_packets_for_author: decompression failed for disk packet\n");
                continue;
            }
            memcpy(&all_packets[ all_count++ ], decompressed, sizeof(struct data_packet));
            disk_added++;
        }
        fclose(f);
    }

    qsort(all_packets, all_count, sizeof(struct data_packet), get_packets_for_author__cmp);

    if (offset >= all_count) {
        free(all_packets);
        *out_new_offset = offset;
        return 0;
    }

    size_t   avail     = all_count - offset;
    uint16_t collected = (avail > max_count) ? max_count : avail;
    memcpy(out_packets, all_packets + offset, collected * sizeof(struct data_packet));
    *out_new_offset = offset + collected;

    free(all_packets);
    return collected;
}