myxce/src/frontend/apiserver/backend/clib.go

//go:build cgo

package backend

/*
#include <stdlib.h>
#include <stdint.h>
#include <errno.h>
#include "client.h"
#include "netio.h"
#include "helper.h"
#include "idgen.h"
#include "controller.capnp.h"
#include "val_helpers.h"

// backend_ctx_t wraps xcecli_ctx_t for use as the C handler's arg.
typedef struct {
    xcecli_ctx_t ctx;
} backend_ctx_t;

// go_push_handler is exported from Go and called from C for every incoming request.
extern int go_push_handler(xce_request_t*, void*);

static int xce_backend_handler(xce_request_t *req, void *arg) {
    return go_push_handler(req, arg);
}

// xce_backend_init initialises the client context, sets the handler,
// starts the network loop, then joins as HostRole_apiserver.
static int xce_backend_init(xce_client_config_t *config, backend_ctx_t *ctx) {
    int rc = xcecli_ctx_init(&(ctx->ctx), config);
    if (rc < 0) return rc;
    xcecli_set_handler(ctx->ctx, (xce_request_handler_t)xce_backend_handler, NULL);
    rc = xcecli_start(ctx->ctx);
    if (rc < 0) return rc;
    return xcecli_join(ctx->ctx, HostRole_apiserver, NULL);
}

// Helpers to extract typed row pointers from a dataset_t.
static xce_job_t    **dataset_jobs(dataset_t *d)    { return (xce_job_t **)d->rows; }
static xce_worker_t **dataset_workers(dataset_t *d) { return (xce_worker_t **)d->rows; }

// CGO cannot access C union fields by name (they appear as [N]byte).
// These thin helpers set the correct union branch from C so Go callers
// don't need unsafe pointer arithmetic.
static void jobctrl_set_siginfo(xce_jobctrl_t *ctrl, xce_signal_t *sig) {
    ctrl->u.siginfo = sig;
}
static void jobctrl_set_ckptoption(xce_jobctrl_t *ctrl, xce_ckpt_options_t *ckpt) {
    ctrl->u.ckptoption = ckpt;
}

// xce_val_t union accessors moved to val_helpers.h (shared with convert.go).
*/
import "C"

import (
	"apiserver/models"
	logger "apiserver/utils"
	"context"
	"fmt"
	"sync"
	"unsafe"

	"go.uber.org/zap"
)

// gCtx holds the C client context for the lifetime of the process.
var gCtx C.backend_ctx_t

// gIDCtx is the C idgen context used to generate job IDs.
var gIDCtx C.idgen_ctx_t
var idCtxOnce sync.Once

// gHandler receives push notifications from the backend.
// Must be set via StartNotificationListener before RegisterToBackend.
var gHandler NotificationHandler

func tryConnect(host string, port int, from, key string) error {
	cAddr := C.CString(host)
	cFrom := C.CString(from)
	cKey := C.CString(key)
	defer func() {
		C.free(unsafe.Pointer(cAddr))
		C.free(unsafe.Pointer(cFrom))
		C.free(unsafe.Pointer(cKey))
	}()

	var cConfig C.xce_client_config_t
	cConfig.nmasters = 1
	cConfig.masters[0] = cAddr
	cConfig.master_ports[0] = C.int(port)
	cConfig.from = cFrom
	cConfig.key = cKey

	rc := C.xce_backend_init(&cConfig, &gCtx)
	if rc < 0 {
		return fmt.Errorf("xce_backend_init failed on %s:%d: %d", host, port, int(rc))
	}
	return nil
}

// RegisterToBackend initialises the XCE client library and joins as an API
// server.  StartNotificationListener should be called first so that push
// events have a handler ready before the first push can arrive.
func RegisterToBackend(config *BackendConfig) error {
	from := config.From
	if from == "" {
		from = "xce_apiserver"
	}
	key := config.Key
	if key == "" {
		key = "keepsecret$"
	}

	if len(config.Peers) > 0 {
		var lastErr error
		for _, peer := range config.Peers {
			host, port, err := parsePeer(peer)
			if err != nil {
				logger.Warn("Skipping invalid peer", zap.String("peer", peer), zap.Error(err))
				lastErr = err
				continue
			}
			if err := tryConnect(host, port, from, key); err != nil {
				logger.Warn("Candidate unreachable, trying next",
					zap.String("host", host), zap.Int("port", port), zap.Error(err))
				lastErr = err
				continue
			}
			logger.Info("Backend connection established (HA)",
				zap.String("host", host), zap.Int("port", port))
			return nil
		}
		return fmt.Errorf("all HA candidates failed: %w", lastErr)
	}

	if err := tryConnect(config.Address, config.Port, from, key); err != nil {
		return err
	}
	logger.Info("Backend connection established",
		zap.String("address", config.Address),
		zap.Int("port", config.Port))
	return nil
}

// StartNotificationListener registers the handler that receives backend push
// events.  Call this before RegisterToBackend.
func StartNotificationListener(handler NotificationHandler) error {
	gHandler = handler
	return nil
}

// initIDGen initialises the C idgen context. Called once via sync.Once.
func initIDGen() {
	C.idgen_ctx_init(&gIDCtx)
}

// GenerateJobID calls the C idgen() to produce a unique job ID.
// clusterID and apiserverID identify the originating instance;
// arraySize=1 for normal jobs, >1 for array jobs.
func GenerateJobID(clusterID, apiserverID int, arraySize int) (uint64, error) {
	idCtxOnce.Do(initIDGen)
	var jobID C.uint64_t
	rc := C.idgen(&gIDCtx, C.size_t(clusterID), C.size_t(apiserverID),
		C.size_t(arraySize), &jobID)
	if rc < 0 {
		return 0, fmt.Errorf("idgen failed: %d", int(rc))
	}
	return uint64(jobID), nil
}

// Cleanup releases C resources. Call at process shutdown.
func Cleanup() {
	C.idgen_ctx_fini(&gIDCtx)
}

// StartHeartbeat launches a goroutine that sends periodic xcecli_heartbeat
// keep-alive messages to the controller. Stops when ctx is cancelled.
func StartHeartbeat(ctx context.Context) {
	go runHeartbeat(ctx, func(secs int) {
		C.xcecli_heartbeat(gCtx.ctx, C.int(secs))
	})
}

// go_push_handler is called by the C library on every incoming request.
// It must be exported so the C trampoline can call it from a C thread.
//
//export go_push_handler
func go_push_handler(req *C.xce_request_t, _ unsafe.Pointer) C.int {
	if req == nil {
		return 0
	}
	defer C.free_Request(req)

	if gHandler == nil {
		return 0
	}

	switch req.opcode {
	case C.OPCode_push:
		return handlePush(req)
	case C.OPCode_notify:
		if req.body_which != C.Request_body_notifymsg {
			return 0
		}
		nd := (*C.xce_notify_data_t)(unsafe.Pointer(C.xce_request_body(req)))
		if nd == nil {
			return 0
		}
		return handleNotifyData(nd)
	}
	return 0
}

// handlePush dispatches an OPCode_push request to the appropriate handler.
func handlePush(req *C.xce_request_t) C.int {
	switch req.body_which {
	case C.Request_body_datamsg:
		dataset := (*C.dataset_t)(unsafe.Pointer(C.xce_request_body(req)))
		return handleDataset(dataset)
	case C.Request_body_configpushmsg:
		cfgdata := (*C.xce_configdata_t)(unsafe.Pointer(C.xce_request_body(req)))
		return handleConfigPush(cfgdata)
	}
	return 0
}

// handleDataset processes a dataset push (jobs or workers).
func handleDataset(d *C.dataset_t) C.int {
	switch d.table {
	case C.Data_data_jobs:
		n := int(d.nrows)
		if n == 0 {
			return 0
		}
		jobs := unsafe.Slice(C.dataset_jobs(d), n)
		for _, jp := range jobs {
			if jp == nil {
				continue
			}
			info := cJobToModel(jp)
			gHandler.OnJobUpdated(info)
		}

	case C.Data_data_workers:
		n := int(d.nrows)
		if n == 0 {
			return 0
		}
		workers := unsafe.Slice(C.dataset_workers(d), n)
		for _, wp := range workers {
			if wp == nil {
				continue
			}
			info := cWorkerToModel(wp)
			gHandler.OnHostUpdated(info)
		}
	}
	return 0
}

// handleConfigPush processes a cluster configuration push.
func handleConfigPush(cfgdata *C.xce_configdata_t) C.int {
	if gHandler == nil {
		return 0
	}
	if cfgdata.data_which != C.ConfigData_data_clusterconfig {
		return 0
	}

	conf := (*C.cluster_config_t)(C.xce_configdata_ptr(cfgdata))
	if conf == nil {
		return 0
	}

	cluster := cClusterConfToClusterInfo(conf)
	if cluster == nil {
		logger.Warn("handleConfigPush: cluster_config_t has empty clustername; skipping OnClusterUpdated")
	}
	dispatchConfigPush(gHandler,
		cClusterConfQueues(conf),
		cClusterConfResources(conf),
		cClusterConfHostConfigs(conf),
		cluster,
	)
	return 0
}

// handleNotifyData processes an OPCode_notify request.
// xce_notify_data_t uses 'event' (not 'data_which') as its union discriminant.
func handleNotifyData(nd *C.xce_notify_data_t) C.int {
	if nd == nil {
		return 0
	}
	switch nd.event {
	case C.NotifyData_data_updateJob:
		jp := (*C.xce_job_t)(C.xce_notify_data_ptr(nd))
		if jp == nil {
			return 0
		}
		gHandler.OnJobUpdated(cJobToModel(jp))

	case C.NotifyData_data_jobFinishInfo:
		fi := (*C.xce_jobfinish_t)(C.xce_notify_data_ptr(nd))
		if fi == nil {
			return 0
		}
		gHandler.OnJobFinished(uint64(fi.jobid), int(fi.exit_code))

	case C.NotifyData_data_jobFailInfo:
		ff := (*C.job_failinfo_t)(C.xce_notify_data_ptr(nd))
		if ff == nil {
			return 0
		}
		gHandler.OnJobFinished(uint64(ff.jobid), int(ff.code))

	case C.NotifyData_data_schedUpdateResvals:
		rvl := (*C.xce_resval_list_t)(C.xce_notify_data_ptr(nd))
		if rvl == nil {
			return 0
		}
		n := int(rvl.count)
		if n == 0 || rvl.vals == nil {
			return 0
		}
		vals := unsafe.Slice(rvl.vals, n)
		for _, vp := range vals {
			if vp == nil {
				continue
			}
			rd := cResvalToDistribution(vp)
			if rd != nil {
				gHandler.OnResourceLocationUpdated(rd)
			}
		}
	}
	return 0
}

// SubmitJob sends a job submission to the backend via xcecli_jobsubmit.
// arraySize=1 for normal jobs, >1 for array jobs.
// The store is updated only when the backend pushes the resulting OnJobCreated
// notification — not immediately here.
func SubmitJob(job *models.JobInfo, arraySize int) error {
	js, sd, si := modelJobToSubmit(job)
	defer freeJobSubmit(js, sd, si)

	js.array_size = C.uint64_t(arraySize)

	var outJobID C.uint64_t
	rc := C.xcecli_jobsubmit(gCtx.ctx, js, &outJobID)
	if rc < 0 {
		return fmt.Errorf("xcecli_jobsubmit failed: %d", int(rc))
	}
	if outJobID != 0 {
		job.Unique_JobID = int64(outJobID)
		job.JobID = int64(ExtractArrayBaseID(uint64(outJobID)))
		job.ArrayIdx = ExtractArrayIdx(uint64(outJobID))
	}
	return nil
}

// ControlJob sends a job control action to the backend via xcecli_jobctrl.
// action must be one of the models.JOB_* constants; content is an optional
// signal name (e.g. "TERM") for kill operations.
func ControlJob(jobid uint64, action int, content string) error {
	var ctrl C.xce_jobctrl_t
	ctrl.jobid = C.uint64_t(jobid)
	ctrl.action = C.int(action)

	// For signal-based actions set signo; for checkpoint set ckpt options.
	// xce_jobctrl_t.u is a C union — we use a static C helper to set the
	// right branch safely.
	switch action {
	case models.JOB_STOP, models.JOB_RESUME, models.JOB_KILL:
		var sig C.xce_signal_t
		sig.signo = C.int(signalNumber(content, action))
		sig.force = 0
		C.jobctrl_set_siginfo(&ctrl, &sig)
	case models.JOB_CHKPNT:
		var ckpt C.xce_ckpt_options_t
		ckpt.kill = 0
		C.jobctrl_set_ckptoption(&ctrl, &ckpt)
	}

	rc := C.xcecli_jobctrl(gCtx.ctx, &ctrl)
	if rc < 0 {
		return fmt.Errorf("xcecli_jobctrl failed: %d", int(rc))
	}
	return nil
}

// ControlQueue sends a queue open/close/activate/deactivate action to the
// backend via xcecli_queuectrl.
func ControlQueue(queuename string, action int) error {
	cQueue := C.CString(queuename)
	defer C.free(unsafe.Pointer(cQueue))

	ctrl := C.xce_queuectrl_t{
		action: C.int(action),
		queue:  cQueue,
	}

	rc := C.xcecli_queuectrl(gCtx.ctx, &ctrl)
	if rc < 0 {
		return fmt.Errorf("xcecli_queuectrl failed: %d", int(rc))
	}
	return nil
}

// ControlHost sends a host open/close action to the backend via xcecli_hostctrl.
func ControlHost(hostname string, action int) error {
	cHost := C.CString(hostname)
	defer C.free(unsafe.Pointer(cHost))

	ctrl := C.xce_hostctrl_t{
		action:   C.int(action),
		hostname: cHost,
	}

	rc := C.xcecli_hostctrl(gCtx.ctx, &ctrl)
	if rc < 0 {
		return fmt.Errorf("xcecli_hostctrl failed: %d", int(rc))
	}
	return nil
}

// signalNumber maps a job action and optional signal name to a POSIX signal
// number, falling back to sensible defaults.
func signalNumber(signame string, action int) int {
	if signame != "" {
		if n, ok := models.JobSignalMapTable[signame]; ok {
			return n
		}
	}
	switch action {
	case models.JOB_KILL:
		return models.JobSignalMapTable["TERM"] // 15
	case models.JOB_STOP:
		return models.JobSignalMapTable["TSTP"] // 20
	case models.JOB_RESUME:
		return models.JobSignalMapTable["CONT"] // 18
	}
	return 15 // SIGTERM as safe default
}