Skip to main content

risingwave_rpc_client/
meta_client.rs

1// Copyright 2022 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::{BTreeMap, HashMap, HashSet};
16use std::fmt::{Debug, Display};
17use std::sync::Arc;
18use std::sync::atomic::AtomicBool;
19use std::sync::atomic::Ordering::Relaxed;
20use std::thread;
21use std::time::{Duration, SystemTime};
22
23use anyhow::{Context, anyhow};
24use async_trait::async_trait;
25use cluster_limit_service_client::ClusterLimitServiceClient;
26use either::Either;
27use futures::stream::BoxStream;
28use list_rate_limits_response::RateLimitInfo;
29use lru::LruCache;
30use replace_job_plan::ReplaceJob;
31use risingwave_common::catalog::{
32    AlterDatabaseParam, FunctionId, IndexId, ObjectId, SecretId, TableId,
33};
34use risingwave_common::config::{MAX_CONNECTION_WINDOW_SIZE, MetaConfig};
35use risingwave_common::hash::WorkerSlotMapping;
36use risingwave_common::id::{
37    ConnectionId, DatabaseId, JobId, SchemaId, SinkId, SubscriptionId, UserId, ViewId, WorkerId,
38};
39use risingwave_common::monitor::EndpointExt;
40use risingwave_common::system_param::AdaptiveParallelismStrategy;
41use risingwave_common::system_param::reader::SystemParamsReader;
42use risingwave_common::telemetry::report::TelemetryInfoFetcher;
43use risingwave_common::util::addr::HostAddr;
44use risingwave_common::util::meta_addr::MetaAddressStrategy;
45use risingwave_common::util::resource_util::cpu::total_cpu_available;
46use risingwave_common::util::resource_util::hostname;
47use risingwave_common::util::resource_util::memory::system_memory_available_bytes;
48use risingwave_common::util::retry::exponential_backoff;
49use risingwave_common::util::version::current_rw_version;
50use risingwave_error::bail;
51use risingwave_error::tonic::ErrorIsFromTonicServerImpl;
52use risingwave_hummock_sdk::change_log::{TableChangeLog, TableChangeLogs};
53use risingwave_hummock_sdk::version::{HummockVersion, HummockVersionDelta};
54use risingwave_hummock_sdk::{
55    CompactionGroupId, HummockEpoch, HummockVersionId, ObjectIdRange, SyncResult,
56};
57use risingwave_pb::backup_service::backup_service_client::BackupServiceClient;
58use risingwave_pb::backup_service::*;
59use risingwave_pb::catalog::{
60    Connection, PbComment, PbDatabase, PbFunction, PbIndex, PbSchema, PbSink, PbSource,
61    PbSubscription, PbTable, PbView, Table,
62};
63use risingwave_pb::cloud_service::cloud_service_client::CloudServiceClient;
64use risingwave_pb::cloud_service::*;
65use risingwave_pb::common::worker_node::Property;
66use risingwave_pb::common::{HostAddress, OptionalUint32, OptionalUint64, WorkerNode, WorkerType};
67use risingwave_pb::configured_monitor_service_client;
68use risingwave_pb::connector_service::sink_coordination_service_client::SinkCoordinationServiceClient;
69use risingwave_pb::ddl_service::alter_owner_request::Object;
70use risingwave_pb::ddl_service::create_iceberg_table_request::{PbSinkJobInfo, PbTableJobInfo};
71use risingwave_pb::ddl_service::ddl_service_client::DdlServiceClient;
72use risingwave_pb::ddl_service::*;
73use risingwave_pb::hummock::compact_task::TaskStatus;
74use risingwave_pb::hummock::get_compaction_score_response::PickerInfo;
75use risingwave_pb::hummock::get_table_change_logs_request::PbTableFilter;
76use risingwave_pb::hummock::hummock_manager_service_client::HummockManagerServiceClient;
77use risingwave_pb::hummock::rise_ctl_update_compaction_config_request::mutable_config::MutableConfig;
78use risingwave_pb::hummock::subscribe_compaction_event_request::Register;
79use risingwave_pb::hummock::write_limits::WriteLimit;
80use risingwave_pb::hummock::*;
81use risingwave_pb::iceberg_compaction::subscribe_iceberg_compaction_event_request::Register as IcebergRegister;
82use risingwave_pb::iceberg_compaction::{
83    SubscribeIcebergCompactionEventRequest, SubscribeIcebergCompactionEventResponse,
84    subscribe_iceberg_compaction_event_request,
85};
86use risingwave_pb::id::{ActorId, FragmentId, HummockSstableId, IcebergCompactionTaskId, SourceId};
87use risingwave_pb::meta::alter_connector_props_request::{
88    AlterConnectorPropsObject, AlterIcebergTableIds, ExtraOptions,
89};
90use risingwave_pb::meta::cancel_creating_jobs_request::PbJobs;
91use risingwave_pb::meta::cluster_service_client::ClusterServiceClient;
92use risingwave_pb::meta::event_log_service_client::EventLogServiceClient;
93use risingwave_pb::meta::heartbeat_service_client::HeartbeatServiceClient;
94use risingwave_pb::meta::hosted_iceberg_catalog_service_client::HostedIcebergCatalogServiceClient;
95use risingwave_pb::meta::list_actor_splits_response::ActorSplit;
96use risingwave_pb::meta::list_actor_states_response::ActorState;
97use risingwave_pb::meta::list_cdc_progress_response::PbCdcProgress;
98use risingwave_pb::meta::list_iceberg_tables_response::IcebergTable;
99use risingwave_pb::meta::list_refresh_table_states_response::RefreshTableState;
100use risingwave_pb::meta::list_streaming_job_states_response::StreamingJobState;
101use risingwave_pb::meta::list_table_fragments_response::TableFragmentInfo;
102use risingwave_pb::meta::meta_member_service_client::MetaMemberServiceClient;
103use risingwave_pb::meta::notification_service_client::NotificationServiceClient;
104use risingwave_pb::meta::scale_service_client::ScaleServiceClient;
105use risingwave_pb::meta::serving_service_client::ServingServiceClient;
106use risingwave_pb::meta::session_param_service_client::SessionParamServiceClient;
107use risingwave_pb::meta::stream_manager_service_client::StreamManagerServiceClient;
108use risingwave_pb::meta::system_params_service_client::SystemParamsServiceClient;
109use risingwave_pb::meta::telemetry_info_service_client::TelemetryInfoServiceClient;
110use risingwave_pb::meta::{FragmentDistribution, *};
111use risingwave_pb::monitor_service::monitor_service_client::MonitorServiceClient;
112use risingwave_pb::monitor_service::stack_trace_request::ActorTracesFormat;
113use risingwave_pb::monitor_service::{StackTraceRequest, StackTraceResponse};
114use risingwave_pb::secret::PbSecretRef;
115use risingwave_pb::stream_plan::StreamFragmentGraph;
116use risingwave_pb::user::alter_default_privilege_request::Operation as AlterDefaultPrivilegeOperation;
117use risingwave_pb::user::update_user_request::UpdateField;
118use risingwave_pb::user::user_service_client::UserServiceClient;
119use risingwave_pb::user::*;
120use thiserror_ext::AsReport;
121use tokio::sync::mpsc::{Receiver, UnboundedSender, unbounded_channel};
122use tokio::sync::oneshot::Sender;
123use tokio::sync::{RwLock, mpsc, oneshot};
124use tokio::task::JoinHandle;
125use tokio::time::{self};
126use tokio_retry::strategy::jitter;
127use tokio_stream::wrappers::UnboundedReceiverStream;
128use tonic::transport::Endpoint;
129use tonic::{Code, Request, Streaming};
130
131use crate::channel::{Channel, WrappedChannelExt};
132use crate::error::{Result, RpcError};
133use crate::hummock_meta_client::{
134    CompactionEventItem, HummockMetaClient, HummockMetaClientChangeLogInfo,
135    IcebergCompactionEventItem,
136};
137use crate::meta_rpc_client_method_impl;
138
139/// Client to meta server. Cloning the instance is lightweight.
140#[derive(Clone, Debug)]
141pub struct MetaClient {
142    worker_id: WorkerId,
143    worker_type: WorkerType,
144    host_addr: HostAddr,
145    inner: GrpcMetaClient,
146    meta_config: Arc<MetaConfig>,
147    cluster_id: String,
148    shutting_down: Arc<AtomicBool>,
149}
150
151impl MetaClient {
152    pub fn worker_id(&self) -> WorkerId {
153        self.worker_id
154    }
155
156    pub fn host_addr(&self) -> &HostAddr {
157        &self.host_addr
158    }
159
160    pub fn worker_type(&self) -> WorkerType {
161        self.worker_type
162    }
163
164    pub fn cluster_id(&self) -> &str {
165        &self.cluster_id
166    }
167
168    /// Subscribe to notification from meta.
169    pub async fn subscribe(
170        &self,
171        subscribe_type: SubscribeType,
172    ) -> Result<Streaming<SubscribeResponse>> {
173        let request = SubscribeRequest {
174            subscribe_type: subscribe_type as i32,
175            host: Some(self.host_addr.to_protobuf()),
176            worker_id: self.worker_id(),
177        };
178
179        let retry_strategy = GrpcMetaClient::retry_strategy_to_bound(
180            Duration::from_secs(self.meta_config.max_heartbeat_interval_secs as u64),
181            true,
182        );
183
184        tokio_retry::Retry::spawn(retry_strategy, || async {
185            let request = request.clone();
186            self.inner.subscribe(request).await
187        })
188        .await
189    }
190
191    pub async fn create_connection(
192        &self,
193        connection_name: String,
194        database_id: DatabaseId,
195        schema_id: SchemaId,
196        owner_id: UserId,
197        req: create_connection_request::Payload,
198    ) -> Result<WaitVersion> {
199        let request = CreateConnectionRequest {
200            name: connection_name,
201            database_id,
202            schema_id,
203            owner_id,
204            payload: Some(req),
205        };
206        let resp = self.inner.create_connection(request).await?;
207        Ok(resp
208            .version
209            .ok_or_else(|| anyhow!("wait version not set"))?)
210    }
211
212    pub async fn create_secret(
213        &self,
214        secret_name: String,
215        database_id: DatabaseId,
216        schema_id: SchemaId,
217        owner_id: UserId,
218        value: Vec<u8>,
219    ) -> Result<WaitVersion> {
220        let request = CreateSecretRequest {
221            name: secret_name,
222            database_id,
223            schema_id,
224            owner_id,
225            value,
226        };
227        let resp = self.inner.create_secret(request).await?;
228        Ok(resp
229            .version
230            .ok_or_else(|| anyhow!("wait version not set"))?)
231    }
232
233    pub async fn list_connections(&self, _name: Option<&str>) -> Result<Vec<Connection>> {
234        let request = ListConnectionsRequest {};
235        let resp = self.inner.list_connections(request).await?;
236        Ok(resp.connections)
237    }
238
239    pub async fn drop_connection(
240        &self,
241        connection_id: ConnectionId,
242        cascade: bool,
243    ) -> Result<WaitVersion> {
244        let request = DropConnectionRequest {
245            connection_id,
246            cascade,
247        };
248        let resp = self.inner.drop_connection(request).await?;
249        Ok(resp
250            .version
251            .ok_or_else(|| anyhow!("wait version not set"))?)
252    }
253
254    pub async fn drop_secret(&self, secret_id: SecretId, cascade: bool) -> Result<WaitVersion> {
255        let request = DropSecretRequest { secret_id, cascade };
256        let resp = self.inner.drop_secret(request).await?;
257        Ok(resp
258            .version
259            .ok_or_else(|| anyhow!("wait version not set"))?)
260    }
261
262    /// Register the current node to the cluster and set the corresponding worker id.
263    ///
264    /// Retry if there's connection issue with the meta node. Exit the process if the registration fails.
265    pub async fn register_new(
266        addr_strategy: MetaAddressStrategy,
267        worker_type: WorkerType,
268        addr: &HostAddr,
269        property: Property,
270        meta_config: Arc<MetaConfig>,
271    ) -> (Self, SystemParamsReader) {
272        let ret =
273            Self::register_new_inner(addr_strategy, worker_type, addr, property, meta_config).await;
274
275        match ret {
276            Ok(ret) => ret,
277            Err(err) => {
278                tracing::error!(error = %err.as_report(), "failed to register worker, exiting...");
279                std::process::exit(1);
280            }
281        }
282    }
283
284    async fn register_new_inner(
285        addr_strategy: MetaAddressStrategy,
286        worker_type: WorkerType,
287        addr: &HostAddr,
288        property: Property,
289        meta_config: Arc<MetaConfig>,
290    ) -> Result<(Self, SystemParamsReader)> {
291        tracing::info!("register meta client using strategy: {}", addr_strategy);
292
293        // Retry until reaching `max_heartbeat_interval_secs`
294        let retry_strategy = GrpcMetaClient::retry_strategy_to_bound(
295            Duration::from_secs(meta_config.max_heartbeat_interval_secs as u64),
296            true,
297        );
298
299        let init_result: Result<_> = tokio_retry::RetryIf::spawn(
300            retry_strategy,
301            || async {
302                let grpc_meta_client =
303                    GrpcMetaClient::new(&addr_strategy, meta_config.clone()).await?;
304
305                let add_worker_resp = grpc_meta_client
306                    .add_worker_node(AddWorkerNodeRequest {
307                        worker_type: worker_type as i32,
308                        host: Some(addr.to_protobuf()),
309                        property: Some(property.clone()),
310                        resource: Some(risingwave_pb::common::worker_node::Resource {
311                            rw_version: current_rw_version(),
312                            total_memory_bytes: system_memory_available_bytes() as _,
313                            total_cpu_cores: total_cpu_available() as _,
314                            hostname: hostname(),
315                        }),
316                    })
317                    .await
318                    .context("failed to add worker node")?;
319
320                let system_params_resp = grpc_meta_client
321                    .get_system_params(GetSystemParamsRequest {})
322                    .await
323                    .context("failed to get initial system params")?;
324
325                Ok((add_worker_resp, system_params_resp, grpc_meta_client))
326            },
327            // Only retry if there's any transient connection issue.
328            // If the error is from our implementation or business, do not retry it.
329            |e: &RpcError| !e.is_from_tonic_server_impl(),
330        )
331        .await;
332
333        let (add_worker_resp, system_params_resp, grpc_meta_client) = init_result?;
334        let worker_id = add_worker_resp
335            .node_id
336            .expect("AddWorkerNodeResponse::node_id is empty");
337
338        let meta_client = Self {
339            worker_id,
340            worker_type,
341            host_addr: addr.clone(),
342            inner: grpc_meta_client,
343            meta_config: meta_config.clone(),
344            cluster_id: add_worker_resp.cluster_id,
345            shutting_down: Arc::new(false.into()),
346        };
347
348        static REPORT_PANIC: std::sync::Once = std::sync::Once::new();
349        REPORT_PANIC.call_once(|| {
350            let meta_client_clone = meta_client.clone();
351            std::panic::update_hook(move |default_hook, info| {
352                // Try to report panic event to meta node.
353                meta_client_clone.try_add_panic_event_blocking(info, None);
354                default_hook(info);
355            });
356        });
357
358        Ok((meta_client, system_params_resp.params.unwrap().into()))
359    }
360
361    /// Activate the current node in cluster to confirm it's ready to serve.
362    pub async fn activate(&self, addr: &HostAddr) -> Result<()> {
363        let request = ActivateWorkerNodeRequest {
364            host: Some(addr.to_protobuf()),
365            node_id: self.worker_id,
366        };
367        let retry_strategy = GrpcMetaClient::retry_strategy_to_bound(
368            Duration::from_secs(self.meta_config.max_heartbeat_interval_secs as u64),
369            true,
370        );
371        tokio_retry::Retry::spawn(retry_strategy, || async {
372            let request = request.clone();
373            self.inner.activate_worker_node(request).await
374        })
375        .await?;
376
377        Ok(())
378    }
379
380    /// Send heartbeat signal to meta service.
381    pub async fn send_heartbeat(&self) -> Result<()> {
382        let request = HeartbeatRequest {
383            node_id: self.worker_id,
384        };
385        let resp = self.inner.heartbeat(request).await?;
386        if let Some(status) = resp.status
387            && status.code() == risingwave_pb::common::status::Code::UnknownWorker
388        {
389            // Ignore the error if we're already shutting down.
390            // Otherwise, exit the process.
391            if !self.shutting_down.load(Relaxed) {
392                tracing::error!(message = status.message, "worker expired");
393                std::process::exit(1);
394            }
395        }
396        Ok(())
397    }
398
399    pub async fn create_database(&self, db: PbDatabase) -> Result<WaitVersion> {
400        let request = CreateDatabaseRequest { db: Some(db) };
401        let resp = self.inner.create_database(request).await?;
402        // TODO: handle error in `resp.status` here
403        Ok(resp
404            .version
405            .ok_or_else(|| anyhow!("wait version not set"))?)
406    }
407
408    pub async fn create_schema(&self, schema: PbSchema) -> Result<WaitVersion> {
409        let request = CreateSchemaRequest {
410            schema: Some(schema),
411        };
412        let resp = self.inner.create_schema(request).await?;
413        // TODO: handle error in `resp.status` here
414        Ok(resp
415            .version
416            .ok_or_else(|| anyhow!("wait version not set"))?)
417    }
418
419    pub async fn create_materialized_view(
420        &self,
421        table: PbTable,
422        graph: StreamFragmentGraph,
423        dependencies: HashSet<ObjectId>,
424        resource_type: streaming_job_resource_type::ResourceType,
425        if_not_exists: bool,
426        refresh_interval_sec: Option<u64>,
427    ) -> Result<WaitVersion> {
428        let request = CreateMaterializedViewRequest {
429            materialized_view: Some(table),
430            fragment_graph: Some(graph),
431            resource_type: Some(PbStreamingJobResourceType {
432                resource_type: Some(resource_type),
433            }),
434            dependencies: dependencies.into_iter().collect(),
435            if_not_exists,
436            refresh_interval_sec,
437        };
438        let resp = self.inner.create_materialized_view(request).await?;
439        // TODO: handle error in `resp.status` here
440        Ok(resp
441            .version
442            .ok_or_else(|| anyhow!("wait version not set"))?)
443    }
444
445    pub async fn drop_materialized_view(
446        &self,
447        table_id: TableId,
448        cascade: bool,
449    ) -> Result<WaitVersion> {
450        let request = DropMaterializedViewRequest { table_id, cascade };
451
452        let resp = self.inner.drop_materialized_view(request).await?;
453        Ok(resp
454            .version
455            .ok_or_else(|| anyhow!("wait version not set"))?)
456    }
457
458    pub async fn create_source(
459        &self,
460        source: PbSource,
461        graph: Option<StreamFragmentGraph>,
462        if_not_exists: bool,
463    ) -> Result<WaitVersion> {
464        let request = CreateSourceRequest {
465            source: Some(source),
466            fragment_graph: graph,
467            if_not_exists,
468        };
469
470        let resp = self.inner.create_source(request).await?;
471        Ok(resp
472            .version
473            .ok_or_else(|| anyhow!("wait version not set"))?)
474    }
475
476    pub async fn create_sink(
477        &self,
478        sink: PbSink,
479        graph: StreamFragmentGraph,
480        dependencies: HashSet<ObjectId>,
481        resource_type: streaming_job_resource_type::ResourceType,
482        if_not_exists: bool,
483        since_timestamp_epoch: Option<u64>,
484    ) -> Result<WaitVersion> {
485        let request = CreateSinkRequest {
486            sink: Some(sink),
487            fragment_graph: Some(graph),
488            dependencies: dependencies.into_iter().collect(),
489            if_not_exists,
490            resource_type: Some(PbStreamingJobResourceType {
491                resource_type: Some(resource_type),
492            }),
493            since_timestamp_epoch,
494        };
495
496        let resp = self.inner.create_sink(request).await?;
497        Ok(resp
498            .version
499            .ok_or_else(|| anyhow!("wait version not set"))?)
500    }
501
502    pub async fn create_subscription(&self, subscription: PbSubscription) -> Result<WaitVersion> {
503        let request = CreateSubscriptionRequest {
504            subscription: Some(subscription),
505        };
506
507        let resp = self.inner.create_subscription(request).await?;
508        Ok(resp
509            .version
510            .ok_or_else(|| anyhow!("wait version not set"))?)
511    }
512
513    pub async fn create_function(&self, function: PbFunction) -> Result<WaitVersion> {
514        let request = CreateFunctionRequest {
515            function: Some(function),
516        };
517        let resp = self.inner.create_function(request).await?;
518        Ok(resp
519            .version
520            .ok_or_else(|| anyhow!("wait version not set"))?)
521    }
522
523    pub async fn create_table(
524        &self,
525        source: Option<PbSource>,
526        table: PbTable,
527        graph: StreamFragmentGraph,
528        job_type: PbTableJobType,
529        if_not_exists: bool,
530        dependencies: HashSet<ObjectId>,
531    ) -> Result<WaitVersion> {
532        let request = CreateTableRequest {
533            materialized_view: Some(table),
534            fragment_graph: Some(graph),
535            source,
536            job_type: job_type as _,
537            if_not_exists,
538            dependencies: dependencies.into_iter().collect(),
539        };
540        let resp = self.inner.create_table(request).await?;
541        // TODO: handle error in `resp.status` here
542        Ok(resp
543            .version
544            .ok_or_else(|| anyhow!("wait version not set"))?)
545    }
546
547    pub async fn comment_on(&self, comment: PbComment) -> Result<WaitVersion> {
548        let request = CommentOnRequest {
549            comment: Some(comment),
550        };
551        let resp = self.inner.comment_on(request).await?;
552        Ok(resp
553            .version
554            .ok_or_else(|| anyhow!("wait version not set"))?)
555    }
556
557    pub async fn alter_name(
558        &self,
559        object: alter_name_request::Object,
560        name: &str,
561    ) -> Result<WaitVersion> {
562        let request = AlterNameRequest {
563            object: Some(object),
564            new_name: name.to_owned(),
565        };
566        let resp = self.inner.alter_name(request).await?;
567        Ok(resp
568            .version
569            .ok_or_else(|| anyhow!("wait version not set"))?)
570    }
571
572    pub async fn alter_owner(&self, object: Object, owner_id: UserId) -> Result<WaitVersion> {
573        let request = AlterOwnerRequest {
574            object: Some(object),
575            owner_id,
576        };
577        let resp = self.inner.alter_owner(request).await?;
578        Ok(resp
579            .version
580            .ok_or_else(|| anyhow!("wait version not set"))?)
581    }
582
583    pub async fn alter_subscription_retention(
584        &self,
585        subscription_id: SubscriptionId,
586        retention_seconds: u64,
587        definition: String,
588    ) -> Result<WaitVersion> {
589        let request = AlterSubscriptionRetentionRequest {
590            subscription_id,
591            retention_seconds,
592            definition,
593        };
594        let resp = self.inner.alter_subscription_retention(request).await?;
595        Ok(resp
596            .version
597            .ok_or_else(|| anyhow!("wait version not set"))?)
598    }
599
600    pub async fn alter_database_param(
601        &self,
602        database_id: DatabaseId,
603        param: AlterDatabaseParam,
604    ) -> Result<WaitVersion> {
605        let request = match param {
606            AlterDatabaseParam::BarrierIntervalMs(barrier_interval_ms) => {
607                let barrier_interval_ms = OptionalUint32 {
608                    value: barrier_interval_ms,
609                };
610                AlterDatabaseParamRequest {
611                    database_id,
612                    param: Some(alter_database_param_request::Param::BarrierIntervalMs(
613                        barrier_interval_ms,
614                    )),
615                }
616            }
617            AlterDatabaseParam::CheckpointFrequency(checkpoint_frequency) => {
618                let checkpoint_frequency = OptionalUint64 {
619                    value: checkpoint_frequency,
620                };
621                AlterDatabaseParamRequest {
622                    database_id,
623                    param: Some(alter_database_param_request::Param::CheckpointFrequency(
624                        checkpoint_frequency,
625                    )),
626                }
627            }
628        };
629        let resp = self.inner.alter_database_param(request).await?;
630        Ok(resp
631            .version
632            .ok_or_else(|| anyhow!("wait version not set"))?)
633    }
634
635    pub async fn alter_set_schema(
636        &self,
637        object: alter_set_schema_request::Object,
638        new_schema_id: SchemaId,
639    ) -> Result<WaitVersion> {
640        let request = AlterSetSchemaRequest {
641            new_schema_id,
642            object: Some(object),
643        };
644        let resp = self.inner.alter_set_schema(request).await?;
645        Ok(resp
646            .version
647            .ok_or_else(|| anyhow!("wait version not set"))?)
648    }
649
650    pub async fn alter_source(&self, source: PbSource) -> Result<WaitVersion> {
651        let request = AlterSourceRequest {
652            source: Some(source),
653        };
654        let resp = self.inner.alter_source(request).await?;
655        Ok(resp
656            .version
657            .ok_or_else(|| anyhow!("wait version not set"))?)
658    }
659
660    pub async fn alter_parallelism(
661        &self,
662        job_id: JobId,
663        parallelism: PbTableParallelism,
664        adaptive_parallelism_strategy: Option<AdaptiveParallelismStrategy>,
665        deferred: bool,
666    ) -> Result<()> {
667        let request = AlterParallelismRequest {
668            table_id: job_id,
669            parallelism: Some(parallelism),
670            deferred,
671            adaptive_parallelism_strategy: adaptive_parallelism_strategy
672                .map(|strategy| strategy.to_string()),
673        };
674
675        self.inner.alter_parallelism(request).await?;
676        Ok(())
677    }
678
679    pub async fn alter_backfill_parallelism(
680        &self,
681        job_id: JobId,
682        parallelism: Option<PbTableParallelism>,
683        adaptive_parallelism_strategy: Option<AdaptiveParallelismStrategy>,
684        deferred: bool,
685    ) -> Result<()> {
686        let request = AlterBackfillParallelismRequest {
687            table_id: job_id,
688            parallelism,
689            deferred,
690            adaptive_parallelism_strategy: adaptive_parallelism_strategy
691                .map(|strategy| strategy.to_string()),
692        };
693
694        self.inner.alter_backfill_parallelism(request).await?;
695        Ok(())
696    }
697
698    pub async fn alter_streaming_job_config(
699        &self,
700        job_id: JobId,
701        entries_to_add: HashMap<String, String>,
702        keys_to_remove: Vec<String>,
703    ) -> Result<()> {
704        let request = AlterStreamingJobConfigRequest {
705            job_id,
706            entries_to_add,
707            keys_to_remove,
708        };
709
710        self.inner.alter_streaming_job_config(request).await?;
711        Ok(())
712    }
713
714    pub async fn alter_fragment_parallelism(
715        &self,
716        fragment_ids: Vec<FragmentId>,
717        parallelism: Option<PbTableParallelism>,
718    ) -> Result<()> {
719        let request = AlterFragmentParallelismRequest {
720            fragment_ids,
721            parallelism,
722        };
723
724        self.inner.alter_fragment_parallelism(request).await?;
725        Ok(())
726    }
727
728    pub async fn alter_cdc_table_backfill_parallelism(
729        &self,
730        table_id: JobId,
731        parallelism: PbTableParallelism,
732    ) -> Result<()> {
733        let request = AlterCdcTableBackfillParallelismRequest {
734            table_id,
735            parallelism: Some(parallelism),
736        };
737        self.inner
738            .alter_cdc_table_backfill_parallelism(request)
739            .await?;
740        Ok(())
741    }
742
743    pub async fn alter_resource_group(
744        &self,
745        job_id: JobId,
746        resource_group: Option<String>,
747        deferred: bool,
748    ) -> Result<()> {
749        let request = AlterResourceGroupRequest {
750            job_id,
751            resource_group,
752            deferred,
753        };
754
755        self.inner.alter_resource_group(request).await?;
756        Ok(())
757    }
758
759    pub async fn alter_database_resource_group(
760        &self,
761        database_id: DatabaseId,
762        resource_group: Option<String>,
763        deferred: bool,
764    ) -> Result<WaitVersion> {
765        let request = AlterDatabaseResourceGroupRequest {
766            database_id,
767            resource_group,
768            deferred,
769        };
770
771        let resp = self.inner.alter_database_resource_group(request).await?;
772        Ok(resp
773            .version
774            .ok_or_else(|| anyhow!("wait version not set"))?)
775    }
776
777    pub async fn alter_swap_rename(
778        &self,
779        object: alter_swap_rename_request::Object,
780    ) -> Result<WaitVersion> {
781        let request = AlterSwapRenameRequest {
782            object: Some(object),
783        };
784        let resp = self.inner.alter_swap_rename(request).await?;
785        Ok(resp
786            .version
787            .ok_or_else(|| anyhow!("wait version not set"))?)
788    }
789
790    pub async fn alter_secret(
791        &self,
792        secret_id: SecretId,
793        secret_name: String,
794        database_id: DatabaseId,
795        schema_id: SchemaId,
796        owner_id: UserId,
797        value: Vec<u8>,
798    ) -> Result<WaitVersion> {
799        let request = AlterSecretRequest {
800            secret_id,
801            name: secret_name,
802            database_id,
803            schema_id,
804            owner_id,
805            value,
806        };
807        let resp = self.inner.alter_secret(request).await?;
808        Ok(resp
809            .version
810            .ok_or_else(|| anyhow!("wait version not set"))?)
811    }
812
813    pub async fn replace_job(
814        &self,
815        graph: StreamFragmentGraph,
816        replace_job: ReplaceJob,
817    ) -> Result<WaitVersion> {
818        let request = ReplaceJobPlanRequest {
819            plan: Some(ReplaceJobPlan {
820                fragment_graph: Some(graph),
821                replace_job: Some(replace_job),
822            }),
823        };
824        let resp = self.inner.replace_job_plan(request).await?;
825        // TODO: handle error in `resp.status` here
826        Ok(resp
827            .version
828            .ok_or_else(|| anyhow!("wait version not set"))?)
829    }
830
831    pub async fn auto_schema_change(&self, schema_change: SchemaChangeEnvelope) -> Result<()> {
832        let request = AutoSchemaChangeRequest {
833            schema_change: Some(schema_change),
834        };
835        let _ = self.inner.auto_schema_change(request).await?;
836        Ok(())
837    }
838
839    /// Block until the pk-index sink `sink_id`'s database has committed through `epoch`, returning
840    /// the coordinator's committed iceberg snapshot id (`None` if no snapshot committed yet).
841    pub async fn wait_iceberg_pk_index_sink_epoch(
842        &self,
843        sink_id: SinkId,
844        epoch: u64,
845    ) -> Result<Option<i64>> {
846        let request = WaitIcebergPkIndexSinkEpochRequest { sink_id, epoch };
847        let resp = self.inner.wait_iceberg_pk_index_sink_epoch(request).await?;
848        Ok(resp.snapshot_id)
849    }
850
851    pub async fn create_view(
852        &self,
853        view: PbView,
854        dependencies: HashSet<ObjectId>,
855    ) -> Result<WaitVersion> {
856        let request = CreateViewRequest {
857            view: Some(view),
858            dependencies: dependencies.into_iter().collect(),
859        };
860        let resp = self.inner.create_view(request).await?;
861        // TODO: handle error in `resp.status` here
862        Ok(resp
863            .version
864            .ok_or_else(|| anyhow!("wait version not set"))?)
865    }
866
867    pub async fn create_index(
868        &self,
869        index: PbIndex,
870        table: PbTable,
871        graph: StreamFragmentGraph,
872        resource_type: streaming_job_resource_type::ResourceType,
873        if_not_exists: bool,
874    ) -> Result<WaitVersion> {
875        let request = CreateIndexRequest {
876            index: Some(index),
877            index_table: Some(table),
878            fragment_graph: Some(graph),
879            if_not_exists,
880            resource_type: Some(PbStreamingJobResourceType {
881                resource_type: Some(resource_type),
882            }),
883        };
884        let resp = self.inner.create_index(request).await?;
885        // TODO: handle error in `resp.status` here
886        Ok(resp
887            .version
888            .ok_or_else(|| anyhow!("wait version not set"))?)
889    }
890
891    pub async fn drop_table(
892        &self,
893        source_id: Option<SourceId>,
894        table_id: TableId,
895        cascade: bool,
896    ) -> Result<WaitVersion> {
897        let request = DropTableRequest {
898            source_id: source_id.map(risingwave_pb::ddl_service::drop_table_request::SourceId::Id),
899            table_id,
900            cascade,
901        };
902
903        let resp = self.inner.drop_table(request).await?;
904        Ok(resp
905            .version
906            .ok_or_else(|| anyhow!("wait version not set"))?)
907    }
908
909    pub async fn compact_iceberg_table(&self, sink_id: SinkId) -> Result<IcebergCompactionTaskId> {
910        let request = CompactIcebergTableRequest { sink_id };
911        let resp = self.inner.compact_iceberg_table(request).await?;
912        Ok(resp.task_id)
913    }
914
915    pub async fn rewrite_iceberg_table_manifests(&self, sink_id: SinkId) -> Result<()> {
916        let request = RewriteIcebergTableManifestsRequest { sink_id };
917        let _resp = self.inner.rewrite_iceberg_table_manifests(request).await?;
918        Ok(())
919    }
920
921    pub async fn expire_iceberg_table_snapshots(&self, sink_id: SinkId) -> Result<()> {
922        let request = ExpireIcebergTableSnapshotsRequest { sink_id };
923        let _resp = self.inner.expire_iceberg_table_snapshots(request).await?;
924        Ok(())
925    }
926
927    pub async fn drop_view(&self, view_id: ViewId, cascade: bool) -> Result<WaitVersion> {
928        let request = DropViewRequest { view_id, cascade };
929        let resp = self.inner.drop_view(request).await?;
930        Ok(resp
931            .version
932            .ok_or_else(|| anyhow!("wait version not set"))?)
933    }
934
935    pub async fn drop_source(&self, source_id: SourceId, cascade: bool) -> Result<WaitVersion> {
936        let request = DropSourceRequest { source_id, cascade };
937        let resp = self.inner.drop_source(request).await?;
938        Ok(resp
939            .version
940            .ok_or_else(|| anyhow!("wait version not set"))?)
941    }
942
943    pub async fn reset_source(&self, source_id: SourceId) -> Result<WaitVersion> {
944        let request = ResetSourceRequest { source_id };
945        let resp = self.inner.reset_source(request).await?;
946        Ok(resp
947            .version
948            .ok_or_else(|| anyhow!("wait version not set"))?)
949    }
950
951    pub async fn drop_sink(&self, sink_id: SinkId, cascade: bool) -> Result<WaitVersion> {
952        let request = DropSinkRequest { sink_id, cascade };
953        let resp = self.inner.drop_sink(request).await?;
954        Ok(resp
955            .version
956            .ok_or_else(|| anyhow!("wait version not set"))?)
957    }
958
959    pub async fn drop_subscription(
960        &self,
961        subscription_id: SubscriptionId,
962        cascade: bool,
963    ) -> Result<WaitVersion> {
964        let request = DropSubscriptionRequest {
965            subscription_id,
966            cascade,
967        };
968        let resp = self.inner.drop_subscription(request).await?;
969        Ok(resp
970            .version
971            .ok_or_else(|| anyhow!("wait version not set"))?)
972    }
973
974    pub async fn drop_index(&self, index_id: IndexId, cascade: bool) -> Result<WaitVersion> {
975        let request = DropIndexRequest { index_id, cascade };
976        let resp = self.inner.drop_index(request).await?;
977        Ok(resp
978            .version
979            .ok_or_else(|| anyhow!("wait version not set"))?)
980    }
981
982    pub async fn drop_function(
983        &self,
984        function_id: FunctionId,
985        cascade: bool,
986    ) -> Result<WaitVersion> {
987        let request = DropFunctionRequest {
988            function_id,
989            cascade,
990        };
991        let resp = self.inner.drop_function(request).await?;
992        Ok(resp
993            .version
994            .ok_or_else(|| anyhow!("wait version not set"))?)
995    }
996
997    pub async fn drop_database(&self, database_id: DatabaseId) -> Result<WaitVersion> {
998        let request = DropDatabaseRequest { database_id };
999        let resp = self.inner.drop_database(request).await?;
1000        Ok(resp
1001            .version
1002            .ok_or_else(|| anyhow!("wait version not set"))?)
1003    }
1004
1005    pub async fn drop_schema(&self, schema_id: SchemaId, cascade: bool) -> Result<WaitVersion> {
1006        let request = DropSchemaRequest { schema_id, cascade };
1007        let resp = self.inner.drop_schema(request).await?;
1008        Ok(resp
1009            .version
1010            .ok_or_else(|| anyhow!("wait version not set"))?)
1011    }
1012
1013    // TODO: using UserInfoVersion instead as return type.
1014    pub async fn create_user(&self, user: UserInfo) -> Result<u64> {
1015        let request = CreateUserRequest { user: Some(user) };
1016        let resp = self.inner.create_user(request).await?;
1017        Ok(resp.version)
1018    }
1019
1020    pub async fn drop_user(&self, user_id: UserId) -> Result<u64> {
1021        let request = DropUserRequest { user_id };
1022        let resp = self.inner.drop_user(request).await?;
1023        Ok(resp.version)
1024    }
1025
1026    pub async fn update_user(
1027        &self,
1028        user: UserInfo,
1029        update_fields: Vec<UpdateField>,
1030    ) -> Result<u64> {
1031        let request = UpdateUserRequest {
1032            user: Some(user),
1033            update_fields: update_fields
1034                .into_iter()
1035                .map(|field| field as i32)
1036                .collect::<Vec<_>>(),
1037        };
1038        let resp = self.inner.update_user(request).await?;
1039        Ok(resp.version)
1040    }
1041
1042    pub async fn grant_privilege(
1043        &self,
1044        user_ids: Vec<UserId>,
1045        privileges: Vec<GrantPrivilege>,
1046        with_grant_option: bool,
1047        granted_by: UserId,
1048    ) -> Result<u64> {
1049        let request = GrantPrivilegeRequest {
1050            user_ids,
1051            privileges,
1052            with_grant_option,
1053            granted_by,
1054        };
1055        let resp = self.inner.grant_privilege(request).await?;
1056        Ok(resp.version)
1057    }
1058
1059    pub async fn revoke_privilege(
1060        &self,
1061        user_ids: Vec<UserId>,
1062        privileges: Vec<GrantPrivilege>,
1063        granted_by: UserId,
1064        revoke_by: UserId,
1065        revoke_grant_option: bool,
1066        cascade: bool,
1067    ) -> Result<u64> {
1068        let request = RevokePrivilegeRequest {
1069            user_ids,
1070            privileges,
1071            granted_by,
1072            revoke_by,
1073            revoke_grant_option,
1074            cascade,
1075        };
1076        let resp = self.inner.revoke_privilege(request).await?;
1077        Ok(resp.version)
1078    }
1079
1080    pub async fn alter_default_privilege(
1081        &self,
1082        user_ids: Vec<UserId>,
1083        database_id: DatabaseId,
1084        schema_ids: Vec<SchemaId>,
1085        operation: AlterDefaultPrivilegeOperation,
1086        granted_by: UserId,
1087    ) -> Result<()> {
1088        let request = AlterDefaultPrivilegeRequest {
1089            user_ids,
1090            database_id,
1091            schema_ids,
1092            operation: Some(operation),
1093            granted_by,
1094        };
1095        self.inner.alter_default_privilege(request).await?;
1096        Ok(())
1097    }
1098
1099    /// Unregister the current node from the cluster.
1100    pub async fn unregister(&self) -> Result<()> {
1101        let request = DeleteWorkerNodeRequest {
1102            host: Some(self.host_addr.to_protobuf()),
1103        };
1104        self.inner.delete_worker_node(request).await?;
1105        self.shutting_down.store(true, Relaxed);
1106        Ok(())
1107    }
1108
1109    /// Try to unregister the current worker from the cluster with best effort. Log the result.
1110    pub async fn try_unregister(&self) {
1111        match self.unregister().await {
1112            Ok(_) => {
1113                tracing::info!(
1114                    worker_id = %self.worker_id(),
1115                    "successfully unregistered from meta service",
1116                )
1117            }
1118            Err(e) => {
1119                tracing::warn!(
1120                    error = %e.as_report(),
1121                    worker_id = %self.worker_id(),
1122                    "failed to unregister from meta service",
1123                );
1124            }
1125        }
1126    }
1127
1128    pub async fn list_worker_nodes(
1129        &self,
1130        worker_type: Option<WorkerType>,
1131    ) -> Result<Vec<WorkerNode>> {
1132        let request = ListAllNodesRequest {
1133            worker_type: worker_type.map(Into::into),
1134            include_starting_nodes: true,
1135        };
1136        let resp = self.inner.list_all_nodes(request).await?;
1137        Ok(resp.nodes)
1138    }
1139
1140    /// Starts a heartbeat worker.
1141    pub fn start_heartbeat_loop(
1142        meta_client: MetaClient,
1143        min_interval: Duration,
1144    ) -> (JoinHandle<()>, Sender<()>) {
1145        let (shutdown_tx, mut shutdown_rx) = tokio::sync::oneshot::channel();
1146        let join_handle = tokio::spawn(async move {
1147            let mut min_interval_ticker = tokio::time::interval(min_interval);
1148            loop {
1149                tokio::select! {
1150                    biased;
1151                    // Shutdown
1152                    _ = &mut shutdown_rx => {
1153                        tracing::info!("Heartbeat loop is stopped");
1154                        return;
1155                    }
1156                    // Wait for interval
1157                    _ = min_interval_ticker.tick() => {},
1158                }
1159                tracing::debug!(target: "events::meta::client_heartbeat", "heartbeat");
1160                match tokio::time::timeout(
1161                    // TODO: decide better min_interval for timeout
1162                    min_interval * 3,
1163                    meta_client.send_heartbeat(),
1164                )
1165                .await
1166                {
1167                    Ok(Ok(_)) => {}
1168                    Ok(Err(err)) => {
1169                        tracing::warn!(error = %err.as_report(), "Failed to send_heartbeat");
1170                    }
1171                    Err(_) => {
1172                        tracing::warn!("Failed to send_heartbeat: timeout");
1173                    }
1174                }
1175            }
1176        });
1177        (join_handle, shutdown_tx)
1178    }
1179
1180    pub async fn risectl_list_state_tables(&self) -> Result<Vec<PbTable>> {
1181        let request = RisectlListStateTablesRequest {};
1182        let resp = self.inner.risectl_list_state_tables(request).await?;
1183        Ok(resp.tables)
1184    }
1185
1186    pub async fn risectl_resume_backfill(
1187        &self,
1188        request: RisectlResumeBackfillRequest,
1189    ) -> Result<()> {
1190        self.inner.risectl_resume_backfill(request).await?;
1191        Ok(())
1192    }
1193
1194    pub async fn flush(&self, database_id: DatabaseId) -> Result<HummockVersionId> {
1195        let request = FlushRequest { database_id };
1196        let resp = self.inner.flush(request).await?;
1197        Ok(resp.hummock_version_id)
1198    }
1199
1200    pub async fn wait(&self, job_id: Option<JobId>) -> Result<WaitVersion> {
1201        let request = WaitRequest { job_id };
1202        let resp = self.inner.wait(request).await?;
1203        Ok(resp
1204            .version
1205            .ok_or_else(|| anyhow!("wait version not set"))?)
1206    }
1207
1208    pub async fn recover(&self) -> Result<()> {
1209        let request = RecoverRequest {};
1210        self.inner.recover(request).await?;
1211        Ok(())
1212    }
1213
1214    pub async fn cancel_creating_jobs(&self, jobs: PbJobs) -> Result<Vec<u32>> {
1215        let request = CancelCreatingJobsRequest { jobs: Some(jobs) };
1216        let resp = self.inner.cancel_creating_jobs(request).await?;
1217        Ok(resp.canceled_jobs)
1218    }
1219
1220    pub async fn list_table_fragments(
1221        &self,
1222        job_ids: &[JobId],
1223    ) -> Result<HashMap<JobId, TableFragmentInfo>> {
1224        let request = ListTableFragmentsRequest {
1225            table_ids: job_ids.to_vec(),
1226        };
1227        let resp = self.inner.list_table_fragments(request).await?;
1228        Ok(resp.table_fragments)
1229    }
1230
1231    pub async fn list_streaming_job_states(&self) -> Result<Vec<StreamingJobState>> {
1232        let resp = self
1233            .inner
1234            .list_streaming_job_states(ListStreamingJobStatesRequest {})
1235            .await?;
1236        Ok(resp.states)
1237    }
1238
1239    pub async fn list_fragment_distributions(
1240        &self,
1241        include_node: bool,
1242    ) -> Result<Vec<FragmentDistribution>> {
1243        let resp = self
1244            .inner
1245            .list_fragment_distribution(ListFragmentDistributionRequest {
1246                include_node: Some(include_node),
1247            })
1248            .await?;
1249        Ok(resp.distributions)
1250    }
1251
1252    pub async fn list_creating_fragment_distribution(&self) -> Result<Vec<FragmentDistribution>> {
1253        let resp = self
1254            .inner
1255            .list_creating_fragment_distribution(ListCreatingFragmentDistributionRequest {
1256                include_node: Some(true),
1257            })
1258            .await?;
1259        Ok(resp.distributions)
1260    }
1261
1262    pub async fn get_fragment_by_id(
1263        &self,
1264        fragment_id: FragmentId,
1265    ) -> Result<Option<FragmentDistribution>> {
1266        let resp = self
1267            .inner
1268            .get_fragment_by_id(GetFragmentByIdRequest { fragment_id })
1269            .await?;
1270        Ok(resp.distribution)
1271    }
1272
1273    pub async fn get_fragment_vnodes(
1274        &self,
1275        fragment_id: FragmentId,
1276    ) -> Result<Vec<(ActorId, Vec<u32>)>> {
1277        let resp = self
1278            .inner
1279            .get_fragment_vnodes(GetFragmentVnodesRequest { fragment_id })
1280            .await?;
1281        Ok(resp
1282            .actor_vnodes
1283            .into_iter()
1284            .map(|actor| (actor.actor_id, actor.vnode_indices))
1285            .collect())
1286    }
1287
1288    pub async fn get_actor_vnodes(&self, actor_id: ActorId) -> Result<Vec<u32>> {
1289        let resp = self
1290            .inner
1291            .get_actor_vnodes(GetActorVnodesRequest { actor_id })
1292            .await?;
1293        Ok(resp.vnode_indices)
1294    }
1295
1296    pub async fn list_actor_states(&self) -> Result<Vec<ActorState>> {
1297        let resp = self
1298            .inner
1299            .list_actor_states(ListActorStatesRequest {})
1300            .await?;
1301        Ok(resp.states)
1302    }
1303
1304    pub async fn list_actor_splits(&self) -> Result<Vec<ActorSplit>> {
1305        let resp = self
1306            .inner
1307            .list_actor_splits(ListActorSplitsRequest {})
1308            .await?;
1309
1310        Ok(resp.actor_splits)
1311    }
1312
1313    pub async fn pause(&self) -> Result<PauseResponse> {
1314        let request = PauseRequest {};
1315        let resp = self.inner.pause(request).await?;
1316        Ok(resp)
1317    }
1318
1319    pub async fn resume(&self) -> Result<ResumeResponse> {
1320        let request = ResumeRequest {};
1321        let resp = self.inner.resume(request).await?;
1322        Ok(resp)
1323    }
1324
1325    pub async fn apply_throttle(
1326        &self,
1327        throttle_target: PbThrottleTarget,
1328        throttle_type: risingwave_pb::common::PbThrottleType,
1329        id: u32,
1330        rate: Option<u32>,
1331    ) -> Result<ApplyThrottleResponse> {
1332        let request = ApplyThrottleRequest {
1333            throttle_target: throttle_target as i32,
1334            throttle_type: throttle_type as i32,
1335            id,
1336            rate,
1337        };
1338        let resp = self.inner.apply_throttle(request).await?;
1339        Ok(resp)
1340    }
1341
1342    pub async fn get_cluster_recovery_status(&self) -> Result<RecoveryStatus> {
1343        let resp = self
1344            .inner
1345            .get_cluster_recovery_status(GetClusterRecoveryStatusRequest {})
1346            .await?;
1347        Ok(resp.get_status().unwrap())
1348    }
1349
1350    pub async fn get_cluster_info(&self) -> Result<GetClusterInfoResponse> {
1351        let request = GetClusterInfoRequest {};
1352        let resp = self.inner.get_cluster_info(request).await?;
1353        Ok(resp)
1354    }
1355
1356    pub async fn reschedule(
1357        &self,
1358        worker_reschedules: HashMap<u32, PbWorkerReschedule>,
1359        revision: u64,
1360        resolve_no_shuffle_upstream: bool,
1361    ) -> Result<(bool, u64)> {
1362        let request = RescheduleRequest {
1363            revision,
1364            resolve_no_shuffle_upstream,
1365            worker_reschedules,
1366        };
1367        let resp = self.inner.reschedule(request).await?;
1368        Ok((resp.success, resp.revision))
1369    }
1370
1371    pub async fn risectl_get_pinned_versions_summary(
1372        &self,
1373    ) -> Result<RiseCtlGetPinnedVersionsSummaryResponse> {
1374        let request = RiseCtlGetPinnedVersionsSummaryRequest {};
1375        self.inner
1376            .rise_ctl_get_pinned_versions_summary(request)
1377            .await
1378    }
1379
1380    pub async fn risectl_get_checkpoint_hummock_version(
1381        &self,
1382    ) -> Result<RiseCtlGetCheckpointVersionResponse> {
1383        let request = RiseCtlGetCheckpointVersionRequest {};
1384        self.inner.rise_ctl_get_checkpoint_version(request).await
1385    }
1386
1387    pub async fn risectl_pause_hummock_version_checkpoint(
1388        &self,
1389    ) -> Result<RiseCtlPauseVersionCheckpointResponse> {
1390        let request = RiseCtlPauseVersionCheckpointRequest {};
1391        self.inner.rise_ctl_pause_version_checkpoint(request).await
1392    }
1393
1394    pub async fn risectl_resume_hummock_version_checkpoint(
1395        &self,
1396    ) -> Result<RiseCtlResumeVersionCheckpointResponse> {
1397        let request = RiseCtlResumeVersionCheckpointRequest {};
1398        self.inner.rise_ctl_resume_version_checkpoint(request).await
1399    }
1400
1401    pub async fn init_metadata_for_replay(
1402        &self,
1403        tables: Vec<PbTable>,
1404        compaction_groups: Vec<CompactionGroupInfo>,
1405    ) -> Result<()> {
1406        let req = InitMetadataForReplayRequest {
1407            tables,
1408            compaction_groups,
1409        };
1410        let _resp = self.inner.init_metadata_for_replay(req).await?;
1411        Ok(())
1412    }
1413
1414    pub async fn replay_version_delta(
1415        &self,
1416        version_delta: HummockVersionDelta,
1417    ) -> Result<(HummockVersion, Vec<CompactionGroupId>)> {
1418        let req = ReplayVersionDeltaRequest {
1419            version_delta: Some(version_delta.into()),
1420        };
1421        let resp = self.inner.replay_version_delta(req).await?;
1422        Ok((
1423            HummockVersion::from_rpc_protobuf(&resp.version.unwrap()),
1424            resp.modified_compaction_groups,
1425        ))
1426    }
1427
1428    pub async fn list_version_deltas(
1429        &self,
1430        start_id: HummockVersionId,
1431        num_limit: u32,
1432        committed_epoch_limit: HummockEpoch,
1433    ) -> Result<Vec<HummockVersionDelta>> {
1434        let req = ListVersionDeltasRequest {
1435            start_id,
1436            num_limit,
1437            committed_epoch_limit,
1438        };
1439        Ok(self
1440            .inner
1441            .list_version_deltas(req)
1442            .await?
1443            .version_deltas
1444            .unwrap()
1445            .version_deltas
1446            .iter()
1447            .map(HummockVersionDelta::from_rpc_protobuf)
1448            .collect())
1449    }
1450
1451    pub async fn trigger_compaction_deterministic(
1452        &self,
1453        version_id: HummockVersionId,
1454        compaction_groups: Vec<CompactionGroupId>,
1455    ) -> Result<()> {
1456        let req = TriggerCompactionDeterministicRequest {
1457            version_id,
1458            compaction_groups,
1459        };
1460        self.inner.trigger_compaction_deterministic(req).await?;
1461        Ok(())
1462    }
1463
1464    pub async fn disable_commit_epoch(&self) -> Result<HummockVersion> {
1465        let req = DisableCommitEpochRequest {};
1466        Ok(HummockVersion::from_rpc_protobuf(
1467            &self
1468                .inner
1469                .disable_commit_epoch(req)
1470                .await?
1471                .current_version
1472                .unwrap(),
1473        ))
1474    }
1475
1476    pub async fn get_assigned_compact_task_num(&self) -> Result<usize> {
1477        let req = GetAssignedCompactTaskNumRequest {};
1478        let resp = self.inner.get_assigned_compact_task_num(req).await?;
1479        Ok(resp.num_tasks as usize)
1480    }
1481
1482    pub async fn risectl_list_compaction_group(&self) -> Result<Vec<CompactionGroupInfo>> {
1483        let req = RiseCtlListCompactionGroupRequest {};
1484        let resp = self.inner.rise_ctl_list_compaction_group(req).await?;
1485        Ok(resp.compaction_groups)
1486    }
1487
1488    pub async fn risectl_update_compaction_config(
1489        &self,
1490        compaction_groups: &[CompactionGroupId],
1491        configs: &[MutableConfig],
1492    ) -> Result<()> {
1493        let req = RiseCtlUpdateCompactionConfigRequest {
1494            compaction_group_ids: compaction_groups.to_vec(),
1495            configs: configs
1496                .iter()
1497                .map(
1498                    |c| rise_ctl_update_compaction_config_request::MutableConfig {
1499                        mutable_config: Some(c.clone()),
1500                    },
1501                )
1502                .collect(),
1503        };
1504        let _resp = self.inner.rise_ctl_update_compaction_config(req).await?;
1505        Ok(())
1506    }
1507
1508    pub async fn backup_meta(&self, remarks: Option<String>) -> Result<u64> {
1509        let req = BackupMetaRequest { remarks };
1510        let resp = self.inner.backup_meta(req).await?;
1511        Ok(resp.job_id)
1512    }
1513
1514    pub async fn get_backup_job_status(&self, job_id: u64) -> Result<(BackupJobStatus, String)> {
1515        let req = GetBackupJobStatusRequest { job_id };
1516        let resp = self.inner.get_backup_job_status(req).await?;
1517        Ok((resp.job_status(), resp.message))
1518    }
1519
1520    pub async fn delete_meta_snapshot(&self, snapshot_ids: &[u64]) -> Result<()> {
1521        let req = DeleteMetaSnapshotRequest {
1522            snapshot_ids: snapshot_ids.to_vec(),
1523        };
1524        let _resp = self.inner.delete_meta_snapshot(req).await?;
1525        Ok(())
1526    }
1527
1528    pub async fn get_meta_snapshot_manifest(&self) -> Result<MetaSnapshotManifest> {
1529        let req = GetMetaSnapshotManifestRequest {};
1530        let resp = self.inner.get_meta_snapshot_manifest(req).await?;
1531        Ok(resp.manifest.expect("should exist"))
1532    }
1533
1534    pub async fn get_telemetry_info(&self) -> Result<TelemetryInfoResponse> {
1535        let req = GetTelemetryInfoRequest {};
1536        let resp = self.inner.get_telemetry_info(req).await?;
1537        Ok(resp)
1538    }
1539
1540    pub async fn get_meta_store_endpoint(&self) -> Result<String> {
1541        let req = GetMetaStoreInfoRequest {};
1542        let resp = self.inner.get_meta_store_info(req).await?;
1543        Ok(resp.meta_store_endpoint)
1544    }
1545
1546    pub async fn alter_sink_props(
1547        &self,
1548        sink_id: SinkId,
1549        changed_props: BTreeMap<String, String>,
1550        changed_secret_refs: BTreeMap<String, PbSecretRef>,
1551        connector_conn_ref: Option<ConnectionId>,
1552    ) -> Result<()> {
1553        let req = AlterConnectorPropsRequest {
1554            object_id: sink_id.as_raw_id(),
1555            changed_props: changed_props.into_iter().collect(),
1556            changed_secret_refs: changed_secret_refs.into_iter().collect(),
1557            connector_conn_ref,
1558            object_type: AlterConnectorPropsObject::Sink as i32,
1559            extra_options: None,
1560        };
1561        let _resp = self.inner.alter_connector_props(req).await?;
1562        Ok(())
1563    }
1564
1565    pub async fn alter_iceberg_table_props(
1566        &self,
1567        table_id: TableId,
1568        sink_id: SinkId,
1569        source_id: SourceId,
1570        changed_props: BTreeMap<String, String>,
1571        changed_secret_refs: BTreeMap<String, PbSecretRef>,
1572        connector_conn_ref: Option<ConnectionId>,
1573    ) -> Result<()> {
1574        let req = AlterConnectorPropsRequest {
1575            object_id: table_id.as_raw_id(),
1576            changed_props: changed_props.into_iter().collect(),
1577            changed_secret_refs: changed_secret_refs.into_iter().collect(),
1578            connector_conn_ref,
1579            object_type: AlterConnectorPropsObject::IcebergTable as i32,
1580            extra_options: Some(ExtraOptions::AlterIcebergTableIds(AlterIcebergTableIds {
1581                sink_id,
1582                source_id,
1583            })),
1584        };
1585        let _resp = self.inner.alter_connector_props(req).await?;
1586        Ok(())
1587    }
1588
1589    pub async fn alter_source_connector_props(
1590        &self,
1591        source_id: SourceId,
1592        changed_props: BTreeMap<String, String>,
1593        changed_secret_refs: BTreeMap<String, PbSecretRef>,
1594        connector_conn_ref: Option<ConnectionId>,
1595    ) -> Result<()> {
1596        let req = AlterConnectorPropsRequest {
1597            object_id: source_id.as_raw_id(),
1598            changed_props: changed_props.into_iter().collect(),
1599            changed_secret_refs: changed_secret_refs.into_iter().collect(),
1600            connector_conn_ref,
1601            object_type: AlterConnectorPropsObject::Source as i32,
1602            extra_options: None,
1603        };
1604        let _resp = self.inner.alter_connector_props(req).await?;
1605        Ok(())
1606    }
1607
1608    pub async fn alter_connection_connector_props(
1609        &self,
1610        connection_id: u32,
1611        changed_props: BTreeMap<String, String>,
1612        changed_secret_refs: BTreeMap<String, PbSecretRef>,
1613    ) -> Result<()> {
1614        let req = AlterConnectorPropsRequest {
1615            object_id: connection_id,
1616            changed_props: changed_props.into_iter().collect(),
1617            changed_secret_refs: changed_secret_refs.into_iter().collect(),
1618            connector_conn_ref: None, // Connections don't reference other connections
1619            object_type: AlterConnectorPropsObject::Connection as i32,
1620            extra_options: None,
1621        };
1622        let _resp = self.inner.alter_connector_props(req).await?;
1623        Ok(())
1624    }
1625
1626    /// Orchestrated source property update with pause/update/resume workflow.
1627    /// This is the "safe" version that pauses sources before updating and resumes after.
1628    pub async fn alter_source_properties_safe(
1629        &self,
1630        source_id: SourceId,
1631        changed_props: BTreeMap<String, String>,
1632        changed_secret_refs: BTreeMap<String, PbSecretRef>,
1633        reset_splits: bool,
1634    ) -> Result<()> {
1635        let req = AlterSourcePropertiesSafeRequest {
1636            source_id: source_id.as_raw_id(),
1637            changed_props: changed_props.into_iter().collect(),
1638            changed_secret_refs: changed_secret_refs.into_iter().collect(),
1639            options: Some(PropertyUpdateOptions { reset_splits }),
1640        };
1641        let _resp = self.inner.alter_source_properties_safe(req).await?;
1642        Ok(())
1643    }
1644
1645    /// Reset source split assignments (UNSAFE - admin only).
1646    /// This clears persisted split metadata and triggers re-discovery.
1647    pub async fn reset_source_splits(&self, source_id: SourceId) -> Result<()> {
1648        let req = ResetSourceSplitsRequest {
1649            source_id: source_id.as_raw_id(),
1650        };
1651        let _resp = self.inner.reset_source_splits(req).await?;
1652        Ok(())
1653    }
1654
1655    /// Inject specific offsets into source splits (UNSAFE - admin only).
1656    /// This can cause data duplication or loss depending on the correctness of the provided offsets.
1657    pub async fn inject_source_offsets(
1658        &self,
1659        source_id: SourceId,
1660        split_offsets: HashMap<String, String>,
1661    ) -> Result<Vec<String>> {
1662        let req = InjectSourceOffsetsRequest {
1663            source_id: source_id.as_raw_id(),
1664            split_offsets,
1665        };
1666        let resp = self.inner.inject_source_offsets(req).await?;
1667        Ok(resp.applied_split_ids)
1668    }
1669
1670    pub async fn set_system_param(
1671        &self,
1672        param: String,
1673        value: Option<String>,
1674    ) -> Result<Option<SystemParamsReader>> {
1675        let req = SetSystemParamRequest { param, value };
1676        let resp = self.inner.set_system_param(req).await?;
1677        Ok(resp.params.map(SystemParamsReader::from))
1678    }
1679
1680    pub async fn get_session_params(&self) -> Result<String> {
1681        let req = GetSessionParamsRequest {};
1682        let resp = self.inner.get_session_params(req).await?;
1683        Ok(resp.params)
1684    }
1685
1686    pub async fn set_session_param(&self, param: String, value: Option<String>) -> Result<String> {
1687        let req = SetSessionParamRequest { param, value };
1688        let resp = self.inner.set_session_param(req).await?;
1689        Ok(resp.param)
1690    }
1691
1692    pub async fn get_ddl_progress(&self) -> Result<Vec<DdlProgress>> {
1693        let req = GetDdlProgressRequest {};
1694        let resp = self.inner.get_ddl_progress(req).await?;
1695        Ok(resp.ddl_progress)
1696    }
1697
1698    pub async fn split_compaction_group(
1699        &self,
1700        group_id: CompactionGroupId,
1701        table_ids_to_new_group: &[TableId],
1702        partition_vnode_count: u32,
1703    ) -> Result<CompactionGroupId> {
1704        let req = SplitCompactionGroupRequest {
1705            group_id,
1706            table_ids: table_ids_to_new_group.to_vec(),
1707            partition_vnode_count,
1708        };
1709        let resp = self.inner.split_compaction_group(req).await?;
1710        Ok(resp.new_group_id)
1711    }
1712
1713    pub async fn get_tables(
1714        &self,
1715        table_ids: Vec<TableId>,
1716        include_dropped_tables: bool,
1717    ) -> Result<HashMap<TableId, Table>> {
1718        let req = GetTablesRequest {
1719            table_ids,
1720            include_dropped_tables,
1721        };
1722        let resp = self.inner.get_tables(req).await?;
1723        Ok(resp.tables)
1724    }
1725
1726    pub async fn list_serving_vnode_mappings(
1727        &self,
1728    ) -> Result<HashMap<FragmentId, (JobId, WorkerSlotMapping)>> {
1729        let req = GetServingVnodeMappingsRequest {};
1730        let resp = self.inner.get_serving_vnode_mappings(req).await?;
1731        let mappings = resp
1732            .worker_slot_mappings
1733            .into_iter()
1734            .map(|p| {
1735                (
1736                    p.fragment_id,
1737                    (
1738                        resp.fragment_to_table
1739                            .get(&p.fragment_id)
1740                            .cloned()
1741                            .unwrap_or_default(),
1742                        WorkerSlotMapping::from_protobuf(p.mapping.as_ref().unwrap()),
1743                    ),
1744                )
1745            })
1746            .collect();
1747        Ok(mappings)
1748    }
1749
1750    pub async fn risectl_list_compaction_status(
1751        &self,
1752    ) -> Result<(
1753        Vec<CompactStatus>,
1754        Vec<CompactTaskAssignment>,
1755        Vec<CompactTaskProgress>,
1756    )> {
1757        let req = RiseCtlListCompactionStatusRequest {};
1758        let resp = self.inner.rise_ctl_list_compaction_status(req).await?;
1759        Ok((
1760            resp.compaction_statuses,
1761            resp.task_assignment,
1762            resp.task_progress,
1763        ))
1764    }
1765
1766    pub async fn get_compaction_score(
1767        &self,
1768        compaction_group_id: CompactionGroupId,
1769    ) -> Result<Vec<PickerInfo>> {
1770        let req = GetCompactionScoreRequest {
1771            compaction_group_id,
1772        };
1773        let resp = self.inner.get_compaction_score(req).await?;
1774        Ok(resp.scores)
1775    }
1776
1777    pub async fn risectl_rebuild_table_stats(&self) -> Result<()> {
1778        let req = RiseCtlRebuildTableStatsRequest {};
1779        let _resp = self.inner.rise_ctl_rebuild_table_stats(req).await?;
1780        Ok(())
1781    }
1782
1783    pub async fn list_branched_object(&self) -> Result<Vec<BranchedObject>> {
1784        let req = ListBranchedObjectRequest {};
1785        let resp = self.inner.list_branched_object(req).await?;
1786        Ok(resp.branched_objects)
1787    }
1788
1789    pub async fn list_active_write_limit(&self) -> Result<HashMap<CompactionGroupId, WriteLimit>> {
1790        let req = ListActiveWriteLimitRequest {};
1791        let resp = self.inner.list_active_write_limit(req).await?;
1792        Ok(resp.write_limits)
1793    }
1794
1795    pub async fn list_hummock_meta_config(&self) -> Result<HashMap<String, String>> {
1796        let req = ListHummockMetaConfigRequest {};
1797        let resp = self.inner.list_hummock_meta_config(req).await?;
1798        Ok(resp.configs)
1799    }
1800
1801    pub async fn get_table_change_logs(
1802        &self,
1803        epoch_only: bool,
1804        start_epoch_inclusive: Option<u64>,
1805        end_epoch_inclusive: Option<u64>,
1806        table_ids: Option<HashSet<TableId>>,
1807        exclude_empty: bool,
1808        limit: Option<u32>,
1809    ) -> Result<TableChangeLogs> {
1810        let req = GetTableChangeLogsRequest {
1811            epoch_only,
1812            start_epoch_inclusive,
1813            end_epoch_inclusive,
1814            table_ids: table_ids.map(|iter| PbTableFilter {
1815                table_ids: iter.into_iter().collect::<Vec<_>>(),
1816            }),
1817            exclude_empty,
1818            limit,
1819        };
1820        let resp = self.inner.get_table_change_logs(req).await?;
1821        Ok(resp
1822            .table_change_logs
1823            .into_iter()
1824            .map(|(id, change_log)| (TableId::new(id), TableChangeLog::from_protobuf(&change_log)))
1825            .collect())
1826    }
1827
1828    pub async fn delete_worker_node(&self, worker: HostAddress) -> Result<()> {
1829        let _resp = self
1830            .inner
1831            .delete_worker_node(DeleteWorkerNodeRequest { host: Some(worker) })
1832            .await?;
1833
1834        Ok(())
1835    }
1836
1837    pub async fn rw_cloud_validate_source(
1838        &self,
1839        source_type: SourceType,
1840        source_config: HashMap<String, String>,
1841    ) -> Result<RwCloudValidateSourceResponse> {
1842        let req = RwCloudValidateSourceRequest {
1843            source_type: source_type.into(),
1844            source_config,
1845        };
1846        let resp = self.inner.rw_cloud_validate_source(req).await?;
1847        Ok(resp)
1848    }
1849
1850    pub async fn sink_coordinate_client(&self) -> SinkCoordinationRpcClient {
1851        self.inner.core.read().await.sink_coordinate_client.clone()
1852    }
1853
1854    pub async fn list_compact_task_assignment(&self) -> Result<Vec<CompactTaskAssignment>> {
1855        let req = ListCompactTaskAssignmentRequest {};
1856        let resp = self.inner.list_compact_task_assignment(req).await?;
1857        Ok(resp.task_assignment)
1858    }
1859
1860    pub async fn list_event_log(&self) -> Result<Vec<EventLog>> {
1861        let req = ListEventLogRequest::default();
1862        let resp = self.inner.list_event_log(req).await?;
1863        Ok(resp.event_logs)
1864    }
1865
1866    pub async fn list_compact_task_progress(&self) -> Result<Vec<CompactTaskProgress>> {
1867        let req = ListCompactTaskProgressRequest {};
1868        let resp = self.inner.list_compact_task_progress(req).await?;
1869        Ok(resp.task_progress)
1870    }
1871
1872    #[cfg(madsim)]
1873    pub fn try_add_panic_event_blocking(
1874        &self,
1875        panic_info: impl Display,
1876        timeout_millis: Option<u64>,
1877    ) {
1878    }
1879
1880    /// If `timeout_millis` is None, default is used.
1881    #[cfg(not(madsim))]
1882    pub fn try_add_panic_event_blocking(
1883        &self,
1884        panic_info: impl Display,
1885        timeout_millis: Option<u64>,
1886    ) {
1887        let event = event_log::EventWorkerNodePanic {
1888            worker_id: self.worker_id,
1889            worker_type: self.worker_type.into(),
1890            host_addr: Some(self.host_addr.to_protobuf()),
1891            panic_info: format!("{panic_info}"),
1892        };
1893        let grpc_meta_client = self.inner.clone();
1894        let _ = thread::spawn(move || {
1895            let rt = tokio::runtime::Builder::new_current_thread()
1896                .enable_all()
1897                .build()
1898                .unwrap();
1899            let req = AddEventLogRequest {
1900                event: Some(add_event_log_request::Event::WorkerNodePanic(event)),
1901            };
1902            rt.block_on(async {
1903                let _ = tokio::time::timeout(
1904                    Duration::from_millis(timeout_millis.unwrap_or(1000)),
1905                    grpc_meta_client.add_event_log(req),
1906                )
1907                .await;
1908            });
1909        })
1910        .join();
1911    }
1912
1913    pub async fn add_sink_fail_evet(
1914        &self,
1915        sink_id: SinkId,
1916        sink_name: String,
1917        connector: String,
1918        error: String,
1919    ) -> Result<()> {
1920        let event = event_log::EventSinkFail {
1921            sink_id,
1922            sink_name,
1923            connector,
1924            error,
1925        };
1926        let req = AddEventLogRequest {
1927            event: Some(add_event_log_request::Event::SinkFail(event)),
1928        };
1929        self.inner.add_event_log(req).await?;
1930        Ok(())
1931    }
1932
1933    pub async fn add_cdc_auto_schema_change_fail_event(
1934        &self,
1935        source_id: SourceId,
1936        table_name: String,
1937        cdc_table_id: String,
1938        upstream_ddl: String,
1939        fail_info: String,
1940    ) -> Result<()> {
1941        let event = event_log::EventAutoSchemaChangeFail {
1942            table_id: source_id.as_cdc_table_id(),
1943            table_name,
1944            cdc_table_id,
1945            upstream_ddl,
1946            fail_info,
1947        };
1948        let req = AddEventLogRequest {
1949            event: Some(add_event_log_request::Event::AutoSchemaChangeFail(event)),
1950        };
1951        self.inner.add_event_log(req).await?;
1952        Ok(())
1953    }
1954
1955    pub async fn cancel_compact_task(&self, task_id: u64, task_status: TaskStatus) -> Result<bool> {
1956        let req = CancelCompactTaskRequest {
1957            task_id,
1958            task_status: task_status as _,
1959        };
1960        let resp = self.inner.cancel_compact_task(req).await?;
1961        Ok(resp.ret)
1962    }
1963
1964    pub async fn get_version_by_epoch(
1965        &self,
1966        epoch: HummockEpoch,
1967        table_id: TableId,
1968    ) -> Result<PbHummockVersion> {
1969        let req = GetVersionByEpochRequest { epoch, table_id };
1970        let resp = self.inner.get_version_by_epoch(req).await?;
1971        Ok(resp.version.unwrap())
1972    }
1973
1974    pub async fn get_cluster_limits(
1975        &self,
1976    ) -> Result<Vec<risingwave_common::util::cluster_limit::ClusterLimit>> {
1977        let req = GetClusterLimitsRequest {};
1978        let resp = self.inner.get_cluster_limits(req).await?;
1979        Ok(resp.active_limits.into_iter().map(|l| l.into()).collect())
1980    }
1981
1982    pub async fn merge_compaction_group(
1983        &self,
1984        left_group_id: CompactionGroupId,
1985        right_group_id: CompactionGroupId,
1986    ) -> Result<()> {
1987        let req = MergeCompactionGroupRequest {
1988            left_group_id,
1989            right_group_id,
1990        };
1991        self.inner.merge_compaction_group(req).await?;
1992        Ok(())
1993    }
1994
1995    /// List all rate limits for sources and backfills
1996    pub async fn list_rate_limits(&self) -> Result<Vec<RateLimitInfo>> {
1997        let request = ListRateLimitsRequest {};
1998        let resp = self.inner.list_rate_limits(request).await?;
1999        Ok(resp.rate_limits)
2000    }
2001
2002    pub async fn list_cdc_progress(&self) -> Result<HashMap<JobId, PbCdcProgress>> {
2003        let request = ListCdcProgressRequest {};
2004        let resp = self.inner.list_cdc_progress(request).await?;
2005        Ok(resp.cdc_progress)
2006    }
2007
2008    pub async fn list_refresh_table_states(&self) -> Result<Vec<RefreshTableState>> {
2009        let request = ListRefreshTableStatesRequest {};
2010        let resp = self.inner.list_refresh_table_states(request).await?;
2011        Ok(resp.states)
2012    }
2013
2014    pub async fn list_iceberg_compaction_status(
2015        &self,
2016    ) -> Result<Vec<list_iceberg_compaction_status_response::IcebergCompactionStatus>> {
2017        let request = ListIcebergCompactionStatusRequest {};
2018        let resp = self.inner.list_iceberg_compaction_status(request).await?;
2019        Ok(resp.statuses)
2020    }
2021
2022    pub async fn list_sink_log_store_tables(
2023        &self,
2024    ) -> Result<Vec<list_sink_log_store_tables_response::SinkLogStoreTable>> {
2025        let request = ListSinkLogStoreTablesRequest {};
2026        let resp = self.inner.list_sink_log_store_tables(request).await?;
2027        Ok(resp.tables)
2028    }
2029
2030    pub async fn create_iceberg_table(
2031        &self,
2032        table_job_info: PbTableJobInfo,
2033        sink_job_info: PbSinkJobInfo,
2034        iceberg_source: PbSource,
2035        if_not_exists: bool,
2036    ) -> Result<WaitVersion> {
2037        let request = CreateIcebergTableRequest {
2038            table_info: Some(table_job_info),
2039            sink_info: Some(sink_job_info),
2040            iceberg_source: Some(iceberg_source),
2041            if_not_exists,
2042        };
2043
2044        let resp = Box::pin(self.inner.create_iceberg_table(request)).await?;
2045        Ok(resp
2046            .version
2047            .ok_or_else(|| anyhow!("wait version not set"))?)
2048    }
2049
2050    pub async fn list_hosted_iceberg_tables(&self) -> Result<Vec<IcebergTable>> {
2051        let request = ListIcebergTablesRequest {};
2052        let resp = self.inner.list_iceberg_tables(request).await?;
2053        Ok(resp.iceberg_tables)
2054    }
2055
2056    /// Get the await tree of all nodes in the cluster.
2057    pub async fn get_cluster_stack_trace(
2058        &self,
2059        actor_traces_format: ActorTracesFormat,
2060    ) -> Result<StackTraceResponse> {
2061        let request = StackTraceRequest {
2062            actor_traces_format: actor_traces_format as i32,
2063        };
2064        let resp = self.inner.stack_trace(request).await?;
2065        Ok(resp)
2066    }
2067
2068    pub async fn set_sync_log_store_aligned(&self, job_id: JobId, aligned: bool) -> Result<()> {
2069        let request = SetSyncLogStoreAlignedRequest { job_id, aligned };
2070        self.inner.set_sync_log_store_aligned(request).await?;
2071        Ok(())
2072    }
2073
2074    pub async fn refresh(&self, request: RefreshRequest) -> Result<RefreshResponse> {
2075        self.inner.refresh(request).await
2076    }
2077
2078    pub async fn list_unmigrated_tables(&self) -> Result<HashMap<TableId, String>> {
2079        let request = ListUnmigratedTablesRequest {};
2080        let resp = self.inner.list_unmigrated_tables(request).await?;
2081        Ok(resp
2082            .tables
2083            .into_iter()
2084            .map(|table| (table.table_id, table.table_name))
2085            .collect())
2086    }
2087}
2088
2089#[async_trait]
2090impl HummockMetaClient for MetaClient {
2091    async fn unpin_version_before(&self, unpin_version_before: HummockVersionId) -> Result<()> {
2092        let req = UnpinVersionBeforeRequest {
2093            context_id: self.worker_id(),
2094            unpin_version_before,
2095        };
2096        self.inner.unpin_version_before(req).await?;
2097        Ok(())
2098    }
2099
2100    async fn get_current_version(&self) -> Result<HummockVersion> {
2101        let req = GetCurrentVersionRequest::default();
2102        Ok(HummockVersion::from_rpc_protobuf(
2103            &self
2104                .inner
2105                .get_current_version(req)
2106                .await?
2107                .current_version
2108                .unwrap(),
2109        ))
2110    }
2111
2112    async fn get_new_object_ids(&self, number: u32) -> Result<ObjectIdRange> {
2113        let resp = self
2114            .inner
2115            .get_new_object_ids(GetNewObjectIdsRequest { number })
2116            .await?;
2117        Ok(ObjectIdRange::new(resp.start_id, resp.end_id))
2118    }
2119
2120    async fn commit_epoch_with_change_log(
2121        &self,
2122        _epoch: HummockEpoch,
2123        _sync_result: SyncResult,
2124        _change_log_info: Option<HummockMetaClientChangeLogInfo>,
2125    ) -> Result<()> {
2126        panic!("Only meta service can commit_epoch in production.")
2127    }
2128
2129    async fn trigger_manual_compaction(
2130        &self,
2131        compaction_group_id: CompactionGroupId,
2132        table_id: JobId,
2133        level: u32,
2134        target_level: Option<u32>,
2135        sst_ids: Vec<HummockSstableId>,
2136        exclusive: bool,
2137    ) -> Result<bool> {
2138        // TODO: support key_range parameter
2139        let req = TriggerManualCompactionRequest {
2140            compaction_group_id,
2141            table_id,
2142            // if table_id not exist, manual_compaction will include all the sst
2143            // without check internal_table_id
2144            level,
2145            target_level,
2146            sst_ids,
2147            exclusive: Some(exclusive),
2148            ..Default::default()
2149        };
2150
2151        let resp = self.inner.trigger_manual_compaction(req).await?;
2152        Ok(resp.should_retry.unwrap_or(false))
2153    }
2154
2155    async fn trigger_full_gc(
2156        &self,
2157        sst_retention_time_sec: u64,
2158        prefix: Option<String>,
2159    ) -> Result<()> {
2160        self.inner
2161            .trigger_full_gc(TriggerFullGcRequest {
2162                sst_retention_time_sec,
2163                prefix,
2164            })
2165            .await?;
2166        Ok(())
2167    }
2168
2169    async fn subscribe_compaction_event(
2170        &self,
2171    ) -> Result<(
2172        UnboundedSender<SubscribeCompactionEventRequest>,
2173        BoxStream<'static, CompactionEventItem>,
2174    )> {
2175        let (request_sender, request_receiver) =
2176            unbounded_channel::<SubscribeCompactionEventRequest>();
2177        request_sender
2178            .send(SubscribeCompactionEventRequest {
2179                event: Some(subscribe_compaction_event_request::Event::Register(
2180                    Register {
2181                        context_id: self.worker_id(),
2182                    },
2183                )),
2184                create_at: SystemTime::now()
2185                    .duration_since(SystemTime::UNIX_EPOCH)
2186                    .expect("Clock may have gone backwards")
2187                    .as_millis() as u64,
2188            })
2189            .context("Failed to subscribe compaction event")?;
2190
2191        let stream = self
2192            .inner
2193            .subscribe_compaction_event(Request::new(UnboundedReceiverStream::new(
2194                request_receiver,
2195            )))
2196            .await?;
2197
2198        Ok((request_sender, Box::pin(stream)))
2199    }
2200
2201    async fn get_version_by_epoch(
2202        &self,
2203        epoch: HummockEpoch,
2204        table_id: TableId,
2205    ) -> Result<PbHummockVersion> {
2206        self.get_version_by_epoch(epoch, table_id).await
2207    }
2208
2209    async fn subscribe_iceberg_compaction_event(
2210        &self,
2211    ) -> Result<(
2212        UnboundedSender<SubscribeIcebergCompactionEventRequest>,
2213        BoxStream<'static, IcebergCompactionEventItem>,
2214    )> {
2215        let (request_sender, request_receiver) =
2216            unbounded_channel::<SubscribeIcebergCompactionEventRequest>();
2217        request_sender
2218            .send(SubscribeIcebergCompactionEventRequest {
2219                event: Some(subscribe_iceberg_compaction_event_request::Event::Register(
2220                    IcebergRegister {
2221                        context_id: self.worker_id(),
2222                    },
2223                )),
2224                create_at: SystemTime::now()
2225                    .duration_since(SystemTime::UNIX_EPOCH)
2226                    .expect("Clock may have gone backwards")
2227                    .as_millis() as u64,
2228            })
2229            .context("Failed to subscribe compaction event")?;
2230
2231        let stream = self
2232            .inner
2233            .subscribe_iceberg_compaction_event(Request::new(UnboundedReceiverStream::new(
2234                request_receiver,
2235            )))
2236            .await?;
2237
2238        Ok((request_sender, Box::pin(stream)))
2239    }
2240
2241    async fn get_table_change_logs(
2242        &self,
2243        epoch_only: bool,
2244        start_epoch_inclusive: Option<u64>,
2245        end_epoch_inclusive: Option<u64>,
2246        table_ids: Option<HashSet<TableId>>,
2247        exclude_empty: bool,
2248        limit: Option<u32>,
2249    ) -> Result<TableChangeLogs> {
2250        self.get_table_change_logs(
2251            epoch_only,
2252            start_epoch_inclusive,
2253            end_epoch_inclusive,
2254            table_ids,
2255            exclude_empty,
2256            limit,
2257        )
2258        .await
2259    }
2260}
2261
2262#[async_trait]
2263impl TelemetryInfoFetcher for MetaClient {
2264    async fn fetch_telemetry_info(&self) -> std::result::Result<Option<String>, String> {
2265        let resp = self
2266            .get_telemetry_info()
2267            .await
2268            .map_err(|e| e.to_report_string())?;
2269        let tracking_id = resp.get_tracking_id().ok();
2270        Ok(tracking_id.map(|id| id.to_owned()))
2271    }
2272}
2273
2274pub type SinkCoordinationRpcClient = SinkCoordinationServiceClient<Channel>;
2275
2276#[derive(Debug, Clone)]
2277struct GrpcMetaClientCore {
2278    cluster_client: ClusterServiceClient<Channel>,
2279    meta_member_client: MetaMemberServiceClient<Channel>,
2280    heartbeat_client: HeartbeatServiceClient<Channel>,
2281    ddl_client: DdlServiceClient<Channel>,
2282    hummock_client: HummockManagerServiceClient<Channel>,
2283    notification_client: NotificationServiceClient<Channel>,
2284    stream_client: StreamManagerServiceClient<Channel>,
2285    user_client: UserServiceClient<Channel>,
2286    scale_client: ScaleServiceClient<Channel>,
2287    backup_client: BackupServiceClient<Channel>,
2288    telemetry_client: TelemetryInfoServiceClient<Channel>,
2289    system_params_client: SystemParamsServiceClient<Channel>,
2290    session_params_client: SessionParamServiceClient<Channel>,
2291    serving_client: ServingServiceClient<Channel>,
2292    cloud_client: CloudServiceClient<Channel>,
2293    sink_coordinate_client: SinkCoordinationRpcClient,
2294    event_log_client: EventLogServiceClient<Channel>,
2295    cluster_limit_client: ClusterLimitServiceClient<Channel>,
2296    hosted_iceberg_catalog_service_client: HostedIcebergCatalogServiceClient<Channel>,
2297    monitor_client: MonitorServiceClient<Channel>,
2298}
2299
2300impl GrpcMetaClientCore {
2301    pub(crate) fn new(channel: Channel) -> Self {
2302        let cluster_client = ClusterServiceClient::new(channel.clone());
2303        let meta_member_client = MetaMemberClient::new(channel.clone());
2304        let heartbeat_client = HeartbeatServiceClient::new(channel.clone());
2305        let ddl_client =
2306            DdlServiceClient::new(channel.clone()).max_decoding_message_size(usize::MAX);
2307        let hummock_client =
2308            HummockManagerServiceClient::new(channel.clone()).max_decoding_message_size(usize::MAX);
2309        let notification_client =
2310            NotificationServiceClient::new(channel.clone()).max_decoding_message_size(usize::MAX);
2311        let stream_client =
2312            StreamManagerServiceClient::new(channel.clone()).max_decoding_message_size(usize::MAX);
2313        let user_client = UserServiceClient::new(channel.clone());
2314        let scale_client =
2315            ScaleServiceClient::new(channel.clone()).max_decoding_message_size(usize::MAX);
2316        let backup_client = BackupServiceClient::new(channel.clone());
2317        let telemetry_client =
2318            TelemetryInfoServiceClient::new(channel.clone()).max_decoding_message_size(usize::MAX);
2319        let system_params_client = SystemParamsServiceClient::new(channel.clone());
2320        let session_params_client = SessionParamServiceClient::new(channel.clone());
2321        let serving_client = ServingServiceClient::new(channel.clone());
2322        let cloud_client = CloudServiceClient::new(channel.clone());
2323        let sink_coordinate_client = SinkCoordinationServiceClient::new(channel.clone())
2324            .max_decoding_message_size(usize::MAX);
2325        let event_log_client = EventLogServiceClient::new(channel.clone());
2326        let cluster_limit_client = ClusterLimitServiceClient::new(channel.clone());
2327        let hosted_iceberg_catalog_service_client =
2328            HostedIcebergCatalogServiceClient::new(channel.clone());
2329        let monitor_client = configured_monitor_service_client(MonitorServiceClient::new(channel));
2330
2331        GrpcMetaClientCore {
2332            cluster_client,
2333            meta_member_client,
2334            heartbeat_client,
2335            ddl_client,
2336            hummock_client,
2337            notification_client,
2338            stream_client,
2339            user_client,
2340            scale_client,
2341            backup_client,
2342            telemetry_client,
2343            system_params_client,
2344            session_params_client,
2345            serving_client,
2346            cloud_client,
2347            sink_coordinate_client,
2348            event_log_client,
2349            cluster_limit_client,
2350            hosted_iceberg_catalog_service_client,
2351            monitor_client,
2352        }
2353    }
2354}
2355
2356/// Client to meta server. Cloning the instance is lightweight.
2357///
2358/// It is a wrapper of tonic client. See [`crate::meta_rpc_client_method_impl`].
2359#[derive(Debug, Clone)]
2360struct GrpcMetaClient {
2361    member_monitor_event_sender: mpsc::Sender<Sender<Result<()>>>,
2362    core: Arc<RwLock<GrpcMetaClientCore>>,
2363}
2364
2365type MetaMemberClient = MetaMemberServiceClient<Channel>;
2366
2367struct MetaMemberGroup {
2368    members: LruCache<http::Uri, Option<MetaMemberClient>>,
2369}
2370
2371struct MetaMemberManagement {
2372    core_ref: Arc<RwLock<GrpcMetaClientCore>>,
2373    members: Either<MetaMemberClient, MetaMemberGroup>,
2374    current_leader: http::Uri,
2375    meta_config: Arc<MetaConfig>,
2376}
2377
2378impl MetaMemberManagement {
2379    const META_MEMBER_REFRESH_PERIOD: Duration = Duration::from_secs(5);
2380
2381    fn host_address_to_uri(addr: HostAddress) -> http::Uri {
2382        format!("http://{}:{}", addr.host, addr.port)
2383            .parse()
2384            .unwrap()
2385    }
2386
2387    async fn recreate_core(&self, channel: Channel) {
2388        let mut core = self.core_ref.write().await;
2389        *core = GrpcMetaClientCore::new(channel);
2390    }
2391
2392    async fn refresh_members(&mut self) -> Result<()> {
2393        let leader_addr = match self.members.as_mut() {
2394            Either::Left(client) => {
2395                let resp = client
2396                    .to_owned()
2397                    .members(MembersRequest {})
2398                    .await
2399                    .map_err(RpcError::from_meta_status)?;
2400                let resp = resp.into_inner();
2401                resp.members.into_iter().find(|member| member.is_leader)
2402            }
2403            Either::Right(member_group) => {
2404                let mut fetched_members = None;
2405
2406                for (addr, client) in &mut member_group.members {
2407                    let members: Result<_> = try {
2408                        let mut client = match client {
2409                            Some(cached_client) => cached_client.to_owned(),
2410                            None => {
2411                                let endpoint = GrpcMetaClient::addr_to_endpoint(addr.clone());
2412                                let channel = GrpcMetaClient::connect_to_endpoint(endpoint)
2413                                    .await
2414                                    .context("failed to create client")
2415                                    .map_err(RpcError::from)?;
2416                                let new_client: MetaMemberClient =
2417                                    MetaMemberServiceClient::new(channel);
2418                                *client = Some(new_client.clone());
2419
2420                                new_client
2421                            }
2422                        };
2423
2424                        let resp = client
2425                            .members(MembersRequest {})
2426                            .await
2427                            .context("failed to fetch members")
2428                            .map_err(RpcError::from)?;
2429
2430                        resp.into_inner().members
2431                    };
2432
2433                    let fetched = members.is_ok();
2434                    fetched_members = Some(members);
2435                    if fetched {
2436                        break;
2437                    }
2438                }
2439
2440                let members = fetched_members
2441                    .context("no member available in the list")?
2442                    .context("could not refresh members")?;
2443
2444                // find new leader
2445                let mut leader = None;
2446                for member in members {
2447                    if member.is_leader {
2448                        leader = Some(member.clone());
2449                    }
2450
2451                    let addr = Self::host_address_to_uri(member.address.unwrap());
2452                    // We don't clean any expired addrs here to deal with some extreme situations.
2453                    if !member_group.members.contains(&addr) {
2454                        tracing::info!("new meta member joined: {}", addr);
2455                        member_group.members.put(addr, None);
2456                    }
2457                }
2458
2459                leader
2460            }
2461        };
2462
2463        if let Some(leader) = leader_addr {
2464            let discovered_leader = Self::host_address_to_uri(leader.address.unwrap());
2465
2466            if discovered_leader != self.current_leader {
2467                tracing::info!("new meta leader {} discovered", discovered_leader);
2468
2469                let retry_strategy = GrpcMetaClient::retry_strategy_to_bound(
2470                    Duration::from_secs(self.meta_config.meta_leader_lease_secs),
2471                    false,
2472                );
2473
2474                let channel = tokio_retry::Retry::spawn(retry_strategy, || async {
2475                    let endpoint = GrpcMetaClient::addr_to_endpoint(discovered_leader.clone());
2476                    GrpcMetaClient::connect_to_endpoint(endpoint).await
2477                })
2478                .await?;
2479
2480                self.recreate_core(channel).await;
2481                self.current_leader = discovered_leader;
2482            }
2483        }
2484
2485        Ok(())
2486    }
2487}
2488
2489impl GrpcMetaClient {
2490    // See `Endpoint::http2_keep_alive_interval`
2491    const ENDPOINT_KEEP_ALIVE_INTERVAL_SEC: u64 = 60;
2492    // See `Endpoint::keep_alive_timeout`
2493    const ENDPOINT_KEEP_ALIVE_TIMEOUT_SEC: u64 = 60;
2494    // Retry base interval in ms for connecting to meta server.
2495    const INIT_RETRY_BASE_INTERVAL_MS: u64 = 10;
2496    // Max retry times for connecting to meta server.
2497    const INIT_RETRY_MAX_INTERVAL_MS: u64 = 2000;
2498
2499    fn start_meta_member_monitor(
2500        &self,
2501        init_leader_addr: http::Uri,
2502        members: Either<MetaMemberClient, MetaMemberGroup>,
2503        force_refresh_receiver: Receiver<Sender<Result<()>>>,
2504        meta_config: Arc<MetaConfig>,
2505    ) -> Result<()> {
2506        let core_ref: Arc<RwLock<GrpcMetaClientCore>> = self.core.clone();
2507        let current_leader = init_leader_addr;
2508
2509        let enable_period_tick = matches!(members, Either::Right(_));
2510
2511        let member_management = MetaMemberManagement {
2512            core_ref,
2513            members,
2514            current_leader,
2515            meta_config,
2516        };
2517
2518        let mut force_refresh_receiver = force_refresh_receiver;
2519
2520        tokio::spawn(async move {
2521            let mut member_management = member_management;
2522            let mut ticker = time::interval(MetaMemberManagement::META_MEMBER_REFRESH_PERIOD);
2523
2524            loop {
2525                let event: Option<Sender<Result<()>>> = if enable_period_tick {
2526                    tokio::select! {
2527                        _ = ticker.tick() => None,
2528                        result_sender = force_refresh_receiver.recv() => {
2529                            if result_sender.is_none() {
2530                                break;
2531                            }
2532
2533                            result_sender
2534                        },
2535                    }
2536                } else {
2537                    let result_sender = force_refresh_receiver.recv().await;
2538
2539                    if result_sender.is_none() {
2540                        break;
2541                    }
2542
2543                    result_sender
2544                };
2545
2546                let tick_result = member_management.refresh_members().await;
2547                if let Err(e) = tick_result.as_ref() {
2548                    tracing::warn!(error = %e.as_report(),  "refresh meta member client failed");
2549                }
2550
2551                if let Some(sender) = event {
2552                    // ignore resp
2553                    let _resp = sender.send(tick_result);
2554                }
2555            }
2556        });
2557
2558        Ok(())
2559    }
2560
2561    async fn force_refresh_leader(&self) -> Result<()> {
2562        let (sender, receiver) = oneshot::channel();
2563
2564        self.member_monitor_event_sender
2565            .send(sender)
2566            .await
2567            .map_err(|e| anyhow!(e))?;
2568
2569        receiver.await.map_err(|e| anyhow!(e))?
2570    }
2571
2572    /// Connect to the meta server from `addrs`.
2573    pub async fn new(strategy: &MetaAddressStrategy, config: Arc<MetaConfig>) -> Result<Self> {
2574        let (channel, addr) = match strategy {
2575            MetaAddressStrategy::LoadBalance(addr) => {
2576                Self::try_build_rpc_channel(vec![addr.clone()]).await
2577            }
2578            MetaAddressStrategy::List(addrs) => Self::try_build_rpc_channel(addrs.clone()).await,
2579        }?;
2580        let (force_refresh_sender, force_refresh_receiver) = mpsc::channel(1);
2581        let client = GrpcMetaClient {
2582            member_monitor_event_sender: force_refresh_sender,
2583            core: Arc::new(RwLock::new(GrpcMetaClientCore::new(channel))),
2584        };
2585
2586        let meta_member_client = client.core.read().await.meta_member_client.clone();
2587        let members = match strategy {
2588            MetaAddressStrategy::LoadBalance(_) => Either::Left(meta_member_client),
2589            MetaAddressStrategy::List(addrs) => {
2590                let mut members = LruCache::new(20.try_into().unwrap());
2591                for addr in addrs {
2592                    members.put(addr.clone(), None);
2593                }
2594                members.put(addr.clone(), Some(meta_member_client));
2595
2596                Either::Right(MetaMemberGroup { members })
2597            }
2598        };
2599
2600        client.start_meta_member_monitor(addr, members, force_refresh_receiver, config.clone())?;
2601
2602        client.force_refresh_leader().await?;
2603
2604        Ok(client)
2605    }
2606
2607    fn addr_to_endpoint(addr: http::Uri) -> Endpoint {
2608        Endpoint::from(addr).initial_connection_window_size(MAX_CONNECTION_WINDOW_SIZE)
2609    }
2610
2611    pub(crate) async fn try_build_rpc_channel(
2612        addrs: impl IntoIterator<Item = http::Uri>,
2613    ) -> Result<(Channel, http::Uri)> {
2614        let endpoints: Vec<_> = addrs
2615            .into_iter()
2616            .map(|addr| (Self::addr_to_endpoint(addr.clone()), addr))
2617            .collect();
2618
2619        let mut last_error = None;
2620
2621        for (endpoint, addr) in endpoints {
2622            match Self::connect_to_endpoint(endpoint).await {
2623                Ok(channel) => {
2624                    tracing::info!("Connect to meta server {} successfully", addr);
2625                    return Ok((channel, addr));
2626                }
2627                Err(e) => {
2628                    tracing::warn!(
2629                        error = %e.as_report(),
2630                        "Failed to connect to meta server {}, trying again",
2631                        addr,
2632                    );
2633                    last_error = Some(e);
2634                }
2635            }
2636        }
2637
2638        if let Some(last_error) = last_error {
2639            Err(anyhow::anyhow!(last_error)
2640                .context("failed to connect to all meta servers")
2641                .into())
2642        } else {
2643            bail!("no meta server address provided")
2644        }
2645    }
2646
2647    async fn connect_to_endpoint(endpoint: Endpoint) -> Result<Channel> {
2648        let channel = endpoint
2649            .http2_keep_alive_interval(Duration::from_secs(Self::ENDPOINT_KEEP_ALIVE_INTERVAL_SEC))
2650            .keep_alive_timeout(Duration::from_secs(Self::ENDPOINT_KEEP_ALIVE_TIMEOUT_SEC))
2651            .connect_timeout(Duration::from_secs(5))
2652            .monitored_connect("grpc-meta-client", Default::default())
2653            .await?
2654            .wrapped();
2655
2656        Ok(channel)
2657    }
2658
2659    pub(crate) fn retry_strategy_to_bound(
2660        high_bound: Duration,
2661        exceed: bool,
2662    ) -> impl Iterator<Item = Duration> {
2663        let iter = exponential_backoff(
2664            Duration::from_millis(Self::INIT_RETRY_BASE_INTERVAL_MS),
2665            Self::INIT_RETRY_BASE_INTERVAL_MS,
2666            Duration::from_millis(Self::INIT_RETRY_MAX_INTERVAL_MS),
2667        )
2668        .map(jitter);
2669
2670        let mut sum = Duration::default();
2671
2672        iter.take_while(move |duration| {
2673            sum += *duration;
2674
2675            if exceed {
2676                sum < high_bound + *duration
2677            } else {
2678                sum < high_bound
2679            }
2680        })
2681    }
2682}
2683
2684macro_rules! for_all_meta_rpc {
2685    ($macro:ident) => {
2686        $macro! {
2687             { cluster_client, add_worker_node, AddWorkerNodeRequest, AddWorkerNodeResponse }
2688            ,{ cluster_client, activate_worker_node, ActivateWorkerNodeRequest, ActivateWorkerNodeResponse }
2689            ,{ cluster_client, delete_worker_node, DeleteWorkerNodeRequest, DeleteWorkerNodeResponse }
2690            ,{ cluster_client, list_all_nodes, ListAllNodesRequest, ListAllNodesResponse }
2691            ,{ cluster_client, get_cluster_recovery_status, GetClusterRecoveryStatusRequest, GetClusterRecoveryStatusResponse }
2692            ,{ cluster_client, get_meta_store_info, GetMetaStoreInfoRequest, GetMetaStoreInfoResponse }
2693            ,{ heartbeat_client, heartbeat, HeartbeatRequest, HeartbeatResponse }
2694            ,{ stream_client, flush, FlushRequest, FlushResponse }
2695            ,{ stream_client, pause, PauseRequest, PauseResponse }
2696            ,{ stream_client, resume, ResumeRequest, ResumeResponse }
2697            ,{ stream_client, apply_throttle, ApplyThrottleRequest, ApplyThrottleResponse }
2698            ,{ stream_client, cancel_creating_jobs, CancelCreatingJobsRequest, CancelCreatingJobsResponse }
2699            ,{ stream_client, list_table_fragments, ListTableFragmentsRequest, ListTableFragmentsResponse }
2700            ,{ stream_client, list_streaming_job_states, ListStreamingJobStatesRequest, ListStreamingJobStatesResponse }
2701            ,{ stream_client, list_fragment_distribution, ListFragmentDistributionRequest, ListFragmentDistributionResponse }
2702            ,{ stream_client, list_creating_fragment_distribution, ListCreatingFragmentDistributionRequest, ListCreatingFragmentDistributionResponse }
2703            ,{ stream_client, list_sink_log_store_tables, ListSinkLogStoreTablesRequest, ListSinkLogStoreTablesResponse }
2704            ,{ stream_client, list_actor_states, ListActorStatesRequest, ListActorStatesResponse }
2705            ,{ stream_client, list_actor_splits, ListActorSplitsRequest, ListActorSplitsResponse }
2706            ,{ stream_client, recover, RecoverRequest, RecoverResponse }
2707            ,{ stream_client, list_rate_limits, ListRateLimitsRequest, ListRateLimitsResponse }
2708            ,{ stream_client, list_cdc_progress, ListCdcProgressRequest, ListCdcProgressResponse }
2709            ,{ stream_client, list_refresh_table_states, ListRefreshTableStatesRequest, ListRefreshTableStatesResponse }
2710            ,{ stream_client, list_iceberg_compaction_status, ListIcebergCompactionStatusRequest, ListIcebergCompactionStatusResponse }
2711            ,{ stream_client, alter_connector_props, AlterConnectorPropsRequest, AlterConnectorPropsResponse }
2712            ,{ stream_client, alter_source_properties_safe, AlterSourcePropertiesSafeRequest, AlterSourcePropertiesSafeResponse }
2713            ,{ stream_client, reset_source_splits, ResetSourceSplitsRequest, ResetSourceSplitsResponse }
2714            ,{ stream_client, inject_source_offsets, InjectSourceOffsetsRequest, InjectSourceOffsetsResponse }
2715            ,{ stream_client, get_fragment_by_id, GetFragmentByIdRequest, GetFragmentByIdResponse }
2716            ,{ stream_client, get_fragment_vnodes, GetFragmentVnodesRequest, GetFragmentVnodesResponse }
2717            ,{ stream_client, get_actor_vnodes, GetActorVnodesRequest, GetActorVnodesResponse }
2718            ,{ stream_client, set_sync_log_store_aligned, SetSyncLogStoreAlignedRequest, SetSyncLogStoreAlignedResponse }
2719            ,{ stream_client, refresh, RefreshRequest, RefreshResponse }
2720            ,{ stream_client, list_unmigrated_tables, ListUnmigratedTablesRequest, ListUnmigratedTablesResponse }
2721            ,{ ddl_client, create_table, CreateTableRequest, CreateTableResponse }
2722            ,{ ddl_client, alter_name, AlterNameRequest, AlterNameResponse }
2723            ,{ ddl_client, alter_owner, AlterOwnerRequest, AlterOwnerResponse }
2724            ,{ ddl_client, alter_subscription_retention, AlterSubscriptionRetentionRequest, AlterSubscriptionRetentionResponse }
2725            ,{ ddl_client, alter_set_schema, AlterSetSchemaRequest, AlterSetSchemaResponse }
2726            ,{ ddl_client, alter_parallelism, AlterParallelismRequest, AlterParallelismResponse }
2727            ,{ ddl_client, alter_backfill_parallelism, AlterBackfillParallelismRequest, AlterBackfillParallelismResponse }
2728            ,{ ddl_client, alter_streaming_job_config, AlterStreamingJobConfigRequest, AlterStreamingJobConfigResponse }
2729            ,{ ddl_client, alter_fragment_parallelism, AlterFragmentParallelismRequest, AlterFragmentParallelismResponse }
2730            ,{ ddl_client, alter_cdc_table_backfill_parallelism, AlterCdcTableBackfillParallelismRequest, AlterCdcTableBackfillParallelismResponse }
2731            ,{ ddl_client, alter_resource_group, AlterResourceGroupRequest, AlterResourceGroupResponse }
2732            ,{ ddl_client, alter_database_resource_group, AlterDatabaseResourceGroupRequest, AlterDatabaseResourceGroupResponse }
2733            ,{ ddl_client, alter_database_param, AlterDatabaseParamRequest, AlterDatabaseParamResponse }
2734            ,{ ddl_client, create_materialized_view, CreateMaterializedViewRequest, CreateMaterializedViewResponse }
2735            ,{ ddl_client, create_view, CreateViewRequest, CreateViewResponse }
2736            ,{ ddl_client, create_source, CreateSourceRequest, CreateSourceResponse }
2737            ,{ ddl_client, create_sink, CreateSinkRequest, CreateSinkResponse }
2738            ,{ ddl_client, create_subscription, CreateSubscriptionRequest, CreateSubscriptionResponse }
2739            ,{ ddl_client, create_schema, CreateSchemaRequest, CreateSchemaResponse }
2740            ,{ ddl_client, create_database, CreateDatabaseRequest, CreateDatabaseResponse }
2741            ,{ ddl_client, create_secret, CreateSecretRequest, CreateSecretResponse }
2742            ,{ ddl_client, create_index, CreateIndexRequest, CreateIndexResponse }
2743            ,{ ddl_client, create_function, CreateFunctionRequest, CreateFunctionResponse }
2744            ,{ ddl_client, drop_table, DropTableRequest, DropTableResponse }
2745            ,{ ddl_client, drop_materialized_view, DropMaterializedViewRequest, DropMaterializedViewResponse }
2746            ,{ ddl_client, drop_view, DropViewRequest, DropViewResponse }
2747            ,{ ddl_client, drop_source, DropSourceRequest, DropSourceResponse }
2748            ,{ ddl_client, reset_source, ResetSourceRequest, ResetSourceResponse }
2749            ,{ ddl_client, drop_secret, DropSecretRequest, DropSecretResponse}
2750            ,{ ddl_client, drop_sink, DropSinkRequest, DropSinkResponse }
2751            ,{ ddl_client, drop_subscription, DropSubscriptionRequest, DropSubscriptionResponse }
2752            ,{ ddl_client, drop_database, DropDatabaseRequest, DropDatabaseResponse }
2753            ,{ ddl_client, drop_schema, DropSchemaRequest, DropSchemaResponse }
2754            ,{ ddl_client, drop_index, DropIndexRequest, DropIndexResponse }
2755            ,{ ddl_client, drop_function, DropFunctionRequest, DropFunctionResponse }
2756            ,{ ddl_client, replace_job_plan, ReplaceJobPlanRequest, ReplaceJobPlanResponse }
2757            ,{ ddl_client, alter_source, AlterSourceRequest, AlterSourceResponse }
2758            ,{ ddl_client, risectl_list_state_tables, RisectlListStateTablesRequest, RisectlListStateTablesResponse }
2759            ,{ ddl_client, risectl_resume_backfill, RisectlResumeBackfillRequest, RisectlResumeBackfillResponse }
2760            ,{ ddl_client, get_ddl_progress, GetDdlProgressRequest, GetDdlProgressResponse }
2761            ,{ ddl_client, create_connection, CreateConnectionRequest, CreateConnectionResponse }
2762            ,{ ddl_client, list_connections, ListConnectionsRequest, ListConnectionsResponse }
2763            ,{ ddl_client, drop_connection, DropConnectionRequest, DropConnectionResponse }
2764            ,{ ddl_client, comment_on, CommentOnRequest, CommentOnResponse }
2765            ,{ ddl_client, get_tables, GetTablesRequest, GetTablesResponse }
2766            ,{ ddl_client, wait, WaitRequest, WaitResponse }
2767            ,{ ddl_client, auto_schema_change, AutoSchemaChangeRequest, AutoSchemaChangeResponse }
2768            ,{ ddl_client, alter_swap_rename, AlterSwapRenameRequest, AlterSwapRenameResponse }
2769            ,{ ddl_client, alter_secret, AlterSecretRequest, AlterSecretResponse }
2770            ,{ ddl_client, compact_iceberg_table, CompactIcebergTableRequest, CompactIcebergTableResponse }
2771            ,{ ddl_client, rewrite_iceberg_table_manifests, RewriteIcebergTableManifestsRequest, RewriteIcebergTableManifestsResponse }
2772            ,{ ddl_client, expire_iceberg_table_snapshots, ExpireIcebergTableSnapshotsRequest, ExpireIcebergTableSnapshotsResponse }
2773            ,{ ddl_client, create_iceberg_table, CreateIcebergTableRequest, CreateIcebergTableResponse }
2774            ,{ ddl_client, wait_iceberg_pk_index_sink_epoch, WaitIcebergPkIndexSinkEpochRequest, WaitIcebergPkIndexSinkEpochResponse }
2775            ,{ hummock_client, unpin_version_before, UnpinVersionBeforeRequest, UnpinVersionBeforeResponse }
2776            ,{ hummock_client, get_current_version, GetCurrentVersionRequest, GetCurrentVersionResponse }
2777            ,{ hummock_client, replay_version_delta, ReplayVersionDeltaRequest, ReplayVersionDeltaResponse }
2778            ,{ hummock_client, list_version_deltas, ListVersionDeltasRequest, ListVersionDeltasResponse }
2779            ,{ hummock_client, get_assigned_compact_task_num, GetAssignedCompactTaskNumRequest, GetAssignedCompactTaskNumResponse }
2780            ,{ hummock_client, trigger_compaction_deterministic, TriggerCompactionDeterministicRequest, TriggerCompactionDeterministicResponse }
2781            ,{ hummock_client, disable_commit_epoch, DisableCommitEpochRequest, DisableCommitEpochResponse }
2782            ,{ hummock_client, get_new_object_ids, GetNewObjectIdsRequest, GetNewObjectIdsResponse }
2783            ,{ hummock_client, trigger_manual_compaction, TriggerManualCompactionRequest, TriggerManualCompactionResponse }
2784            ,{ hummock_client, trigger_full_gc, TriggerFullGcRequest, TriggerFullGcResponse }
2785            ,{ hummock_client, rise_ctl_get_pinned_versions_summary, RiseCtlGetPinnedVersionsSummaryRequest, RiseCtlGetPinnedVersionsSummaryResponse }
2786            ,{ hummock_client, rise_ctl_list_compaction_group, RiseCtlListCompactionGroupRequest, RiseCtlListCompactionGroupResponse }
2787            ,{ hummock_client, rise_ctl_update_compaction_config, RiseCtlUpdateCompactionConfigRequest, RiseCtlUpdateCompactionConfigResponse }
2788            ,{ hummock_client, rise_ctl_get_checkpoint_version, RiseCtlGetCheckpointVersionRequest, RiseCtlGetCheckpointVersionResponse }
2789            ,{ hummock_client, rise_ctl_pause_version_checkpoint, RiseCtlPauseVersionCheckpointRequest, RiseCtlPauseVersionCheckpointResponse }
2790            ,{ hummock_client, rise_ctl_resume_version_checkpoint, RiseCtlResumeVersionCheckpointRequest, RiseCtlResumeVersionCheckpointResponse }
2791            ,{ hummock_client, init_metadata_for_replay, InitMetadataForReplayRequest, InitMetadataForReplayResponse }
2792            ,{ hummock_client, split_compaction_group, SplitCompactionGroupRequest, SplitCompactionGroupResponse }
2793            ,{ hummock_client, rise_ctl_list_compaction_status, RiseCtlListCompactionStatusRequest, RiseCtlListCompactionStatusResponse }
2794            ,{ hummock_client, get_compaction_score, GetCompactionScoreRequest, GetCompactionScoreResponse }
2795            ,{ hummock_client, rise_ctl_rebuild_table_stats, RiseCtlRebuildTableStatsRequest, RiseCtlRebuildTableStatsResponse }
2796            ,{ hummock_client, subscribe_compaction_event, impl tonic::IntoStreamingRequest<Message = SubscribeCompactionEventRequest>, Streaming<SubscribeCompactionEventResponse> }
2797            ,{ hummock_client, subscribe_iceberg_compaction_event, impl tonic::IntoStreamingRequest<Message = SubscribeIcebergCompactionEventRequest>, Streaming<SubscribeIcebergCompactionEventResponse> }
2798            ,{ hummock_client, list_branched_object, ListBranchedObjectRequest, ListBranchedObjectResponse }
2799            ,{ hummock_client, list_active_write_limit, ListActiveWriteLimitRequest, ListActiveWriteLimitResponse }
2800            ,{ hummock_client, list_hummock_meta_config, ListHummockMetaConfigRequest, ListHummockMetaConfigResponse }
2801            ,{ hummock_client, list_compact_task_assignment, ListCompactTaskAssignmentRequest, ListCompactTaskAssignmentResponse }
2802            ,{ hummock_client, list_compact_task_progress, ListCompactTaskProgressRequest, ListCompactTaskProgressResponse }
2803            ,{ hummock_client, cancel_compact_task, CancelCompactTaskRequest, CancelCompactTaskResponse}
2804            ,{ hummock_client, get_version_by_epoch, GetVersionByEpochRequest, GetVersionByEpochResponse }
2805            ,{ hummock_client, merge_compaction_group, MergeCompactionGroupRequest, MergeCompactionGroupResponse }
2806            ,{ hummock_client, get_table_change_logs, GetTableChangeLogsRequest, GetTableChangeLogsResponse }
2807            ,{ user_client, create_user, CreateUserRequest, CreateUserResponse }
2808            ,{ user_client, update_user, UpdateUserRequest, UpdateUserResponse }
2809            ,{ user_client, drop_user, DropUserRequest, DropUserResponse }
2810            ,{ user_client, grant_privilege, GrantPrivilegeRequest, GrantPrivilegeResponse }
2811            ,{ user_client, revoke_privilege, RevokePrivilegeRequest, RevokePrivilegeResponse }
2812            ,{ user_client, alter_default_privilege, AlterDefaultPrivilegeRequest, AlterDefaultPrivilegeResponse }
2813            ,{ scale_client, get_cluster_info, GetClusterInfoRequest, GetClusterInfoResponse }
2814            ,{ scale_client, reschedule, RescheduleRequest, RescheduleResponse }
2815            ,{ notification_client, subscribe, SubscribeRequest, Streaming<SubscribeResponse> }
2816            ,{ backup_client, backup_meta, BackupMetaRequest, BackupMetaResponse }
2817            ,{ backup_client, get_backup_job_status, GetBackupJobStatusRequest, GetBackupJobStatusResponse }
2818            ,{ backup_client, delete_meta_snapshot, DeleteMetaSnapshotRequest, DeleteMetaSnapshotResponse}
2819            ,{ backup_client, get_meta_snapshot_manifest, GetMetaSnapshotManifestRequest, GetMetaSnapshotManifestResponse}
2820            ,{ telemetry_client, get_telemetry_info, GetTelemetryInfoRequest, TelemetryInfoResponse}
2821            ,{ system_params_client, get_system_params, GetSystemParamsRequest, GetSystemParamsResponse }
2822            ,{ system_params_client, set_system_param, SetSystemParamRequest, SetSystemParamResponse }
2823            ,{ session_params_client, get_session_params, GetSessionParamsRequest, GetSessionParamsResponse }
2824            ,{ session_params_client, set_session_param, SetSessionParamRequest, SetSessionParamResponse }
2825            ,{ serving_client, get_serving_vnode_mappings, GetServingVnodeMappingsRequest, GetServingVnodeMappingsResponse }
2826            ,{ cloud_client, rw_cloud_validate_source, RwCloudValidateSourceRequest, RwCloudValidateSourceResponse }
2827            ,{ event_log_client, list_event_log, ListEventLogRequest, ListEventLogResponse }
2828            ,{ event_log_client, add_event_log, AddEventLogRequest, AddEventLogResponse }
2829            ,{ cluster_limit_client, get_cluster_limits, GetClusterLimitsRequest, GetClusterLimitsResponse }
2830            ,{ hosted_iceberg_catalog_service_client, list_iceberg_tables, ListIcebergTablesRequest, ListIcebergTablesResponse }
2831            ,{ monitor_client, stack_trace, StackTraceRequest, StackTraceResponse }
2832        }
2833    };
2834}
2835
2836impl GrpcMetaClient {
2837    async fn refresh_client_if_needed(&self, code: Code) {
2838        if matches!(
2839            code,
2840            Code::Unknown | Code::Unimplemented | Code::Unavailable
2841        ) {
2842            tracing::debug!("matching tonic code {}", code);
2843            let (result_sender, result_receiver) = oneshot::channel();
2844            if self
2845                .member_monitor_event_sender
2846                .try_send(result_sender)
2847                .is_ok()
2848            {
2849                if let Ok(Err(e)) = result_receiver.await {
2850                    tracing::warn!(error = %e.as_report(), "force refresh meta client failed");
2851                }
2852            } else {
2853                tracing::debug!("skipping the current refresh, somewhere else is already doing it")
2854            }
2855        }
2856    }
2857}
2858
2859impl GrpcMetaClient {
2860    for_all_meta_rpc! { meta_rpc_client_method_impl }
2861}