refactor: use Weak ref to break cyclic references (#17743)

MetaServiceImpl is not dropped when there is an alive connection.
Thus it can not hold a Strong reference to MetaNode.
Otherwise when stopping, `MetaNode` is not dropped.
This might prevent releasing other resources held by MetaNode.

Author: drmingdrmer

src/meta/client/src/grpc_client.rs

@@ -988,6 +988,8 @@ impl MetaGrpcClient {
 
         let mut client = self.get_established_client().await?;
 
+        // Since 1.2.677, initial_flush is added to WatchRequest.
+        // If the server is not upto date to support this feature, return an error.
         if watch_request.initial_flush {
             let server_version = client.server_protocol_version();
             let least_server_version = 1002677;

src/meta/raft-store/src/applier.rs

@@ -128,6 +128,7 @@ where SM: StateMachineApi + 'static
         // Send queued change events to subscriber
         if let Some(sender) = self.sm.event_sender() {
             for event in self.changes.drain(..) {
+                debug!("send to EventSender: {:?}", event);
                 sender.send(event);
             }
         }

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

@@ -15,6 +15,7 @@
 use std::io;
 use std::pin::Pin;
 use std::sync::Arc;
+use std::sync::Weak;
 
 use arrow_flight::BasicAuth;
 use databend_common_base::base::tokio::sync::mpsc;
@@ -74,9 +75,11 @@ use watcher::key_range::build_key_range;
 use watcher::util::new_watch_sink;
 use watcher::util::try_forward;
 use watcher::watch_stream::WatchStream;
+use watcher::watch_stream::WatchStreamSender;
 
 use crate::message::ForwardRequest;
 use crate::message::ForwardRequestBody;
+use crate::meta_service::watcher::DispatcherHandle;
 use crate::meta_service::watcher::WatchTypes;
 use crate::meta_service::MetaNode;
 use crate::metrics::network_metrics;
@@ -88,17 +91,36 @@ use crate::version::MIN_METACLI_SEMVER;
 
 pub struct MetaServiceImpl {
     token: GrpcToken,
-    pub(crate) meta_node: Arc<MetaNode>,
+    /// MetaServiceImpl is not dropped if there is an alive connection.
+    ///
+    /// Thus make the reference to [`MetaNode`] a Weak reference so that it does not prevent [`MetaNode`] to be dropped
+    pub(crate) meta_node: Weak<MetaNode>,
+}
+
+impl Drop for MetaServiceImpl {
+    fn drop(&mut self) {
+        if let Some(meta_node) = self.meta_node.upgrade() {
+            debug!("MetaServiceImpl::drop: id={}", meta_node.raft_store.id);
+        } else {
+            debug!("MetaServiceImpl::drop: inner MetaNode already dropped");
+        }
+    }
 }
 
 impl MetaServiceImpl {
-    pub fn create(meta_node: Arc<MetaNode>) -> Self {
+    pub fn create(meta_node: Weak<MetaNode>) -> Self {
         Self {
             token: GrpcToken::create(),
             meta_node,
         }
     }
 
+    pub fn try_get_meta_node(&self) -> Result<Arc<MetaNode>, Status> {
+        self.meta_node.upgrade().ok_or_else(|| {
+            Status::internal("MetaNode is already dropped, can not serve new requests")
+        })
+    }
+
     fn check_token(&self, metadata: &MetadataMap) -> Result<GrpcClaim, Status> {
         let token = metadata
             .get_bin("auth-token-bin")
@@ -117,7 +139,7 @@ impl MetaServiceImpl {
         let req: MetaGrpcReq = request.try_into()?;
         debug!("{}: Received MetaGrpcReq: {:?}", func_name!(), req);
 
-        let m = &self.meta_node;
+        let m = self.try_get_meta_node()?;
         let reply = match &req {
             MetaGrpcReq::UpsertKV(a) => {
                 let res = m
@@ -144,8 +166,9 @@ impl MetaServiceImpl {
 
         let req = ForwardRequest::new(1, req);
 
-        let res = self
-            .meta_node
+        let meta_node = self.try_get_meta_node()?;
+
+        let res = meta_node
             .handle_forwardable_request::<MetaGrpcReadReq>(req.clone())
             .log_elapsed_info(format!("ReadRequest: {:?}", req))
             .await
@@ -168,8 +191,9 @@ impl MetaServiceImpl {
 
         let forward_req = ForwardRequest::new(1, ForwardRequestBody::Write(ent));
 
-        let forward_res = self
-            .meta_node
+        let meta_node = self.try_get_meta_node()?;
+
+        let forward_res = meta_node
             .handle_forwardable_request(forward_req)
             .log_elapsed_info(format!("TxnRequest: {}", txn))
             .await;
@@ -342,7 +366,8 @@ impl MetaService for MetaServiceImpl {
     ) -> Result<Response<Self::ExportStream>, Status> {
         let _guard = RequestInFlight::guard();
 
-        let meta_node = &self.meta_node;
+        let meta_node = self.try_get_meta_node()?;
+
         let strm = meta_node.raft_store.inner().export();
 
         let chunk_size = 32;
@@ -370,7 +395,8 @@ impl MetaService for MetaServiceImpl {
     ) -> Result<Response<Self::ExportV1Stream>, Status> {
         let _guard = RequestInFlight::guard();
 
-        let meta_node = &self.meta_node;
+        let meta_node = self.try_get_meta_node()?;
+
         let strm = meta_node.raft_store.inner().export();
 
         let chunk_size = request.get_ref().chunk_size.unwrap_or(32) as usize;
@@ -400,18 +426,13 @@ impl MetaService for MetaServiceImpl {
 
         let (tx, rx) = mpsc::channel(4);
 
-        let mn = &self.meta_node;
+        let mn = self.try_get_meta_node()?;
 
-        let sender = mn.add_watcher(watch, tx.clone()).await?;
+        let weak_sender = mn.add_watcher(watch, tx.clone()).await?;
 
-        let dispatcher_handle = mn.dispatcher_handle.clone();
+        let weak_handle = Arc::downgrade(&mn.dispatcher_handle);
         let on_drop = move || {
-            let Some(sender) = sender.upgrade() else {
-                return;
-            };
-            dispatcher_handle.request(move |dispatcher| {
-                dispatcher.remove_watcher(sender);
-            })
+            try_remove_sender(weak_sender, weak_handle, "on-drop-WatchStream");
         };
 
         let stream = WatchStream::new(rx, Box::new(on_drop));
@@ -453,7 +474,8 @@ impl MetaService for MetaServiceImpl {
 
         let _guard = RequestInFlight::guard();
 
-        let meta_node = &self.meta_node;
+        let meta_node = self.try_get_meta_node()?;
+
         let members = meta_node.get_grpc_advertise_addrs().await;
 
         let resp = MemberListReply { data: members };
@@ -467,8 +489,10 @@ impl MetaService for MetaServiceImpl {
         _request: Request<Empty>,
     ) -> Result<Response<ClusterStatus>, Status> {
         let _guard = RequestInFlight::guard();
-        let status = self
-            .meta_node
+
+        let meta_node = self.try_get_meta_node()?;
+
+        let status = meta_node
             .get_status()
             .await
             .map_err(|e| Status::internal(format!("get meta node status failed: {}", e)))?;
@@ -543,3 +567,38 @@ fn thread_tracking_guard<T>(req: &tonic::Request<T>) -> Option<TrackingGuard> {
 
     None
 }
+
+/// Try to remove a [`WatchStream`] from the subscriber.
+///
+/// This function receives two weak references: one to the sender and one to the dispatcher handle.
+/// Using weak references prevents memory leaks by avoiding cyclic references when these are captured in closures.
+/// If either reference can't be upgraded, the function logs the situation and returns early.
+fn try_remove_sender(
+    weak_sender: Weak<WatchStreamSender<WatchTypes>>,
+    weak_handle: Weak<DispatcherHandle>,
+    ctx: &str,
+) {
+    debug!("{ctx}: try removing:  {:?}", weak_sender);
+
+    let Some(d) = weak_handle.upgrade() else {
+        debug!(
+            "{ctx}: when try removing {:?}; dispatcher is already dropped",
+            weak_sender
+        );
+        return;
+    };
+
+    let Some(sender) = weak_sender.upgrade() else {
+        debug!(
+            "on_drop is called for WatchStream {:?}; but sender is already dropped",
+            weak_sender
+        );
+        return;
+    };
+
+    debug!("{ctx}: request is sent to remove watcher: {}", sender);
+
+    d.request(move |dispatcher| {
+        dispatcher.remove_watcher(sender);
+    })
+}

src/meta/service/src/api/grpc_server.rs

@@ -26,6 +26,8 @@ use databend_common_meta_types::protobuf::FILE_DESCRIPTOR_SET;
 use databend_common_meta_types::GrpcConfig;
 use databend_common_meta_types::MetaNetworkError;
 use fastrace::prelude::*;
+use futures::future::select;
+use futures::future::Either;
 use log::info;
 use tonic::transport::Identity;
 use tonic::transport::Server;
@@ -34,14 +36,24 @@ use tonic::transport::ServerTlsConfig;
 use crate::api::grpc::grpc_service::MetaServiceImpl;
 use crate::configs::Config;
 use crate::meta_service::MetaNode;
+use crate::util::DropDebug;
 
 pub struct GrpcServer {
     conf: Config,
+    /// GrpcServer is the main container of the gRPC service.
+    /// [`MetaNode`] should never be dropped while [`GrpcServer`] is alive.
+    /// Therefore, it is held by a strong reference (Arc) to ensure proper lifetime management.
     pub(crate) meta_node: Arc<MetaNode>,
     join_handle: Option<JoinHandle<()>>,
     stop_grpc_tx: Option<Sender<()>>,
 }
 
+impl Drop for GrpcServer {
+    fn drop(&mut self) {
+        info!("GrpcServer::drop: id={}", self.conf.raft_config.id);
+    }
+}
+
 impl GrpcServer {
     pub fn create(conf: Config, meta_node: Arc<MetaNode>) -> Self {
         Self {
@@ -88,13 +100,17 @@ impl GrpcServer {
 
         info!("start gRPC listening: {}", addr);
 
-        let grpc_impl = MetaServiceImpl::create(meta_node.clone());
+        let grpc_impl = MetaServiceImpl::create(Arc::downgrade(&meta_node));
         let grpc_srv = MetaServiceServer::new(grpc_impl)
             .max_decoding_message_size(GrpcConfig::MAX_DECODING_SIZE)
             .max_encoding_message_size(GrpcConfig::MAX_ENCODING_SIZE);
 
+        let id = conf.raft_config.id;
+
         let j = databend_common_base::runtime::spawn(
             async move {
+                let _d = DropDebug::new(format!("GrpcServer(id={}) spawned service task", id));
+
                 let res = builder
                     .add_service(reflect_srv)
                     .add_service(grpc_srv)
@@ -128,20 +144,38 @@ impl GrpcServer {
     }
 
     async fn do_stop(&mut self, _force: Option<tokio::sync::broadcast::Receiver<()>>) {
+        let id = self.meta_node.raft_store.id;
+        let ctx = format!("gRPC-task(id={id})");
+
         if let Some(stop_grpc_tx) = self.stop_grpc_tx.take() {
-            info!("Sending stop signal to gRPC server");
+            info!("Sending stop signal to {ctx}");
             let _ = stop_grpc_tx.send(());
         }
 
-        if let Some(j) = self.join_handle.take() {
-            info!("Waiting for gRPC server stop");
-            let x = tokio::time::timeout(Duration::from_millis(1_000), j).await;
-            info!("Done: waiting for grpc stop: res: {:?}", x);
+        if let Some(jh) = self.join_handle.take() {
+            info!("Waiting for {ctx} to stop");
+
+            let timeout = Duration::from_millis(1_000);
+            let sleep = tokio::time::sleep(timeout);
+
+            let slp = std::pin::pin!(sleep);
+
+            match select(slp, jh).await {
+                Either::Left((_v1, jh)) => {
+                    info!("{ctx} stop timeout after {:?}; force abort", timeout);
+                    jh.abort();
+                    jh.await.ok();
+                    info!("Done: waiting for {ctx} force stop");
+                }
+                Either::Right((_v2, _slp)) => {
+                    info!("Done: waiting for {ctx} normal stop");
+                }
+            }
         }
 
-        info!("Waiting for meta_node stop");
+        info!("Waiting for meta_node(id={id}) to stop");
         let x = tokio::time::timeout(Duration::from_millis(1_000), self.meta_node.stop()).await;
-        info!("Done: waiting for meta_node stop: res: {:?}", x);
+        info!("Done: waiting for meta_node(id={id}) stop: res: {:?}", x);
     }
 
     async fn tls_config(conf: &Config) -> Result<Option<ServerTlsConfig>, std::io::Error> {

src/meta/service/src/lib.rs

@@ -26,3 +26,5 @@ pub mod raft_client;
 pub(crate) mod request_handling;
 pub mod store;
 pub mod version;
+
+pub mod util;

src/meta/service/src/meta_service/meta_node.rs

@@ -111,6 +111,9 @@ pub type MetaRaft = Raft<TypeConfig>;
 /// MetaNode is the container of metadata related components and threads, such as storage, the raft node and a raft-state monitor.
 pub struct MetaNode {
     pub raft_store: RaftStore,
+    /// MetaNode hold a strong reference to the dispatcher handle.
+    ///
+    /// Other components should keep a weak one.
     pub dispatcher_handle: Arc<DispatcherHandle>,
     pub raft: MetaRaft,
     pub running_tx: watch::Sender<()>,

src/meta/service/src/util.rs

@@ -0,0 +1,36 @@
+// Copyright 2021 Datafuse 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 log::debug;
+
+/// A struct that implements the Drop trait to log a message when it is dropped.
+///
+/// It can be used to track the lifetime. For example, use it a struct field of as a local variable of a closure.
+pub struct DropDebug {
+    message: String,
+}
+
+impl Drop for DropDebug {
+    fn drop(&mut self) {
+        debug!("DropDebug: {}", self.message);
+    }
+}
+
+impl DropDebug {
+    pub fn new(m: impl ToString) -> Self {
+        Self {
+            message: m.to_string(),
+        }
+    }
+}

src/meta/service/tests/it/grpc/metasrv_grpc_kv_api_restart_cluster.rs

@@ -190,6 +190,8 @@ async fn test_kv_api_restart_cluster_token_expired() -> anyhow::Result<()> {
         stopped_tcs
     };
 
+    sleep(Duration::from_secs(1)).await;
+
     info!("--- restart the cluster");
     let tcs = {
         let mut tcs = vec![];

src/meta/service/tests/it/tests/service.rs

@@ -96,7 +96,7 @@ pub fn make_grpc_client(addresses: Vec<String>) -> Result<Arc<ClientHandle>, Met
         addresses,
         "root",
         "xxx",
-        None,
+        Some(Duration::from_secs(2)), // timeout
         Some(Duration::from_secs(10)),
         None,
     )?;