myxce/src/controller/worker.c

#include <errno.h>
#include <stdbool.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

#include "cluster.h"
#include "controller.capnp.h"
#include "persist.h"
#include "wobj.h"
#include "worker.h"
#include "xce/cluster_def.h"
#include "xce/errnox.h"
#include "xstring.h"

static int hostmatched(const char* pattern, const char* name)
{
    if (strcmp(pattern, name) == 0) {
        return 1;
    }

    return 0;
}

static xce_hostconfig_t* find_matched_hostconfig(cluster_t* c, const char* name)
{
    xce_hostconfig_t*  defaultconfig = NULL;
    xce_hostconfig_t** hconfigs      = NULL;
    int                i, nhconfigs;

    if (c->myconf == NULL) {
        errno = EINVAL;
        return NULL;
    }

    if (c->myconf->cluster_config == NULL) {
        errno = EINVAL;
        return NULL;
    }

    if (c->myconf->cluster_config->hconfigs == NULL) {
        return NULL;
    }

    nhconfigs = c->myconf->cluster_config->nhconfigs;
    hconfigs  = c->myconf->cluster_config->hconfigs;

    for (i = 0; i < nhconfigs; i++) {
        xce_hostconfig_t* config = hconfigs[i];

        if (strcmp(config->name, "default") == 0) {
            defaultconfig = config;
            continue;
        }

        if (hostmatched(config->name, name)) {
            return config;
        }
    }

    return defaultconfig;
}

static int update_val_res(cluster_t* c, worker_t* w, xce_resource_t* r,
                          xce_val_t* val)
{
    int        idx;
    zarray_t*  array = NULL;
    xce_val_t* v;

    idx = r->idx;

    if (r->flags & XCE_RESF_SHARED) {
        /* update shared resource */
        array = &(c->sharedinst);
    } else {
        /* update host based resource */
        array = &(w->host.resvals);
    }

    if (idx < array->caps) {
        v = (xce_val_t*)array->data[idx];
    } else {
        v = NULL;
    }

    if (v == NULL) {
        int rc;

        v  = (xce_val_t*)calloc(1, sizeof(*v));
        rc = zarray_put(array, idx, v);
        if (rc < 0) {
            FREEUP(v);
            return rc;
        }
    }

    memcpy(v, val, sizeof(*v));

    return 0;
}

static int update_resval(cluster_t* c, worker_t* w, xce_resval_t* val)
{
    xce_resource_t* r = NULL;
    xce_val_t*      v = NULL;
    int             rc;

    lock(c->resmgr, RTTI_LOCK_READ);
    rc = key_find(c->resmgr, (void*)val->name, (void**)&r);
    unlock(c->resmgr);

    if (rc < 0) {
        return rc;
    }
    if (r == NULL) {
        ERROR("resource `%s' can not be found", val->name);
        return -errno;
    }

    if (r->idx < 0) {
        ERROR("invalid resource definition for %s", val->name);
        return -EINVAL;
    }

    v = val->val;

    rc = update_val_res(c, w, r, v);

    return rc;
}

int update_resvals(cluster_t* c, worker_t* w, int nres, xce_resval_t* vals)
{
    int rc = 0;
    int i;

    for (i = 0; i < nres; i++) {
        rc = update_resval(c, w, &(vals[i]));
        if (rc < 0) {
            PERROR("update_resval");
        }
    }

    return 0;
}

int update_resval_list(cluster_t* c, worker_t* w, xce_resval_list_t* l)
{
    int rc = 0;
    int i;

    for (i = 0; i < l->count; i++) {
        rc = update_resval(c, w, l->vals[i]);
        if (rc < 0) {
            PERROR("update_resval");
        }
    }

    return 0;
}

static int mark_host_res(cluster_t* c, worker_t* w, xce_config_t* config)
{
    xce_resource_t* r = NULL;
    char*           n = config->name;
    int             rc;
    /* clang-format-off */
    xce_val_t val;
    /* clang-format-on */

    lock(c->resmgr, RTTI_LOCK_READ);
    rc = key_find(c->resmgr, (void*)n, (void**)&r);
    unlock(c->resmgr);

    if (rc < 0) {
        return rc;
    }

    if (r == NULL) {
        ERROR("resource `%s' can not be found", n);
        return -ENOENT;
    }

    if (r->idx < 0) {
        ERROR("invalid resource definition for `%s'", n);
        return -EINVAL;
    }

    if (r->flags & XCE_RESF_SHARED) {
        ERROR("resource `%s' should not be a shared resource", n);
        return -EINVAL;
    }

    switch (r->valtype) {
    case XCE_VAL_STRING:
        val.which = Value_u_svalue;
        if (config->value == NULL) {
            val.u.svalue = "";
        } else {
            val.u.svalue = config->value;
        }

        break;
    case XCE_VAL_NUMERIC:
        val.which = Value_u_fvalue;
        if (config->value == NULL) {
            val.u.fvalue = 0.0;
        } else {
            val.u.fvalue = atof(config->value);
        }
        break;
    case XCE_VAL_BOOLEAN:
        val.which    = Value_u_bvalue;
        val.u.bvalue = 1;
        break;
    }
    rc = update_val_res(c, w, r, &val);

    return rc;
}

static int add_hostattrs_resvals(cluster_t* c, worker_t* w)
{
#define N_HOST_ATTRS (8)
    xce_val_t values[N_HOST_ATTRS] = {
        { .which = Value_u_svalue, .u = { .svalue = w->host.attrs->hosttype } },
        { .which = Value_u_svalue,
         .u     = { .svalue = w->host.attrs->hostmodel }                      },
        { .which = Value_u_fvalue,
         .u     = { .fvalue = w->host.attrs->cpufactor }                      },
        { .which = Value_u_fvalue, .u = { .fvalue = w->host.attrs->maxcpus }  },
        { .which = Value_u_fvalue, .u = { .fvalue = w->host.attrs->maxmem }   },
        { .which = Value_u_fvalue, .u = { .fvalue = w->host.attrs->maxswap }  },
        { .which = Value_u_fvalue, .u = { .fvalue = w->host.attrs->maxtmp }   },
        { .which = Value_u_fvalue, .u = { .fvalue = w->host.attrs->ndisks }   }
    };
    xce_resval_t resvals[N_HOST_ATTRS] = {
        { .name = "type",    .val = &values[0] },
        { .name = "model",   .val = &values[1] },
        { .name = "cpuf",    .val = &values[2] },
        { .name = "maxcpus", .val = &values[3] },
        { .name = "maxmem",  .val = &values[4] },
        { .name = "maxswap", .val = &values[5] },
        { .name = "maxtmp",  .val = &values[6] },
        { .name = "ndisks",  .val = &values[7] }
    };

    return update_resvals(c, w, N_HOST_ATTRS, &(resvals[0]));
}

static int set_host_slots(cluster_t* c, worker_t* w, int nslots)
{
    /* clang-format off */
    xce_val_t values[1] = {
        {
            .which = Value_u_fvalue,
            .u = {
                .fvalue = (float)nslots
            }
        }
    };
    xce_resval_t resvals[1] = {
        {
            .name = "slots",
            .val = &values[0]
        }
    };
    /* clang-format on */

    return update_resvals(c, w, 1, &(resvals[0]));
}

int add_worker(void* cls, endpoint_t* ep, const char* host, hostattrs_t* attrs)
{
    cluster_t* c = (cluster_t*)cls;
    /* clang-format off */
    worker_t   worker = {
        .host = (xce_worker_t){
            .name = (char*)host,
        }
    };
    /* clang-format on */
    worker_t*         w = NULL;
    int               idx, i;
    xce_hostconfig_t* hconfig = NULL;
    int               nslots;

    if ((cls == NULL) || (host == NULL)) {
        return -EINVAL;
    }

    w = calloc(1, sizeof(*w));
    if (w == NULL) {
        return -ENOMEM;
    }

    hconfig = find_matched_hostconfig(c, host);
    if (hconfig == NULL) {
        return -errno;
    }

    worker.host.attrs = attrs;

    w = makeWorker(c->workermgr, &(worker.host));
    if (w == NULL) {
        return -errno;
    }

    /* we should hold the update lock */
    /* If admin makes the host in locked or closed, we need keep this */

    lock(w, RTTI_LOCK_WRITE);

    w->host.status &= ~XCE_WORKER_UNAVAIL; /* status OK */
    w->ep = ep;

    /* override host type? */
    if (strcmp(hconfig->hosttype, "!") != 0) {
        dereference(w->host.attrs->hosttype);
        w->host.attrs->hosttype = makeString(hconfig->hosttype);
    }

    /* override host model */
    if (strcmp(hconfig->hostmodel, "!") != 0) {
        dereference(w->host.attrs->hostmodel);
        w->host.attrs->hostmodel = makeString(hconfig->hostmodel);
    }
    add_hostattrs_resvals(c, w);
    idx = w->idx;

    /* using auto detected or pre configured MXJ? */
    nslots = hconfig->mxj;
    if (nslots < 0) {
        nslots = w->host.attrs->maxcpus;
    }

    set_host_slots(c, w, nslots);

    if (w->host.dispatch_window_config != NULL) {
        free_timewindow_config(w->host.dispatch_window_config);
        w->host.dispatch_window_config = NULL;
    }

    if (hconfig->dispatchwindow != NULL) {
        w->host.dispatch_window_config
            = dup_timewindow_config(hconfig->dispatchwindow);
    }

    if (w->host.thresholds != NULL) {
        free_resource_thresholds(w->host.thresholds, w->host.nthresholds);
        w->host.nthresholds = 0;
        w->host.thresholds  = NULL;
    }

    if (hconfig->thresholds != NULL) {
        w->host.thresholds  = dup_resource_thresholds(hconfig->thresholds,
                                                      hconfig->nthresholds);
        w->host.nthresholds = hconfig->nthresholds;
    }

    /* add bool resources to host */
    for (i = 0; i < hconfig->nresources; i++) {
        mark_host_res(c, w, hconfig->resources[i]);
    }

    unlock(w);

    return idx;
}

int restore_worker(void* cls, xce_worker_t host)
{
    cluster_t* c   = (cluster_t*)cls;
    int        idx = 0;
    worker_t*  w   = NULL;

    w = calloc(1, sizeof(*w));
    if (w == NULL) {
        return -ENOMEM;
    }
    w = makeWorker(c->workermgr, &host);
    if (w == NULL) {
        return -errno;
    }

    lock(w, RTTI_LOCK_WRITE);

    add_hostattrs_resvals(c, w);
    idx = w->idx;
    w->host.status |= XCE_WORKER_UNAVAIL;

    unlock(w);

    return idx;
}

struct worker_config_publish_args {
    cluster_t*  c;
    endpoint_t* ep;
};

static int check_conf4worker_(xce_config_t* c)
{
    if (c != NULL) {
        return 1;
    }

    return 0;
}

static int worker_push_config(cluster_t* c, endpoint_t* ep)
{
    xce_configdata_t configdata;
    int              rc = 0;
    xlist_t          clist;

    memset(&configdata, 0x0, sizeof(configdata));

    configdata.data.conf4worker = calloc(1, sizeof(xce_conf4worker_t));
    if (configdata.data.conf4worker == NULL) {
        return -ENOMEM;
    }

    xlist_init(&clist, 256);

    collect_confvars(c, (collect_check_fn_t)check_conf4worker_, &clist);

    configdata.data.conf4worker->n_confvars = clist.count;
    configdata.data.conf4worker->confvars   = (xce_config_t**)clist.data;
    configdata.data_which                   = ConfigData_data_workerconfig;

    configdata.data.conf4worker->clustername
        = c->myconf->cluster_config->clustername;

    rc = xce_pushconfig_req(ep, c->myhost, &configdata);

    FREEUP(configdata.data.conf4worker);

    xlist_fini(&clist);

    return rc;
}

/* consider to push jobs on hosts to */
int on_worker_join(void* cls, endpoint_t* ep)
{
    cluster_t* c  = (cluster_t*)cls;
    int        rc = 0;

    rc = worker_push_config(c, ep);

    return rc;
}

int get_resval(cluster_t* c, char* resname, char* host, xce_val_t* val)
{
    xce_resource_t* r     = NULL;
    zarray_t*       array = NULL;
    xce_val_t*      v     = NULL;
    worker_t*       w     = NULL;
    int             rc    = 0;

    lock(c->resmgr, RTTI_LOCK_READ);
    rc = key_find(c->resmgr, (void*)resname, (void**)&r);
    unlock(c->resmgr);

    if (rc < 0) {
        return rc;
    }

    if (r == NULL) {
        return -ENOENT;
    }

    if ((r->flags & XCE_RESF_SHARED) == 0) {
        if ((host == NULL) || (strlen(host) == 0)) {

            return -EINVAL;
        }

        lock(c->workermgr, RTTI_LOCK_READ);
        rc = key_find(c->workermgr, (void*)host, (void**)&w);
        if (w != NULL) {
            reference(w);
        }
        unlock(c->workermgr);
        if (rc < 0) {
            return rc;
        }
        if (w == NULL) {
            return -ENOENT;
        }

        array = &(w->host.resvals);
    } else {
        array = &(c->sharedinst);
    }

    if (r->idx < array->caps) {
        v = (xce_val_t*)array->data[r->idx];
    }

    if (v == NULL) {
        if (w != NULL) {
            dereference(w);
        }

        return -ENOENT;
    }

    memcpy(val, v, sizeof(*val));

    if (w != NULL) {
        dereference(w);
    }

    return 1;
}

int mark_worker_unavail(cluster_t* c, worker_t* w)
{
    xce_run_hooks(&(c->hooksys), XCE_CONTROLLER_HOOK_WORKER, HOOK_BEFORE_EVENT,
                  HOOK_DETAIL_WORKER_DEAD, c, w);

    w->host.status |= XCE_WORKER_UNAVAIL;
    if (w->host.dead_since == 0)
        w->host.dead_since = time(NULL);

    xce_run_hooks(&(c->hooksys), XCE_CONTROLLER_HOOK_WORKER, HOOK_AFTER_EVENT,
                  HOOK_DETAIL_WORKER_DEAD, c, w);
    return 0;
}

int mark_worker_back(cluster_t* c, worker_t* w)
{
    xce_run_hooks(&(c->hooksys), XCE_CONTROLLER_HOOK_WORKER, HOOK_BEFORE_EVENT,
                  HOOK_DETAIL_WORKER_UPDATE, c, w);

    w->host.status &= ~XCE_WORKER_UNAVAIL;

    xce_run_hooks(&(c->hooksys), XCE_CONTROLLER_HOOK_WORKER, HOOK_AFTER_EVENT,
                  HOOK_DETAIL_WORKER_UPDATE, c, w);
    return 0;
}

int handle_worker_leave(cluster_t* c, int idx)
{
    worker_t* w = NULL;
    int       rc;

    TRACE("%s worker %d broken connection", __func__, idx);

    if (c == NULL) {
        SET_ERROR(EINVAL);
        return -errno;
    }

    if (idx < 0) {
        SET_ERROR(EINVAL);
        return -errno;
    }

    rc = vget(c->workermgr, idx, (void**)&w);
    if (rc < 0) {
        return rc;
    }

    if (w == NULL) {
        ERROR("fail to find worker by id %d", idx);
        SET_ERROR(E_XCE_BAD_WORKER);
        return -errno;
    }

    lock(w, RTTI_LOCK_WRITE);

    w->ep = NULL;

    mark_worker_unavail(c, w);

    unlock(w);

    return 0;
}

int worker_ctrl(cluster_t* c, xce_hostctrl_t* ctrl)
{
    worker_t* w  = NULL;
    int       rc = 0;

    if (!c || !ctrl || !ctrl->hostname || *ctrl->hostname == '\0') {
        TRACE("%s invalid arguments.", __func__);
        return -EINVAL;
    }
    TRACE("%s worker_ctrl worker %s", __func__, ctrl->hostname);

    lock(c->workermgr, RTTI_LOCK_READ);
    rc = key_find(c->workermgr, (void*)ctrl->hostname, (void**)&w);
    if (w != NULL) {
        reference(w);
    }
    unlock(c->workermgr);

    if (rc < 0) {
        return rc;
    }

    if (w == NULL) {
        PERROR("workermgr_lock");
        return -ENOENT;
    }

    lock(w, RTTI_LOCK_WRITE);
    switch (ctrl->action) {
    case HOST_OPEN:
        if (!(w->host.status & XCE_WORKER_CLOSED)) {
            goto clean;
        }
        w->host.status &= ~XCE_WORKER_CLOSED;
        break;
    case HOST_CLOSE:
        if (w->host.status & XCE_WORKER_CLOSED) {
            goto clean;
        }
        w->host.status |= XCE_WORKER_CLOSED;
        break;
    default:
        rc = -1;
        break;
    }
    if (rc == 0) {
        xce_run_hooks(&(c->hooksys), XCE_CONTROLLER_HOOK_WORKER,
                      HOOK_AFTER_EVENT, HOOK_DETAIL_WORKER_UPDATE, c, w);
    }

clean:
    unlock(w);
    dereference(w);
    return rc;
}

int do_hostcontrol_op(cluster_t* c, endpoint_t* ep, xce_request_t* request)
{
    int         rc = 0;
    xce_reply_t reply;

    init_reply_by_req(&reply, request);

    rc               = worker_ctrl(c, request->body.hostctrlmsg);
    reply.status     = rc;
    reply.data_which = RequestReply_data_nodata;

    xce_reply(ep, c->myhost, &reply);
    return rc;
}

/* Reconstruct a worker from a WAL_WORKER_CTRL Cap'n Proto payload.
 * Payload layout: [uint32_t detail][Cap'n Proto xce_worker_t].
 * Returns -EEXIST if the worker is already registered (idempotent replay). */
int worker_from_capnp(cluster_t* c, const void* capnp_payload, size_t len)
{
    if (!capnp_payload || len <= sizeof(uint32_t))
        return -EINVAL;

    /* Skip the 4-byte detail discriminator written by persist_worker_hook_fun */
    const uint8_t* data     = (const uint8_t*)capnp_payload + sizeof(uint32_t);
    size_t         data_len = len - sizeof(uint32_t);

    struct capn   c_ctx;
    capn_init_mem(&c_ctx, data, data_len, 0);
    Worker_ptr wp;
    wp.p = capn_getp(capn_root(&c_ctx), 0, 1);
    struct Worker cw;
    xce_worker_t  host;
    memset(&host, 0, sizeof(host));
    read_Worker(&cw, wp);
    decode_Worker(&host, &cw);
    capn_free(&c_ctx);

    /* Idempotency guard: skip if already present */
    worker_t* existing = NULL;
    lock(c->workermgr, RTTI_LOCK_READ);
    int found = key_find(c->workermgr, (void*)host.name, (void**)&existing);
    unlock(c->workermgr);
    if (found > 0 && existing != NULL)
        return -EEXIST;

    return restore_worker(c, host);
}

/* worker_should_retain lives in persist.c —
 * it only depends on xce_worker_t fields and persist_config_t,
 * and must be accessible from the unit test binary which links persist.c. */