g3tiles: add backend alive-connection rpc method

Author: zh-jq-b

g3tiles/proto/build.rs

@@ -20,6 +20,7 @@ fn main() {
         .file("schema/types.capnp")
         .file("schema/proc.capnp")
         .file("schema/server.capnp")
+        .file("schema/backend.capnp")
         .run()
         .unwrap();
 }

g3tiles/proto/schema/backend.capnp

@@ -0,0 +1,5 @@
+@0xb3f076019fd2cc68;
+
+interface BackendControl {
+  aliveConnection @0 () -> (count :UInt64);
+}

g3tiles/proto/schema/proc.capnp

@@ -4,6 +4,8 @@ using Types = import "types.capnp";
 
 using Server = import "server.capnp";
 
+using Backend = import "backend.capnp";
+
 interface ProcControl {
   #
 
@@ -24,4 +26,5 @@ interface ProcControl {
 
   reloadBackend @9 (name :Text) -> (result :Types.OperationResult);
   listBackend @10 () -> (result :List(Text));
+  getBackend @13 (name: Text) -> (backend :Types.FetchResult(Backend.BackendControl));
 }

g3tiles/proto/src/lib.rs

@@ -28,3 +28,7 @@ pub mod proc_capnp {
 pub mod server_capnp {
     include!(concat!(env!("OUT_DIR"), "/server_capnp.rs"));
 }
+
+pub mod backend_capnp {
+    include!(concat!(env!("OUT_DIR"), "/backend_capnp.rs"));
+}

g3tiles/src/backend/dummy_close.rs

@@ -85,6 +85,10 @@ impl Backend for DummyCloseBackend {
         Ok(())
     }
 
+    fn alive_connection(&self) -> u64 {
+        0
+    }
+
     async fn stream_connect(&self, _task_notes: &ServerTaskNotes) -> StreamConnectResult {
         Err(StreamConnectError::UpstreamNotResolved)
     }

g3tiles/src/backend/keyless_quic/mod.rs

@@ -196,6 +196,10 @@ impl Backend for KeylessQuicBackend {
         Ok(())
     }
 
+    fn alive_connection(&self) -> u64 {
+        self.pool_handle.alive_connection()
+    }
+
     async fn keyless(&self, req: KeylessRequest) -> KeylessResponse {
         let err = KeylessInternalErrorResponse::new(req.header());
         if !self.config.wait_new_channel && self.stats.alive_channel() <= 0 {

g3tiles/src/backend/keyless_tcp/mod.rs

@@ -207,6 +207,10 @@ impl Backend for KeylessTcpBackend {
         Ok(())
     }
 
+    fn alive_connection(&self) -> u64 {
+        self.pool_handle.alive_connection()
+    }
+
     async fn keyless(&self, req: KeylessRequest) -> KeylessResponse {
         let err = KeylessInternalErrorResponse::new(req.header());
         if !self.config.wait_new_channel && self.stats.alive_channel() <= 0 {

g3tiles/src/backend/mod.rs

@@ -35,7 +35,7 @@ mod stream_tcp;
 
 mod ops;
 pub use ops::load_all;
-pub(crate) use ops::{reload, update_dependency_to_discover};
+pub(crate) use ops::{get_backend, reload, update_dependency_to_discover};
 
 mod registry;
 pub(crate) use registry::{get_names, get_or_insert_default};
@@ -53,6 +53,8 @@ pub(crate) trait Backend {
     fn discover(&self) -> &NodeName;
     fn update_discover(&self) -> anyhow::Result<()>;
 
+    fn alive_connection(&self) -> u64;
+
     async fn stream_connect(&self, _task_notes: &ServerTaskNotes) -> StreamConnectResult {
         Err(StreamConnectError::UpstreamNotResolved) // TODO
     }

g3tiles/src/backend/ops.rs

@@ -69,6 +69,13 @@ pub async fn load_all() -> anyhow::Result<()> {
     Ok(())
 }
 
+pub(crate) fn get_backend(name: &NodeName) -> anyhow::Result<ArcBackend> {
+    match registry::get(name) {
+        Some(backend) => Ok(backend),
+        None => Err(anyhow!("no backend named {name} found")),
+    }
+}
+
 pub(crate) async fn reload(
     name: &NodeName,
     position: Option<YamlDocPosition>,

g3tiles/src/backend/stream_tcp.rs

@@ -179,6 +179,11 @@ impl Backend for StreamTcpBackend {
         Ok(())
     }
 
+    fn alive_connection(&self) -> u64 {
+        // TODO add alive connection stats
+        0
+    }
+
     async fn stream_connect(&self, task_notes: &ServerTaskNotes) -> StreamConnectResult {
         let Some(next_addr) = self.select_peer(task_notes) else {
             return Err(StreamConnectError::UpstreamNotResolved);

g3tiles/src/control/capnp/backend.rs

@@ -0,0 +1,47 @@
+/*
+ * Copyright 2025 ByteDance and/or its affiliates.
+ *
+ * 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 capnp::capability::Promise;
+
+use g3_types::metrics::NodeName;
+
+use g3tiles_proto::backend_capnp::backend_control;
+
+use crate::backend::ArcBackend;
+
+pub(super) struct BackendControlImpl {
+    backend: ArcBackend,
+}
+
+impl BackendControlImpl {
+    pub(super) fn new_client(name: &str) -> anyhow::Result<backend_control::Client> {
+        let name = unsafe { NodeName::new_unchecked(name) };
+        let backend = crate::backend::get_backend(&name)?;
+        Ok(capnp_rpc::new_client(BackendControlImpl { backend }))
+    }
+}
+
+impl backend_control::Server for BackendControlImpl {
+    fn alive_connection(
+        &mut self,
+        _params: backend_control::AliveConnectionParams,
+        mut results: backend_control::AliveConnectionResults,
+    ) -> Promise<(), capnp::Error> {
+        let alive_count = self.backend.alive_connection();
+        results.get().set_count(alive_count);
+        Promise::ok(())
+    }
+}

g3tiles/src/control/capnp/mod.rs

@@ -20,6 +20,7 @@ mod common;
 use common::set_operation_result;
 mod proc;
 
+mod backend;
 mod server;
 
 pub fn stop_working_thread() {

g3tiles/src/control/capnp/proc.rs

@@ -19,6 +19,7 @@ use capnp_rpc::pry;
 
 use g3_types::metrics::NodeName;
 
+use g3tiles_proto::backend_capnp::backend_control;
 use g3tiles_proto::proc_capnp::proc_control;
 use g3tiles_proto::server_capnp::server_control;
 use g3tiles_proto::types_capnp::fetch_result;
@@ -185,6 +186,19 @@ impl proc_control::Server for ProcControlImpl {
         }
         Promise::ok(())
     }
+
+    fn get_backend(
+        &mut self,
+        params: proc_control::GetBackendParams,
+        mut results: proc_control::GetBackendResults,
+    ) -> Promise<(), capnp::Error> {
+        let backend = pry!(pry!(pry!(params.get()).get_name()).to_str());
+        pry!(set_fetch_result::<backend_control::Owned>(
+            results.get().init_backend(),
+            super::backend::BackendControlImpl::new_client(backend),
+        ));
+        Promise::ok(())
+    }
 }
 
 fn set_fetch_result<'a, T>(

g3tiles/src/module/keyless/backend/pool.rs

@@ -55,6 +55,7 @@ enum KeylessPoolCmd {
 #[derive(Clone)]
 pub(crate) struct KeylessConnectionPoolHandle {
     cmd_sender: mpsc::Sender<KeylessPoolCmd>,
+    stats: Arc<ConnectionPoolStats>,
 }
 
 impl KeylessConnectionPoolHandle {
@@ -69,6 +70,10 @@ impl KeylessConnectionPoolHandle {
     pub(crate) fn request_new_connection(&self) {
         let _ = self.cmd_sender.try_send(KeylessPoolCmd::NewConnection);
     }
+
+    pub(crate) fn alive_connection(&self) -> u64 {
+        self.stats.alive_count() as u64
+    }
 }
 
 pub(crate) struct KeylessConnectionPool<C: KeylessUpstreamConnection> {
@@ -123,11 +128,12 @@ where
             keyless_request_receiver,
             graceful_close_wait,
         );
+        let stats = pool.stats.clone();
         let (cmd_sender, cmd_receiver) = mpsc::channel(CMD_CHANNEL_SIZE);
         tokio::spawn(async move {
             pool.into_running(cmd_receiver).await;
         });
-        KeylessConnectionPoolHandle { cmd_sender }
+        KeylessConnectionPoolHandle { cmd_sender, stats }
     }
 
     async fn into_running(mut self, mut cmd_receiver: mpsc::Receiver<KeylessPoolCmd>) {

g3tiles/utils/ctl/src/backend.rs

@@ -0,0 +1,58 @@
+/*
+ * Copyright 2025 ByteDance and/or its affiliates.
+ *
+ * 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 clap::{Arg, ArgMatches, Command};
+use futures_util::future::TryFutureExt;
+
+use g3_ctl::CommandResult;
+
+use g3tiles_proto::backend_capnp::backend_control;
+use g3tiles_proto::proc_capnp::proc_control;
+
+pub const COMMAND: &str = "backend";
+
+const COMMAND_ARG_NAME: &str = "name";
+
+const SUBCOMMAND_ALIVE_CONNECTION: &str = "alive-connection";
+
+pub fn command() -> Command {
+    Command::new(COMMAND)
+        .arg(Arg::new(COMMAND_ARG_NAME).required(true).num_args(1))
+        .subcommand_required(true)
+        .subcommand(Command::new(SUBCOMMAND_ALIVE_CONNECTION))
+}
+
+async fn alive_connection(client: &backend_control::Client) -> CommandResult<()> {
+    let req = client.alive_connection_request();
+    let rsp = req.send().promise.await?;
+    let count = rsp.get()?.get_count();
+    println!("{count}");
+    Ok(())
+}
+
+pub async fn run(client: &proc_control::Client, args: &ArgMatches) -> CommandResult<()> {
+    let name = args.get_one::<String>(COMMAND_ARG_NAME).unwrap();
+
+    let (subcommand, _) = args.subcommand().unwrap();
+    match subcommand {
+        SUBCOMMAND_ALIVE_CONNECTION => {
+            super::proc::get_backend(client, name)
+                .and_then(|backend| async move { alive_connection(&backend).await })
+                .await
+        }
+        _ => unreachable!(),
+    }
+}

g3tiles/utils/ctl/src/main.rs

@@ -24,6 +24,7 @@ use g3tiles_proto::proc_capnp::proc_control;
 mod common;
 mod proc;
 
+mod backend;
 mod server;
 
 fn build_cli_args() -> Command {
@@ -39,6 +40,7 @@ fn build_cli_args() -> Command {
         .subcommand(proc::commands::reload_discover())
         .subcommand(proc::commands::reload_backend())
         .subcommand(server::command())
+        .subcommand(backend::command())
 }
 
 #[tokio::main(flavor = "current_thread")]
@@ -74,6 +76,7 @@ async fn main() -> anyhow::Result<()> {
                 proc::COMMAND_RELOAD_DISCOVER => proc::reload_discover(&proc_control, args).await,
                 proc::COMMAND_RELOAD_BACKEND => proc::reload_backend(&proc_control, args).await,
                 server::COMMAND => server::run(&proc_control, args).await,
+                backend::COMMAND => backend::run(&proc_control, args).await,
                 _ => Err(CommandError::Cli(anyhow!(
                     "unsupported command {subcommand}"
                 ))),

g3tiles/utils/ctl/src/proc.rs

@@ -18,6 +18,7 @@ use clap::ArgMatches;
 
 use g3_ctl::CommandResult;
 
+use g3tiles_proto::backend_capnp::backend_control;
 use g3tiles_proto::proc_capnp::proc_control;
 use g3tiles_proto::server_capnp::server_control;
 
@@ -199,3 +200,13 @@ pub(crate) async fn get_server(
     let rsp = req.send().promise.await?;
     parse_fetch_result(rsp.get()?.get_server()?)
 }
+
+pub(crate) async fn get_backend(
+    client: &proc_control::Client,
+    name: &str,
+) -> CommandResult<backend_control::Client> {
+    let mut req = client.get_backend_request();
+    req.get().set_name(name);
+    let rsp = req.send().promise.await?;
+    parse_fetch_result(rsp.get()?.get_backend()?)
+}