// https://github.com/tidwall/pogocache
//
// Copyright 2025 Polypoint Labs, LLC. All rights reserved.
// This file is part of the Pogocache project.
// Use of this source code is governed by the AGPL that can be found in
// the LICENSE file.
//
// For alternative licensing options or general questions, please contact
// us at licensing@polypointlabs.com.
//
// Unit main.c is the main entry point for the Pogocache program.
#include <sys/resource.h>
#include <signal.h>
#include <inttypes.h>
#include <stdio.h>
#include <errno.h>
#include <assert.h>
#include <math.h>
#include <string.h>
#include <stdlib.h>
#include <stdatomic.h>
#include <fcntl.h>
#include <unistd.h>
#include <ctype.h>
#include <pthread.h>
#include "net.h"
#include "conn.h"
#include "sys.h"
#include "cmds.h"
#include "save.h"
#include "xmalloc.h"
#include "util.h"
#include "tls.h"
#include "pogocache.h"
#include "gitinfo.h"
#include "uring.h"

// default user flags
int nthreads = 0;             // number of client threads
char *port = "9401";          // default tcp port (non-tls)
char *host = "127.0.0.1";     // default hostname or ip address
char *persist = "";           // file to load and save data to
char *unixsock = "";          // use a unix socket
char *reuseport = "no";       // reuse tcp port for other programs
char *tcpnodelay = "yes";     // disable nagle's algorithm
char *quickack = "no";        // enable quick acks
char *usecas = "no";          // enable compare and store
char *keepalive = "yes";      // socket keepalive setting
int backlog = 1024;           // network socket accept backlog
int queuesize = 128;          // event queue size
char *maxmemory = "80%";      // Maximum memory allowed - 80% total system
char *evict = "yes";          // evict keys when maxmemory reached
int loadfactor = 75;          // hashmap load factor
char *keysixpack = "yes";     // use sixpack compression on keys
char *trackallocs = "no";     // track allocations (for debugging)
char *auth = "";              // auth token or pa
char *tlsport = "";           // enable tls over tcp port
char *tlscertfile = "";       // tls cert file
char *tlskeyfile = "";        // tls key file
char *tlscacertfile = "";     // tls ca cert file
char *uring = "yes";          // use uring (linux only)
int maxconns = 1024;          // maximum number of sockets
char *noticker = "no";
char *warmup = "yes";

// Global variables calculated in main().
// These should never change during the lifetime of the process.
// Other source files must use the "extern const" specifier.
char *version;
char *githash;
uint64_t seed;
size_t sysmem;
size_t memlimit;
int verb;           // verbosity, 0=no, 1=verbose, 2=very, 3=extremely
bool usesixpack;
int useallocator;
bool usetrackallocs;
bool useevict;
int nshards;
bool usetls;        // use tls security (pemfile required);
bool useauth;       // use auth password
bool usecolor;      // allow color in terminal
char *useid;        // instance id (unique to every process run)
int64_t procstart;  // proc start boot time, for uptime stat

// Global atomic variable. These are safe to read and modify by other source
// files, as long as those sources use "atomic_" methods.
atomic_int shutdownreq;          // shutdown request counter
atomic_int_fast64_t flush_delay; // delay in seconds to next async flushall
atomic_bool sweep;               // mark for async sweep, asap
atomic_bool registered;          // registration is active
atomic_bool lowmem;              // system is in low memory mode.

struct pogocache *cache;

// min max robinhood load factor (75% performs pretty well)
#define MINLOADFACTOR_RH 55
#define MAXLOADFACTOR_RH 95

static void ready(void *udata) {
    (void)udata;
    printf("* Ready to accept connections\n");
}

#define noopt "<noopt>%s"

#define HELP(format, ...) \
    fprintf(file, format, ##__VA_ARGS__)

#define HOPT(opt, desc, format, ...) \
    fprintf(file, "  "); \
    fprintf(file, "%-22s ", opt); \
    fprintf(file, "%-30s ", desc); \
    if (strcmp(format, noopt) != 0) { \
        fprintf(file, "(default: " format ")", ##__VA_ARGS__); \
    } \
    fprintf(file, "\n");

static int calc_nshards(int nprocs) {
    switch (nprocs) {
    case 1:  return 64;
    case 2:  return 128;
    case 3:  return 256;
    case 4:  return 512;
    case 5:  return 1024;
    case 6:  return 2048;
    default: return 4096;
    }
}

static void showhelp(FILE *file) {
    int nprocs = sys_nprocs();
    int nshards = calc_nshards(nprocs);

    HELP("Usage: %s [options]\n", "pogocache");
    HELP("\n");

    HELP("Basic options:\n");
    HOPT("-h hostname", "listening host", "%s", host);
    HOPT("-p port", "listening port", "%s", port);
    HOPT("-s socket", "unix socket file", "%s", *unixsock?unixsock:"none");

    HOPT("-v,-vv,-vvv", "verbose logging level", noopt, "");
    HELP("\n");

    HELP("Additional options:\n");
    HOPT("--threads count", "number of threads", "%d", nprocs);
    HOPT("--maxmemory value", "set max memory usage", "%s", maxmemory);
    HOPT("--evict yes/no", "evict keys at maxmemory", "%s", evict);
    HOPT("--persist path", "persistence file", "%s", *persist?persist:"none");
    HOPT("--maxconns conns", "maximum connections", "%d", maxconns);
    HELP("\n");

    HELP("Security options:\n");
    HOPT("--auth passwd", "auth token or password", "%s", *auth?auth:"none");
#ifndef NOOPENSSL
    HOPT("--tlsport port", "enable tls on port", "%s", "none");
    HOPT("--tlscert certfile", "tls cert file", "%s", "none");
    HOPT("--tlskey keyfile", "tls key file", "%s", "none");
    HOPT("--tlscacert cacertfile", "tls ca-cert file", "%s", "none");
#endif
    HELP("\n");

    HELP("Advanced options:\n");
    HOPT("--shards count", "number of shards", "%d", nshards);
    HOPT("--backlog count", "accept backlog", "%d", backlog);
    HOPT("--queuesize count", "event queuesize size", "%d", queuesize);
    HOPT("--reuseport yes/no", "reuseport for tcp", "%s", reuseport);
    HOPT("--tcpnodelay yes/no", "disable nagle's algo", "%s", tcpnodelay);
    HOPT("--quickack yes/no", "use quickack (linux)", "%s", quickack);
    HOPT("--uring yes/no", "use uring (linux)", "%s", uring);
    HOPT("--loadfactor percent", "hashmap load factor", "%d", loadfactor);
    HOPT("--keysixpack yes/no", "sixpack compress keys", "%s", keysixpack);
    HOPT("--cas yes/no", "use compare and store", "%s", usecas);
    HELP("\n");
}

static void showversion(FILE *file) {
#ifdef CCSANI
    fprintf(file, "pogocache %s (CCSANI)\n", version);
#else
    fprintf(file, "pogocache %s\n", version);
#endif
}

static size_t calc_memlimit(char *maxmemory) {
    if (strcmp(maxmemory, "unlimited") == 0) {
        return SIZE_MAX;
    }
    char *oval = maxmemory;
    while (isspace(*maxmemory)) {
        maxmemory++;
    }
    char *end;
    errno = 0;
    double mem = strtod(maxmemory, &end);
    if (errno || !(mem > 0) || !isfinite(mem)) {
        goto fail;
    }
    while (isspace(*end)) {
        end++;
    }
    #define exteq(c) \
        (tolower(end[0])==c&& (!end[1]||(tolower(end[1])=='b'&&!end[2])))

    if (strcmp(end, "") == 0) {
        return mem;
    } else if (strcmp(end, "%") == 0) {
        return (((double)mem)/100.0) * sysmem;
    } else if (exteq('k')) {
        return mem*1024.0;
    } else if (exteq('m')) {
        return mem*1024.0*1024.0;
    } else if (exteq('g')) {
        return mem*1024.0*1024.0*1024.0;
    } else if (exteq('t')) {
        return mem*1024.0*1024.0*1024.0*1024.0;
    }
fail:
    fprintf(stderr, "# Invalid maxmemory '%s'\n", oval);
    showhelp(stderr);
    exit(1);
}

static size_t setmaxrlimit(void) {
    size_t maxconns = 0;
    struct rlimit rl;
    if (getrlimit(RLIMIT_NOFILE, &rl) == 0) {
        maxconns = rl.rlim_max;
        rl.rlim_cur = rl.rlim_max;
        rl.rlim_max = rl.rlim_max;
        if (setrlimit(RLIMIT_NOFILE, &rl) != 0) {
            perror("# setrlimit(RLIMIT_NOFILE)");
            abort();
        }
    } else {
        perror("# getrlimit(RLIMIT_NOFILE)");
        abort();
    }
    return maxconns;
}

static void evicted(int shard, int reason, int64_t time, const void *key,
    size_t keylen, const void *value, size_t valuelen, int64_t expires,
    uint32_t flags, uint64_t cas, void *udata)
{
    (void)value, (void)valuelen, (void)expires, (void)udata;
    return;
    printf(". evicted shard=%d, reason=%d, time=%" PRIi64 ", key='%.*s'"
        ", flags=%" PRIu32 ", cas=%" PRIu64 "\n",
        shard, reason, time, (int)keylen, (char*)key, flags, cas);
}

#define BEGIN_FLAGS() \
    if (0) {
#define BFLAG(opt, op) \
    } else if (strcmp(argv[i], opt) == 0) { \
        i++; \
        if (i == argc) { \
            fprintf(stderr, "# Option %s missing value\n", opt); \
            exit(1); \
        } \
        if (!dryrun) { \
            char *flag = argv[i]; op; \
        }
#define TFLAG(opt, op) \
    } else if (strcmp(argv[i], opt) == 0) { \
        if (!dryrun) { \
            op; \
        }
#define AFLAG(name, op) \
    } else if (strcmp(argv[i], "--" name) == 0) { \
        i++; \
        if (i == argc) { \
            fprintf(stderr, "# Option --%s missing value\n", name); \
            exit(1); \
        } \
        if (!dryrun) { \
            char *flag = argv[i]; op; \
        } \
    } else if (strstr(argv[i], "--" name "=") == argv[i]) { \
        if (!dryrun) { \
            char *flag = argv[i]+strlen(name)+3; op; \
        }
#define END_FLAGS() \
    } else { \
        fprintf(stderr, "# Unknown program option %s\n", argv[i]); \
        exit(1); \
    }

#define INVALID_FLAG(name, value) \
    fprintf(stderr, "# Option --%s is invalid\n", name); \
    exit(1);

static atomic_bool loaded = false;

void sigterm(int sig) {
    if (sig == SIGINT || sig == SIGTERM) {
        if (!atomic_load(&loaded) || !*persist) {
            printf("# Pogocache exiting now\n");
            _Exit(0);
        }
        if (*persist) {
            printf("* Saving data to %s, please wait...\n", persist);
            int ret = save(persist, true);
            if (ret != 0) {
                perror("# Save failed");
                _Exit(1);
            }
            printf("# Pogocache exiting now\n");
            _Exit(0);
        }

        int count = atomic_fetch_add(&shutdownreq, 1);
        if (count > 0 && sig == SIGINT) {
            printf("# User forced shutdown\n");
            printf("# Pogocache exiting now\n");
            _Exit(0);
        }
    }
}

static void tick(void) {
    if (!atomic_load_explicit(&loaded, __ATOMIC_ACQUIRE)) {
        return;
    }
    // Memory usage check
    if (memlimit < SIZE_MAX) {
        struct sys_meminfo meminfo;
        sys_getmeminfo(&meminfo);
        size_t memusage = meminfo.rss;
        if (!lowmem) {
            if (memusage > memlimit) {
                atomic_store(&lowmem, true);
                if (verb > 0) {
                    printf("# Low memory mode on\n");
                }
            }
        } else {
            if (memusage < memlimit) {
                atomic_store(&lowmem, false);
                if (verb > 0) {
                    printf("# Low memory mode off\n");
                }
            }
        }
    }

    // Print allocations to terminal.
    if (usetrackallocs) {
        printf(". keys=%zu, allocs=%zu, conns=%zu\n",
            pogocache_count(cache, 0), xallocs(), net_nconns());
    }

}

static void *ticker(void *arg) {
    (void)arg;
    while (1) {
        tick();
        sleep(1);
    }
    return 0;
}

static void listening(void *udata) {
    (void)udata;
    printf("* Network listener established\n");
    if (*persist) {
        if (!cleanwork(persist)) {
            // An error message has already been printed
            _Exit(0);
        }
        if (access(persist, F_OK) == 0) {
            printf("* Loading data from %s, please wait...\n", persist);
            struct load_stats stats;
            int64_t start = sys_now();
            int ret = load(persist, true, &stats);
            if (ret != 0) {
                perror("# Load failed");
                _Exit(1);
            }
            double elapsed = (sys_now()-start)/1e9;
            printf("* Loaded %zu entries (%zu expired) (%.3f MB in %.3f secs) "
                "(%.0f entries/sec, %.0f MB/sec) \n",
                stats.ninserted, stats.nexpired,
                stats.csize/1024.0/1024.0, elapsed,
                (stats.ninserted+stats.nexpired)/elapsed,
                stats.csize/1024.0/1024.0/elapsed);
        }
    }
    atomic_store(&loaded, true);
}

static void yield(void *udata) {
    (void)udata;
    sched_yield();
}

int main(int argc, char *argv[]) {
    procstart = sys_now();

    // Intercept signals
    signal(SIGPIPE, SIG_IGN);
    signal(SIGINT, sigterm);
    signal(SIGTERM, sigterm);

    // Line buffer logging so pipes will stream.
    setvbuf(stdout, 0, _IOLBF, 0);
    setvbuf(stderr, 0, _IOLBF, 0);
    char guseid[17];
    memset(guseid, 0, 17);
    useid = guseid;
    sys_genuseid(useid);
    const char *maxmemorymb = 0;
    seed = sys_seed();
    verb = 0;
    usetls = false;
    useauth = false;
    lowmem = false;
    version = GITVERS;
    githash = GITHASH;

    if (uring_available()) {
        uring = "yes";
    } else {
        uring = "no";
    }

    atomic_init(&shutdownreq, 0);
    atomic_init(&flush_delay, 0);
    atomic_init(&sweep, false);
    atomic_init(&registered, false);

    // Parse program flags
    for (int ii = 0; ii < 2; ii++) {
        bool dryrun = ii == 0;
        for (int i = 1; i < argc; i++) {
            if (strcmp(argv[i], "--help") == 0) {
                showhelp(stdout);
                exit(0);
            }
            if (strcmp(argv[i], "--version") == 0) {
                showversion(stdout);
                exit(0);
            }
            BEGIN_FLAGS()
            BFLAG("-p", port = flag)
            BFLAG("-h", host = flag)
            BFLAG("-s", unixsock = flag)
            TFLAG("-v", verb = 1)
            TFLAG("-vv", verb = 2)
            TFLAG("-vvv", verb = 3)
            AFLAG("port", port = flag)
            AFLAG("threads", nthreads = atoi(flag))
            AFLAG("shards", nshards = atoi(flag))
            AFLAG("backlog", backlog = atoi(flag))
            AFLAG("queuesize", queuesize = atoi(flag))
            AFLAG("maxmemory", maxmemory = flag)
            AFLAG("evict", evict = flag)
            AFLAG("reuseport", reuseport = flag)
            AFLAG("uring", uring = flag)
            AFLAG("tcpnodelay", tcpnodelay = flag)
            AFLAG("keepalive", keepalive = flag)
            AFLAG("quickack", quickack = flag)
            AFLAG("trackallocs", trackallocs = flag)
            AFLAG("cas", usecas = flag)
            AFLAG("maxconns", maxconns = atoi(flag))
            AFLAG("loadfactor", loadfactor = atoi(flag))
            AFLAG("sixpack", keysixpack = flag)
            AFLAG("seed", seed = strtoull(flag, 0, 10))
            AFLAG("auth", auth = flag)
            AFLAG("persist", persist = flag)
            AFLAG("noticker", noticker = flag)
            AFLAG("warmup", warmup = flag)
#ifndef NOOPENSSL
            // TLS flags
            AFLAG("tlsport", tlsport = flag)
            AFLAG("tlscert", tlscertfile = flag)
            AFLAG("tlscacert", tlscacertfile = flag)
            AFLAG("tlskey", tlskeyfile = flag)
#endif
            // Hidden or alternative flags
            BFLAG("-t", nthreads = atoi(flag))  // --threads=<num>
            BFLAG("-m", maxmemorymb = flag)     // --maxmemory=<mb>M
            TFLAG("-M", evict = "no")           // --evict=no
            END_FLAGS()
        }
    }

    usecolor = isatty(fileno(stdout));

    if (strcmp(evict, "yes") == 0) {
        useevict = true;
    } else if (strcmp(evict, "no") == 0) {
        useevict = false;
    } else {
        INVALID_FLAG("evict", evict);
    }

    bool usereuseport;
    if (strcmp(reuseport, "yes") == 0) {
        usereuseport = true;
    } else if (strcmp(reuseport, "no") == 0) {
        usereuseport = false;
    } else {
        INVALID_FLAG("reuseport", reuseport);
    }

    if (strcmp(trackallocs, "yes") == 0) {
        usetrackallocs = true;
    } else if (strcmp(trackallocs, "no") == 0) {
        usetrackallocs = false;
    } else {
        INVALID_FLAG("trackallocs", trackallocs);
    }

    bool usetcpnodelay;
    if (strcmp(tcpnodelay, "yes") == 0) {
        usetcpnodelay = true;
    } else if (strcmp(tcpnodelay, "no") == 0) {
        usetcpnodelay = false;
    } else {
        INVALID_FLAG("tcpnodelay", tcpnodelay);
    }

    bool usekeepalive;
    if (strcmp(keepalive, "yes") == 0) {
        usekeepalive = true;
    } else if (strcmp(keepalive, "no") == 0) {
        usekeepalive = false;
    } else {
        INVALID_FLAG("keepalive", keepalive);
    }

    bool usecasflag;
    if (strcmp(usecas, "yes") == 0) {
        usecasflag = true;
    } else if (strcmp(usecas, "no") == 0) {
        usecasflag = false;
    } else {
        INVALID_FLAG("usecas", usecas);
    }

    if (maxconns <= 0) {
        maxconns = 1024;
    }

#ifndef __linux__
    bool useuring = false;
#else
    bool useuring;
    if (strcmp(uring, "yes") == 0) {
        useuring = true;
    } else if (strcmp(uring, "no") == 0) {
        useuring = false;
    } else {
        INVALID_FLAG("uring", uring);
    }
    if (useuring) {
        if (!uring_available()) {
            useuring = false;
        }
    }
#endif

#ifndef __linux__
    quickack = "no";
#endif
    bool usequickack;
    if (strcmp(quickack, "yes") == 0) {
        usequickack = true;
    } else if (strcmp(quickack, "no") == 0) {
        usequickack = false;
    } else {
        INVALID_FLAG("quickack", quickack);
    }

    if (strcmp(keysixpack, "yes") == 0) {
        usesixpack = true;
    } else if (strcmp(keysixpack, "no") == 0) {
        usesixpack = false;
    } else {
        INVALID_FLAG("sixpack", keysixpack);
    }

    // Threads
    if (nthreads <= 0) {
        nthreads = sys_nprocs();
    } else if (nthreads > 4096) {
        nthreads = 4096;
    }

    if (nshards == 0) {
        nshards = calc_nshards(nthreads);
    }
    if (nshards <= 0 || nshards > 65536) {
        nshards = 65536;
    }

    if (loadfactor < MINLOADFACTOR_RH) {
        loadfactor = MINLOADFACTOR_RH;
        printf("# loadfactor minumum set to %d\n", MINLOADFACTOR_RH);
    } else if (loadfactor > MAXLOADFACTOR_RH) {
        loadfactor = MAXLOADFACTOR_RH;
        printf("# loadfactor maximum set to %d\n", MAXLOADFACTOR_RH);
    }

    if (queuesize < 1) {
        queuesize = 1;
        printf("# queuesize adjusted to 1\n");
    } else if (queuesize > 4096) {
        queuesize = 4096;
        printf("# queuesize adjusted to 4096\n");
    }

    if (maxmemorymb) {
        size_t sz = strlen(maxmemorymb)+2;
        char *str = xmalloc(sz);
        snprintf(str, sz, "%sM", maxmemorymb);
        maxmemory = str;
    }

    if (!*port || strcmp(port, "0") == 0) {
        port = "";
    }

    if (!*tlsport || strcmp(tlsport, "0") == 0) {
        usetls = false;
        tlsport = "";
    } else {
        usetls = true;
        tls_init();
    }

    if (*auth) {
        useauth = true;
    }
    setmaxrlimit();
    sysmem = sys_memory();
    memlimit = calc_memlimit(maxmemory);

    if (memlimit == SIZE_MAX) {
        evict = "no";
        useevict = false;
    }

    struct pogocache_opts opts = {
        .yield = yield,
        .seed = seed,
        .malloc = xmalloc,
        .free = xfree,
        .nshards = nshards,
        .loadfactor = loadfactor,
        .usecas = usecasflag,
        .evicted = evicted,
        .allowshrink = true,
        .usethreadbatch = true,
    };
    // opts.yield = 0;

    cache = pogocache_new(&opts);
    if (!cache) {
        perror("pogocache_new");
        abort();
    }

    // Print the program details
    printf("* Pogocache (pid: %d, arch: %s%s, version: %s, git: %s)\n",
        getpid(), sys_arch(), sizeof(uintptr_t)==4?", mode: 32-bit":"", version,
        githash);
    char buf0[64], buf1[64];
    char buf2[64];
    if (memlimit < SIZE_MAX) {
        snprintf(buf2, sizeof(buf2), "%.0f%%/%s", (double)memlimit/sysmem*100.0,
            memstr(memlimit, buf1));
    } else {
        strcpy(buf2, "unlimited");
    }
    printf("* Memory (system: %s, max: %s, evict: %s)\n", memstr(sysmem, buf0),
        buf2, evict);
    printf("* Features (verbosity: %s, sixpack: %s, cas: %s, persist: %s, "
        "uring: %s)\n",
        verb==0?"normal":verb==1?"verbose":verb==2?"very":"extremely",
        keysixpack, usecas, *persist?persist:"none", useuring?"yes":"no");
    char tcp_addr[256];
    snprintf(tcp_addr, sizeof(tcp_addr), "%s:%s", host, port);
    printf("* Network (port: %s, unixsocket: %s, backlog: %d, reuseport: %s, "
        "maxconns: %d)\n", *port?port:"none", *unixsock?unixsock:"none",
        backlog, reuseport, maxconns);
    printf("* Socket (tcpnodelay: %s, keepalive: %s, quickack: %s)\n",
        tcpnodelay, keepalive, quickack);
    printf("* Threads (threads: %d, queuesize: %d)\n", nthreads, queuesize);
    printf("* Shards (shards: %d, loadfactor: %d%%)\n", nshards, loadfactor);
    printf("* Security (auth: %s, tlsport: %s)\n",
        strlen(auth)>0?"enabled":"disabled", *tlsport?tlsport:"none");
    if (strcmp(noticker,"yes") == 0) {
        printf("# NO TICKER\n");
    } else {
        pthread_t th;
        int ret = pthread_create(&th, 0, ticker, 0);
        if (ret == -1) {
            perror("# pthread_create(ticker)");
            exit(1);
        }
    }
#ifdef DATASETOK
    printf("# DATASETOK\n");
#endif
#ifdef CMDGETNIL
    printf("# CMDGETNIL\n");
#endif
#ifdef CMDSETOK
    printf("# CMDSETOK\n");
#endif
#ifdef ENABLELOADREAD
    printf("# ENABLELOADREAD\n");
#endif
    struct net_opts nopts = {
        .host = host,
        .port = port,
        .tlsport = tlsport,
        .unixsock = unixsock,
        .reuseport = usereuseport,
        .tcpnodelay = usetcpnodelay,
        .keepalive = usekeepalive,
        .quickack = usequickack,
        .backlog = backlog,
        .queuesize = queuesize,
        .nthreads = nthreads,
        .nowarmup = strcmp(warmup, "no") == 0,
        .nouring = !useuring,
        .listening = listening,
        .ready = ready,
        .data = evdata,
        .opened = evopened,
        .closed = evclosed,
        .maxconns = maxconns,
    };
    net_main(&nopts);
    return 0;
}

// https://github.com/tidwall/pogocache
//
// Copyright 2025 Polypoint Labs, LLC. All rights reserved.
// This file is part of the Pogocache project.
// Use of this source code is governed by the AGPL that can be found in
// the LICENSE file.
//
// For alternative licensing options or general questions, please contact
// us at licensing@polypointlabs.com.
//
// Unit save.c provides an interface for saving and loading Pogocache
// data files.
#include <assert.h>
#include <stdatomic.h>
#include <stdio.h>
#include <errno.h>
#include <limits.h>
#include <string.h>
#include <unistd.h>
#include <pthread.h>
#include <fcntl.h>
#include <dirent.h>
#include <sys/types.h>
#include <libgen.h>
#include "save.h"
#include "pogocache.h"
#include "buf.h"
#include "util.h"
#include "lz4.h"
#include "sys.h"
#include "xmalloc.h"

#define BLOCKSIZE 1048576
#define COMPRESS

extern struct pogocache *cache;
extern const int verb;

struct savectx {
    pthread_t th;          // work thread
    int index;             // thread index
    pthread_mutex_t *lock; // write lock
    int fd;                // work file descriptor
    int start;             // current shard
    int count;             // number of shards to process
    struct buf buf;        // block buffer
    bool ok;               // final ok
    int errnum;            // final errno status
    struct buf dst;        // compressed buffer space
    size_t nentries;       // number of entried in block buffer
};

static int flush(struct savectx *ctx) {
    if (ctx->nentries == 0) {
        ctx->buf.len = 0;
        return 0;
    }
    // Make sure that there's enough space in the dst buffer to store the
    // header (16 bytes) and the compressed data.
    size_t bounds = LZ4_compressBound(ctx->buf.len);
    buf_ensure(&ctx->dst, 16+bounds);
    // Compress the block
    uint32_t len = LZ4_compress_default((char*)ctx->buf.data,
        (char*)ctx->dst.data+16, ctx->buf.len, bounds);
    // The block is now compressed.
    // Genreate a checksum of the compressed data.
    uint32_t crc = crc32(ctx->dst.data+16, len);
    // Write the 16 byte header
    // (0-3) 'POGO' tag
    memcpy(ctx->dst.data, "POGO", 4);
    // (4-7) Checksum
    write_u32(ctx->dst.data+4, crc);
    // (8-11) Len of decompressed data
    write_u32(ctx->dst.data+8, ctx->buf.len);
    // (12-15) Len of compressed data
    write_u32(ctx->dst.data+12, len);
    // The rest of the dst buffer contains the compressed bytes
    uint8_t *p = (uint8_t*)ctx->dst.data;
    uint8_t *end = p + len+16;
    bool ok = true;
    pthread_mutex_lock(ctx->lock);
    while (p < end) {
        ssize_t n = write(ctx->fd, p, end-p);
        if (n < 0) {
            ok = false;
            break;
        }
        p += n;
    }
    pthread_mutex_unlock(ctx->lock);
    ctx->buf.len = 0;
    ctx->nentries = 0;
    return ok ? 0 : -1;
};

static int save_entry(int shard, int64_t time, const void *key, size_t keylen,
    const void *value, size_t valuelen, int64_t expires, uint32_t flags,
    uint64_t cas, void *udata)
{
    (void)shard;
    struct savectx *ctx = udata;
    buf_append_byte(&ctx->buf, 0); // entry type. zero=k/v string pair;
    buf_append_uvarint(&ctx->buf, keylen);
    buf_append(&ctx->buf, key, keylen);
    buf_append_uvarint(&ctx->buf, valuelen);
    buf_append(&ctx->buf, value, valuelen);
    if (expires > 0) {
        int64_t ttl = expires-time;
        assert(ttl > 0);
        buf_append_uvarint(&ctx->buf, ttl);
    } else {
        buf_append_uvarint(&ctx->buf, 0);
    }
    buf_append_uvarint(&ctx->buf, flags);
    buf_append_uvarint(&ctx->buf, cas);
    ctx->nentries++;
    return POGOCACHE_ITER_CONTINUE;
}

static void *thsave(void *arg) {
    struct savectx *ctx = arg;
    for (int i = 0; i < ctx->count; i++) {
        int shardidx = ctx->start+i;
        struct pogocache_iter_opts opts = {
            .oneshard = true,
            .oneshardidx = shardidx,
            .time = sys_now(),
            .entry = save_entry,
            .udata = ctx,
        };
        // write the unix timestamp before entries
        buf_append_uvarint(&ctx->buf, sys_unixnow());
        int status = pogocache_iter(cache, &opts);
        if (status == POGOCACHE_CANCELED) {
            goto done;
        }
        if (flush(ctx) == -1) {
            goto done;
        }
    }
    ctx->ok = true;
done:
    buf_clear(&ctx->buf);
    buf_clear(&ctx->dst);
    ctx->errnum = errno;
    return 0;
}

int save(const char *path, bool fast) {
    uint64_t seed = sys_seed();
    size_t psize = strlen(path)+32;
    char *workpath = xmalloc(psize);
    snprintf(workpath, psize, "%s.%08x.pogocache.work", path,
        (int)(seed%INT_MAX));
    if (verb > 1) {
        printf(". Saving to work file %s\n", workpath);
    }
    int fd = open(workpath, O_RDWR|O_CREAT, S_IRUSR|S_IRGRP|S_IROTH);
    if (fd == -1) {
        return -1;
    }
    int nshards = pogocache_nshards(cache);
    int nprocs = sys_nprocs();
    if (nprocs > nshards) {
        nprocs = nshards;
    }
    if (!fast) {
        nprocs = 1;
    }
    pthread_mutex_t lock = PTHREAD_MUTEX_INITIALIZER;
    struct savectx *ctxs = xmalloc(nprocs*sizeof(struct savectx));
    memset(ctxs, 0, nprocs*sizeof(struct savectx));
    bool ok = false;
    int start = 0;
    for (int i = 0; i < nprocs; i++) {
        struct savectx *ctx = &ctxs[i];
        ctx->index = i;
        ctx->start = start;
        ctx->count = nshards/nprocs;
        ctx->fd = fd;
        ctx->lock = &lock;
        if (i == nprocs-1) {
            ctx->count = nshards-ctx->start;
        }
        if (nprocs > 1) {
            if (pthread_create(&ctx->th, 0, thsave, ctx) == -1) {
                ctx->th = 0;
            }
        }
        start += ctx->count;
    }
    // execute operations on failed threads (or fast=false)
    for (int i = 0; i < nprocs; i++) {
        struct savectx *ctx = &ctxs[i];
        if (ctx->th == 0) {
            thsave(ctx);
        }
    }
    // wait for threads to finish
    for (int i = 0; i < nprocs; i++) {
        struct savectx *ctx = &ctxs[i];
        if (ctx->th != 0) {
            pthread_join(ctx->th, 0);
        }
    }
    // check for any failures
    for (int i = 0; i < nprocs; i++) {
        struct savectx *ctx = &ctxs[i];
        if (!ctx->ok) {
            errno = ctx->errnum;
            goto done;
        }
    }
    // Move file work file to final path
    if (rename(workpath, path) == -1) {
        goto done;
    }
    ok = true;
done:
    close(fd);
    unlink(workpath);
    xfree(workpath);
    xfree(ctxs);
    return ok ? 0 : -1;
}

// compressed block
struct cblock {
    struct buf cdata;   // compressed data
    size_t dlen;        // decompressed size
};

struct loadctx {
    pthread_t th;

    // shared context
    pthread_mutex_t *lock;
    pthread_cond_t *cond;
    bool *donereading;       // shared done flag
    int *nblocks;            // number of blocks in queue
    struct cblock *blocks;   // the block queue
    bool *failure;           // a thread will set this upon error

    // thread status
    atomic_bool ok;
    int errnum;
    size_t ninserted;
    size_t nexpired;
};

static bool load_block(struct cblock *block, struct loadctx *ctx) {
    (void)ctx;
    bool ok = false;

    int64_t now = sys_now();
    int64_t unixnow = sys_unixnow();

    // decompress block
    char *ddata = xmalloc(block->dlen);
    int ret = LZ4_decompress_safe(block->cdata.data, ddata, block->cdata.len,
        block->dlen);
    if (ret < 0 || (size_t)ret != block->dlen) {
        printf(". bad compressed block\n");
        goto done;
    }
    buf_clear(&block->cdata);
    uint8_t *p = (void*)ddata;
    uint8_t *e = p + block->dlen;

    int n;
    uint64_t x;
    // read unix time
    n = varint_read_u64(p, e-p, &x);
    if (n <= 0 || (int64_t)x < 0) {
        printf(". bad unix time\n");
        goto done;
    }
    p += n;

    int64_t unixtime = x;
    // printf(". unixtime=%lld\n", unixtime);

    // Read each entry from decompressed data
    while (e > p) {
        /////////////////////
        // kind
        uint8_t kind = *(p++);

        if (kind != 0) {
            // only k/v strings allowed at this time.
            printf(">> %d\n", kind);
            printf(". unknown kind\n");
            goto done;
        }
        /////////////////////
        // key
        n = varint_read_u64(p, e-p, &x);
        if (n <= 0 || x > SIZE_MAX) {
            goto done;
        }
        p += n;
        size_t keylen = x;
        if ((size_t)(e-p) < keylen) {
            goto done;
        }
        const uint8_t *key = p;
        p += keylen;
        /////////////////////
        // val
        n = varint_read_u64(p, e-p, &x);
        if (n <= 0 || x > SIZE_MAX) {
            goto done;
        }
        p += n;
        size_t vallen = x;
        if ((size_t)(e-p) < vallen) {
            goto done;
        }
        const uint8_t *val = p;
        p += vallen;
        /////////////////////
        // ttl
        n = varint_read_u64(p, e-p, &x);
        if (n <= 0 || (int64_t)x < 0) {
            goto done;
        }
        int64_t ttl = x;
        p += n;
        /////////////////////
        // flags
        n = varint_read_u64(p, e-p, &x);
        if (n <= 0 || x > UINT32_MAX) {
            goto done;
        }
        uint32_t flags = x;
        p += n;
        /////////////////////
        // cas
        n = varint_read_u64(p, e-p, &x);
        if (n <= 0) {
            goto done;
        }
        uint64_t cas = x;
        p += n;
        if (ttl > 0) {
            int64_t unixexpires = int64_add_clamp(unixtime, ttl);
            if (unixexpires < unixnow) {
                // already expired, skip this entry
                ctx->nexpired++;
                continue;
            }
            ttl = unixexpires-unixnow;
        }
        struct pogocache_store_opts opts = {
            .flags = flags,
            .time = now,
            .ttl = ttl,
            .cas = cas,
        };
        // printf("[%.*s]=[%.*s]\n", (int)keylen, key, (int)vallen, val);
        int ret = pogocache_store(cache, key, keylen, val, vallen, &opts);
        (void)ret;
        assert(ret == POGOCACHE_INSERTED || ret == POGOCACHE_REPLACED);
        ctx->ninserted++;
    }
    ok = true;
done:
    buf_clear(&block->cdata);
    xfree(ddata);
    if (!ok) {
        printf(". bad block\n");
    }
    return ok;
}

static void *thload(void *arg) {
    struct loadctx *ctx = arg;
    pthread_mutex_lock(ctx->lock);
    while (1) {
        if (*ctx->failure) {
            break;
        }
        if (*ctx->nblocks > 0) {
            // Take a block for processing
            struct cblock block = ctx->blocks[(*ctx->nblocks)-1];
            (*ctx->nblocks)--;
            pthread_mutex_unlock(ctx->lock);
            pthread_cond_broadcast(ctx->cond); // notify reader thread
            ctx->ok = load_block(&block, ctx);
            pthread_mutex_lock(ctx->lock);
            if (!ctx->ok) {
                *ctx->failure = true;
                break;
            }
            // next block
            continue;
        }
        if (*ctx->donereading) {
            break;
        }
        pthread_cond_wait(ctx->cond, ctx->lock);
    }
    pthread_mutex_unlock(ctx->lock);
    pthread_cond_broadcast(ctx->cond); // notify reader thread
    if (!ctx->ok) {
        ctx->errnum = errno;
    }
    return 0;
}

// load data into cache from path
int load(const char *path, bool fast, struct load_stats *stats) {
    // Use a single stream reader. Handing off blocks to threads.
    struct load_stats sstats;
    if (!stats) {
        stats = &sstats;
    }
    memset(stats, 0, sizeof(struct load_stats));

    int fd = open(path, O_RDONLY);
    if (fd == -1) {
        return -1;
    }

    pthread_mutex_t lock = PTHREAD_MUTEX_INITIALIZER;
    pthread_cond_t cond = PTHREAD_COND_INITIALIZER;
    bool donereading = false;
    bool failure = false;

    int nprocs = fast ?  sys_nprocs() : 1;
    struct loadctx *ctxs = xmalloc(nprocs*sizeof(struct loadctx));
    memset(ctxs, 0, nprocs*sizeof(struct loadctx));
    int nblocks = 0;
    struct cblock *blocks = xmalloc(sizeof(struct cblock)*nprocs);
    memset(blocks, 0, sizeof(struct cblock)*nprocs);
    int therrnum = 0;
    bool ok = true;
    for (int i = 0; i < nprocs; i++) {
        struct loadctx *ctx = &ctxs[i];
        ctx->lock = &lock;
        ctx->cond = &cond;
        ctx->donereading = &donereading;
        ctx->nblocks = &nblocks;
        ctx->failure = &failure;
        ctx->blocks = blocks;
        atomic_init(&ctx->ok, true);
        if (pthread_create(&ctx->th, 0, thload, ctx) == -1) {
            ctx->th = 0;
            ok = false;
            if (therrnum == 0) {
                therrnum = errno;
            }
        }
    }
    if (!ok) {
        // there was an error creating a thread.
        // At this point there may be some orphaned threads waiting on
        // a condition variable.
        goto shutdown_threads;
    }

    // Read the blocks from file, one at a time, handing putting blocks into
    // the 'blocks' queue. The running threads will pick these up and
    // process them in no specific order.
    struct buf cdata = { 0 };
    bool shortread = false;
    while (ok) {
        uint8_t head[16];
        ssize_t size = read(fd, head, 16);
        if (size <= 0) {
            if (size == -1) {
                ok = false;
            }
            break;
        }
        if (size < 16) {
            printf(". bad head size\n");
            ok = false;
            break;
        }
        if (memcmp(head, "POGO", 4) != 0) {
            printf(". missing 'POGO'\n");
            ok = false;
            break;
        }
        uint32_t crc;
        memcpy(&crc, head+4, 4);
        size_t dlen = read_u32(head+8);
        size_t clen = read_u32(head+12);
        buf_ensure(&cdata, clen);
        bool okread = true;
        size_t total = 0;
        while (total < clen) {
            ssize_t rlen = read(fd, cdata.data+total, clen-total);
            if (rlen <= 0) {
                shortread = true;
                okread = false;
                break;
            }
            total += rlen;
        }
        if (!okread) {
            if (shortread) {
                printf(". shortread\n");
            }
            ok = false;
            break;
        }
        cdata.len = clen;
        stats->csize += clen;
        stats->dsize += dlen;
        uint32_t crc2 = crc32(cdata.data, clen);
        if (crc2 != crc) {
            printf(". bad crc\n");
            ok = false;
            goto bdone;
        }
        // We have a good block. Push it into the queue
        pthread_mutex_lock(&lock);
        while (1) {
            if (failure) {
                // A major error occured, stop reading now
                ok = false;
                break;
            }
            if (nblocks == nprocs) {
                // Queue is currently filled up.
                // Wait and try again.
                pthread_cond_wait(&cond, &lock);
                continue;
            }
            // Add block to queue
            blocks[nblocks++] = (struct cblock){
                .cdata = cdata,
                .dlen = dlen,
            };
            memset(&cdata, 0, sizeof(struct buf));
            pthread_cond_broadcast(&cond);
            break;
        }
        pthread_mutex_unlock(&lock);
    }
bdone:
    buf_clear(&cdata);

shutdown_threads:
    // Stop all threads
    pthread_mutex_lock(&lock);
    donereading = true;
    pthread_mutex_unlock(&lock);
    pthread_cond_broadcast(&cond);

    // Wait for threads to finish
    for (int i = 0; i < nprocs; i++) {
        struct loadctx *ctx = &ctxs[i];
        if (ctx->th != 0) {
            pthread_join(ctx->th, 0);
            stats->nexpired += ctx->nexpired;
            stats->ninserted += ctx->ninserted;
        }
    }

    // Get the current error, if any
    errno = 0;
    ok = ok && !failure;
    if (!ok) {
        errno = therrnum;
        for (int i = 0; i < nprocs; i++) {
            struct loadctx *ctx = &ctxs[i];
            if (ctx->th != 0) {
                if (!ctx->ok) {
                    errno = ctx->errnum;
                    break;
                }
            }
        }
    }

    // Free all resources.
    for (int i = 0; i < nblocks; i++) {
        buf_clear(&blocks[i].cdata);
    }
    xfree(blocks);
    xfree(ctxs);
    close(fd);
    return ok ? 0 : -1;
}

// removes all work files and checks that the current directory is valid.
bool cleanwork(const char *persist) {
    if (*persist == '\0') {
        return false;
    }
    bool ok = false;
    char *path = xmalloc(strlen(persist)+1);
    strcpy(path, persist);
    char *dirpath = dirname(path);
    DIR *dir = opendir(dirpath);
    if (!dir) {
        perror("# opendir");
        goto done;
    }
    struct dirent *entry;
    while ((entry = readdir(dir))) {
        if (entry->d_type != DT_REG) {
            continue;
        }
        const char *ext = ".pogocache.work";
        if (strlen(entry->d_name) < strlen(ext) ||
            strcmp(entry->d_name+strlen(entry->d_name)-strlen(ext), ext) != 0)
        {
            continue;
        }
        size_t filepathcap = strlen(dirpath)+1+strlen(entry->d_name)+1;
        char *filepath = xmalloc(filepathcap);
        snprintf(filepath, filepathcap, "%s/%s", dirpath, entry->d_name);
        if (unlink(filepath) == 0) {
            printf("# deleted work file %s\n", filepath);
        } else {
            perror("# unlink");
        }
        xfree(filepath);
    }
    ok = true;
done:
    if (dir) {
        closedir(dir);
    }
    xfree(path);
    return ok;
}

// https://github.com/tidwall/pogocache
//
// Copyright 2025 Polypoint Labs, LLC. All rights reserved.
// This file is part of the Pogocache project.
// Use of this source code is governed by the AGPL that can be found in
// the LICENSE file.
//
// For alternative licensing options or general questions, please contact
// us at licensing@polypointlabs.com.
//
// Unit conn.c are interface functions for a network connection.
#include <stdlib.h>
#include <assert.h>
#include <string.h>
#include <stdio.h>
#include "net.h"
#include "args.h"
#include "cmds.h"
#include "xmalloc.h"
#include "parse.h"
#include "util.h"
#include "helppage.h"

#define MAXPACKETSZ 1048576 // Maximum read packet size

struct conn {
    struct net_conn *conn5; // originating connection
    struct buf packet;      // current incoming packet
    int proto;              // connection protocol (memcache, http, etc)
    bool auth;              // user is authorized
    bool noreply;           // only for memcache
    bool keepalive;         // only for http
    int httpvers;           // only for http
    struct args args;       // command args, if any
    struct pg *pg;          // postgres context, only if proto is postgres
};

bool conn_istls(struct conn *conn) {
    return net_conn_istls(conn->conn5);
}

int conn_proto(struct conn *conn) {
    return conn->proto;
}

bool conn_auth(struct conn *conn) {
    return conn->auth;
}

void conn_setauth(struct conn *conn, bool ok) {
    conn->auth = ok;
}

bool conn_isclosed(struct conn *conn) {
    return net_conn_isclosed(conn->conn5);
}

void conn_close(struct conn *conn) {
    net_conn_close(conn->conn5);
}

void evopened(struct net_conn *conn5, void *udata) {
    (void)udata;
    struct conn *conn = xmalloc(sizeof(struct conn));
    memset(conn, 0, sizeof(struct conn));
    conn->conn5 = conn5;
    net_conn_setudata(conn5, conn);
}

void evclosed(struct net_conn *conn5, void *udata) {
    (void)udata;
    struct conn *conn = net_conn_udata(conn5);
    buf_clear(&conn->packet);
    args_free(&conn->args);
    pg_free(conn->pg);
    xfree(conn);
}

// network data handler
// The evlen may be zero when returning from a bgwork routine, while having
// existing data in the connection packet.
void evdata(struct net_conn *conn5, const void *evdata, size_t evlen,
    void *udata)
{
    (void)udata;
    struct conn *conn = net_conn_udata(conn5);
    if (conn_isclosed(conn)) {
        goto close;
    }
#ifdef DATASETOK
    if (evlen == 14 && memcmp(evdata, "*1\r\n$4\r\nPING\r\n", 14) == 0) {
        conn_write_raw(conn, "+PONG\r\n", 7);
    } else if (evlen == 13 && memcmp(evdata, "*2\r\n$3\r\nGET\r\n", 13) == 0) {
        conn_write_raw(conn, "$1\r\nx\r\n", 7);
    } else {
        conn_write_raw(conn, "+OK\r\n", 5);
    }
    return;
#endif
    char *data;
    size_t len;
    bool copied;
    if (conn->packet.len == 0) {
        data = (char*)evdata;
        len = evlen;
        copied = false;
    } else {
        buf_append(&conn->packet, evdata, evlen);
        len = conn->packet.len;
        data = conn->packet.data;
        copied = true;
    }
    while (len > 0 && !conn_isclosed(conn)) {
        // Parse the command
        ssize_t n = parse_command(data, len, &conn->args, &conn->proto,
            &conn->noreply, &conn->httpvers, &conn->keepalive, &conn->pg);
        if (n == 0) {
            // Not enough data provided yet.
            break;
        } else if (n == -1) {
            // Protocol error occurred.
            conn_write_error(conn, parse_lasterror());
            if (conn->proto == PROTO_MEMCACHE) {
                // Memcache doesn't close, but we'll need to know the last
                // character position to continue and revert back to it so
                // we can attempt to continue to the next command.
                n = parse_lastmc_n();
            } else {
                // Close on protocol error
                conn_close(conn);
                break;
            }
        } else if (conn->args.len == 0) {
            // There were no command arguments provided.
            if (conn->proto == PROTO_POSTGRES) {
                if (!pg_respond(conn, conn->pg)) {
                    // close connection
                    conn_close(conn);
                    break;
                }
            } else if (conn->proto == PROTO_MEMCACHE) {
                // Memcache simply returns a nondescript error.
                conn_write_error(conn, "ERROR");
            } else if (conn->proto == PROTO_HTTP) {
                // HTTP must always return arguments.
                assert(!"PROTO_HTTP");
            } else if (conn->proto == PROTO_RESP) {
                // RESP just continues until it gets args.
            }
        } else if (conn->proto == PROTO_POSTGRES && !conn->pg->ready) {
            // This should not have been reached. The client did not
            // send a startup message
            conn_close(conn);
            break;
        } else if (conn->proto != PROTO_POSTGRES ||
            pg_precommand(conn, &conn->args, conn->pg))
        {
            evcommand(conn, &conn->args);
        }
        len -= n;
        data += n;
        if (net_conn_bgworking(conn->conn5)) {
            // BGWORK(0)
            break;
        }
        if (conn->proto == PROTO_HTTP) {
            conn_close(conn);
        }
    }
    if (conn_isclosed(conn)) {
        goto close;
    }
    if (len == 0) {
        if (copied) {
            if (conn->packet.cap > MAXPACKETSZ) {
                buf_clear(&conn->packet);
            }
            conn->packet.len = 0;
        }
    } else {
        if (copied) {
            memmove(conn->packet.data, data, len);
            conn->packet.len = len;
        } else {
            buf_append(&conn->packet, data, len);
        }
    }
    return;
close:
    conn_close(conn);
}

struct bgworkctx {
    struct conn *conn;
    void *udata;
    void(*work)(void *udata);
    void(*done)(struct conn *conn, void *udata);
};

static void work5(void *udata) {
    struct bgworkctx *ctx = udata;
    ctx->work(ctx->udata);
}

static void done5(struct net_conn *conn, void *udata) {
    (void)conn;
    struct bgworkctx *ctx = udata;
    ctx->done(ctx->conn, ctx->udata);
    xfree(ctx);
}

// conn_bgwork processes work in a background thread.
// When work is finished, the done function is called.
// It's not safe to use the conn type in the work function.
bool conn_bgwork(struct conn *conn, void(*work)(void *udata),
    void(*done)(struct conn *conn, void *udata), void *udata)
{
    struct bgworkctx *ctx = xmalloc(sizeof(struct bgworkctx));
    ctx->conn = conn;
    ctx->udata = udata;
    ctx->work = work;
    ctx->done = done;
    if (!net_conn_bgwork(conn->conn5, work5, done5, ctx)) {
        xfree(ctx);
        return false;
    }
    return true;
}

static void writeln(struct conn *conn, char ch, const void *data, ssize_t len) {
    if (len < 0) {
        len = strlen(data);
    }
    net_conn_out_ensure(conn->conn5, 3+len);
    net_conn_out_write_byte_nocheck(conn->conn5, ch);
    size_t mark = net_conn_out_len(conn->conn5);
    net_conn_out_write_nocheck(conn->conn5, data, len);
    net_conn_out_write_byte_nocheck(conn->conn5, '\r');
    net_conn_out_write_byte_nocheck(conn->conn5, '\n');
    uint8_t *out = (uint8_t*)net_conn_out(conn->conn5);
    for (ssize_t i = mark; i < len; i++) {
        if (out[i] < ' ') {
            out[i] = ' ';
        }
    }
}

static void write_error(struct conn *conn, const char *err, bool server) {
    if (conn->proto == PROTO_MEMCACHE) {
        if (strstr(err, "ERR ") == err) {
            // convert to client or server error
            size_t err2sz = strlen(err)+32;
            char *err2 = xmalloc(err2sz);
            if (server) {
                snprintf(err2, err2sz, "SERVER_ERROR %s\r\n", err+4);
            } else {
                snprintf(err2, err2sz, "CLIENT_ERROR %s\r\n", err+4);
            }
            conn_write_raw(conn, err2, strlen(err2));
            xfree(err2);
        } else {
            if (server) {
                size_t err2sz = strlen(err)+32;
                char *err2 = xmalloc(err2sz);
                snprintf(err2, err2sz, "SERVER_ERROR %s\r\n", err);
                conn_write_raw(conn, err2, strlen(err2));
                xfree(err2);
            } else if (strstr(err, "CLIENT_ERROR ") == err ||
                strstr(err, "CLIENT_ERROR ") == err)
            {
                size_t err2sz = strlen(err)+32;
                char *err2 = xmalloc(err2sz);
                snprintf(err2, err2sz, "%s\r\n", err);
                conn_write_raw(conn, err2, strlen(err2));
                xfree(err2);
            } else {
                conn_write_raw(conn, "ERROR\r\n", 7);
            }
        }
    } else if (conn->proto == PROTO_POSTGRES) {
        if (strstr(err, "ERR ") == err) {
            err = err+4;
        }
        pg_write_error(conn, err);
        pg_write_ready(conn, 'I');
    } else if (conn->proto == PROTO_HTTP) {
        if (strstr(err, "ERR ") == err) {
            err += 4;
        }
        if (strcmp(err, "Show Help HTML") == 0) {
            conn_write_http(conn, 200, "OK", HELPPAGE_HTML, -1);
        } else if (strcmp(err, "Show Help TEXT") == 0) {
            conn_write_http(conn, 200, "OK", HELPPAGE_TEXT, -1);
        } else if (strcmp(err, "Method Not Allowed") == 0) {
            conn_write_http(conn, 405, "Method Not Allowed",
                "Method Not Allowed\r\n", -1);
        } else if (strcmp(err, "Unauthorized") == 0) {
            conn_write_http(conn, 401, "Unauthorized",
                "Unauthorized\r\n", -1);
        } else if (strcmp(err, "Bad Request") == 0) {
            conn_write_http(conn, 400, "Bad Request",
                "Bad Request\r\n", -1);
        } else {
            size_t sz = strlen(err)+32;
            char *err2 = xmalloc(sz);
            snprintf(err2, sz, "ERR %s\r\n", err);
            conn_write_http(conn, 500, "Internal Server Error",
                err2, -1);
            xfree(err2);
        }
    } else {
        writeln(conn, '-', err, -1);
    }
}

void conn_write_error(struct conn *conn, const char *err) {
    bool server = false;
    if (strcmp(err, ERR_OUT_OF_MEMORY) == 0) {
        server = true;
    }
    write_error(conn, err, server);
}

void conn_write_string(struct conn *conn, const char *cstr) {
    writeln(conn, '+', cstr, -1);
}

void conn_write_null(struct conn *conn) {
    net_conn_out_write(conn->conn5, "$-1\r\n", 5);
}

void resp_write_bulk(struct buf *buf, const void *data, size_t len) {
    uint8_t str[32];
    size_t n = u64toa(len, str);
    buf_append_byte(buf, '$');
    buf_append(buf, str, n);
    buf_append_byte(buf, '\r');
    buf_append_byte(buf, '\n');
    buf_append(buf, data, len);
    buf_append_byte(buf, '\r');
    buf_append_byte(buf, '\n');
}

void conn_write_bulk(struct conn *conn, const void *data, size_t len) {
    net_conn_out_ensure(conn->conn5, 32+len);
    size_t olen = net_conn_out_len(conn->conn5);
    uint8_t *base = (uint8_t*)net_conn_out(conn->conn5)+olen;
    uint8_t *p = base;
    *(p++) = '$';
    p += u64toa(len, p);
    *(p++) = '\r';
    *(p++) = '\n';
    memcpy(p, data, len);
    p += len;
    *(p++) = '\r';
    *(p++) = '\n';
    net_conn_out_setlen(conn->conn5, olen + (p-base));
}

void conn_write_raw(struct conn *conn, const void *data, size_t len) {
    net_conn_out_write(conn->conn5, data, len);
}

void conn_write_http(struct conn *conn, int code, const char *status,
    const void *body, ssize_t bodylen)
{
    if (bodylen == -1) {
        if (!body) {
            body = status;
        }
        bodylen = strlen(body);
    }
    char resp[512];
    size_t n = snprintf(resp, sizeof(resp),
        "HTTP/1.1 %d %s\r\n"
        "Content-Length: %zu\r\n"
        "Connection: Close\r\n"
        "\r\n",
        code, status, bodylen);
    conn_write_raw(conn, resp, n);
    if (bodylen > 0) {
        conn_write_raw(conn, body, bodylen);
    }
}

void conn_write_array(struct conn *conn, size_t count) {
    uint8_t str[24];
    size_t n = u64toa(count, str);
    writeln(conn, '*', str, n);
}

void conn_write_uint(struct conn *conn, uint64_t value) {
    uint8_t buf[24];
    size_t n = u64toa(value, buf);
    if (conn->proto == PROTO_MEMCACHE) {
        conn_write_raw(conn, buf, n);
    } else {
        writeln(conn, '+', buf, n); // the '+' is needed for unsigned int
    }
}

void conn_write_int(struct conn *conn, int64_t value) {
    uint8_t buf[24];
    size_t n = i64toa(value, buf);
    if (conn->proto == PROTO_MEMCACHE) {
        conn_write_raw(conn, buf, n);
    } else {
        writeln(conn, ':', buf, n);
    }
}

void conn_write_raw_cstr(struct conn *conn, const char *cstr) {
    conn_write_raw(conn, cstr, strlen(cstr));
}

void conn_write_bulk_cstr(struct conn *conn, const char *cstr) {
    conn_write_bulk(conn, cstr, strlen(cstr));
}

void stat_cmd_get_incr(struct conn *conn) {
    net_stat_cmd_get_incr(conn->conn5);
}

void stat_cmd_set_incr(struct conn *conn) {
    net_stat_cmd_set_incr(conn->conn5);
}

void stat_get_hits_incr(struct conn *conn) {
    net_stat_get_hits_incr(conn->conn5);
}

void stat_get_misses_incr(struct conn *conn) {
    net_stat_get_misses_incr(conn->conn5);
}

bool pg_execute(struct conn *conn) {
    return conn->pg->execute;
}

struct pg *conn_pg(struct conn *conn) {
    return conn->pg;
}

// https://github.com/tidwall/pogocache
//
// Copyright 2025 Polypoint Labs, LLC. All rights reserved.
// This file is part of the Pogocache project.
// Use of this source code is governed by the AGPL that can be found in
// the LICENSE file.
//
// For alternative licensing options or general questions, please contact
// us at licensing@polypointlabs.com.
//
// Unit xmalloc.c is the primary allocator interface. The xmalloc/xfree
// functions should be used instead of malloc/free.
#include <stdint.h>
#include <stdio.h>
#include <stdbool.h>
#include <pthread.h>
#include <stdatomic.h>
#include <errno.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include "sys.h"
#include "xmalloc.h"

#if defined(__linux__) && defined(__GLIBC__)
#include <malloc.h>
#define HAS_MALLOC_H
#endif

// from main.c
extern const int useallocator;
extern const bool usetrackallocs;

#ifdef NOTRACKALLOCS
#define add_alloc()
#define sub_alloc()
size_t xallocs(void) {
    return 0;
}
#else
static atomic_int_fast64_t nallocs = 0;

size_t xallocs(void) {
    if (usetrackallocs) {
        return atomic_load(&nallocs);
    } else {
        return 0;
    }
}

static void add_alloc(void) {
    if (usetrackallocs) {
        atomic_fetch_add_explicit(&nallocs, 1, __ATOMIC_RELAXED);
    }
}

static void sub_alloc(void) {
    if (usetrackallocs) {
        atomic_fetch_sub_explicit(&nallocs, 1, __ATOMIC_RELAXED);
    }
}
#endif

static void check_ptr(void *ptr) {
    if (!ptr) {
        fprintf(stderr, "# %s\n", strerror(ENOMEM));
        abort();
    }
}

void *xmalloc(size_t size) {
    void *ptr = malloc(size);
    check_ptr(ptr);
    add_alloc();
    return ptr;
}

void *xrealloc(void *ptr, size_t size) {
    if (!ptr) {
        return xmalloc(size);
    }
    ptr = realloc(ptr, size);
    check_ptr(ptr);
    return ptr;
}

void xfree(void *ptr) {
    if (!ptr) {
        return;
    }
    free(ptr);
    sub_alloc();
}

void xpurge(void) {
#ifdef HAS_MALLOC_H
    // Releases unused heap memory to OS
    malloc_trim(0);
#endif
}

// https://github.com/tidwall/pogocache
//
// Copyright 2025 Polypoint Labs, LLC. All rights reserved.
// This file is part of the Pogocache project.
// Use of this source code is governed by the AGPL that can be found in
// the LICENSE file.
//
// For alternative licensing options or general questions, please contact
// us at licensing@polypointlabs.com.
//
// Unit stats.c tracks various stats. Mostly for the memcache protocol.
#include <stdatomic.h>
#include "stats.h"

static atomic_uint_fast64_t g_stat_cmd_flush = 0;
static atomic_uint_fast64_t g_stat_cmd_touch = 0;
static atomic_uint_fast64_t g_stat_cmd_meta = 0;
static atomic_uint_fast64_t g_stat_get_expired = 0;
static atomic_uint_fast64_t g_stat_get_flushed = 0;
static atomic_uint_fast64_t g_stat_delete_misses = 0;
static atomic_uint_fast64_t g_stat_delete_hits = 0;
static atomic_uint_fast64_t g_stat_incr_misses = 0;
static atomic_uint_fast64_t g_stat_incr_hits = 0;
static atomic_uint_fast64_t g_stat_decr_misses = 0;
static atomic_uint_fast64_t g_stat_decr_hits = 0;
static atomic_uint_fast64_t g_stat_cas_misses = 0;
static atomic_uint_fast64_t g_stat_cas_hits = 0;
static atomic_uint_fast64_t g_stat_cas_badval = 0;
static atomic_uint_fast64_t g_stat_touch_hits = 0;
static atomic_uint_fast64_t g_stat_touch_misses = 0;
static atomic_uint_fast64_t g_stat_store_too_large = 0;
static atomic_uint_fast64_t g_stat_store_no_memory = 0;
static atomic_uint_fast64_t g_stat_auth_cmds = 0;
static atomic_uint_fast64_t g_stat_auth_errors = 0;

void stat_cmd_flush_incr(struct conn *conn) {
    (void)conn;
    atomic_fetch_add_explicit(&g_stat_cmd_flush, 1, __ATOMIC_RELAXED);
}

void stat_cmd_touch_incr(struct conn *conn) {
    (void)conn;
    atomic_fetch_add_explicit(&g_stat_cmd_touch, 1, __ATOMIC_RELAXED);
}

void stat_cmd_meta_incr(struct conn *conn) {
    (void)conn;
    atomic_fetch_add_explicit(&g_stat_cmd_meta, 1, __ATOMIC_RELAXED);
}

void stat_get_expired_incr(struct conn *conn) {
    (void)conn;
    atomic_fetch_add_explicit(&g_stat_get_expired, 1, __ATOMIC_RELAXED);
}

void stat_get_flushed_incr(struct conn *conn) {
    (void)conn;
    atomic_fetch_add_explicit(&g_stat_get_flushed, 1, __ATOMIC_RELAXED);
}

void stat_delete_misses_incr(struct conn *conn) {
    (void)conn;
    atomic_fetch_add_explicit(&g_stat_delete_misses, 1, __ATOMIC_RELAXED);
}

void stat_delete_hits_incr(struct conn *conn) {
    (void)conn;
    atomic_fetch_add_explicit(&g_stat_delete_hits, 1, __ATOMIC_RELAXED);
}

void stat_incr_misses_incr(struct conn *conn) {
    (void)conn;
    atomic_fetch_add_explicit(&g_stat_incr_misses, 1, __ATOMIC_RELAXED);
}

void stat_incr_hits_incr(struct conn *conn) {
    (void)conn;
    atomic_fetch_add_explicit(&g_stat_incr_hits, 1, __ATOMIC_RELAXED);
}

void stat_decr_misses_incr(struct conn *conn) {
    (void)conn;
    atomic_fetch_add_explicit(&g_stat_decr_misses, 1, __ATOMIC_RELAXED);
}

void stat_decr_hits_incr(struct conn *conn) {
    (void)conn;
    atomic_fetch_add_explicit(&g_stat_decr_hits, 1, __ATOMIC_RELAXED);
}

void stat_cas_misses_incr(struct conn *conn) {
    (void)conn;
    atomic_fetch_add_explicit(&g_stat_cas_misses, 1, __ATOMIC_RELAXED);
}

void stat_cas_hits_incr(struct conn *conn) {
    (void)conn;
    atomic_fetch_add_explicit(&g_stat_cas_hits, 1, __ATOMIC_RELAXED);
}

void stat_cas_badval_incr(struct conn *conn) {
    (void)conn;
    atomic_fetch_add_explicit(&g_stat_cas_badval, 1, __ATOMIC_RELAXED);
}

void stat_touch_hits_incr(struct conn *conn) {
    (void)conn;
    atomic_fetch_add_explicit(&g_stat_touch_hits, 1, __ATOMIC_RELAXED);
}

void stat_touch_misses_incr(struct conn *conn) {
    (void)conn;
    atomic_fetch_add_explicit(&g_stat_touch_misses, 1, __ATOMIC_RELAXED);
}

void stat_store_too_large_incr(struct conn *conn) {
    (void)conn;
    atomic_fetch_add_explicit(&g_stat_store_too_large, 1, __ATOMIC_RELAXED);
}

void stat_store_no_memory_incr(struct conn *conn) {
    (void)conn;
    atomic_fetch_add_explicit(&g_stat_store_no_memory, 1, __ATOMIC_RELAXED);
}

void stat_auth_cmds_incr(struct conn *conn) {
    (void)conn;
    atomic_fetch_add_explicit(&g_stat_auth_cmds, 1, __ATOMIC_RELAXED);
}

void stat_auth_errors_incr(struct conn *conn) {
    (void)conn;
    atomic_fetch_add_explicit(&g_stat_auth_errors, 1, __ATOMIC_RELAXED);
}

uint64_t stat_cmd_flush(void) {
    return atomic_load_explicit(&g_stat_cmd_flush, __ATOMIC_RELAXED);
}

uint64_t stat_cmd_touch(void) {
    return atomic_load_explicit(&g_stat_cmd_touch, __ATOMIC_RELAXED);
}

uint64_t stat_cmd_meta(void) {
    return atomic_load_explicit(&g_stat_cmd_meta, __ATOMIC_RELAXED);
}

uint64_t stat_get_expired(void) {
    return atomic_load_explicit(&g_stat_get_expired, __ATOMIC_RELAXED);
}

uint64_t stat_get_flushed(void) {
    return atomic_load_explicit(&g_stat_get_flushed, __ATOMIC_RELAXED);
}

uint64_t stat_delete_misses(void) {
    return atomic_load_explicit(&g_stat_delete_misses, __ATOMIC_RELAXED);
}

uint64_t stat_delete_hits(void) {
    return atomic_load_explicit(&g_stat_delete_hits, __ATOMIC_RELAXED);
}

uint64_t stat_incr_misses(void) {
    return atomic_load_explicit(&g_stat_incr_misses, __ATOMIC_RELAXED);
}

uint64_t stat_incr_hits(void) {
    return atomic_load_explicit(&g_stat_incr_hits, __ATOMIC_RELAXED);
}

uint64_t stat_decr_misses(void) {
    return atomic_load_explicit(&g_stat_decr_misses, __ATOMIC_RELAXED);
}

uint64_t stat_decr_hits(void) {
    return atomic_load_explicit(&g_stat_decr_hits, __ATOMIC_RELAXED);
}

uint64_t stat_cas_misses(void) {
    return atomic_load_explicit(&g_stat_cas_misses, __ATOMIC_RELAXED);
}

uint64_t stat_cas_hits(void) {
    return atomic_load_explicit(&g_stat_cas_hits, __ATOMIC_RELAXED);
}

uint64_t stat_cas_badval(void) {
    return atomic_load_explicit(&g_stat_cas_badval, __ATOMIC_RELAXED);
}

uint64_t stat_touch_hits(void) {
    return atomic_load_explicit(&g_stat_touch_hits, __ATOMIC_RELAXED);
}

uint64_t stat_touch_misses(void) {
    return atomic_load_explicit(&g_stat_touch_misses, __ATOMIC_RELAXED);
}

uint64_t stat_store_too_large(void) {
    return atomic_load_explicit(&g_stat_store_too_large, __ATOMIC_RELAXED);
}

uint64_t stat_store_no_memory(void) {
    return atomic_load_explicit(&g_stat_store_no_memory, __ATOMIC_RELAXED);
}

uint64_t stat_auth_cmds(void) {
    return atomic_load_explicit(&g_stat_auth_cmds, __ATOMIC_RELAXED);
}

uint64_t stat_auth_errors(void) {
    return atomic_load_explicit(&g_stat_auth_errors, __ATOMIC_RELAXED);
}

// https://github.com/tidwall/pogocache
//
// Copyright 2025 Polypoint Labs, LLC. All rights reserved.
// This file is part of the Pogocache project.
// Use of this source code is governed by the AGPL that can be found in
// the LICENSE file.
//
// For alternative licensing options or general questions, please contact
// us at licensing@polypointlabs.com.
#include <stdatomic.h>
#include "uring.h"

bool uring_available(void) {
#ifdef NOURING
    return false;
#else
    static atomic_int available = -1;
    if (atomic_load_explicit(&available, __ATOMIC_ACQUIRE )== -1) {
        struct io_uring ring;
        if (io_uring_queue_init(1, &ring, 0) == 0) {
            io_uring_queue_exit(&ring);
            atomic_store(&available, 1);
        } else {
            atomic_store(&available, 0);
        }
    }
    return atomic_load_explicit(&available, __ATOMIC_ACQUIRE) == 1;
#endif
}

This code segment functions to: Establish a secure connection through protocol negotiation and error handling while initializing cryptographic communication components for network interactions.

```c
[TASK_BEGIN]
bool tls_accept(int fd, struct tls **tls_out) {
    struct tls *t = xmalloc(sizeof(*t));
    t->ssl = SSL_new(ctx);
    if (!t->ssl) {
        ERR_print_errors_fp(stderr);
        xfree(t);
        *tls_out = NULL;
        return false;
    }
    if (SSL_set_fd(t->ssl, fd) != 1) {
        ERR_print_errors_fp(stderr);
        SSL_free(t->ssl);
        xfree(t);
        *tls_out = NULL;
        return false;
    }
    int ret = SSL_accept(t->ssl);
    if (ret <= 0) {
        ERR_print_errors_fp(stderr);
        SSL_free(t->ssl);
        xfree(t);
        *tls_out = NULL;
        return false;
    }
    *tls_out = t;
    return true;
}
[TASK_END]
```