diff --git a/include/scheduler.h b/include/scheduler.h new file mode 100644 index 0000000..e48fcb3 --- /dev/null +++ b/include/scheduler.h @@ -0,0 +1,33 @@ +#ifndef RPC_SCHEDULER_H +#define RPC_SCHEDULER_H + +#include "rpc/server.h" + +#include + +typedef struct rpc_scheduler rpc_scheduler; +typedef struct rpc_call rpc_call; +typedef void (*rpc_call_done_fn)(rpc_call *call, void *user_data); + +struct rpc_ctx { + uint64_t call_id; + uint32_t proc_id; + void *server; +}; + +int rpc_scheduler_init(rpc_scheduler **out); +void rpc_scheduler_destroy(rpc_scheduler *scheduler); +int rpc_scheduler_submit(rpc_scheduler *scheduler, uint64_t call_id, + uint32_t proc_id, rpc_handler_fn handler, + void *handler_data, const uint8_t *payload, + size_t payload_len, rpc_call_done_fn done, + void *done_data); +void rpc_scheduler_run_ready(rpc_scheduler *scheduler); + +int rpc_call_result(const rpc_call *call); +const rpc_writer *rpc_call_response(const rpc_call *call); +const char *rpc_call_error(const rpc_call *call); +uint64_t rpc_call_id(const rpc_call *call); +uint32_t rpc_call_proc_id(const rpc_call *call); + +#endif diff --git a/src/scheduler.c b/src/scheduler.c new file mode 100644 index 0000000..568207e --- /dev/null +++ b/src/scheduler.c @@ -0,0 +1,238 @@ +#include "scheduler.h" + +#include "arena.h" + +#define MCO_USE_VMEM_ALLOCATOR +#define MCO_ZERO_MEMORY +#define MCO_DEFAULT_STACK_SIZE (1024 * 1024) +#define MINICORO_IMPL +#include "minicoro.h" + +#include +#include +#include + +#define RPC_CALL_ARENA_CAPACITY 262144u + +struct rpc_scheduler { + rpc_call *head; + rpc_call *tail; + rpc_fixed_arena call_arena; +}; + +struct rpc_call { + mco_coro *co; + rpc_ctx ctx; + rpc_handler_fn handler; + void *handler_data; + uint8_t *payload; + size_t payload_len; + rpc_value *args; + size_t argc; + rpc_writer response; + int result; + char error[160]; + int completed; + rpc_call_done_fn done; + void *done_data; + rpc_call *next_ready; +}; + +static void call_free(rpc_scheduler *scheduler, rpc_call *call) { + if (!call) { + return; + } + if (call->co) { + (void)mco_destroy(call->co); + } + rpc_values_free(call->args); + rpc_writer_free(&call->response); + free(call->payload); + rpc_fixed_arena_free(&scheduler->call_arena, call); +} + +static int enqueue(rpc_scheduler *scheduler, rpc_call *call) { + call->next_ready = NULL; + if (scheduler->tail) { + scheduler->tail->next_ready = call; + } else { + scheduler->head = call; + } + scheduler->tail = call; + return 0; +} + +static rpc_call *dequeue(rpc_scheduler *scheduler) { + rpc_call *call = scheduler->head; + if (!call) { + return NULL; + } + scheduler->head = call->next_ready; + if (!scheduler->head) { + scheduler->tail = NULL; + } + call->next_ready = NULL; + return call; +} + +static void call_entry(mco_coro *co) { + rpc_call *call = mco_get_user_data(co); + call->result = call->handler(&call->ctx, call->args, call->argc, + &call->response, call->handler_data); + call->completed = 1; + if (call->result != 0 && call->error[0] == '\0') { + snprintf(call->error, sizeof(call->error), "procedure returned %d", + call->result); + } +} + +int rpc_scheduler_init(rpc_scheduler **out) { + if (!out) { + return -1; + } + rpc_scheduler *scheduler = calloc(1, sizeof(*scheduler)); + if (!scheduler) { + return -1; + } + if (rpc_fixed_arena_init(&scheduler->call_arena, sizeof(rpc_call), + RPC_CALL_ARENA_CAPACITY) != 0) { + free(scheduler); + return -1; + } + *out = scheduler; + return 0; +} + +void rpc_scheduler_destroy(rpc_scheduler *scheduler) { + if (!scheduler) { + return; + } + rpc_call *call = NULL; + while ((call = dequeue(scheduler)) != NULL) { + call_free(scheduler, call); + } + rpc_fixed_arena_destroy(&scheduler->call_arena); + free(scheduler); +} + +int rpc_scheduler_submit(rpc_scheduler *scheduler, uint64_t call_id, + uint32_t proc_id, rpc_handler_fn handler, + void *handler_data, const uint8_t *payload, + size_t payload_len, rpc_call_done_fn done, + void *done_data) { + if (!scheduler || !handler || (!payload && payload_len > 0)) { + return -1; + } + + rpc_call *call = rpc_fixed_arena_alloc(&scheduler->call_arena); + if (!call) { + return -1; + } + call->ctx.call_id = call_id; + call->ctx.proc_id = proc_id; + call->handler = handler; + call->handler_data = handler_data; + call->payload_len = payload_len; + call->done = done; + call->done_data = done_data; + rpc_writer_init(&call->response); + + if (payload_len > 0) { + call->payload = malloc(payload_len); + if (!call->payload) { + call_free(scheduler, call); + return -1; + } + memcpy(call->payload, payload, payload_len); + } + + if (rpc_payload_decode(call->payload, call->payload_len, &call->args, + &call->argc) != 0) { + snprintf(call->error, sizeof(call->error), "malformed payload"); + call->result = -1; + call->completed = 1; + done(call, done_data); + call_free(scheduler, call); + return 0; + } + + mco_desc desc = mco_desc_init(call_entry, 0); + desc.user_data = call; + mco_result rc = mco_create(&call->co, &desc); + if (rc != MCO_SUCCESS) { + snprintf(call->error, sizeof(call->error), "coroutine create failed: %s", + mco_result_description(rc)); + call->result = -1; + call->completed = 1; + done(call, done_data); + call_free(scheduler, call); + return 0; + } + + if (enqueue(scheduler, call) != 0) { + call_free(scheduler, call); + return -1; + } + return 0; +} + +void rpc_scheduler_run_ready(rpc_scheduler *scheduler) { + if (!scheduler) { + return; + } + + rpc_call *call = NULL; + while ((call = dequeue(scheduler)) != NULL) { + mco_result rc = mco_resume(call->co); + if (rc != MCO_SUCCESS) { + snprintf(call->error, sizeof(call->error), "coroutine resume failed: %s", + mco_result_description(rc)); + call->result = -1; + call->completed = 1; + } + + if (!call->completed && mco_status(call->co) == MCO_SUSPENDED) { + if (enqueue(scheduler, call) == 0) { + continue; + } + snprintf(call->error, sizeof(call->error), "scheduler enqueue failed"); + call->result = -1; + call->completed = 1; + } + + if (call->done) { + call->done(call, call->done_data); + } + call_free(scheduler, call); + } +} + +int rpc_call_result(const rpc_call *call) { return call ? call->result : -1; } + +const rpc_writer *rpc_call_response(const rpc_call *call) { + return call ? &call->response : NULL; +} + +const char *rpc_call_error(const rpc_call *call) { + return call && call->error[0] ? call->error : "procedure failed"; +} + +uint64_t rpc_call_id(const rpc_call *call) { + return call ? call->ctx.call_id : 0; +} + +uint32_t rpc_call_proc_id(const rpc_call *call) { + return call ? call->ctx.proc_id : 0; +} + +uint64_t rpc_ctx_call_id(const rpc_ctx *ctx) { return ctx ? ctx->call_id : 0; } + +uint32_t rpc_ctx_proc_id(const rpc_ctx *ctx) { return ctx ? ctx->proc_id : 0; } + +void rpc_ctx_yield(rpc_ctx *ctx) { + (void)ctx; + mco_coro *co = mco_running(); + if (co) { + (void)mco_yield(co); + } +} diff --git a/tests/test_scheduler.c b/tests/test_scheduler.c new file mode 100644 index 0000000..b2b1e49 --- /dev/null +++ b/tests/test_scheduler.c @@ -0,0 +1,47 @@ +#include "scheduler.h" + +#include + +static int done_count; + +static int yielding_handler(rpc_ctx *ctx, const rpc_value *args, size_t argc, + rpc_writer *out, void *user_data) { + int *seen = user_data; + assert(argc == 1); + assert(args[0].type == RPC_TYPE_I64); + *seen += 1; + rpc_ctx_yield(ctx); + *seen += 1; + return rpc_writer_i64(out, args[0].as.i64 + 1); +} + +static void on_done(rpc_call *call, void *user_data) { + (void)user_data; + done_count++; + assert(rpc_call_result(call) == 0); + const rpc_writer *w = rpc_call_response(call); + rpc_value *values = NULL; + size_t count = 0; + assert(rpc_payload_decode(w->data, w->len, &values, &count) == 0); + assert(count == 1); + assert(values[0].type == RPC_TYPE_I64); + assert(values[0].as.i64 == 42); + rpc_values_free(values); +} + +int main(void) { + rpc_scheduler *scheduler = NULL; + int seen = 0; + rpc_writer payload; + rpc_writer_init(&payload); + assert(rpc_writer_i64(&payload, 41) == 0); + assert(rpc_scheduler_init(&scheduler) == 0); + assert(rpc_scheduler_submit(scheduler, 7, 9, yielding_handler, &seen, + payload.data, payload.len, on_done, NULL) == 0); + rpc_scheduler_run_ready(scheduler); + assert(seen == 2); + assert(done_count == 1); + rpc_scheduler_destroy(scheduler); + rpc_writer_free(&payload); + return 0; +}