Add background_shutdown_runtime dependency to EtcdElectionClient …

…and `SourceManager` (#123)

src/meta/src/rpc/election/etcd.rs

@@ -20,43 +20,23 @@ use std::time::Duration;
 use anyhow::anyhow;
 use etcd_client::{ConnectOptions, Error, GetOptions, LeaderKey, ResignOptions};
 use risingwave_common::bail;
+use risingwave_common::util::runtime::BackgroundShutdownRuntime;
 use serde::Serialize;
 use tokio::runtime::Runtime;
 use tokio::sync::watch::Receiver;
 use tokio::sync::{oneshot, watch};
 use tokio::time;
 use tokio_stream::StreamExt;
 
+use crate::rpc::election::META_ELECTION_KEY;
 use crate::storage::WrappedEtcdClient;
-use crate::MetaResult;
-
-const META_ELECTION_KEY: &str = "__meta_election_";
-
-#[derive(Debug, Serialize)]
-pub struct ElectionMember {
-    pub id: String,
-    pub is_leader: bool,
-}
-
-#[async_trait::async_trait]
-pub trait ElectionClient: Send + Sync + 'static {
-    fn id(&self) -> MetaResult<String>;
-    async fn run_once(&self, ttl: i64, stop: Receiver<()>) -> MetaResult<()>;
-    fn subscribe(&self) -> Receiver<bool>;
-    async fn leader(&self) -> MetaResult<Option<ElectionMember>>;
-    async fn get_members(&self) -> MetaResult<Vec<ElectionMember>>;
-    async fn is_leader(&self) -> bool;
-
-    fn runtime_ref(&self) -> Option<Arc<Runtime>> {
-        None
-    }
-}
+use crate::{ElectionClient, ElectionMember, MetaResult};
 
 pub struct EtcdElectionClient {
     id: String,
     is_leader_sender: watch::Sender<bool>,
     client: WrappedEtcdClient,
-    runtime: Arc<Runtime>,
+    runtime: Arc<BackgroundShutdownRuntime>,
 }
 
 #[async_trait::async_trait]
@@ -355,18 +335,19 @@ impl EtcdElectionClient {
 
         let client = WrappedEtcdClient::connect(endpoints, options, auth_enabled).await?;
 
-        let runtime = Runtime::new().map_err(|e| {
-            anyhow!(
-                "create runtime for election client failed {}",
-                e.to_string()
-            )
-        })?;
+        let runtime = tokio::runtime::Builder::new_multi_thread()
+            .thread_name("risingwave-election-client")
+            .enable_all()
+            .build()
+            .map_err(|e| anyhow!(e))?;
+
+        let runtime: Arc<BackgroundShutdownRuntime> = Arc::new(runtime.into());
 
         Ok(Self {
             id,
             is_leader_sender: sender,
             client,
-            runtime: Arc::new(runtime),
+            runtime,
         })
     }
 }

src/meta/src/rpc/election/mod.rs

@@ -14,6 +14,9 @@
 pub mod etcd;
 pub mod sql;
 
+use std::sync::Arc;
+
+use risingwave_common::util::runtime::BackgroundShutdownRuntime;
 use serde::Serialize;
 use tokio::sync::watch::Receiver;
 
@@ -35,4 +38,8 @@ pub trait ElectionClient: Send + Sync + 'static {
     async fn leader(&self) -> MetaResult<Option<ElectionMember>>;
     async fn get_members(&self) -> MetaResult<Vec<ElectionMember>>;
     async fn is_leader(&self) -> bool;
+
+    fn runtime_ref(&self) -> Option<Arc<BackgroundShutdownRuntime>> {
+        None
+    }
 }

src/meta/src/stream/source_manager.rs

@@ -24,6 +24,7 @@ use anyhow::anyhow;
 use itertools::Itertools;
 use risingwave_common::catalog::TableId;
 use risingwave_connector::dispatch_source_prop;
+use risingwave_common::util::runtime::BackgroundShutdownRuntime;
 use risingwave_connector::source::{
     ConnectorProperties, SourceEnumeratorContext, SourceEnumeratorInfo, SourceProperties,
     SplitEnumerator, SplitId, SplitImpl, SplitMetaData,
@@ -52,6 +53,7 @@ pub struct SourceManager {
     barrier_scheduler: BarrierScheduler,
     core: Mutex<SourceManagerCore>,
     metrics: Arc<MetaMetrics>,
+    runtime: Arc<BackgroundShutdownRuntime>,
 }
 
 const MAX_FAIL_CNT: u32 = 10;
@@ -496,11 +498,20 @@ impl SourceManager {
         fragment_manager: FragmentManagerRef,
         metrics: Arc<MetaMetrics>,
     ) -> MetaResult<Self> {
+        let runtime = tokio::runtime::Builder::new_multi_thread()
+            .thread_name("risingwave-source-manager")
+            .enable_all()
+            .build()
+            .map_err(|e| anyhow!(e))?;
+
+        let runtime: Arc<BackgroundShutdownRuntime> = Arc::new(runtime.into());
+
         let mut managed_sources = HashMap::new();
         {
             let sources = catalog_manager.list_sources().await;
             for source in sources {
                 Self::create_source_worker_async(
+                    runtime.clone(),
                     env.connector_client(),
                     source,
                     &mut managed_sources,
@@ -534,6 +545,7 @@ impl SourceManager {
             core,
             paused: Mutex::new(()),
             metrics,
+            runtime,
         })
     }
 
@@ -675,6 +687,7 @@ impl SourceManager {
             tracing::warn!("source {} already registered", source.get_id());
         } else {
             Self::create_source_worker(
+                self.runtime.clone(),
                 self.env.connector_client(),
                 source,
                 &mut core.managed_sources,
@@ -687,6 +700,7 @@ impl SourceManager {
     }
 
     fn create_source_worker_async(
+        runtime: Arc<BackgroundShutdownRuntime>,
         connector_client: Option<ConnectorClient>,
         source: Source,
         managed_sources: &mut HashMap<SourceId, ConnectorSourceWorkerHandle>,
@@ -701,8 +715,7 @@ impl SourceManager {
         let source_id = source.id;
 
         let connector_properties = extract_prop_from_source(&source)?;
-
-        let handle = tokio::spawn(async move {
+        let handle = runtime.spawn(async move {
             let mut ticker = time::interval(Self::DEFAULT_SOURCE_TICK_INTERVAL);
             ticker.set_missed_tick_behavior(MissedTickBehavior::Skip);
 
@@ -745,6 +758,7 @@ impl SourceManager {
     }
 
     async fn create_source_worker(
+        runtime: Arc<BackgroundShutdownRuntime>,
         connector_client: Option<ConnectorClient>,
         source: &Source,
         managed_sources: &mut HashMap<SourceId, ConnectorSourceWorkerHandle>,
@@ -785,7 +799,7 @@ impl SourceManager {
                     })??;
             }
 
-            tokio::spawn(async move { worker.run(sync_call_rx).await })
+            runtime.spawn(async move { worker.run(sync_call_rx).await })
         });
 
         managed_sources.insert(