goosefs_sdk/client/
worker_manager.rs

1// Copyright (C) 2026 Tencent. All rights reserved.
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
15//! Goosefs Worker Manager client for worker discovery.
16//!
17//! Wraps `WorkerManagerMasterClientService` (Master:9200) to fetch
18//! the list of live workers and their addresses.
19//!
20//! ## HA / Multi-Master Support
21//!
22//! When multiple Master addresses are configured, uses
23//! [`MasterInquireClient`] to discover the Primary Master.
24
25use std::sync::Arc;
26
27use tonic::service::interceptor::InterceptedService;
28use tonic::transport::Channel;
29use tracing::{debug, instrument};
30
31use crate::auth::{ChannelAuthenticator, ChannelIdInterceptor, SaslStreamGuard};
32use crate::client::master_inquire::{create_master_inquire_client, MasterInquireClient};
33use crate::config::GoosefsConfig;
34use crate::error::{Error, Result};
35use crate::proto::grpc::block::{
36    worker_manager_master_client_service_client::WorkerManagerMasterClientServiceClient,
37    GetWorkerInfoListPOptions, WorkerInfo,
38};
39
40/// Type alias for the authenticated WorkerManager gRPC client.
41type AuthenticatedWorkerMgrClient =
42    WorkerManagerMasterClientServiceClient<InterceptedService<Channel, ChannelIdInterceptor>>;
43
44/// Client for `WorkerManagerMasterClientService` (Master:9200).
45///
46/// Used to discover the live worker list for block routing.
47#[derive(Clone)]
48pub struct WorkerManagerClient {
49    inner: AuthenticatedWorkerMgrClient,
50    /// Keeps the SASL authentication stream alive for the channel's lifetime.
51    _sasl_guard: Arc<Option<SaslStreamGuard>>,
52}
53
54impl WorkerManagerClient {
55    /// Connect to the Goosefs Master for worker management.
56    ///
57    /// In HA mode, discovers the Primary Master first via the inquire client.
58    pub async fn connect(config: &GoosefsConfig) -> Result<Self> {
59        let inquire_client = create_master_inquire_client(config);
60        Self::connect_with_inquire(config, inquire_client).await
61    }
62
63    /// Connect using an externally-provided [`MasterInquireClient`].
64    ///
65    /// This allows sharing the same inquire client with `MasterClient`,
66    /// avoiding redundant Primary discovery.
67    pub async fn connect_with_inquire(
68        config: &GoosefsConfig,
69        inquire_client: Arc<dyn MasterInquireClient>,
70    ) -> Result<Self> {
71        let primary_addr = inquire_client.get_primary_rpc_address().await?;
72        let endpoint_uri = format!("http://{}", primary_addr);
73
74        let endpoint = Channel::from_shared(endpoint_uri)
75            .map_err(|e| Error::ConfigError {
76                message: format!("invalid master endpoint: {}", e),
77            })?
78            .connect_timeout(config.connect_timeout)
79            .timeout(config.request_timeout);
80
81        let channel = endpoint.connect().await?;
82
83        // Perform SASL authentication based on the configured auth type
84        let authenticator =
85            ChannelAuthenticator::new(config.auth_type, config.auth_username.clone(), None)
86                .with_auth_timeout(config.auth_timeout);
87
88        let mut auth_channel = authenticator.authenticate(channel).await?;
89        let sasl_guard = auth_channel.take_sasl_guard();
90        debug!(addr = %primary_addr, auth_type = %config.auth_type, "connected to WorkerManagerMasterClientService");
91
92        Ok(Self {
93            inner: WorkerManagerMasterClientServiceClient::new(auth_channel.channel),
94            _sasl_guard: Arc::new(sasl_guard),
95        })
96    }
97
98    /// Create from an existing tonic channel.
99    ///
100    /// **Note**: This bypasses authentication.
101    pub fn from_channel(channel: Channel) -> Self {
102        let interceptor = ChannelIdInterceptor::new("test-no-auth".to_string());
103        let intercepted = InterceptedService::new(channel, interceptor);
104        Self {
105            inner: WorkerManagerMasterClientServiceClient::new(intercepted),
106            _sasl_guard: Arc::new(None),
107        }
108    }
109
110    /// Fetch the full list of workers from the Master.
111    #[instrument(skip(self))]
112    pub async fn get_worker_info_list(&self) -> Result<Vec<WorkerInfo>> {
113        let req = GetWorkerInfoListPOptions {};
114
115        let resp = self.inner.clone().get_worker_info_list(req).await?;
116
117        Ok(resp.into_inner().worker_infos)
118    }
119}
goosefs_sdk/client/worker_manager.rs

goosefs_sdk/client/
worker_manager.rs
