mirror of
https://github.com/valkey-io/valkey.git
synced 2026-07-31 21:54:23 -04:00
Jacob MurphyandGitHub
2b7d5f7760
As discussed in https://github.com/valkey-io/valkey/issues/2579 Notably, I am exposing this feature as "Atomic Slot Migration" to modules. If we want to call it something else, we should consider that now (e.g.. "Replication-Based Slot Migration"?) Also note: I am exposing both target and source node in the event. It is not guaranteed that either target or source would be the node the event fires on (e.g. replicas will fire the event after replica containment is introduced). Even though it could be potentially inferred from CLUSTER SLOTS, it should help modules parse it this way. Modules should be able to infer whether it is occurring on primary/replica from `ctx` flags, so not duplicating that here. Closes #2579 --------- Signed-off-by: Jacob Murphy <[email protected]>
414 lines
14 KiB
C
414 lines
14 KiB
C
#include "valkeymodule.h"
|
|
|
|
#include <string.h>
|
|
#include <assert.h>
|
|
|
|
/* Module configuration, save aux or not? */
|
|
#define CONF_AUX_OPTION_NO_AUX 0
|
|
#define CONF_AUX_OPTION_SAVE2 1 << 0
|
|
#define CONF_AUX_OPTION_BEFORE_KEYSPACE 1 << 1
|
|
#define CONF_AUX_OPTION_AFTER_KEYSPACE 1 << 2
|
|
#define CONF_AUX_OPTION_NO_DATA 1 << 3
|
|
long long conf_aux_count = 0;
|
|
|
|
/* Registered type */
|
|
ValkeyModuleType *testrdb_type = NULL;
|
|
|
|
/* Global values to store and persist to aux */
|
|
ValkeyModuleString *before_str = NULL;
|
|
ValkeyModuleString *after_str = NULL;
|
|
|
|
/* Global values used to keep aux from db being loaded (in case of async_loading) */
|
|
ValkeyModuleString *before_str_temp = NULL;
|
|
ValkeyModuleString *after_str_temp = NULL;
|
|
|
|
/* Indicates whether there is an async replication in progress.
|
|
* We control this value from ValkeyModuleEvent_ReplAsyncLoad events. */
|
|
int async_loading = 0;
|
|
|
|
int n_aux_load_called = 0;
|
|
|
|
void replAsyncLoadCallback(ValkeyModuleCtx *ctx, ValkeyModuleEvent e, uint64_t sub, void *data)
|
|
{
|
|
VALKEYMODULE_NOT_USED(e);
|
|
VALKEYMODULE_NOT_USED(data);
|
|
|
|
switch (sub) {
|
|
case VALKEYMODULE_SUBEVENT_REPL_ASYNC_LOAD_STARTED:
|
|
assert(async_loading == 0);
|
|
async_loading = 1;
|
|
break;
|
|
case VALKEYMODULE_SUBEVENT_REPL_ASYNC_LOAD_ABORTED:
|
|
/* Discard temp aux */
|
|
if (before_str_temp)
|
|
ValkeyModule_FreeString(ctx, before_str_temp);
|
|
if (after_str_temp)
|
|
ValkeyModule_FreeString(ctx, after_str_temp);
|
|
before_str_temp = NULL;
|
|
after_str_temp = NULL;
|
|
|
|
async_loading = 0;
|
|
break;
|
|
case VALKEYMODULE_SUBEVENT_REPL_ASYNC_LOAD_COMPLETED:
|
|
if (before_str)
|
|
ValkeyModule_FreeString(ctx, before_str);
|
|
if (after_str)
|
|
ValkeyModule_FreeString(ctx, after_str);
|
|
before_str = before_str_temp;
|
|
after_str = after_str_temp;
|
|
|
|
before_str_temp = NULL;
|
|
after_str_temp = NULL;
|
|
|
|
async_loading = 0;
|
|
break;
|
|
default:
|
|
assert(0);
|
|
}
|
|
}
|
|
|
|
void *testrdb_type_load(ValkeyModuleIO *rdb, int encver) {
|
|
int count = ValkeyModule_LoadSigned(rdb);
|
|
ValkeyModuleString *str = ValkeyModule_LoadString(rdb);
|
|
float f = ValkeyModule_LoadFloat(rdb);
|
|
long double ld = ValkeyModule_LoadLongDouble(rdb);
|
|
/* Context creation is part of the test. Creating the context will force the
|
|
* core to allocate a context, which needs to be cleaned up when the
|
|
* ValkeyModuleIO is destructed. */
|
|
ValkeyModuleCtx *ctx = ValkeyModule_GetContextFromIO(rdb);
|
|
if (ValkeyModule_IsIOError(rdb)) {
|
|
if (str)
|
|
ValkeyModule_FreeString(ctx, str);
|
|
return NULL;
|
|
}
|
|
/* Using the values only after checking for io errors. */
|
|
assert(count==1);
|
|
assert(encver==1);
|
|
assert(f==1.5f);
|
|
assert(ld==0.333333333333333333L);
|
|
return str;
|
|
}
|
|
|
|
void testrdb_type_save(ValkeyModuleIO *rdb, void *value) {
|
|
ValkeyModuleString *str = (ValkeyModuleString*)value;
|
|
ValkeyModule_SaveSigned(rdb, 1);
|
|
ValkeyModule_SaveString(rdb, str);
|
|
ValkeyModule_SaveFloat(rdb, 1.5);
|
|
ValkeyModule_SaveLongDouble(rdb, 0.333333333333333333L);
|
|
}
|
|
|
|
void testrdb_aux_save(ValkeyModuleIO *rdb, int when) {
|
|
if (!(conf_aux_count & CONF_AUX_OPTION_BEFORE_KEYSPACE)) assert(when == VALKEYMODULE_AUX_AFTER_RDB);
|
|
if (!(conf_aux_count & CONF_AUX_OPTION_AFTER_KEYSPACE)) assert(when == VALKEYMODULE_AUX_BEFORE_RDB);
|
|
assert(conf_aux_count!=CONF_AUX_OPTION_NO_AUX);
|
|
/* Context creation is part of the test. Creating the context will force the
|
|
* core to allocate a context, which needs to be cleaned up when the
|
|
* ValkeyModuleIO is destructed. */
|
|
ValkeyModuleCtx *ctx = ValkeyModule_GetContextFromIO(rdb);
|
|
VALKEYMODULE_NOT_USED(ctx);
|
|
if (when == VALKEYMODULE_AUX_BEFORE_RDB) {
|
|
if (before_str) {
|
|
ValkeyModule_SaveSigned(rdb, 1);
|
|
ValkeyModule_SaveString(rdb, before_str);
|
|
} else {
|
|
ValkeyModule_SaveSigned(rdb, 0);
|
|
}
|
|
} else {
|
|
if (after_str) {
|
|
ValkeyModule_SaveSigned(rdb, 1);
|
|
ValkeyModule_SaveString(rdb, after_str);
|
|
} else {
|
|
ValkeyModule_SaveSigned(rdb, 0);
|
|
}
|
|
}
|
|
}
|
|
|
|
int testrdb_aux_load(ValkeyModuleIO *rdb, int encver, int when) {
|
|
assert(encver == 1);
|
|
if (!(conf_aux_count & CONF_AUX_OPTION_BEFORE_KEYSPACE)) assert(when == VALKEYMODULE_AUX_AFTER_RDB);
|
|
if (!(conf_aux_count & CONF_AUX_OPTION_AFTER_KEYSPACE)) assert(when == VALKEYMODULE_AUX_BEFORE_RDB);
|
|
assert(conf_aux_count!=CONF_AUX_OPTION_NO_AUX);
|
|
ValkeyModuleCtx *ctx = ValkeyModule_GetContextFromIO(rdb);
|
|
if (when == VALKEYMODULE_AUX_BEFORE_RDB) {
|
|
if (async_loading == 0) {
|
|
if (before_str)
|
|
ValkeyModule_FreeString(ctx, before_str);
|
|
before_str = NULL;
|
|
int count = ValkeyModule_LoadSigned(rdb);
|
|
if (ValkeyModule_IsIOError(rdb))
|
|
return VALKEYMODULE_ERR;
|
|
if (count)
|
|
before_str = ValkeyModule_LoadString(rdb);
|
|
} else {
|
|
if (before_str_temp)
|
|
ValkeyModule_FreeString(ctx, before_str_temp);
|
|
before_str_temp = NULL;
|
|
int count = ValkeyModule_LoadSigned(rdb);
|
|
if (ValkeyModule_IsIOError(rdb))
|
|
return VALKEYMODULE_ERR;
|
|
if (count)
|
|
before_str_temp = ValkeyModule_LoadString(rdb);
|
|
}
|
|
} else {
|
|
if (async_loading == 0) {
|
|
if (after_str)
|
|
ValkeyModule_FreeString(ctx, after_str);
|
|
after_str = NULL;
|
|
int count = ValkeyModule_LoadSigned(rdb);
|
|
if (ValkeyModule_IsIOError(rdb))
|
|
return VALKEYMODULE_ERR;
|
|
if (count)
|
|
after_str = ValkeyModule_LoadString(rdb);
|
|
} else {
|
|
if (after_str_temp)
|
|
ValkeyModule_FreeString(ctx, after_str_temp);
|
|
after_str_temp = NULL;
|
|
int count = ValkeyModule_LoadSigned(rdb);
|
|
if (ValkeyModule_IsIOError(rdb))
|
|
return VALKEYMODULE_ERR;
|
|
if (count)
|
|
after_str_temp = ValkeyModule_LoadString(rdb);
|
|
}
|
|
}
|
|
|
|
if (ValkeyModule_IsIOError(rdb))
|
|
return VALKEYMODULE_ERR;
|
|
return VALKEYMODULE_OK;
|
|
}
|
|
|
|
void testrdb_type_free(void *value) {
|
|
if (value)
|
|
ValkeyModule_FreeString(NULL, (ValkeyModuleString*)value);
|
|
}
|
|
|
|
int testrdb_set_before(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int argc)
|
|
{
|
|
if (argc != 2) {
|
|
ValkeyModule_WrongArity(ctx);
|
|
return VALKEYMODULE_OK;
|
|
}
|
|
|
|
if (before_str)
|
|
ValkeyModule_FreeString(ctx, before_str);
|
|
before_str = argv[1];
|
|
ValkeyModule_RetainString(ctx, argv[1]);
|
|
ValkeyModule_ReplyWithLongLong(ctx, 1);
|
|
return VALKEYMODULE_OK;
|
|
}
|
|
|
|
int testrdb_get_before(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int argc)
|
|
{
|
|
VALKEYMODULE_NOT_USED(argv);
|
|
if (argc != 1){
|
|
ValkeyModule_WrongArity(ctx);
|
|
return VALKEYMODULE_OK;
|
|
}
|
|
if (before_str)
|
|
ValkeyModule_ReplyWithString(ctx, before_str);
|
|
else
|
|
ValkeyModule_ReplyWithStringBuffer(ctx, "", 0);
|
|
return VALKEYMODULE_OK;
|
|
}
|
|
|
|
/* For purpose of testing module events, expose variable state during async_loading. */
|
|
int testrdb_async_loading_get_before(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int argc)
|
|
{
|
|
VALKEYMODULE_NOT_USED(argv);
|
|
if (argc != 1){
|
|
ValkeyModule_WrongArity(ctx);
|
|
return VALKEYMODULE_OK;
|
|
}
|
|
if (before_str_temp)
|
|
ValkeyModule_ReplyWithString(ctx, before_str_temp);
|
|
else
|
|
ValkeyModule_ReplyWithStringBuffer(ctx, "", 0);
|
|
return VALKEYMODULE_OK;
|
|
}
|
|
|
|
int testrdb_set_after(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int argc)
|
|
{
|
|
if (argc != 2){
|
|
ValkeyModule_WrongArity(ctx);
|
|
return VALKEYMODULE_OK;
|
|
}
|
|
|
|
if (after_str)
|
|
ValkeyModule_FreeString(ctx, after_str);
|
|
after_str = argv[1];
|
|
ValkeyModule_RetainString(ctx, argv[1]);
|
|
ValkeyModule_ReplyWithLongLong(ctx, 1);
|
|
return VALKEYMODULE_OK;
|
|
}
|
|
|
|
int testrdb_get_after(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int argc)
|
|
{
|
|
VALKEYMODULE_NOT_USED(argv);
|
|
if (argc != 1){
|
|
ValkeyModule_WrongArity(ctx);
|
|
return VALKEYMODULE_OK;
|
|
}
|
|
if (after_str)
|
|
ValkeyModule_ReplyWithString(ctx, after_str);
|
|
else
|
|
ValkeyModule_ReplyWithStringBuffer(ctx, "", 0);
|
|
return VALKEYMODULE_OK;
|
|
}
|
|
|
|
int testrdb_set_key(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int argc)
|
|
{
|
|
if (argc != 3){
|
|
ValkeyModule_WrongArity(ctx);
|
|
return VALKEYMODULE_OK;
|
|
}
|
|
|
|
ValkeyModuleKey *key = ValkeyModule_OpenKey(ctx, argv[1], VALKEYMODULE_WRITE);
|
|
ValkeyModuleString *str = ValkeyModule_ModuleTypeGetValue(key);
|
|
if (str)
|
|
ValkeyModule_FreeString(ctx, str);
|
|
ValkeyModule_ModuleTypeSetValue(key, testrdb_type, argv[2]);
|
|
ValkeyModule_RetainString(ctx, argv[2]);
|
|
ValkeyModule_CloseKey(key);
|
|
ValkeyModule_ReplyWithLongLong(ctx, 1);
|
|
return VALKEYMODULE_OK;
|
|
}
|
|
|
|
int testrdb_get_key(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int argc)
|
|
{
|
|
if (argc != 2){
|
|
ValkeyModule_WrongArity(ctx);
|
|
return VALKEYMODULE_OK;
|
|
}
|
|
|
|
ValkeyModuleKey *key = ValkeyModule_OpenKey(ctx, argv[1], VALKEYMODULE_READ);
|
|
ValkeyModuleString *str = ValkeyModule_ModuleTypeGetValue(key);
|
|
ValkeyModule_CloseKey(key);
|
|
ValkeyModule_ReplyWithString(ctx, str);
|
|
return VALKEYMODULE_OK;
|
|
}
|
|
|
|
int testrdb_get_n_aux_load_called(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int argc)
|
|
{
|
|
VALKEYMODULE_NOT_USED(ctx);
|
|
VALKEYMODULE_NOT_USED(argv);
|
|
VALKEYMODULE_NOT_USED(argc);
|
|
ValkeyModule_ReplyWithLongLong(ctx, n_aux_load_called);
|
|
return VALKEYMODULE_OK;
|
|
}
|
|
|
|
int test2rdb_aux_load(ValkeyModuleIO *rdb, int encver, int when) {
|
|
VALKEYMODULE_NOT_USED(rdb);
|
|
VALKEYMODULE_NOT_USED(encver);
|
|
VALKEYMODULE_NOT_USED(when);
|
|
n_aux_load_called++;
|
|
return VALKEYMODULE_OK;
|
|
}
|
|
|
|
void test2rdb_aux_save(ValkeyModuleIO *rdb, int when) {
|
|
VALKEYMODULE_NOT_USED(rdb);
|
|
VALKEYMODULE_NOT_USED(when);
|
|
}
|
|
|
|
int ValkeyModule_OnLoad(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int argc) {
|
|
VALKEYMODULE_NOT_USED(argv);
|
|
VALKEYMODULE_NOT_USED(argc);
|
|
if (ValkeyModule_Init(ctx,"testrdb",1,VALKEYMODULE_APIVER_1) == VALKEYMODULE_ERR)
|
|
return VALKEYMODULE_ERR;
|
|
|
|
ValkeyModule_SetModuleOptions(ctx, VALKEYMODULE_OPTIONS_HANDLE_IO_ERRORS | VALKEYMODULE_OPTIONS_HANDLE_REPL_ASYNC_LOAD | VALKEYMODULE_OPTIONS_HANDLE_ATOMIC_SLOT_MIGRATION);
|
|
|
|
if (argc > 0)
|
|
ValkeyModule_StringToLongLong(argv[0], &conf_aux_count);
|
|
|
|
if (conf_aux_count==CONF_AUX_OPTION_NO_AUX) {
|
|
ValkeyModuleTypeMethods datatype_methods = {
|
|
.version = 1,
|
|
.rdb_load = testrdb_type_load,
|
|
.rdb_save = testrdb_type_save,
|
|
.aof_rewrite = NULL,
|
|
.digest = NULL,
|
|
.free = testrdb_type_free,
|
|
};
|
|
|
|
testrdb_type = ValkeyModule_CreateDataType(ctx, "test__rdb", 1, &datatype_methods);
|
|
if (testrdb_type == NULL)
|
|
return VALKEYMODULE_ERR;
|
|
} else if (!(conf_aux_count & CONF_AUX_OPTION_NO_DATA)) {
|
|
ValkeyModuleTypeMethods datatype_methods = {
|
|
.version = VALKEYMODULE_TYPE_METHOD_VERSION,
|
|
.rdb_load = testrdb_type_load,
|
|
.rdb_save = testrdb_type_save,
|
|
.aof_rewrite = NULL,
|
|
.digest = NULL,
|
|
.free = testrdb_type_free,
|
|
.aux_load = testrdb_aux_load,
|
|
.aux_save = testrdb_aux_save,
|
|
.aux_save_triggers = ((conf_aux_count & CONF_AUX_OPTION_BEFORE_KEYSPACE) ? VALKEYMODULE_AUX_BEFORE_RDB : 0) |
|
|
((conf_aux_count & CONF_AUX_OPTION_AFTER_KEYSPACE) ? VALKEYMODULE_AUX_AFTER_RDB : 0)
|
|
};
|
|
|
|
if (conf_aux_count & CONF_AUX_OPTION_SAVE2) {
|
|
datatype_methods.aux_save2 = testrdb_aux_save;
|
|
}
|
|
|
|
testrdb_type = ValkeyModule_CreateDataType(ctx, "test__rdb", 1, &datatype_methods);
|
|
if (testrdb_type == NULL)
|
|
return VALKEYMODULE_ERR;
|
|
} else {
|
|
|
|
/* Used to verify that aux_save2 api without any data, saves nothing to the RDB. */
|
|
ValkeyModuleTypeMethods datatype_methods = {
|
|
.version = VALKEYMODULE_TYPE_METHOD_VERSION,
|
|
.aux_load = test2rdb_aux_load,
|
|
.aux_save = test2rdb_aux_save,
|
|
.aux_save_triggers = ((conf_aux_count & CONF_AUX_OPTION_BEFORE_KEYSPACE) ? VALKEYMODULE_AUX_BEFORE_RDB : 0) |
|
|
((conf_aux_count & CONF_AUX_OPTION_AFTER_KEYSPACE) ? VALKEYMODULE_AUX_AFTER_RDB : 0)
|
|
};
|
|
if (conf_aux_count & CONF_AUX_OPTION_SAVE2) {
|
|
datatype_methods.aux_save2 = test2rdb_aux_save;
|
|
}
|
|
|
|
ValkeyModule_CreateDataType(ctx, "test__rdb", 1, &datatype_methods);
|
|
}
|
|
|
|
if (ValkeyModule_CreateCommand(ctx,"testrdb.set.before", testrdb_set_before,"deny-oom",0,0,0) == VALKEYMODULE_ERR)
|
|
return VALKEYMODULE_ERR;
|
|
|
|
if (ValkeyModule_CreateCommand(ctx,"testrdb.get.before", testrdb_get_before,"",0,0,0) == VALKEYMODULE_ERR)
|
|
return VALKEYMODULE_ERR;
|
|
|
|
if (ValkeyModule_CreateCommand(ctx,"testrdb.async_loading.get.before", testrdb_async_loading_get_before,"",0,0,0) == VALKEYMODULE_ERR)
|
|
return VALKEYMODULE_ERR;
|
|
|
|
if (ValkeyModule_CreateCommand(ctx,"testrdb.set.after", testrdb_set_after,"deny-oom",0,0,0) == VALKEYMODULE_ERR)
|
|
return VALKEYMODULE_ERR;
|
|
|
|
if (ValkeyModule_CreateCommand(ctx,"testrdb.get.after", testrdb_get_after,"",0,0,0) == VALKEYMODULE_ERR)
|
|
return VALKEYMODULE_ERR;
|
|
|
|
if (ValkeyModule_CreateCommand(ctx,"testrdb.set.key", testrdb_set_key,"deny-oom",1,1,1) == VALKEYMODULE_ERR)
|
|
return VALKEYMODULE_ERR;
|
|
|
|
if (ValkeyModule_CreateCommand(ctx,"testrdb.get.key", testrdb_get_key,"",1,1,1) == VALKEYMODULE_ERR)
|
|
return VALKEYMODULE_ERR;
|
|
|
|
if (ValkeyModule_CreateCommand(ctx,"testrdb.get.n_aux_load_called", testrdb_get_n_aux_load_called,"",1,1,1) == VALKEYMODULE_ERR)
|
|
return VALKEYMODULE_ERR;
|
|
|
|
ValkeyModule_SubscribeToServerEvent(ctx,
|
|
ValkeyModuleEvent_ReplAsyncLoad, replAsyncLoadCallback);
|
|
|
|
return VALKEYMODULE_OK;
|
|
}
|
|
|
|
int ValkeyModule_OnUnload(ValkeyModuleCtx *ctx) {
|
|
if (before_str)
|
|
ValkeyModule_FreeString(ctx, before_str);
|
|
if (after_str)
|
|
ValkeyModule_FreeString(ctx, after_str);
|
|
if (before_str_temp)
|
|
ValkeyModule_FreeString(ctx, before_str_temp);
|
|
if (after_str_temp)
|
|
ValkeyModule_FreeString(ctx, after_str_temp);
|
|
return VALKEYMODULE_OK;
|
|
}
|