redict/tests/modules/fork.c

// SPDX-FileCopyrightText: 2024 Redict Contributors
// SPDX-FileCopyrightText: 2024 Salvatore Sanfilippo <antirez at gmail dot com>
//
// SPDX-License-Identifier: BSD-3-Clause
// SPDX-License-Identifier: LGPL-3.0-only

/* define macros for having usleep */
#define _BSD_SOURCE
#define _DEFAULT_SOURCE

#include "redictmodule.h"
#include <string.h>
#include <assert.h>
#include <unistd.h>

#define UNUSED(V) ((void) V)

int child_pid = -1;
int exitted_with_code = -1;

void done_handler(int exitcode, int bysignal, void *user_data) {
    child_pid = -1;
    exitted_with_code = exitcode;
    assert(user_data==(void*)0xdeadbeef);
    UNUSED(bysignal);
}

int fork_create(RedictModuleCtx *ctx, RedictModuleString **argv, int argc)
{
    long long code_to_exit_with;
    long long usleep_us;
    if (argc != 3) {
        RedictModule_WrongArity(ctx);
        return REDICTMODULE_OK;
    }

    if(!RMAPI_FUNC_SUPPORTED(RedictModule_Fork)){
        RedictModule_ReplyWithError(ctx, "Fork api is not supported in the current redict version");
        return REDICTMODULE_OK;
    }

    RedictModule_StringToLongLong(argv[1], &code_to_exit_with);
    RedictModule_StringToLongLong(argv[2], &usleep_us);
    exitted_with_code = -1;
    int fork_child_pid = RedictModule_Fork(done_handler, (void*)0xdeadbeef);
    if (fork_child_pid < 0) {
        RedictModule_ReplyWithError(ctx, "Fork failed");
        return REDICTMODULE_OK;
    } else if (fork_child_pid > 0) {
        /* parent */
        child_pid = fork_child_pid;
        RedictModule_ReplyWithLongLong(ctx, child_pid);
        return REDICTMODULE_OK;
    }

    /* child */
    RedictModule_Log(ctx, "notice", "fork child started");
    usleep(usleep_us);
    RedictModule_Log(ctx, "notice", "fork child exiting");
    RedictModule_ExitFromChild(code_to_exit_with);
    /* unreachable */
    return 0;
}

int fork_exitcode(RedictModuleCtx *ctx, RedictModuleString **argv, int argc)
{
    UNUSED(argv);
    UNUSED(argc);
    RedictModule_ReplyWithLongLong(ctx, exitted_with_code);
    return REDICTMODULE_OK;
}

int fork_kill(RedictModuleCtx *ctx, RedictModuleString **argv, int argc)
{
    UNUSED(argv);
    UNUSED(argc);
    if (RedictModule_KillForkChild(child_pid) != REDICTMODULE_OK)
        RedictModule_ReplyWithError(ctx, "KillForkChild failed");
    else
        RedictModule_ReplyWithLongLong(ctx, 1);
    child_pid = -1;
    return REDICTMODULE_OK;
}

int RedictModule_OnLoad(RedictModuleCtx *ctx, RedictModuleString **argv, int argc) {
    UNUSED(argv);
    UNUSED(argc);
    if (RedictModule_Init(ctx,"fork",1,REDICTMODULE_APIVER_1)== REDICTMODULE_ERR)
        return REDICTMODULE_ERR;

    if (RedictModule_CreateCommand(ctx,"fork.create", fork_create,"",0,0,0) == REDICTMODULE_ERR)
        return REDICTMODULE_ERR;

    if (RedictModule_CreateCommand(ctx,"fork.exitcode", fork_exitcode,"",0,0,0) == REDICTMODULE_ERR)
        return REDICTMODULE_ERR;

    if (RedictModule_CreateCommand(ctx,"fork.kill", fork_kill,"",0,0,0) == REDICTMODULE_ERR)
        return REDICTMODULE_ERR;

    return REDICTMODULE_OK;
}