Merge pull request #9334 from drmingdrmer/50-drop-watcher

feat(meta/watch): remove stream when a watch client is dropped

Author: drmingdrmer

src/meta/service/src/api/grpc/grpc_service.rs

@@ -58,6 +58,7 @@ use crate::version::from_digit_ver;
 use crate::version::to_digit_ver;
 use crate::version::METASRV_SEMVER;
 use crate::version::MIN_METACLI_SEMVER;
+use crate::watcher::WatchStream;
 
 pub struct MetaServiceImpl {
     token: GrpcToken,
@@ -239,10 +240,20 @@ impl MetaService for MetaServiceImpl {
     ) -> Result<Response<Self::WatchStream>, Status> {
         let (tx, rx) = mpsc::channel(4);
 
-        self.meta_node.add_watcher(request.into_inner(), tx);
+        let mn = &self.meta_node;
 
-        let output_stream = tokio_stream::wrappers::ReceiverStream::new(rx);
-        Ok(Response::new(Box::pin(output_stream) as Self::WatchStream))
+        let add_res = mn.add_watcher(request.into_inner(), tx).await;
+
+        match add_res {
+            Ok(watcher) => {
+                let stream = WatchStream::new(rx, watcher, mn.dispatcher_handle.clone());
+                Ok(Response::new(Box::pin(stream) as Self::WatchStream))
+            }
+            Err(e) => {
+                // TODO: test error return.
+                Err(Status::invalid_argument(e))
+            }
+        }
     }
 
     async fn transaction(

src/meta/service/src/meta_service/raftmeta.rs

@@ -61,6 +61,7 @@ use common_meta_types::MetaOperationError;
 use common_meta_types::MetaStartupError;
 use common_meta_types::Node;
 use common_meta_types::NodeId;
+use futures::channel::oneshot;
 use itertools::Itertools;
 use openraft::Config;
 use openraft::LogId;
@@ -86,7 +87,7 @@ use crate::store::RaftStoreBare;
 use crate::watcher::DispatcherSender;
 use crate::watcher::EventDispatcher;
 use crate::watcher::EventDispatcherHandle;
-use crate::watcher::WatchEvent;
+use crate::watcher::Watcher;
 use crate::watcher::WatcherSender;
 use crate::Opened;
 
@@ -1085,10 +1086,22 @@ impl MetaNode {
         Ok(resp)
     }
 
-    pub(crate) fn add_watcher(&self, request: WatchRequest, tx: WatcherSender) {
-        let _ = self
-            .dispatcher_handle
-            .tx
-            .send(WatchEvent::AddWatcher((request, tx)));
+    pub(crate) async fn add_watcher(
+        &self,
+        request: WatchRequest,
+        tx: WatcherSender,
+    ) -> Result<Watcher, &'static str> {
+        let (resp_tx, resp_rx) = oneshot::channel();
+
+        self.dispatcher_handle.request(|d: &mut EventDispatcher| {
+            let add_res = d.add_watcher(request, tx);
+            let _ = resp_tx.send(add_res);
+        });
+
+        let recv_res = resp_rx.await;
+        match recv_res {
+            Ok(add_res) => add_res,
+            Err(_e) => Err("dispatcher closed"),
+        }
     }
 }

src/meta/service/src/watcher/mod.rs

@@ -18,8 +18,8 @@ mod watcher_stream;
 pub(crate) use watcher_manager::DispatcherSender;
 pub(crate) use watcher_manager::EventDispatcher;
 pub use watcher_manager::EventDispatcherHandle;
-pub(crate) use watcher_manager::WatchEvent;
 pub use watcher_manager::WatcherId;
 pub use watcher_manager::WatcherSender;
-pub use watcher_stream::WatcherInfo;
-pub use watcher_stream::WatcherStream;
+pub use watcher_stream::WatchStream;
+pub use watcher_stream::WatchStreamHandle;
+pub use watcher_stream::Watcher;

src/meta/service/src/watcher/watcher_manager.rs

@@ -15,7 +15,6 @@ use core::ops::Range;
 
 use common_base::base::tokio;
 use common_base::base::tokio::sync::mpsc;
-use common_base::base::tokio::sync::mpsc::Sender;
 use common_base::base::tokio::sync::oneshot;
 use common_base::rangemap::RangeMap;
 use common_base::rangemap::RangeMapKey;
@@ -31,23 +30,22 @@ use tonic::Status;
 use tracing::info;
 use tracing::warn;
 
-use super::WatcherStream;
+use super::WatchStreamHandle;
 use crate::metrics::network_metrics;
 use crate::metrics::server_metrics;
+use crate::watcher::Watcher;
 
 pub type WatcherId = i64;
 
 /// A sender for dispatcher to send event to interested watchers.
-pub type WatcherSender = Sender<Result<WatchResponse, Status>>;
+pub type WatcherSender = mpsc::Sender<Result<WatchResponse, Status>>;
 
 /// A sender for event source, such as raft state machine, to send event to [`EventDispatcher`].
 #[derive(Clone, Debug)]
 pub(crate) struct DispatcherSender(pub(crate) mpsc::UnboundedSender<WatchEvent>);
 
+/// An event sent to EventDispatcher.
 pub(crate) enum WatchEvent {
-    /// Inform the dispatcher to add a new watcher.
-    AddWatcher((WatchRequest, WatcherSender)),
-
     /// Submit a kv change event to dispatcher
     KVChange(Change<Vec<u8>, String>),
 
@@ -59,6 +57,7 @@ pub(crate) enum WatchEvent {
     },
 }
 
+#[derive(Clone, Debug)]
 pub struct EventDispatcherHandle {
     /// For sending event or command to the dispatcher.
     pub(crate) tx: mpsc::UnboundedSender<WatchEvent>,
@@ -94,7 +93,7 @@ pub struct EventDispatcher {
     event_rx: mpsc::UnboundedReceiver<WatchEvent>,
 
     /// map range to WatcherId
-    watcher_range_map: RangeMap<String, WatcherId, WatcherStream>,
+    watcher_range_map: RangeMap<String, WatcherId, WatchStreamHandle>,
 
     current_watcher_id: WatcherId,
 }
@@ -120,9 +119,6 @@ impl EventDispatcher {
         loop {
             if let Some(event) = self.event_rx.recv().await {
                 match event {
-                    WatchEvent::AddWatcher((req, tx)) => {
-                        self.add_watcher(req, tx).await;
-                    }
                     WatchEvent::KVChange(kv_change) => {
                         self.dispatch_event(kv_change).await;
                     }
@@ -178,57 +174,64 @@ impl EventDispatcher {
                     watcher_id, err
                 );
                 remove_range_keys.push(RangeMapKey::new(
-                    stream.watcher.key.clone()..stream.watcher.key_end.clone(),
+                    stream.watcher.key_range.clone(),
                     watcher_id,
                 ));
             };
         }
 
         // TODO: when a watcher stream is dropped, send a event to remove the watcher explicitly
         for range_key in remove_range_keys {
-            self.remove_watcher(range_key);
+            self.remove_watcher(&range_key);
         }
     }
 
     #[tracing::instrument(level = "debug", skip(self))]
-    pub async fn add_watcher(&mut self, create: WatchRequest, tx: WatcherSender) {
-        info!("create_watcher_stream: {:?}", create);
-
-        let range = match EventDispatcher::get_range_key(create.key.clone(), &create.key_end) {
+    pub fn add_watcher(
+        &mut self,
+        create: WatchRequest,
+        tx: WatcherSender,
+    ) -> Result<Watcher, &'static str> {
+        info!("add_watcher: {:?}", create);
+
+        let range = match EventDispatcher::build_key_range(create.key.clone(), &create.key_end) {
             Ok(range) => range,
-            Err(_) => return,
+            Err(e) => return Err(e),
         };
 
         self.current_watcher_id += 1;
         let watcher_id = self.current_watcher_id;
         let filter: FilterType = create.filter_type();
 
-        let watcher_stream = WatcherStream::new(
-            watcher_id,
-            filter,
-            tx,
-            range.start.clone(),
-            range.end.clone(),
-        );
+        let watcher = Watcher::new(watcher_id, filter, range.clone());
+        let stream_handle = WatchStreamHandle::new(watcher.clone(), tx);
 
         self.watcher_range_map
-            .insert(range, watcher_id, watcher_stream);
+            .insert(range, watcher_id, stream_handle);
 
         server_metrics::incr_watchers(1);
+
+        Ok(watcher)
     }
 
-    #[tracing::instrument(level = "debug", skip(self))]
-    fn remove_watcher(&mut self, key: RangeMapKey<String, WatcherId>) {
-        self.watcher_range_map.remove_by_key(&key);
+    #[tracing::instrument(level = "debug", skip_all)]
+    pub fn remove_watcher(&mut self, key: &RangeMapKey<String, WatcherId>) {
+        info!("remove_watcher: {:?}", key);
+
+        self.watcher_range_map.remove_by_key(key);
 
+        // TODO: decrease it only when the key is actually removed
         server_metrics::incr_watchers(-1);
     }
 
-    fn get_range_key(key: String, key_end: &Option<String>) -> Result<Range<String>, bool> {
+    fn build_key_range(
+        key: String,
+        key_end: &Option<String>,
+    ) -> Result<Range<String>, &'static str> {
         match key_end {
             Some(key_end) => {
                 if &key > key_end {
-                    return Err(false);
+                    return Err("empty range");
                 }
                 Ok(key..key_end.to_string())
             }

src/meta/service/src/watcher/watcher_stream.rs

@@ -12,48 +12,54 @@
 //  See the License for the specific language governing permissions and
 //  limitations under the License.
 
+use std::ops::Range;
+use std::pin::Pin;
+use std::task::Context;
+use std::task::Poll;
+
 use common_base::base::tokio::sync::mpsc::error::SendError;
+use common_base::base::tokio::sync::mpsc::Receiver;
+use common_base::rangemap::RangeMapKey;
 use common_meta_types::protobuf::watch_request::FilterType;
 use common_meta_types::protobuf::WatchResponse;
+use futures::Stream;
 use tonic::Status;
 
 use super::WatcherId;
 use super::WatcherSender;
+use crate::watcher::EventDispatcherHandle;
 
 /// Attributes of a watcher that is interested in kv change events.
-pub struct WatcherInfo {
+#[derive(Clone, Debug)]
+pub struct Watcher {
     pub id: WatcherId,
 
+    /// Defines how to filter keys with `key_range`.
     pub filter_type: FilterType,
 
-    pub key: String,
-
-    pub key_end: String,
+    /// The range of key this watcher is interested in.
+    pub key_range: Range<String>,
 }
 
-pub struct WatcherStream {
-    pub watcher: WatcherInfo,
+impl Watcher {
+    pub fn new(id: WatcherId, filter_type: FilterType, key_range: Range<String>) -> Self {
+        Self {
+            id,
+            filter_type,
+            key_range,
+        }
+    }
+}
 
+/// A handle of a watching stream, for feeding messages to the stream.
+pub struct WatchStreamHandle {
+    pub watcher: Watcher,
     tx: WatcherSender,
 }
 
-impl WatcherStream {
-    pub fn new(
-        id: WatcherId,
-        filter_type: FilterType,
-        tx: WatcherSender,
-        key: String,
-        key_end: String,
-    ) -> Self {
-        WatcherStream {
-            watcher: WatcherInfo {
-                id,
-                filter_type,
-                key,
-                key_end,
-            },
-            tx,
-        }
+impl WatchStreamHandle {
+    pub fn new(watcher: Watcher, tx: WatcherSender) -> Self {
+        WatchStreamHandle { watcher, tx }
     }
 
     pub async fn send(
@@ -63,3 +69,59 @@ impl WatcherStream {
         self.tx.send(Ok(resp)).await
     }
 }
+
+/// A wrapper around [`tokio::sync::mpsc::Receiver`] that implements [`Stream`].
+#[derive(Debug)]
+pub struct WatchStream<T> {
+    inner: Receiver<T>,
+    watcher: Watcher,
+    dispatcher: EventDispatcherHandle,
+}
+
+impl<T> Drop for WatchStream<T> {
+    fn drop(&mut self) {
+        let rng = self.watcher.key_range.clone();
+        let id = self.watcher.id;
+
+        self.dispatcher.request(move |d| {
+            let key = RangeMapKey::new(rng, id);
+            d.remove_watcher(&key)
+        })
+    }
+}
+
+impl<T> WatchStream<T> {
+    /// Create a new `WatcherStream`.
+    pub fn new(rx: Receiver<T>, watcher: Watcher, dispatcher: EventDispatcherHandle) -> Self {
+        Self {
+            inner: rx,
+            watcher,
+            dispatcher,
+        }
+    }
+
+    /// Closes the receiving half of a channel without dropping it.
+    pub fn close(&mut self) {
+        self.inner.close()
+    }
+}
+
+impl<T> Stream for WatchStream<T> {
+    type Item = T;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        self.inner.poll_recv(cx)
+    }
+}
+
+impl<T> AsRef<Receiver<T>> for WatchStream<T> {
+    fn as_ref(&self) -> &Receiver<T> {
+        &self.inner
+    }
+}
+
+impl<T> AsMut<Receiver<T>> for WatchStream<T> {
+    fn as_mut(&mut self) -> &mut Receiver<T> {
+        &mut self.inner
+    }
+}