feat(common): LocalSystemParamManager for worker node

src/common/Cargo.toml

@@ -15,6 +15,7 @@ normal = ["workspace-hack"]
 
 [dependencies]
 anyhow = "1"
+arc-swap = "1"
 arrow-array = "33"
 arrow-schema = "33"
 async-trait = "0.1"

src/common/src/system_param/local_manager.rs

@@ -0,0 +1,93 @@
+// Copyright 2023 RisingWave Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::ops::Deref;
+use std::sync::Arc;
+
+use arc_swap::ArcSwap;
+use risingwave_pb::meta::SystemParams;
+use tokio::sync::broadcast::{channel, Sender};
+
+use super::reader::SystemParamsReader;
+use crate::util::channel_util::broadcast::IgnoreLaggedReceiver;
+
+pub type SystemParamsReaderRef = Arc<ArcSwap<SystemParamsReader>>;
+
+/// The system parameter manager on worker nodes. It provides two methods for other components to
+/// read the latest system parameters:
+/// - `get_params` returns a reference to the latest parameters that is atomically updated.
+/// - `subscribe_params` returns a channel on which calling `recv` will get the latest parameters.
+///   Compared with `get_params`, the caller can be explicitly notifed of parameter change.
+pub struct LocalSystemParamManager {
+    /// The latest parameters.
+    params: SystemParamsReaderRef,
+
+    /// Sender of the latest parameters.
+    tx: Sender<SystemParamsReaderRef>,
+}
+
+impl LocalSystemParamManager {
+    pub fn new(params: SystemParamsReader) -> Self {
+        // Only care about latest params.
+        let (tx, _) = channel(1);
+        Self {
+            params: Arc::new(ArcSwap::from_pointee(params)),
+            tx,
+        }
+    }
+
+    pub fn get_params(&self) -> SystemParamsReaderRef {
+        self.params.clone()
+    }
+
+    pub fn try_set_params(&self, new_params: SystemParams) {
+        let new_params_reader = SystemParamsReader::from(new_params);
+        if self.params.load().deref().deref() != &new_params_reader {
+            self.params.store(Arc::new(new_params_reader.clone()));
+            // Ignore no active receiver.
+            let _ = self.tx.send(self.params.clone());
+        }
+    }
+
+    pub fn subscribe_parmams(&self) -> IgnoreLaggedReceiver<SystemParamsReaderRef> {
+        self.tx.subscribe().into()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_manager() {
+        let p = SystemParams::default().into();
+        let manager = LocalSystemParamManager::new(p);
+        let shared_params = manager.get_params();
+
+        let new_params = SystemParams {
+            sstable_size_mb: Some(1),
+            ..Default::default()
+        };
+
+        let mut params_rx = manager.subscribe_parmams();
+
+        manager.try_set_params(new_params.clone());
+
+        assert_eq!(
+            **params_rx.recv().await.unwrap().load(),
+            new_params.clone().into()
+        );
+        assert_eq!(**shared_params.load(), new_params.into());
+    }
+}

src/common/src/system_param/mod.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+pub mod local_manager;
 pub mod reader;
 
 use std::collections::HashSet;

src/common/src/system_param/reader.rs

@@ -21,7 +21,7 @@ use super::system_params_to_kv;
 /// - Avoid misuse of deprecated fields by hiding their getters.
 /// - Abstract fallback logic for fields that might not be provided by meta service due to backward
 ///   compatibility.
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, PartialEq)]
 pub struct SystemParamsReader {
     prost: ProstSystemParams,
 }

src/common/src/util/channel_util.rs

@@ -0,0 +1,62 @@
+// Copyright 2023 RisingWave Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+pub mod broadcast {
+    use tokio::sync::broadcast::error::RecvError;
+    use tokio::sync::broadcast::Receiver;
+
+    /// A wrapper around `[tokio::sync::broadcast::Receiver]` that instead of returning
+    /// `Err(RecvError::Lagged)` when the capacity is reached, keep trying `recv` and ignoring
+    /// skipped items until success.
+    pub struct IgnoreLaggedReceiver<T> {
+        inner: Receiver<T>,
+    }
+
+    impl<T: Clone> IgnoreLaggedReceiver<T> {
+        pub async fn recv(&mut self) -> Result<T, RecvError> {
+            loop {
+                match self.inner.recv().await {
+                    Ok(v) => break Ok(v),
+                    Err(RecvError::Lagged(_)) => continue,
+                    Err(e) => break Err(e),
+                }
+            }
+        }
+    }
+
+    impl<T> From<Receiver<T>> for IgnoreLaggedReceiver<T> {
+        fn from(inner: Receiver<T>) -> Self {
+            Self { inner }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    mod broadcast {
+        use tokio::sync::broadcast::channel;
+
+        use crate::util::channel_util::broadcast::IgnoreLaggedReceiver;
+
+        #[tokio::test]
+        async fn test_ignore_lagged_receiver() {
+            let (tx, rx) = channel(1);
+            let mut rx = IgnoreLaggedReceiver::from(rx);
+            for i in 0..2 {
+                tx.send(i as u32).unwrap();
+            }
+            assert_eq!(rx.recv().await, Ok(1));
+        }
+    }
+}

src/common/src/util/mod.rs

@@ -20,6 +20,7 @@ use crate::error::ErrorCode::InternalError;
 use crate::error::{Result, RwError};
 
 pub mod addr;
+pub mod channel_util;
 pub mod chunk_coalesce;
 pub mod compress;
 pub mod encoding_for_comparison;