feat(meta-client): add watch_with_initialization() method (#17944)

This method requires the databend-meta server to respond a
`is_initialization` flag in each watch response for the client to
distinguish between initialization events and key-value change events.
This feature is required by the cache mod.

Author: drmingdrmer

Cargo.lock

@@ -31,6 +31,7 @@ use crate::tenant_key::resource::TenantResource;
 /// e.g. TableId, DatabaseId as a value.
 ///
 /// `DataId<R>` can be dereferenced to `u64`.
+/// `DataId<R>` will take place of `Id<T>` in future.
 ///
 /// When an id is used as a key in a key-value store,
 /// it is serialized in another way to keep order.

src/meta/app/src/data_id.rs

@@ -29,6 +29,8 @@ use databend_common_meta_kvapi::kvapi::ListKVReq;
 use databend_common_meta_types::protobuf::ClientInfo;
 use databend_common_meta_types::protobuf::ClusterStatus;
 use databend_common_meta_types::protobuf::StreamItem;
+use databend_common_meta_types::protobuf::WatchRequest;
+use databend_common_meta_types::protobuf::WatchResponse;
 use databend_common_meta_types::ConnectionError;
 use databend_common_meta_types::MetaClientError;
 use databend_common_meta_types::MetaError;
@@ -45,6 +47,7 @@ use crate::grpc_metrics;
 use crate::message;
 use crate::message::Response;
 use crate::ClientWorkerRequest;
+use crate::InitFlag;
 use crate::RequestFor;
 use crate::Streamed;
 
@@ -104,6 +107,22 @@ impl ClientHandle {
         Ok(strm)
     }
 
+    /// Watch without requiring the initialization support
+    pub async fn watch(
+        &self,
+        watch: WatchRequest,
+    ) -> Result<tonic::codec::Streaming<WatchResponse>, MetaClientError> {
+        self.request(watch).await
+    }
+
+    /// Watch with requiring the initialization support
+    pub async fn watch_with_initialization(
+        &self,
+        watch: WatchRequest,
+    ) -> Result<tonic::codec::Streaming<WatchResponse>, MetaClientError> {
+        self.request((watch, InitFlag)).await
+    }
+
     /// Send a request to the internal worker task, which will be running in another runtime.
     #[fastrace::trace]
     #[async_backtrace::framed]

src/meta/client/src/client_handle.rs

@@ -28,6 +28,7 @@ use databend_common_meta_types::protobuf::StreamItem;
 use databend_common_meta_types::protobuf::WatchRequest;
 use databend_common_meta_types::protobuf::WatchResponse;
 use databend_common_meta_types::GrpcHelper;
+use databend_common_meta_types::MetaHandshakeError;
 use databend_common_meta_types::TxnReply;
 use databend_common_meta_types::TxnRequest;
 use log::error;
@@ -44,6 +45,7 @@ use crate::endpoints::Endpoints;
 use crate::grpc_client::AuthInterceptor;
 use crate::grpc_client::RealClient;
 use crate::required::Features;
+use crate::FeatureSpec;
 
 /// Update the client state according to the result of an RPC.
 trait HandleRPCResult<T> {
@@ -148,11 +150,15 @@ impl EstablishedClient {
         self.features.contains_key(feature)
     }
 
-    pub fn ensure_feature(&self, feature: &str) -> Result<(), Status> {
+    pub fn ensure_feature_spec(&self, spec: &FeatureSpec) -> Result<(), MetaHandshakeError> {
+        self.ensure_feature(spec.0)
+    }
+
+    pub fn ensure_feature(&self, feature: &str) -> Result<(), MetaHandshakeError> {
         if self.has_feature(feature) {
             Ok(())
         } else {
-            Err(Status::failed_precondition(format!(
+            Err(MetaHandshakeError::new(format!(
                 "Feature {} is not supported by the server; server:{{version: {}, features: {:?}}}",
                 feature, self.server_protocol_version, self.features
             )))

src/meta/client/src/established_client.rs

@@ -44,6 +44,7 @@ use crate::message::GetClusterStatus;
 use crate::message::GetEndpoints;
 use crate::message::MakeEstablishedClient;
 use crate::message::Streamed;
+use crate::InitFlag;
 
 /// Bind a request type to its corresponding response type.
 pub trait RequestFor: Clone + fmt::Debug {
@@ -185,6 +186,10 @@ impl RequestFor for WatchRequest {
     type Reply = tonic::codec::Streaming<WatchResponse>;
 }
 
+impl RequestFor for (WatchRequest, InitFlag) {
+    type Reply = tonic::codec::Streaming<WatchResponse>;
+}
+
 impl RequestFor for ExportReq {
     type Reply = tonic::codec::Streaming<WatchResponse>;
 }

src/meta/client/src/grpc_action.rs

@@ -88,6 +88,7 @@ use crate::grpc_action::RequestFor;
 use crate::grpc_metrics;
 use crate::message;
 use crate::message::Response;
+use crate::required::features;
 use crate::required::std;
 use crate::required::supported_features;
 use crate::required::Features;
@@ -365,6 +366,10 @@ impl MetaGrpcClient {
                 let resp = self.watch(r).await;
                 Response::Watch(resp)
             }
+            message::Request::WatchWithInitialization((r, _)) => {
+                let resp = self.watch_with_initialization(r).await;
+                Response::WatchWithInitialization(resp)
+            }
             message::Request::Export(r) => {
                 let resp = self.export(r).await;
                 Response::Export(resp)
@@ -726,6 +731,27 @@ impl MetaGrpcClient {
         Ok(res.into_inner())
     }
 
+    /// Create a watching stream that receives KV change events.
+    ///
+    /// This method is similar to `watch`, but it also sends all existing key-values with `is_initialization=true`
+    /// before starting the watch stream.
+    #[fastrace::trace]
+    #[async_backtrace::framed]
+    pub(crate) async fn watch_with_initialization(
+        &self,
+        watch_request: WatchRequest,
+    ) -> Result<tonic::codec::Streaming<WatchResponse>, MetaClientError> {
+        debug!("{}: handle watch request: {:?}", self, watch_request);
+
+        let mut client = self.get_established_client().await?;
+        client.ensure_feature_spec(&features::WATCH)?;
+        client.ensure_feature_spec(&features::WATCH_INITIAL_FLUSH)?;
+        client.ensure_feature_spec(&features::WATCH_INIT_FLAG)?;
+
+        let res = client.watch(watch_request).await?;
+        Ok(res.into_inner())
+    }
+
     /// Export all data in json from metasrv.
     #[fastrace::trace]
     #[async_backtrace::framed]

src/meta/client/src/grpc_client.rs

@@ -39,6 +39,7 @@ pub use grpc_action::MetaGrpcReq;
 pub use grpc_action::RequestFor;
 pub use grpc_client::MetaGrpcClient;
 pub use message::ClientWorkerRequest;
+pub use message::InitFlag;
 pub use message::Streamed;
 pub use required::FeatureSpec;
 pub use required::VersionTuple;

src/meta/client/src/lib.rs

@@ -71,6 +71,10 @@ impl<T> Streamed<T> {
     }
 }
 
+/// A marker type to indicate that the watch response should return a flag event indicating the end of initialization.
+#[derive(Debug, Clone)]
+pub struct InitFlag;
+
 /// Meta-client handle-to-worker request body
 #[derive(Debug, Clone, derive_more::From)]
 pub enum Request {
@@ -89,6 +93,17 @@ pub enum Request {
     /// Watch KV changes, expecting a Stream that reports KV change events
     Watch(WatchRequest),
 
+    /// Watch KV changes with optional initialization flush of all existing values.
+    ///
+    /// When `initial_flush` is set to true in the WatchRequest:
+    /// - First sends all existing key-values with `is_initialization=true`
+    /// - Then sends a special event with no data to mark the end of initialization
+    /// - Finally streams real-time changes as they occur
+    ///
+    /// This provides a reliable way to initialize client-side caches and then
+    /// keep them synchronized with server state.
+    WatchWithInitialization((WatchRequest, InitFlag)),
+
     /// Export all data
     Export(ExportReq),
 
@@ -113,6 +128,7 @@ impl Request {
             Request::Upsert(_) => "Upsert",
             Request::Txn(_) => "Txn",
             Request::Watch(_) => "Watch",
+            Request::WatchWithInitialization(_) => "WatchWithInitialization",
             Request::Export(_) => "Export",
             Request::MakeEstablishedClient(_) => "MakeClient",
             Request::GetEndpoints(_) => "GetEndpoints",
@@ -130,6 +146,7 @@ pub enum Response {
     Upsert(Result<UpsertKVReply, MetaError>),
     Txn(Result<TxnReply, MetaError>),
     Watch(Result<tonic::codec::Streaming<WatchResponse>, MetaClientError>),
+    WatchWithInitialization(Result<tonic::codec::Streaming<WatchResponse>, MetaClientError>),
     Export(Result<tonic::codec::Streaming<ExportedChunk>, MetaError>),
     MakeEstablishedClient(Result<EstablishedClient, MetaClientError>),
     GetEndpoints(Result<Vec<String>, MetaError>),
@@ -155,6 +172,9 @@ impl fmt::Debug for Response {
             Response::Watch(x) => {
                 write!(f, "Watch({:?})", x)
             }
+            Response::WatchWithInitialization(x) => {
+                write!(f, "WatchWithInitialization({:?})", x)
+            }
             Response::Export(x) => {
                 write!(f, "Export({:?})", x)
             }
@@ -197,6 +217,10 @@ impl Response {
                 .as_ref()
                 .err()
                 .map(|x| x as &(dyn std::error::Error + 'static)),
+            Response::WatchWithInitialization(res) => res
+                .as_ref()
+                .err()
+                .map(|x| x as &(dyn std::error::Error + 'static)),
             Response::Export(res) => res
                 .as_ref()
                 .err()

src/meta/client/src/message.rs

@@ -34,6 +34,8 @@ pub mod features {
     pub const EXPORT_V1:           FeatureSpec = ("export_v1",           (1, 2, 315));
     pub const WATCH:               FeatureSpec = ("watch",               (1, 2, 259));
     pub const WATCH_INITIAL_FLUSH: FeatureSpec = ("watch/initial_flush", (1, 2, 677));
+    /// WatchResponse contains a flag to indicate whether the event is initialization event or change event.
+    pub const WATCH_INIT_FLAG:     FeatureSpec = ("watch/init_flag",     (1, 2, 736));
     pub const MEMBER_LIST:         FeatureSpec = ("member_list",         (1, 2, 259));
     pub const GET_CLUSTER_STATUS:  FeatureSpec = ("get_cluster_status",  (1, 2, 259));
     pub const GET_CLIENT_INFO:     FeatureSpec = ("get_client_info",     (1, 2, 259));
@@ -50,6 +52,7 @@ pub fn all() -> &'static [FeatureSpec] {
         features::EXPORT_V1,
         features::WATCH,
         features::WATCH_INITIAL_FLUSH,
+        features::WATCH_INIT_FLAG,
         features::MEMBER_LIST,
         features::GET_CLUSTER_STATUS,
         features::GET_CLIENT_INFO,
@@ -81,6 +84,8 @@ pub fn std() -> &'static [FeatureSpec] {
         // features::EXPORT_V1,
         features::WATCH,
         features::WATCH_INITIAL_FLUSH,
+        // MIN_METASRV_VER does not include this feature, thus it is optional
+        // features::WATCH_INIT_FLAG,
         features::MEMBER_LIST,
         features::GET_CLUSTER_STATUS,
         features::GET_CLIENT_INFO,
@@ -110,6 +115,7 @@ pub fn read_write_watch() -> &'static [FeatureSpec] {
         features::TRANSACTION,
         features::WATCH,
         features::WATCH_INITIAL_FLUSH,
+        features::WATCH_INIT_FLAG,
     ];
 
     REQUIRES

src/meta/client/src/required.rs

@@ -365,7 +365,7 @@ async fn test_watch_initialization_flush() -> anyhow::Result<()> {
             filter_type: FilterType::All.into(),
             initial_flush: true,
         };
-        client.request(watch).await?
+        client.watch_with_initialization(watch).await?
     };
 
     let cache = Arc::new(Mutex::new(BTreeMap::new()));