feat(ctl): list serving fragment mappings

proto/meta.proto

@@ -411,3 +411,14 @@ service SystemParamsService {
   rpc GetSystemParams(GetSystemParamsRequest) returns (GetSystemParamsResponse);
   rpc SetSystemParam(SetSystemParamRequest) returns (SetSystemParamResponse);
 }
+
+message GetServingVnodeMappingsRequest {}
+
+message GetServingVnodeMappingsResponse {
+  repeated FragmentParallelUnitMapping mappings = 1;
+  map<uint32, uint32> fragment_to_table = 2;
+}
+
+service ServingService {
+  rpc GetServingVnodeMappings(GetServingVnodeMappingsRequest) returns (GetServingVnodeMappingsResponse);
+}

src/ctl/src/cmd_impl/meta.rs

@@ -17,9 +17,11 @@ mod cluster_info;
 mod connection;
 mod pause_resume;
 mod reschedule;
+mod serving;
 
 pub use backup_meta::*;
 pub use cluster_info::*;
 pub use connection::*;
 pub use pause_resume::*;
 pub use reschedule::*;
+pub use serving::*;

src/ctl/src/cmd_impl/meta/serving.rs

@@ -0,0 +1,94 @@
+// Copyright 2023 RisingWave Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::collections::HashMap;
+
+use comfy_table::{Row, Table};
+use itertools::Itertools;
+use risingwave_common::hash::{ParallelUnitId, VirtualNode};
+use risingwave_pb::common::{WorkerNode, WorkerType};
+
+use crate::CtlContext;
+
+pub async fn list_serving_fragment_mappings(context: &CtlContext) -> anyhow::Result<()> {
+    let meta_client = context.meta_client().await?;
+    let mappings = meta_client.list_serving_vnode_mappings().await?;
+    let workers = meta_client
+        .list_worker_nodes(WorkerType::ComputeNode)
+        .await?;
+    let mut pu_to_worker: HashMap<ParallelUnitId, &WorkerNode> = HashMap::new();
+    for w in &workers {
+        for pu in &w.parallel_units {
+            pu_to_worker.insert(pu.id, w);
+        }
+    }
+
+    let mut table = Table::new();
+    table.set_header({
+        let mut row = Row::new();
+        row.add_cell("Table Id".into());
+        row.add_cell("Fragment Id".into());
+        row.add_cell("Parallel Unit Id".into());
+        row.add_cell("Virtual Node".into());
+        row.add_cell("Worker".into());
+        row
+    });
+
+    let rows = mappings
+        .iter()
+        .flat_map(|(fragment_id, (table_id, mapping))| {
+            let mut pu_vnodes: HashMap<ParallelUnitId, Vec<VirtualNode>> = HashMap::new();
+            for (vnode, pu) in mapping.iter_with_vnode() {
+                pu_vnodes.entry(pu).or_insert(vec![]).push(vnode);
+            }
+            pu_vnodes.into_iter().map(|(pu_id, vnodes)| {
+                (
+                    *table_id,
+                    *fragment_id,
+                    pu_id,
+                    vnodes,
+                    pu_to_worker.get(&pu_id),
+                )
+            })
+        })
+        .collect_vec();
+    for (table_id, fragment_id, pu_id, vnodes, worker) in
+        rows.into_iter().sorted_by_key(|(t, f, p, ..)| (*t, *f, *p))
+    {
+        let mut row = Row::new();
+        row.add_cell(table_id.into());
+        row.add_cell(fragment_id.into());
+        row.add_cell(pu_id.into());
+        row.add_cell(
+            format!(
+                "{} in total: {}",
+                vnodes.len(),
+                vnodes
+                    .into_iter()
+                    .sorted()
+                    .map(|v| v.to_index().to_string())
+                    .join(",")
+            )
+            .into(),
+        );
+        if let Some(w) = worker && let Some(addr) = w.host.as_ref() {
+            row.add_cell(format!("id: {}; {}:{}", w.id, addr.host, addr.port).into());
+        } else {
+            row.add_cell("".into());
+        }
+        table.add_row(row);
+    }
+    println!("{table}");
+    Ok(())
+}

src/ctl/src/lib.rs

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#![feature(let_chains)]
+
 use anyhow::Result;
 use clap::{Parser, Subcommand};
 use cmd_impl::bench::BenchCommands;
@@ -229,6 +231,9 @@ enum MetaCommands {
 
     /// List all existing connections in the catalog
     ListConnections,
+
+    /// List fragment to parallel units mapping for serving
+    ListServingFragmentMapping,
 }
 
 pub async fn start(opts: CliOpts) -> Result<()> {
@@ -363,6 +368,9 @@ pub async fn start_impl(opts: CliOpts, context: &CtlContext) -> Result<()> {
         Commands::Meta(MetaCommands::ListConnections) => {
             cmd_impl::meta::list_connections(context).await?
         }
+        Commands::Meta(MetaCommands::ListServingFragmentMapping) => {
+            cmd_impl::meta::list_serving_fragment_mappings(context).await?
+        }
         Commands::Trace => cmd_impl::trace::trace(context).await?,
         Commands::Profile { sleep } => cmd_impl::profile::profile(context, sleep).await?,
     }

src/meta/src/lib.rs

@@ -37,14 +37,14 @@
 
 pub mod backup_restore;
 mod barrier;
-pub(crate) mod batch;
 #[cfg(not(madsim))] // no need in simulation test
 mod dashboard;
 mod error;
 pub mod hummock;
 pub mod manager;
 mod model;
 mod rpc;
+pub(crate) mod serving;
 pub mod storage;
 mod stream;
 pub(crate) mod telemetry;

src/meta/src/manager/catalog/fragment.rs

@@ -96,6 +96,10 @@ impl FragmentManagerCore {
             })
             .collect()
     }
+
+    pub fn table_fragments(&self) -> &BTreeMap<TableId, TableFragments> {
+        &self.table_fragments
+    }
 }
 
 /// `FragmentManager` stores definition and status of fragment as well as the actors inside.

src/meta/src/rpc/server.rs

@@ -36,6 +36,7 @@ use risingwave_pb::meta::heartbeat_service_server::HeartbeatServiceServer;
 use risingwave_pb::meta::meta_member_service_server::MetaMemberServiceServer;
 use risingwave_pb::meta::notification_service_server::NotificationServiceServer;
 use risingwave_pb::meta::scale_service_server::ScaleServiceServer;
+use risingwave_pb::meta::serving_service_server::ServingServiceServer;
 use risingwave_pb::meta::stream_manager_service_server::StreamManagerServiceServer;
 use risingwave_pb::meta::system_params_service_server::SystemParamsServiceServer;
 use risingwave_pb::meta::telemetry_info_service_server::TelemetryInfoServiceServer;
@@ -51,10 +52,10 @@ use super::intercept::MetricsMiddlewareLayer;
 use super::service::health_service::HealthServiceImpl;
 use super::service::notification_service::NotificationServiceImpl;
 use super::service::scale_service::ScaleServiceImpl;
+use super::service::serving_service::ServingServiceImpl;
 use super::DdlServiceImpl;
 use crate::backup_restore::BackupManager;
 use crate::barrier::{BarrierScheduler, GlobalBarrierManager};
-use crate::batch::ServingVnodeMapping;
 use crate::hummock::{CompactionScheduler, HummockManager};
 use crate::manager::{
     CatalogManager, ClusterManager, FragmentManager, IdleManager, MetaOpts, MetaSrvEnv,
@@ -72,10 +73,11 @@ use crate::rpc::service::stream_service::StreamServiceImpl;
 use crate::rpc::service::system_params_service::SystemParamsServiceImpl;
 use crate::rpc::service::telemetry_service::TelemetryInfoServiceImpl;
 use crate::rpc::service::user_service::UserServiceImpl;
+use crate::serving::ServingVnodeMapping;
 use crate::storage::{EtcdMetaStore, MemStore, MetaStore, WrappedEtcdClient as EtcdClient};
 use crate::stream::{GlobalStreamManager, SourceManager};
 use crate::telemetry::{MetaReportCreator, MetaTelemetryInfoFetcher};
-use crate::{batch, hummock, MetaError, MetaResult};
+use crate::{hummock, serving, MetaError, MetaResult};
 
 #[derive(Debug)]
 pub enum MetaStoreBackend {
@@ -358,7 +360,7 @@ pub async fn start_service_as_election_leader<S: MetaStore>(
             .unwrap(),
     );
     let serving_vnode_mapping = Arc::new(ServingVnodeMapping::default());
-    batch::on_meta_start(
+    serving::on_meta_start(
         env.notification_manager_ref(),
         cluster_manager.clone(),
         fragment_manager.clone(),
@@ -538,6 +540,8 @@ pub async fn start_service_as_election_leader<S: MetaStore>(
     let backup_srv = BackupServiceImpl::new(backup_manager);
     let telemetry_srv = TelemetryInfoServiceImpl::new(meta_store.clone());
     let system_params_srv = SystemParamsServiceImpl::new(system_params_manager.clone());
+    let serving_srv =
+        ServingServiceImpl::new(serving_vnode_mapping.clone(), fragment_manager.clone());
 
     if let Some(prometheus_addr) = address_info.prometheus_addr {
         MetricsManager::boot_metrics_service(
@@ -579,7 +583,7 @@ pub async fn start_service_as_election_leader<S: MetaStore>(
     sub_tasks.push(SystemParamsManager::start_params_notifier(system_params_manager.clone()).await);
     sub_tasks.push(HummockManager::hummock_timer_task(hummock_manager).await);
     sub_tasks.push(
-        batch::start_serving_vnode_mapping_worker(
+        serving::start_serving_vnode_mapping_worker(
             env.notification_manager_ref(),
             cluster_manager.clone(),
             fragment_manager.clone(),
@@ -686,6 +690,7 @@ pub async fn start_service_as_election_leader<S: MetaStore>(
         .add_service(BackupServiceServer::new(backup_srv))
         .add_service(SystemParamsServiceServer::new(system_params_srv))
         .add_service(TelemetryInfoServiceServer::new(telemetry_srv))
+        .add_service(ServingServiceServer::new(serving_srv))
         .serve_with_shutdown(address_info.listen_addr, async move {
             tokio::select! {
                 res = svc_shutdown_rx.changed() => {

src/meta/src/rpc/service/mod.rs

@@ -21,6 +21,7 @@ pub mod hummock_service;
 pub mod meta_member_service;
 pub mod notification_service;
 pub mod scale_service;
+pub mod serving_service;
 pub mod stream_service;
 pub mod system_params_service;
 pub mod telemetry_service;

src/meta/src/rpc/service/notification_service.rs

@@ -29,12 +29,12 @@ use tokio_stream::wrappers::UnboundedReceiverStream;
 use tonic::{Request, Response, Status};
 
 use crate::backup_restore::BackupManagerRef;
-use crate::batch::ServingVnodeMappingRef;
 use crate::hummock::HummockManagerRef;
 use crate::manager::{
     Catalog, CatalogManagerRef, ClusterManagerRef, FragmentManagerRef, MetaSrvEnv, Notification,
     NotificationVersion, WorkerKey,
 };
+use crate::serving::ServingVnodeMappingRef;
 use crate::storage::MetaStore;
 
 pub struct NotificationServiceImpl<S: MetaStore> {

src/meta/src/rpc/service/serving_service.rs

@@ -0,0 +1,81 @@
+// Copyright 2023 RisingWave Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use itertools::Itertools;
+use risingwave_pb::meta::serving_service_server::ServingService;
+use risingwave_pb::meta::{
+    FragmentParallelUnitMapping, GetServingVnodeMappingsRequest, GetServingVnodeMappingsResponse,
+};
+use tonic::{Request, Response, Status};
+
+use crate::manager::FragmentManagerRef;
+use crate::serving::ServingVnodeMappingRef;
+use crate::storage::MetaStore;
+
+pub struct ServingServiceImpl<S: MetaStore> {
+    serving_vnode_mapping: ServingVnodeMappingRef,
+    fragment_manager: FragmentManagerRef<S>,
+}
+
+impl<S> ServingServiceImpl<S>
+where
+    S: MetaStore,
+{
+    pub fn new(
+        serving_vnode_mapping: ServingVnodeMappingRef,
+        fragment_manager: FragmentManagerRef<S>,
+    ) -> Self {
+        Self {
+            serving_vnode_mapping,
+            fragment_manager,
+        }
+    }
+}
+
+#[async_trait::async_trait]
+impl<S> ServingService for ServingServiceImpl<S>
+where
+    S: MetaStore,
+{
+    async fn get_serving_vnode_mappings(
+        &self,
+        _request: Request<GetServingVnodeMappingsRequest>,
+    ) -> Result<Response<GetServingVnodeMappingsResponse>, Status> {
+        let mappings = self
+            .serving_vnode_mapping
+            .all()
+            .into_iter()
+            .map(|(fragment_id, mapping)| FragmentParallelUnitMapping {
+                fragment_id,
+                mapping: Some(mapping.to_protobuf()),
+            })
+            .collect();
+        let fragment_to_table = {
+            let guard = self.fragment_manager.get_fragment_read_guard().await;
+            guard
+                .table_fragments()
+                .iter()
+                .flat_map(|(table_id, tf)| {
+                    tf.fragment_ids()
+                        .map(|fragment_id| (fragment_id, table_id.table_id))
+                        .collect_vec()
+                })
+                .collect()
+        };
+        Ok(Response::new(GetServingVnodeMappingsResponse {
+            mappings,
+            fragment_to_table,
+        }))
+    }
+}