fix: cache miss after NATs restarts (#6726)

Refreshing in-memory cache in such NATs event trigger by NATs client's
callback closure.

---------

Co-authored-by: Hengfei Yang <hengfei.yang@gmail.com>

src/common/infra/cluster/nats.rs

@@ -23,14 +23,25 @@ use config::{
     utils::json,
 };
 use infra::{
-    db::{NEED_WATCH, get_coordinator},
+    db::{NEED_WATCH, get_coordinator, nats},
     dist_lock,
     errors::{Error, Result},
 };
-use tokio::task;
+use tokio::{sync::mpsc, task};
+
+use crate::common::infra::config::update_cache;
 
 /// Register and keep alive the node to cluster
 pub(crate) async fn register_and_keep_alive() -> Result<()> {
+    // first, init NATs client with channel communicating NATs events
+    let (nats_event_tx, nats_event_rx) = mpsc::channel::<nats::NatsEvent>(10);
+    if let Err(e) = nats::init_nats_client(nats_event_tx).await {
+        log::error!("[CLUSTER] NATs client init failed: {}", e);
+        return Err(e);
+    }
+    // a child task to handle NATs event
+    task::spawn(async move { update_cache(nats_event_rx).await });
+
     if let Err(e) = register().await {
         log::error!("[CLUSTER] register failed: {}", e);
         return Err(e);

src/common/infra/config.rs

@@ -33,7 +33,7 @@ use hashbrown::HashMap;
 use infra::table::short_urls::ShortUrlRecord;
 use once_cell::sync::Lazy;
 use parking_lot::RwLock;
-use tokio::sync::RwLock as TokioRwLock;
+use tokio::sync::{RwLock as TokioRwLock, mpsc};
 use vector_enrichment::TableRegistry;
 
 use crate::{
@@ -101,3 +101,46 @@ pub static WS_SESSIONS: Lazy<RwAHashMap<String, Arc<TokioRwLock<WsSession>>>> =
     Lazy::new(Default::default);
 pub static USER_ROLES_CACHE: Lazy<RwAHashMap<String, CachedUserRoles>> =
     Lazy::new(Default::default);
+
+/// Refreshes in-memory cache in the event of NATs restart.
+///
+/// We should add all in-memory caches listed above that a CacheMiss is not followed by
+/// reading from db, e.g. UserSession, Pipeline
+pub(crate) async fn update_cache(mut nats_event_rx: mpsc::Receiver<infra::db::nats::NatsEvent>) {
+    let mut first_conenction = true;
+    while let Some(event) = nats_event_rx.recv().await {
+        if let infra::db::nats::NatsEvent::Connected = event {
+            // Skip refreshing if it's the first time connecting to the client
+            if first_conenction {
+                first_conenction = false;
+                continue;
+            }
+            log::info!(
+                "[infra::config] received NATs event: {event}, refreshing in-memory cache for Alerts, Pipelines, RealtimeTriggers, Schema, Users, and UserSessions"
+            );
+            if let Err(e) = crate::service::db::alerts::alert::cache().await {
+                log::error!("Error refreshing in-memory cache \"Alerts\": {}", e);
+            }
+            if let Err(e) = crate::service::db::pipeline::cache().await {
+                log::error!("Error refreshing in-memory cache \"Pipelines\": {}", e);
+            }
+            if let Err(e) = crate::service::db::alerts::realtime_triggers::cache().await {
+                log::error!(
+                    "Error refreshing in-memory cache \"RealtimeTriggers\": {}",
+                    e
+                );
+            }
+            if let Err(e) = crate::service::db::schema::cache().await {
+                log::error!("Error refreshing in-memory cache \"Schema\": {}", e);
+            }
+            if let Err(e) = crate::service::db::session::cache().await {
+                log::error!("Error refreshing in-memory cache \"UserSessions\": {}", e);
+            }
+            if let Err(e) = crate::service::db::user::cache().await {
+                log::error!("Error refreshing in-memory cache \"Users\": {}", e);
+            }
+        }
+    }
+
+    log::info!("[infra::config] stops to listen to NATs event to refresh in-memory caches");
+}

src/infra/src/db/nats.rs

@@ -21,6 +21,7 @@ use std::{
     time::Duration,
 };
 
+pub use async_nats::Event as NatsEvent;
 use async_nats::{Client, ServerAddr, jetstream};
 use async_trait::async_trait;
 use bytes::Bytes;
@@ -49,8 +50,22 @@ const SUPER_CLUSTER_PREFIX: &str = "super_cluster_kv_";
 
 static NATS_CLIENT: OnceCell<Client> = OnceCell::const_new();
 
-pub async fn get_nats_client() -> &'static Client {
-    NATS_CLIENT.get_or_init(connect).await
+/// Initialize a global NATS client with a mpsc channel sender for sending NATs events.
+pub async fn init_nats_client(nats_event_sender: mpsc::Sender<async_nats::Event>) -> Result<()> {
+    let client = connect(nats_event_sender).await;
+
+    NATS_CLIENT
+        .set(client)
+        .map_err(|e| Error::Message(format!("[NATS:init] failed to set global client: {}", e)))
+}
+
+pub async fn get_nats_client() -> Result<Client> {
+    NATS_CLIENT
+        .get()
+        .ok_or(Error::Message(
+            "[NATS:get_nats_client] NATs client not initialized".to_string(),
+        ))
+        .cloned()
 }
 
 async fn get_bucket_by_key<'a>(
@@ -58,7 +73,7 @@ async fn get_bucket_by_key<'a>(
     key: &'a str,
 ) -> Result<(jetstream::kv::Store, &'a str)> {
     let cfg = get_config();
-    let client = get_nats_client().await.clone();
+    let client = get_nats_client().await?;
     let jetstream = jetstream::new(client);
     let key = key.trim_start_matches('/');
     let bucket_name = key.split('/').next().unwrap();
@@ -147,7 +162,7 @@ impl super::Db for NatsDb {
     }
 
     async fn stats(&self) -> Result<super::Stats> {
-        let client = get_nats_client().await.clone();
+        let client = get_nats_client().await?;
         let jetstream = async_nats::jetstream::new(client);
         let mut keys_count = 0;
         let mut bytes_len = 0;
@@ -545,7 +560,7 @@ pub async fn create_table() -> Result<()> {
     Ok(())
 }
 
-pub async fn connect() -> async_nats::Client {
+pub async fn connect(nats_event_sender: mpsc::Sender<async_nats::Event>) -> async_nats::Client {
     let cfg = get_config();
     if cfg.common.print_key_config {
         log::info!("Nats init get_config(): {:?}", cfg.nats);
@@ -559,6 +574,17 @@ pub async fn connect() -> async_nats::Client {
     if !cfg.nats.user.is_empty() {
         opts = opts.user_and_password(cfg.nats.user.to_string(), cfg.nats.password.to_string());
     }
+    opts = opts.event_callback(move |event| {
+        let sender = nats_event_sender.clone();
+        async move {
+            if let Err(e) = sender.send(event).await {
+                log::error!(
+                    "NATs client event callback channel failed to send event: {}",
+                    e
+                );
+            }
+        }
+    });
     let addrs = cfg
         .nats
         .addr

src/infra/src/queue/nats.rs

@@ -68,7 +68,7 @@ impl Default for NatsQueue {
 impl super::Queue for NatsQueue {
     async fn create(&self, topic: &str) -> Result<()> {
         let cfg = config::get_config();
-        let client = get_nats_client().await.clone();
+        let client = get_nats_client().await?;
         let jetstream = jetstream::new(client);
         let topic_name = format!("{}{}", self.prefix, topic);
         let config = jetstream::stream::Config {
@@ -85,7 +85,7 @@ impl super::Queue for NatsQueue {
 
     /// you can pub message with the topic or topic.* to match the topic
     async fn publish(&self, topic: &str, value: Bytes) -> Result<()> {
-        let client = get_nats_client().await.clone();
+        let client = get_nats_client().await?;
         let jetstream = jetstream::new(client);
         // Publish a message to the stream
         let topic_name = format!("{}{}", self.prefix, topic);
@@ -100,7 +100,7 @@ impl super::Queue for NatsQueue {
         let consumer_name = self.consumer_name.clone();
         let is_durable = self.is_durable;
         let _task: JoinHandle<Result<()>> = tokio::task::spawn(async move {
-            let client = get_nats_client().await.clone();
+            let client = get_nats_client().await?;
             let jetstream = jetstream::new(client);
             let stream = jetstream.get_stream(&stream_name).await?;
             let config = jetstream::consumer::pull::Config {