From e567f7c3a146ebdbc0bdbc9f3c6ca8304fb331f6 Mon Sep 17 00:00:00 2001 From: Taiming Liu Date: Tue, 6 May 2025 19:33:31 -0700 Subject: [PATCH] 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 --- src/common/infra/cluster/nats.rs | 15 +++++++++-- src/common/infra/config.rs | 45 +++++++++++++++++++++++++++++++- src/infra/src/db/nats.rs | 36 +++++++++++++++++++++---- src/infra/src/queue/nats.rs | 6 ++--- 4 files changed, 91 insertions(+), 11 deletions(-) diff --git a/src/common/infra/cluster/nats.rs b/src/common/infra/cluster/nats.rs index caeaa7be81..fcb0db7b6f 100644 --- a/src/common/infra/cluster/nats.rs +++ b/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::(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); diff --git a/src/common/infra/config.rs b/src/common/infra/config.rs index 9ddd3200b7..13452b6185 100644 --- a/src/common/infra/config.rs +++ b/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>>> = Lazy::new(Default::default); pub static USER_ROLES_CACHE: Lazy> = 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) { + 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"); +} diff --git a/src/infra/src/db/nats.rs b/src/infra/src/db/nats.rs index 022fa0fa94..b8ff7020f3 100644 --- a/src/infra/src/db/nats.rs +++ b/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 = 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) -> 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 { + 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 { - 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::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 diff --git a/src/infra/src/queue/nats.rs b/src/infra/src/queue/nats.rs index 674e4069a4..5d742e628d 100644 --- a/src/infra/src/queue/nats.rs +++ b/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> = 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 {