swapdb_event

Signed-off-by: Dmitry Polyakovsky <[email protected]>

Author: Dmitry Polyakovsky

examples/server_events.rs

@@ -13,7 +13,7 @@ use valkey_module_macros::{
     client_changed_event_handler, config_changed_event_handler, cron_event_handler,
     flush_event_handler, fork_child_event_handler, key_event_handler,
     master_link_change_event_handler, persistence_event_handler, repl_async_load_event_handler,
-    replica_change_event_handler, shutdown_event_handler,
+    replica_change_event_handler, shutdown_event_handler, swapdb_event_handler,
 };
 
 static NUM_FLUSHES: AtomicI64 = AtomicI64::new(0);
@@ -27,6 +27,7 @@ static IS_MASTER_LINK_UP: AtomicBool = AtomicBool::new(false);
 static NUM_FORK_CHILD_EVENTS: AtomicI64 = AtomicI64::new(0);
 static NUM_REPLICA_CHANGE_EVENTS: AtomicI64 = AtomicI64::new(0);
 static NUM_REPL_ASYNC_LOAD_EVENTS: AtomicI64 = AtomicI64::new(0);
+static NUM_SWAP_DB_EVENTS: AtomicI64 = AtomicI64::new(0);
 
 #[flush_event_handler]
 fn flushed_event_handler(_ctx: &Context, flush_event: FlushSubevent) {
@@ -179,6 +180,12 @@ fn repl_async_load_event_handler(ctx: &Context, repl_async_load_subevent: ReplAs
     NUM_REPL_ASYNC_LOAD_EVENTS.fetch_add(1, Ordering::SeqCst);
 }
 
+#[swapdb_event_handler]
+fn swapdb_event_handler(ctx: &Context, _: u64) {
+    NUM_SWAP_DB_EVENTS.fetch_add(1, Ordering::SeqCst);
+    ctx.log_notice("Databases swapped");
+}
+
 fn num_flushed(_ctx: &Context, _args: Vec<ValkeyString>) -> ValkeyResult {
     Ok(ValkeyValue::Integer(NUM_FLUSHES.load(Ordering::SeqCst)))
 }
@@ -235,7 +242,13 @@ fn num_repl_async_load_events(_ctx: &Context, _args: Vec<ValkeyString>) -> Valke
     ))
 }
 
-fn init(ctx: &Context, args: &[ValkeyString]) -> Status {
+fn num_swapdb_events(_ctx: &Context, _args: Vec<ValkeyString>) -> ValkeyResult {
+    Ok(ValkeyValue::Integer(
+        NUM_SWAP_DB_EVENTS.load(Ordering::SeqCst),
+    ))
+}
+
+fn init(ctx: &Context, _args: &[ValkeyString]) -> Status {
     // https://valkey.io/topics/modules-api-ref/#ValkeyModule_SetModuleOptions
     // otherwise you get:  Skipping diskless-load because there are modules that are not aware of async replication.
     // needed for repl_async_load_event_handler
@@ -262,6 +275,7 @@ valkey_module! {
         ["is_master_link_up", is_master_link_up, "readonly", 0, 0, 0],
         ["num_fork_child_events", num_fork_child_events, "readonly", 0, 0, 0],
         ["num_replica_change_events", num_replica_change_events, "readonly", 0, 0, 0],
-        ["num_repl_async_load_events", num_repl_async_load_events, "readonly", 0, 0, 0]
+        ["num_repl_async_load_events", num_repl_async_load_events, "readonly", 0, 0, 0],
+        ["num_swapdb_events", num_swapdb_events, "readonly", 0, 0, 0],
     ]
 }

src/context/server_events.rs

@@ -140,6 +140,9 @@ pub static REPLICA_CHANGE_SERVER_EVENTS_LIST: [fn(&Context, ReplicaChangeSubeven
 #[distributed_slice()]
 pub static REPL_ASYNC_LOAD_SERVER_EVENTS_LIST: [fn(&Context, ReplAsyncLoadSubevent)] = [..];
 
+#[distributed_slice()]
+pub static SWAPDB_SERVER_EVENTS_LIST: [fn(&Context, u64)] = [..];
+
 extern "C" fn cron_callback(
     ctx: *mut raw::RedisModuleCtx,
     _eid: raw::RedisModuleEvent,
@@ -401,6 +404,18 @@ extern "C" fn repl_async_load_event_callback(
         });
 }
 
+extern "C" fn swapdb_callback(
+    ctx: *mut raw::RedisModuleCtx,
+    _eid: raw::RedisModuleEvent,
+    subevent: u64,
+    _data: *mut ::std::os::raw::c_void,
+) {
+    let ctx = Context::new(ctx);
+    SWAPDB_SERVER_EVENTS_LIST.iter().for_each(|callback| {
+        callback(&ctx, subevent);
+    });
+}
+
 fn register_single_server_event_type<T>(
     ctx: &Context,
     callbacks: &[fn(&Context, T)],
@@ -511,5 +526,11 @@ pub fn register_server_events(ctx: &Context) -> Result<(), ValkeyError> {
         raw::REDISMODULE_EVENT_REPL_ASYNC_LOAD,
         Some(repl_async_load_event_callback),
     )?;
+    register_single_server_event_type(
+        ctx,
+        &SWAPDB_SERVER_EVENTS_LIST,
+        raw::REDISMODULE_EVENT_SWAPDB,
+        Some(swapdb_callback),
+    )?;
     Ok(())
 }

tests/integration.rs

@@ -1943,8 +1943,25 @@ fn test_repl_asnc_load_event() -> Result<()> {
 
     // check num_repl_async_load_events on the replica
     let event_count1: i64 = redis::cmd("num_repl_async_load_events").query(&mut replica_con)?;
-    // started and completed fire so result is 2, aborted even tdoes not fire
+    // started and completed fire so result is 2, aborted event does not fire
     assert_eq!(event_count1, 2);
 
     Ok(())
 }
+
+#[test]
+fn test_swapdb_event() -> Result<()> {
+    let port: u16 = 6528;
+    let _guards = vec![start_valkey_server_with_module("server_events", port)
+        .with_context(|| FAILED_TO_START_SERVER)?];
+    let mut con = get_valkey_connection(port).with_context(|| FAILED_TO_CONNECT_TO_SERVER)?;
+
+    // run swapdb between db 0 and db 1
+    let _: () = redis::cmd("swapdb").arg(&["0", "1"]).query(&mut con)?;
+
+    // check swapdb event count
+    let event_count1: i64 = redis::cmd("num_swapdb_events").query(&mut con)?;
+    assert_eq!(event_count1, 1);
+
+    Ok(())
+}

valkeymodule-rs-macros/src/lib.rs

@@ -614,4 +614,26 @@ pub fn repl_async_load_event_handler(_attr: TokenStream, item: TokenStream) -> T
         #ast
     };
     gen.into()
-}
+}
+
+/// Proc macro which is set on a function that need to be called whenever a swapdb event happens.
+/// The function must accept a [Context] and [u64].
+///
+/// Example:
+///
+/// ```rust,no_run,ignore
+/// #[swapdb_event_handler]
+/// fn swapdb_event_handler(ctx: &Context, _: u64) { ... }
+/// ```
+#[proc_macro_attribute]
+pub fn swapdb_event_handler(_attr: TokenStream, item: TokenStream) -> TokenStream {
+    let ast: ItemFn = match syn::parse(item) {
+        Ok(res) => res,
+        Err(e) => return e.to_compile_error().into(),
+    };
+    let gen = quote! {
+        #[linkme::distributed_slice(valkey_module::server_events::SWAPDB_SERVER_EVENTS_LIST)]
+        #ast
+    };
+    gen.into()
+}