/*
 * test_alloc_runner_result.c - scheduler-allocator-runner task 3.3: the
 * runner's result assembly (Req 4.1-4.4) and its out-of-memory paths
 * (Req 5.4-5.7, 7.2), through sched_allocator_run() with a scripted fake
 * plan generator over real workers.
 *
 * Injectable steps, in the order the runner reaches them:
 *   HOST_ARRAY   calloc() of the host-allocation array: a call-counting
 *                test_calloc() replacing calloc in alloc_runner.c only
 *   RECORDS      sched_make_allocs() (stubbed here): host 0's records,
 *                host 1's records, then the shared records
 *   HOST_ALLOC   make_host_alloc() (--wrap): host 0, then host 1
 *   JOB_ALLOC    make_job_alloc() (--wrap)
 * Each is armed with the 1-based call number that fails; 0 = never.
 */
#include <errno.h>
#include <stdlib.h>
#include <string.h>

#include <unity/unity.h>

#include "preason.h"
#include "resobj.h"
#include "resreq.h"
#include "sched/alloc_iter.h"
#include "sched/types.h"
#include "scheduler.h"
#include "utils.h"
#include "xcealloc.h"
#include "xstring.h"
#include "xvec.h"

/* White-box sources under test; calloc() is replaced in the runner only */
static void* test_calloc(size_t n, size_t sz);
#define calloc test_calloc
#include "../src/alloc_runner.c"
#undef calloc
#include "../src/core.c"
#include "../src/host_iter.c"
#include "../src/resource.c"
#include "../src/schedobj.c"

/* ---------------------------------------------------------------------- */
/* Link-satisfying stubs (same convention as test_alloc_runner.c)          */
/* ---------------------------------------------------------------------- */

int  plugin_register(void* s, void* p) { return (void)s, (void)p, 0; }
void plugin_unregister(void* arg) { (void)arg; }

int plugin_job_update(void* s, xce_sched_job_t* j)
{
    return (void)s, (void)j, 0;
}
int plugin_queue_update(void* s, xce_sched_queue_t* q)
{
    return (void)s, (void)q, 0;
}
int plugin_job_rusage_update(void* s, xce_sched_job_t* j, xce_job_rusages_t* r)
{
    return (void)s, (void)j, (void)r, 0;
}

/* ---------------------------------------------------------------------- */
/* Fault injection                                                         */
/* ---------------------------------------------------------------------- */

typedef struct {
    int calls;
    int fail_on; /* 1-based call number; 0 = never */
} inj_t;

enum {
    HOST_ARRAY,
    RECORDS,
    HOST_ALLOC,
    JOB_ALLOC,
    NSTEPS
};
static inj_t g_inj[NSTEPS];

static int inj_fails(int step)
{
    g_inj[step].calls++;
    return (g_inj[step].fail_on > 0)
        && (g_inj[step].calls == g_inj[step].fail_on);
}

static void* test_calloc(size_t n, size_t sz)
{
    if (inj_fails(HOST_ARRAY)) {
        errno = ENOMEM;
        return NULL;
    }
    return calloc(n, sz);
}

/* Stub for scheduler.c's sched_make_allocs(): real make_alloc() records
 * (resmgr name, requested value, the given reservation, duration -1 for a
 * NULL duration array); NULL for nres <= 0 as the real one. */
xce_alloc_t** sched_make_allocs(xce_scheduler_t* sched, char* reserve, int nres,
                                int* idxs, float* vals, int* duration)
{
    xce_alloc_t** allocs;
    int           i;

    if (inj_fails(RECORDS)) {
        errno = ENOMEM;
        return NULL;
    }
    if ((sched == NULL) || (nres <= 0) || (idxs == NULL) || (vals == NULL)) {
        errno = EINVAL;
        return NULL;
    }

    allocs = (xce_alloc_t**)calloc(nres, sizeof(*allocs));
    TEST_ASSERT_NOT_NULL(allocs);
    for (i = 0; i < nres; i++) {
        xce_resource_t* r = NULL;

        vget(sched->core.resmgr, idxs[i], (void**)&r);
        TEST_ASSERT_NOT_NULL(r);
        allocs[i] = make_alloc(r->name, reserve, 0, vals[i],
                               (duration != NULL) ? duration[i] : -1);
        TEST_ASSERT_NOT_NULL(allocs[i]);
    }

    return allocs;
}

xce_host_alloc_t* __real_make_host_alloc(char*, int, xce_alloc_t**);
xce_job_alloc_t*  __real_make_job_alloc(uint64_t, int, xce_host_alloc_t**);

xce_host_alloc_t* __wrap_make_host_alloc(char* host, int nalloc,
                                         xce_alloc_t** allocs)
{
    if (inj_fails(HOST_ALLOC)) {
        errno = ENOMEM;
        return NULL;
    }
    return __real_make_host_alloc(host, nalloc, allocs);
}

xce_job_alloc_t* __wrap_make_job_alloc(uint64_t jobid, int nhost,
                                       xce_host_alloc_t** hallocs)
{
    if (inj_fails(JOB_ALLOC)) {
        errno = ENOMEM;
        return NULL;
    }
    return __real_make_job_alloc(jobid, nhost, hallocs);
}

/* ---------------------------------------------------------------------- */
/* plugin_prepare_allocation() stub                                        */
/* ---------------------------------------------------------------------- */

static int g_slots_idx = -1;
static int g_mem_idx   = -1;
static int g_lic_idx   = -1;
static int g_lic2_idx  = -1;
static int g_shared_on = 0; /* job wants lic 2 and lic2 5 */
static int g_host_none = 0; /* per-host request left empty */

/* Job-level pass: the shared wants. Per-host pass: mem 10 per slot, then
 * the bucket's slot count. */
int plugin_prepare_allocation(xce_sched_core_t* sched, xce_sched_ctx_t* ctx)
{
    (void)sched;
    xvec_idval_fini(&(ctx->alloc_req.host));
    xvec_idval_init(&(ctx->alloc_req.host), 8);
    if (ctx->worker == NULL) {
        xvec_idval_fini(&(ctx->alloc_req.shared));
        xvec_idval_init(&(ctx->alloc_req.shared), 8);
        if (g_shared_on) {
            TEST_ASSERT_TRUE(
                xvec_idval_store(&(ctx->alloc_req.shared), g_lic_idx, 2.0f)
                > 0);
            TEST_ASSERT_TRUE(
                xvec_idval_store(&(ctx->alloc_req.shared), g_lic2_idx, 5.0f)
                > 0);
        }
        return 0;
    }
    if (g_host_none) {
        return 0;
    }

    TEST_ASSERT_TRUE(xvec_idval_store(&(ctx->alloc_req.host), g_mem_idx,
                                      10.0f * (float)ctx->nslots)
                     > 0);
    TEST_ASSERT_TRUE(xvec_idval_store(&(ctx->alloc_req.host), g_slots_idx,
                                      (float)ctx->nslots)
                     > 0);
    return 0;
}

/* ---------------------------------------------------------------------- */
/* Scripted fake plan generator: plan [3,1], or none                       */
/* ---------------------------------------------------------------------- */

static int              g_no_plan    = 0;
static int              g_next_calls = 0;
static int              g_fini_calls = 0;
static int              g_slots[2]   = { 3, 1 };
static xce_alloc_plan_t g_plan;

static int fake_plan_init(xce_alloc_plan_iter_t* iter,
                          xce_primary_resreq_t* resreq, xce_sched_core_t* core)
{
    (void)resreq, (void)core;
    iter->private = NULL;
    return 0;
}

static xce_alloc_plan_t* fake_plan_next(xce_alloc_plan_iter_t* iter)
{
    (void)iter;
    if (g_no_plan || (g_next_calls++ > 0)) {
        return NULL;
    }
    memset(&g_plan, 0, sizeof(g_plan));
    g_plan.nbuckets     = 2;
    g_plan.bucket_slots = g_slots;
    return &g_plan;
}

static void fake_plan_fini(xce_alloc_plan_iter_t* iter)
{
    (void)iter;
    g_fini_calls++;
}

static xce_sched_allocator_iface_t g_iface = {
    .accepts   = NULL,
    .plan_init = fake_plan_init,
    .plan_next = fake_plan_next,
    .plan_fini = fake_plan_fini,
};

void setUp(void)
{
    memset(g_inj, 0, sizeof(g_inj));
    g_shared_on  = 0;
    g_host_none  = 0;
    g_no_plan    = 0;
    g_next_calls = 0;
    g_fini_calls = 0;
}

void tearDown(void) { }

/* ---------------------------------------------------------------------- */
/* Fixture: hosts ra-a and ra-b (4 slots, 100 mem), lic/lic2 10 each       */
/* ---------------------------------------------------------------------- */

typedef struct {
    xce_scheduler_t        sched;
    xce_resource_t *       slots, *mem, *lic, *lic2;
    xce_sched_worker_t *   a, *b;
    xce_sched_plugin_t     plugin;
    xce_primary_resreq_t   req;
    xce_sched_obj_t        obj;
    xce_primary_schedobj_t primary_obj;
    xce_sched_ctx_t        ctx;
} rr_fixture_t;

static xce_resource_t* rr_res(rr_fixture_t* f, const char* name, int flags)
{
    xce_resource_t* r = makeResource(
        f->sched.core.resmgr,
        &(xce_resource_t) { .name = (char*)name, .flags = flags });

    TEST_ASSERT_NOT_NULL(r);
    return r;
}

static xce_sched_worker_t* rr_add_host(rr_fixture_t* f, const char* name)
{
    xce_sched_worker_t* w;
    int                 idx[2]   = { f->slots->idx, f->mem->idx };
    float               total[2] = { 4.0f, 100.0f };
    float               inuse[2] = { 0.0f, 0.0f };

    w = (xce_sched_worker_t*)calloc(1, sizeof(*w));
    TEST_ASSERT_NOT_NULL(w);
    w->name = makeString((char*)name);
    TEST_ASSERT_NOT_NULL(w->name);
    zarray_init(&(w->resvals), 8, (zarray_item_killer_t)kill_sched_resval);
    zarray_init(&(w->resInUse), 8, (zarray_item_killer_t)kill_sched_resval);
    put_resvals(&(w->resvals), 2, idx, total);
    put_resvals(&(w->resInUse), 2, idx, inuse);
    TEST_ASSERT_TRUE(sched_add_worker(&(f->sched.core), w) >= 0);

    return w;
}

static void rr_fixture_init(rr_fixture_t* f)
{
    int   i;
    int   idx[2];
    float total[2] = { 10.0f, 10.0f };

    memset(f, 0, sizeof(*f));
    TEST_ASSERT_EQUAL_INT(0, xce_sched_core_init(&(f->sched.core)));
    /* xce_sched_core_fini() finalizes these; both drivers calloc them */
    for (i = 0; i < 2; i++) {
        f->sched.core.pend_updates[i] = (dataq_t*)calloc(1, sizeof(dataq_t));
        TEST_ASSERT_NOT_NULL(f->sched.core.pend_updates[i]);
        dataq_init(f->sched.core.pend_updates[i], NULL);
    }

    f->slots    = rr_res(f, "slots", 0);
    f->mem      = rr_res(f, "mem", 0);
    f->lic      = rr_res(f, "lic", XCE_RESF_GLOBAL);
    f->lic2     = rr_res(f, "lic2", XCE_RESF_GLOBAL);
    g_slots_idx = f->slots->idx;
    g_mem_idx   = f->mem->idx;
    g_lic_idx   = f->lic->idx;
    g_lic2_idx  = f->lic2->idx;
    idx[0]      = f->lic->idx;
    idx[1]      = f->lic2->idx;
    put_resvals(&(f->sched.core.sharedInst), 2, idx, total);

    f->a = rr_add_host(f, "ra-a");
    f->b = rr_add_host(f, "ra-b");

    f->plugin.name            = (char*)"fake_allocator";
    f->plugin.sched           = &(f->sched);
    f->plugin.core            = &(f->sched.core);
    f->plugin.allocator_iface = &g_iface;

    f->req.minTasks       = 1;
    f->req.maxTasks       = 4;
    f->primary_obj.resreq = &(f->req);
    f->obj.tag            = XCE_RESREQ_PRIMARY;
    f->obj.u.primary      = &(f->primary_obj);
    TEST_ASSERT_EQUAL_INT(
        0, sched_ctx_init(&(f->ctx), &(f->sched.core), &(f->obj)));
    sched_ctx_set_primary(&(f->ctx), &(f->req));
}

static void rr_fixture_fini(rr_fixture_t* f)
{
    free_JobAlloc_ptr(&(f->ctx.alloc));
    free_PendReason_ptr(&(f->ctx.reason));
    sched_ctx_fini(&(f->ctx));
    xce_sched_core_fini(&(f->sched.core));
    xstring_table_fini();
}

static int rr_run(rr_fixture_t* f)
{
    return sched_allocator_run(&(f->plugin), &(f->ctx));
}

static float rr_debit(rr_fixture_t* f, xce_sched_worker_t* w, xce_resource_t* r)
{
    xvec_t* v = (xvec_t*)zarray_get(&(f->ctx.uncomitted.hosts), w->idx);

    return (v == NULL) ? 0.0f : ((float*)v->data)[r->idx];
}

/* Host debit kept (the claim's request) or released (0) */
static void rr_assert_debit(rr_fixture_t* f, xce_sched_worker_t* w, float slots)
{
    TEST_ASSERT_EQUAL_FLOAT(10.0f * slots, rr_debit(f, w, f->mem));
    TEST_ASSERT_EQUAL_FLOAT(slots, rr_debit(f, w, f->slots));
}

static void rr_assert_shared_debit(rr_fixture_t* f)
{
    float* v = (float*)f->ctx.uncomitted.shared->data;

    TEST_ASSERT_EQUAL_FLOAT(2.0f, v[f->lic->idx]);
    TEST_ASSERT_EQUAL_FLOAT(5.0f, v[f->lic2->idx]);
}

static void rr_assert_record(xce_alloc_t* a, const char* resname, float value)
{
    TEST_ASSERT_NOT_NULL(a);
    TEST_ASSERT_EQUAL_STRING(resname, a->resname);
    TEST_ASSERT_EQUAL_STRING("", a->reserve);
    TEST_ASSERT_EQUAL_FLOAT(value, a->value);
    TEST_ASSERT_EQUAL_INT(-1, a->duration);
}

static void rr_assert_host(xce_host_alloc_t* ha, const char* host, float slots)
{
    TEST_ASSERT_NOT_NULL(ha);
    TEST_ASSERT_EQUAL_STRING(host, ha->host);
    TEST_ASSERT_EQUAL_INT(2, ha->nallocs);
    TEST_ASSERT_NOT_NULL(ha->allocs);
    rr_assert_record(ha->allocs[0], "mem", 10.0f * slots);
    rr_assert_record(ha->allocs[1], "slots", slots);
}

/* ---------------------------------------------------------------------- */
/* Success (Req 4.1-4.4, 5.6)                                              */
/* ---------------------------------------------------------------------- */

/* Plan [3,1]: bucket 1 on ra-a, bucket 2 on ra-b; two shared records. */
static void test_success_builds_hosts_and_shared_records(void)
{
    rr_fixture_t     f;
    xce_job_alloc_t* ja;

    g_shared_on = 1;
    rr_fixture_init(&f);

    TEST_ASSERT_EQUAL_INT(1, rr_run(&f));
    ja = f.ctx.alloc;
    TEST_ASSERT_NOT_NULL(ja);
    TEST_ASSERT_TRUE(ja->jobid == 0);
    TEST_ASSERT_EQUAL_INT(2, ja->nhosts);
    rr_assert_host(ja->allocs[0], "ra-a", 3.0f);
    rr_assert_host(ja->allocs[1], "ra-b", 1.0f);
    TEST_ASSERT_EQUAL_INT(2, ja->nres_shared);
    TEST_ASSERT_NOT_NULL(ja->allocs_shared);
    rr_assert_record(ja->allocs_shared[0], "lic", 2.0f);
    rr_assert_record(ja->allocs_shared[1], "lic2", 5.0f);
    TEST_ASSERT_EQUAL_INT(3, g_inj[RECORDS].calls);
    TEST_ASSERT_EQUAL_INT(1, g_inj[HOST_ARRAY].calls);

    /* debits stay with the attempt context */
    rr_assert_debit(&f, f.a, 3.0f);
    rr_assert_debit(&f, f.b, 1.0f);
    rr_assert_shared_debit(&f);
    TEST_ASSERT_EQUAL_INT(1, g_fini_calls);
    rr_fixture_fini(&f);
}

/* No shared request: no shared record, no shared records call. */
static void test_success_without_shared_request(void)
{
    rr_fixture_t f;

    rr_fixture_init(&f);

    TEST_ASSERT_EQUAL_INT(1, rr_run(&f));
    TEST_ASSERT_NOT_NULL(f.ctx.alloc);
    TEST_ASSERT_EQUAL_INT(2, f.ctx.alloc->nhosts);
    TEST_ASSERT_NULL(f.ctx.alloc->allocs_shared);
    TEST_ASSERT_EQUAL_INT(0, f.ctx.alloc->nres_shared);
    TEST_ASSERT_EQUAL_INT(2, g_inj[RECORDS].calls);
    rr_fixture_fini(&f);
}

/* An empty per-host request builds a host allocation with no record. */
static void test_success_with_empty_host_request(void)
{
    rr_fixture_t f;

    g_host_none = 1;
    rr_fixture_init(&f);

    TEST_ASSERT_EQUAL_INT(1, rr_run(&f));
    TEST_ASSERT_NOT_NULL(f.ctx.alloc);
    TEST_ASSERT_EQUAL_INT(2, f.ctx.alloc->nhosts);
    TEST_ASSERT_EQUAL_STRING("ra-a", f.ctx.alloc->allocs[0]->host);
    TEST_ASSERT_EQUAL_INT(0, f.ctx.alloc->allocs[0]->nallocs);
    TEST_ASSERT_NULL(f.ctx.alloc->allocs[0]->allocs);
    TEST_ASSERT_EQUAL_STRING("ra-b", f.ctx.alloc->allocs[1]->host);
    rr_fixture_fini(&f);
}

/* No plan satisfied: pending, nothing on the attempt. */
static void test_pending_leaves_no_alloc(void)
{
    rr_fixture_t f;

    g_no_plan = 1;
    rr_fixture_init(&f);

    TEST_ASSERT_EQUAL_INT(0, rr_run(&f));
    TEST_ASSERT_NULL(f.ctx.alloc);
    TEST_ASSERT_EQUAL_INT(0, g_inj[HOST_ARRAY].calls);
    TEST_ASSERT_EQUAL_INT(1, g_fini_calls);
    rr_fixture_fini(&f);
}

/* ---------------------------------------------------------------------- */
/* Out of memory (Req 5.4-5.7, 7.2)                                        */
/* ---------------------------------------------------------------------- */

/* Fails `step` at call `nth` (shared request on): -ENOMEM, no allocation
 * on the attempt, generator finished once; ra-a / ra-b keep their debit
 * only if their claim became a host allocation (or the failure came after
 * all of them); the shared amounts stay counted. Leaks and double frees
 * are ASAN's to report. */
static void rr_check_oom(int step, int nth, int keep_a, int keep_b)
{
    rr_fixture_t f;

    setUp();
    g_shared_on = 1;
    rr_fixture_init(&f);
    g_inj[step].fail_on = nth;

    TEST_ASSERT_EQUAL_INT(-ENOMEM, rr_run(&f));
    TEST_ASSERT_EQUAL_INT(nth, g_inj[step].calls);
    TEST_ASSERT_NULL(f.ctx.alloc);
    TEST_ASSERT_EQUAL_INT(1, g_fini_calls);
    rr_assert_debit(&f, f.a, keep_a ? 3.0f : 0.0f);
    rr_assert_debit(&f, f.b, keep_b ? 1.0f : 0.0f);
    rr_assert_shared_debit(&f);
    rr_fixture_fini(&f);
}

static void test_oom_host_array(void) { rr_check_oom(HOST_ARRAY, 1, 0, 0); }

static void test_oom_host0_records(void) { rr_check_oom(RECORDS, 1, 0, 0); }

static void test_oom_host1_records(void) { rr_check_oom(RECORDS, 2, 1, 0); }

static void test_oom_host0_host_alloc(void)
{
    rr_check_oom(HOST_ALLOC, 1, 0, 0);
}

static void test_oom_host1_host_alloc(void)
{
    rr_check_oom(HOST_ALLOC, 2, 1, 0);
}

static void test_oom_job_alloc(void) { rr_check_oom(JOB_ALLOC, 1, 1, 1); }

static void test_oom_shared_records(void) { rr_check_oom(RECORDS, 3, 1, 1); }

int main(void)
{
    UNITY_BEGIN();
    RUN_TEST(test_success_builds_hosts_and_shared_records);
    RUN_TEST(test_success_without_shared_request);
    RUN_TEST(test_success_with_empty_host_request);
    RUN_TEST(test_pending_leaves_no_alloc);
    RUN_TEST(test_oom_host_array);
    RUN_TEST(test_oom_host0_records);
    RUN_TEST(test_oom_host1_records);
    RUN_TEST(test_oom_host0_host_alloc);
    RUN_TEST(test_oom_host1_host_alloc);
    RUN_TEST(test_oom_job_alloc);
    RUN_TEST(test_oom_shared_records);
    return UNITY_END();
}