feat: on_init_unsafe and on_post_upgrade_unsafe

src/libs/macros/src/lib.rs

@@ -447,3 +447,61 @@ pub fn on_post_upgrade(attr: TokenStream, item: TokenStream) -> TokenStream {
 pub fn on_init(attr: TokenStream, item: TokenStream) -> TokenStream {
     hook_macro(Hook::OnInit, attr, item)
 }
+
+/// The `on_post_upgrade_unsafe` function is a procedural macro attribute for hooking into the `OnPostUpgrade` event.
+/// Unlike `on_post_upgrade`, this variant **executes synchronously** and is intended for advanced use cases
+/// where immediate execution is required.
+///
+/// This macro should only be used in scenarios where deferred execution (via async hooks) is **not an option**.
+/// Regular users should use `#[on_post_upgrade]` instead.
+///
+/// # Warning ⚠️
+/// - This function **executes immediately** during a satellite upgrade.
+/// - It **bypasses** the usual async execution model, which may lead to **unexpected behavior** if used improperly.
+/// - Any error in this function will cause the upgrade to fail!!!
+/// - If the upgrade fails, the satellite might become unresponsive and lose its data.
+/// - **Developers should prefer `on_post_upgrade` unless they explicitly need synchronous behavior.**
+///
+/// # Example (Restricted Usage)
+///
+/// ```rust
+/// #[on_post_upgrade_unsafe]
+/// fn on_post_upgrade_critical() {
+///     // Perform necessary actions immediately on upgrade
+/// }
+/// ```
+///
+/// **Note:** This function is hidden from public documentation to discourage general usage.
+#[doc(hidden)]
+#[proc_macro_attribute]
+pub fn on_post_upgrade_unsafe(attr: TokenStream, item: TokenStream) -> TokenStream {
+    hook_macro(Hook::OnInitUnsafe, attr, item)
+}
+
+/// The `on_init_unsafe` function is a procedural macro attribute for hooking into the `OnInit` event.
+/// It serves the same purpose as `on_init`, but **executes synchronously**, ensuring that initialization
+/// logic is run **immediately** instead of being deferred.
+///
+/// This function is intended for **special cases only** where an immediate initialization
+/// step is necessary before any asynchronous tasks are triggered.
+///
+/// # Warning ⚠️
+/// - This function **runs immediately** at initialization time.
+/// - It **bypasses deferred execution**, meaning long-running operations **may slow down startup**.
+/// - Regular users **should use `on_init` instead**, unless there's a strict requirement for synchronous behavior.
+///
+/// # Example (Restricted Usage)
+///
+/// ```rust
+/// #[on_init_unsafe]
+/// fn critical_on_init() {
+///     // This runs synchronously before any async operations are triggered
+/// }
+/// ```
+///
+/// **Note:** This function is hidden from public documentation to prevent unintended usage.
+#[doc(hidden)]
+#[proc_macro_attribute]
+pub fn on_init_unsafe(attr: TokenStream, item: TokenStream) -> TokenStream {
+    hook_macro(Hook::OnInitUnsafe, attr, item)
+}

src/libs/macros/src/parser.rs

@@ -25,7 +25,9 @@ pub enum Hook {
     OnDeleteManyAssets,
     OnDeleteFilteredAssets,
     OnInit,
+    OnInitUnsafe,
     OnPostUpgrade,
+    OnPostUpgradeUnsafe,
     AssertSetDoc,
     AssertDeleteDoc,
     AssertUploadAsset,
@@ -47,6 +49,8 @@ fn map_hook_name(hook: Hook) -> String {
         Hook::OnDeleteFilteredAssets => "juno_on_delete_filtered_assets".to_string(),
         Hook::OnInit => "juno_on_init".to_string(),
         Hook::OnPostUpgrade => "juno_on_post_upgrade".to_string(),
+        Hook::OnInitUnsafe => "juno_on_init_unsafe".to_string(),
+        Hook::OnPostUpgradeUnsafe => "juno_on_post_upgrade_unsafe".to_string(),
         Hook::AssertSetDoc => "juno_assert_set_doc".to_string(),
         Hook::AssertDeleteDoc => "juno_assert_delete_doc".to_string(),
         Hook::AssertUploadAsset => "juno_assert_upload_asset".to_string(),
@@ -106,7 +110,9 @@ fn parse_hook(hook: &Hook, attr: TokenStream, item: TokenStream) -> Result<Token
     let hook_fn = Ident::new(&map_hook_name(hook.clone()), proc_macro2::Span::call_site());
 
     match hook {
-        Hook::OnPostUpgrade | Hook::OnInit => parse_lifecycle_hook(&ast, signature, &hook_fn),
+        Hook::OnPostUpgrade | Hook::OnInit | Hook::OnPostUpgradeUnsafe | Hook::OnInitUnsafe => {
+            parse_lifecycle_hook(&ast, signature, &hook_fn)
+        }
         _ => parse_doc_hook(&ast, signature, &hook_fn, hook, attr),
     }
 }

src/libs/satellite/Cargo.toml

@@ -26,6 +26,8 @@ on_delete_many_assets = []
 on_delete_filtered_assets = []
 on_init = []
 on_post_upgrade = []
+on_init_unsafe = []
+on_post_upgrade_unsafe = []
 assert_set_doc = []
 assert_delete_doc = []
 assert_upload_asset = []

src/libs/satellite/src/hooks.rs

@@ -55,6 +55,9 @@ extern "Rust" {
 
     fn juno_on_init();
     fn juno_on_post_upgrade();
+
+    fn juno_on_init_unsafe();
+    fn juno_on_post_upgrade_unsafe();
 }
 
 #[allow(unused_variables)]
@@ -385,6 +388,16 @@ pub fn invoke_on_init() {
     }
 }
 
+#[allow(unused_variables)]
+pub fn invoke_on_init_unsafe() {
+    #[cfg(feature = "on_init_unsafe")]
+    {
+        unsafe {
+            juno_on_init_unsafe();
+        }
+    }
+}
+
 #[allow(unused_variables)]
 pub fn invoke_on_post_upgrade() {
     #[cfg(feature = "on_post_upgrade")]
@@ -397,6 +410,18 @@ pub fn invoke_on_post_upgrade() {
     }
 }
 
+#[allow(unused_variables)]
+pub fn invoke_on_post_upgrade_unsafe() {
+    #[cfg(feature = "on_post_upgrade_unsafe")]
+    {
+        unsafe {
+            set_timer(Duration::ZERO, || {
+                juno_on_post_upgrade_unsafe();
+            });
+        }
+    }
+}
+
 fn should_invoke_doc_hook<T>(
     collections: Option<Vec<String>>,
     context: &HookContext<DocContext<T>>,

src/libs/satellite/src/satellite.rs

@@ -18,8 +18,8 @@ use crate::db::types::state::{Doc, DocContext, DocUpsert};
 use crate::hooks::{
     invoke_on_delete_asset, invoke_on_delete_doc, invoke_on_delete_filtered_assets,
     invoke_on_delete_filtered_docs, invoke_on_delete_many_assets, invoke_on_delete_many_docs,
-    invoke_on_init, invoke_on_post_upgrade, invoke_on_set_doc, invoke_on_set_many_docs,
-    invoke_upload_asset,
+    invoke_on_init, invoke_on_init_unsafe, invoke_on_post_upgrade, invoke_on_post_upgrade_unsafe,
+    invoke_on_set_doc, invoke_on_set_many_docs, invoke_upload_asset,
 };
 use crate::memory::{get_memory_upgrades, init_stable_state, STATE};
 use crate::random::defer_init_random_seed;
@@ -88,6 +88,8 @@ pub fn init() {
         };
     });
 
+    invoke_on_init_unsafe();
+
     invoke_on_init();
 }
 
@@ -111,6 +113,8 @@ pub fn post_upgrade() {
     defer_init_certified_assets();
     defer_init_random_seed();
 
+    invoke_on_post_upgrade_unsafe();
+
     invoke_on_post_upgrade();
 }