goosefs_sdk/context.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//! `FileSystemContext` — shared connection pool and routing context.
16//!
17//! This module implements the **three-layer connection architecture** that
18//! eliminates repeated TCP+SASL handshakes:
19//!
20//! ```text
21//! Layer 3: FileSystemContext — lifecycle manager + unified acquisition API
22//! │
23//! ├── Arc<MasterClient> — persistent Master gRPC channel
24//! ├── Arc<WorkerManagerClient> — persistent WorkerMgr gRPC channel
25//! ├── Arc<WorkerClientPool> — shared Worker connection pool
26//! ├── Arc<WorkerRouter> — shared consistent-hash router
27//! └── Option<Arc<HeartbeatTask>> — periodic metrics heartbeat (when enabled)
28//!
29//! Layer 2: WorkerClientPool / WorkerRouter — connection & routing management
30//!
31//! Layer 1: MasterClient / WorkerManagerClient / WorkerClient — gRPC stubs
32//! ```
33//!
34//! # Before vs After
35//!
36//! | Operation | Before (per-call) | After (shared) |
37//! |-----------|-------------------|----------------|
38//! | `BaseFileSystem::get_status()` | 1 TCP+SASL | 0 (reused) |
39//! | `GoosefsFileInStream::open()` | 2 TCP+SASL | 0 (reused) |
40//! | Reading N blocks | N TCP connects | ~N_workers (pooled) |
41//!
42//! # Usage
43//!
44//! ```rust,no_run
45//! use goosefs_sdk::context::FileSystemContext;
46//! use goosefs_sdk::config::GoosefsConfig;
47//! use goosefs_sdk::fs::FileSystem; // needed to call trait methods
48//!
49//! # async fn example() -> goosefs_sdk::error::Result<()> {
50//! // Build once, share across all operations
51//! let ctx = FileSystemContext::connect(GoosefsConfig::new("127.0.0.1:9200")).await?;
52//!
53//! // Pass ctx into filesystem operations
54//! use goosefs_sdk::fs::BaseFileSystem;
55//! let fs = BaseFileSystem::from_context(ctx.clone());
56//! let status = fs.get_status("/data/file.parquet").await?;
57//! # Ok(())
58//! # }
59//! ```
60
61use std::sync::atomic::{AtomicBool, Ordering};
62use std::sync::Arc;
63use std::time::Duration;
64
65use tokio::sync::Mutex;
66use tracing::{debug, warn};
67
68use crate::block::router::WorkerRouter;
69use crate::cache::CacheManager;
70#[cfg(feature = "page-cache")]
71use crate::cache::LocalCacheManager;
72use crate::client::metrics_master::MetricsClient;
73use crate::client::metrics_master::MetricsMasterClient;
74use crate::client::{
75 create_master_inquire_client, MasterClient, MasterClientPool, MasterInquireClient,
76 WorkerClientPool, WorkerManagerClient,
77};
78use crate::config::{ConfigRefresher, GoosefsConfig, TransparentAccelerationSwitch};
79use crate::error::{Error, Result};
80use crate::metadata_cache::MetadataCache;
81use crate::metrics::heartbeat::{resolve_app_id, HeartbeatTask};
82#[cfg(feature = "metrics-pushgateway")]
83use crate::metrics::pushgateway::{PushgatewayConfig, PushgatewayTask};
84use crate::metrics::reporter::ClientMetricsReporter;
85
86/// How often the background refresh loop checks whether the worker list is stale.
87/// Matches the `DEFAULT_WORKER_REFRESH_TTL` (30s) in `WorkerRouter`.
88const REFRESH_CHECK_INTERVAL: Duration = Duration::from_secs(30);
89
90/// How often the background config refresh loop runs (default 60s).
91///
92/// Mirrors Java's `refreshInterval` (default 60s) in `NamespaceRefreshThread`.
93/// This is intentionally separate from [`REFRESH_CHECK_INTERVAL`] so that
94/// config reloading and worker-list refreshing run on independent cadences.
95const CONFIG_REFRESH_INTERVAL: Duration = Duration::from_secs(60);
96
97/// Shared connection context for Goosefs filesystem operations.
98///
99/// A single `FileSystemContext` instance should be created per Goosefs cluster
100/// and shared across all `BaseFileSystem`, `GoosefsFileInStream`, and
101/// `GoosefsFileWriter` instances that connect to that cluster.
102///
103/// The context owns:
104/// - One persistent gRPC channel to the Master
105/// - One persistent gRPC channel to the WorkerManager service
106/// - One `WorkerClientPool` shared across all readers and writers
107/// - One `WorkerRouter` that tracks live workers and routes block reads
108pub struct FileSystemContext {
109 config: Arc<GoosefsConfig>,
110
111 /// Persistent Master gRPC connection pool (metadata RPCs).
112 ///
113 /// Pool of `config.master_connection_pool_size` channels (default 1).
114 /// Scheduling strategy is controlled by `master_connection_pool_schedule`
115 /// (default `RoundRobin`; set to `P2C` for adaptive load balancing).
116 master_pool: Arc<MasterClientPool>,
117
118 /// Persistent WorkerManager gRPC connection (`GetWorkerInfoList`).
119 ///
120 /// `None` when the remote Master does not expose `WorkerManagerMasterClientService`
121 /// (e.g. older GooseFS versions). In that case, operations that need worker
122 /// discovery will log a warning and fall back to the Master RPC path.
123 worker_manager: Option<Arc<WorkerManagerClient>>,
124
125 /// Worker gRPC connection pool — shared across all readers and writers.
126 worker_pool: Arc<WorkerClientPool>,
127
128 /// Consistent-hash router with TTL refresh and local-first preference.
129 worker_router: Arc<WorkerRouter>,
130
131 /// Client-side local page cache, when `config.client_cache_enabled`.
132 ///
133 /// `None` when the cache is disabled or failed to initialize (the cache is
134 /// best-effort, so an init failure degrades to no-cache rather than
135 /// failing `connect()`). Shared across all readers in this context.
136 cache_manager: Option<Arc<dyn CacheManager>>,
137
138 /// Client-side metadata cache (status + listing + negative).
139 ///
140 /// `None` when `goosefs.user.metadata.cache.enabled` is false (the Java
141 /// default) or when expiration is `<= 0`. When present, `get_status` /
142 /// `list_status` / `exists` / open share this LRU, and the write path
143 /// invalidates path + parent after a successful mutation.
144 metadata_cache: Option<Arc<MetadataCache>>,
145
146 /// HA Master address discovery client (shared between master + wm).
147 inquire_client: Arc<dyn MasterInquireClient>,
148
149 /// Periodic config refresher — reloads `goosefs-site.properties` when
150 /// expired and updates the transparent acceleration switch flags.
151 ///
152 /// Mirrors Java's `ConfigurationUtils.loadIfExpire()` +
153 /// `AbstractCompatibleFileSystem.refreshTransparentAccelerationSwitch()`.
154 config_refresher: Arc<ConfigRefresher>,
155
156 /// Set to `true` once `close()` has been called.
157 closed: Arc<AtomicBool>,
158
159 /// Handle to the background worker-list TTL-refresh task.
160 /// Aborted on `close()`.
161 worker_refresh_task: Mutex<Option<tokio::task::JoinHandle<()>>>,
162
163 /// Handle to the background config refresh task.
164 /// Aborted on `close()`.
165 config_refresh_task: Mutex<Option<tokio::task::JoinHandle<()>>>,
166
167 /// Periodic metrics heartbeat task.
168 /// `None` when `config.metrics_enabled = false`.
169 /// Shut down gracefully (with final flush) in `close()`.
170 metrics_heartbeat: Mutex<Option<Arc<HeartbeatTask>>>,
171
172 /// Prometheus Pushgateway background push task.
173 /// `None` when `config.pushgateway_enabled = false`.
174 /// Shut down gracefully (with final flush) in `close()`.
175 #[cfg(feature = "metrics-pushgateway")]
176 pushgateway_task: Mutex<Option<PushgatewayTask>>,
177}
178
179impl FileSystemContext {
180 // ── Construction ────────────────────────────────────────────────────────
181
182 /// Build a `FileSystemContext` by connecting to the Goosefs cluster.
183 ///
184 /// Establishes persistent connections to the Master and WorkerManager,
185 /// fetches the initial worker list, and starts a background refresh task.
186 ///
187 /// This is the **only** call that performs network I/O. All subsequent
188 /// operations on the context are zero-cost Arc clones.
189 pub async fn connect(config: GoosefsConfig) -> Result<Arc<Self>> {
190 config
191 .validate()
192 .map_err(|message| Error::InvalidArgument { message })?;
193 let config = Arc::new(config);
194
195 // Build a shared inquire client so Master + WorkerManager both use the
196 // same singleflight-deduped HA discovery.
197 let inquire_client = create_master_inquire_client(&config);
198
199 // Connect Master pool (required) and WorkerManager (optional on older clusters).
200 let (pool_res, wm_res) = tokio::join!(
201 MasterClientPool::connect_with_inquire(&config, inquire_client.clone()),
202 WorkerManagerClient::connect_with_inquire(&config, inquire_client.clone()),
203 );
204 let master_pool = Arc::new(pool_res?);
205
206 // Build the router with failure/refresh TTLs.
207 let worker_router = Arc::new(WorkerRouter::with_ttls(
208 Duration::from_secs(60), // failure_ttl
209 Duration::from_secs(30), // worker_refresh_ttl (matches Go SDK)
210 ));
211
212 // WorkerManager is optional — older Master versions may not expose
213 // WorkerManagerMasterClientService / GetWorkerInfoList.
214 let worker_manager = match wm_res {
215 Ok(wm) => {
216 // Try to fetch the initial worker list; if the RPC itself fails
217 // (e.g. Method not found), treat the WorkerManager as unavailable.
218 match wm.get_worker_info_list().await {
219 Ok(workers) => {
220 if workers.is_empty() {
221 warn!("WorkerManager returned empty worker list — proceeding without worker discovery");
222 None
223 } else {
224 debug!(count = workers.len(), "initial worker list fetched");
225 worker_router.update_workers(workers).await;
226 Some(Arc::new(wm))
227 }
228 }
229 Err(e) => {
230 warn!("GetWorkerInfoList failed ({}), proceeding without worker discovery. \
231 Master-only operations (CreateFile, GetStatus, etc.) will still work.", e);
232 None
233 }
234 }
235 }
236 Err(e) => {
237 warn!(
238 "WorkerManager connection failed ({}), proceeding without worker discovery. \
239 Master-only operations (CreateFile, GetStatus, etc.) will still work.",
240 e
241 );
242 None
243 }
244 };
245
246 // Build the shared worker connection pool.
247 let worker_pool = WorkerClientPool::new_shared((*config).clone());
248
249 // Build the client-side local page cache (best-effort).
250 #[cfg(feature = "page-cache")]
251 let cache_manager: Option<Arc<dyn CacheManager>> = if config.client_cache_enabled {
252 match LocalCacheManager::from_config(&config).await {
253 Ok(mgr) => {
254 debug!(
255 page_size = config.client_cache_page_size,
256 dirs = ?config.client_cache_dirs,
257 "client local page cache enabled"
258 );
259 Some(mgr as Arc<dyn CacheManager>)
260 }
261 Err(e) => {
262 warn!(error = %e, "failed to init client page cache; continuing without cache");
263 None
264 }
265 }
266 } else {
267 None
268 };
269 #[cfg(not(feature = "page-cache"))]
270 let cache_manager: Option<Arc<dyn CacheManager>> = None;
271
272 // Java `FileSystem.Factory`: enabled → MetadataCachingBaseFileSystem.
273 // Rust hangs the same LRU on the context instead of swapping types.
274 let metadata_cache = if config.metadata_cache_enabled {
275 MetadataCache::maybe_new(
276 config.metadata_cache_expiration,
277 config.metadata_cache_max_size,
278 )
279 } else {
280 None
281 };
282 if let Some(c) = &metadata_cache {
283 crate::metrics::gauge(crate::metrics::name::CLIENT_METADATA_CACHE_ENABLED).set(1);
284 debug!(
285 ttl_ms = config.metadata_cache_expiration.as_millis(),
286 capacity = config.metadata_cache_max_size,
287 "metadata cache enabled, ttl={:?}",
288 c.ttl(),
289 );
290 }
291
292 let ctx = Arc::new(Self {
293 config: config.clone(),
294 master_pool,
295 worker_manager,
296 worker_pool,
297 worker_router,
298 cache_manager,
299 metadata_cache,
300 inquire_client,
301 config_refresher: Arc::new(ConfigRefresher::from_config(&config)),
302 closed: Arc::new(AtomicBool::new(false)),
303 worker_refresh_task: Mutex::new(None),
304 config_refresh_task: Mutex::new(None),
305 metrics_heartbeat: Mutex::new(None),
306 #[cfg(feature = "metrics-pushgateway")]
307 pushgateway_task: Mutex::new(None),
308 });
309
310 // Start the background worker-list refresh loop.
311 ctx.clone().start_worker_refresh_task().await;
312 // Start the background config refresh loop (separate cadence).
313 ctx.clone().start_config_refresh_task().await;
314 // Start the metrics heartbeat task (no-op when metrics_enabled = false).
315 ctx.clone().start_metrics_heartbeat_task().await?;
316 // Start the Pushgateway push task (no-op when pushgateway_enabled = false).
317 #[cfg(feature = "metrics-pushgateway")]
318 ctx.clone().start_pushgateway_task().await;
319
320 Ok(ctx)
321 }
322
323 // ── Acquisition API ──────────────────────────────────────────────────────
324
325 /// Return a shared `MasterClient` from the pool.
326 ///
327 /// With `master_connection_pool_schedule = RoundRobin` (default) this
328 /// cycles through pooled channels in order. With `P2C` it picks the
329 /// least-loaded connection out of two random candidates. Per-channel
330 /// in-flight counts are tracked inside `MasterClient::with_retry`, so
331 /// the load signal stays accurate even for clients cloned out of the
332 /// pool (e.g. by `GoosefsFileWriter`).
333 pub fn acquire_master(&self) -> Arc<MasterClient> {
334 self.master_pool.pick()
335 }
336
337 /// Return the shared `MasterClientPool` (zero-cost Arc clone).
338 pub fn acquire_master_pool(&self) -> Arc<MasterClientPool> {
339 self.master_pool.clone()
340 }
341
342 /// Return the shared `WorkerManagerClient` (zero-cost Arc clone).
343 ///
344 /// Returns `None` when the Master does not support `GetWorkerInfoList`.
345 pub fn acquire_worker_manager(&self) -> Option<Arc<WorkerManagerClient>> {
346 self.worker_manager.clone()
347 }
348
349 /// Return the shared `WorkerClientPool` (zero-cost Arc clone).
350 pub fn acquire_worker_pool(&self) -> Arc<WorkerClientPool> {
351 self.worker_pool.clone()
352 }
353
354 /// Return the shared `WorkerRouter` (zero-cost Arc clone).
355 pub fn acquire_router(&self) -> Arc<WorkerRouter> {
356 self.worker_router.clone()
357 }
358
359 /// Return the shared client-side page cache, if enabled.
360 ///
361 /// `None` when `config.client_cache_enabled = false` or the cache failed
362 /// to initialize. Readers consult this on the random-read path.
363 pub fn acquire_cache_manager(&self) -> Option<Arc<dyn CacheManager>> {
364 self.cache_manager.clone()
365 }
366
367 /// Return the shared metadata cache, if constructed.
368 ///
369 /// `None` when `metadata_cache_enabled` is false or expiration is `<= 0`.
370 pub fn acquire_metadata_cache(&self) -> Option<Arc<MetadataCache>> {
371 self.metadata_cache.clone()
372 }
373
374 /// Invalidate `path`, and optionally its parent directory listing/status.
375 ///
376 /// Idempotent no-op when the cache is disabled. Write paths must call this
377 /// **after** a successful mutation (Rust is more conservative than Java,
378 /// which invalidates before the RPC).
379 pub fn invalidate_metadata(&self, path: &str, with_parent: bool) {
380 if let Some(cache) = &self.metadata_cache {
381 if with_parent {
382 cache.invalidate_with_parent(path);
383 } else {
384 cache.invalidate(path);
385 }
386 }
387 }
388
389 /// OpenDAL-compatible entry. Equivalent to `invalidate_metadata(path, true)`.
390 ///
391 /// **Signature is frozen** — OpenDAL `core.rs` already calls this.
392 pub fn invalidate_file_info(&self, path: &str) {
393 self.invalidate_metadata(path, true);
394 }
395
396 /// Fetch `FileInfo` for `path`, consulting the metadata cache when present.
397 ///
398 /// Open paths share this with `BaseFileSystem::get_status` so a prior
399 /// `get_status` hit means open issues zero extra getStatus RPCs.
400 /// CheckBlocks enrichment must clone the result (INV-MC-D1).
401 pub(crate) async fn get_file_info_cached(
402 &self,
403 path: &str,
404 ) -> Result<crate::proto::grpc::file::FileInfo> {
405 let master = self.acquire_master();
406 crate::metadata_cache::get_status_through_cache(
407 self.metadata_cache.as_deref(),
408 path,
409 self.config.file_metadata_sync_interval,
410 || {
411 master.get_status_with_load_type(
412 path,
413 Some(self.config.file_metadata_load_type),
414 Some(self.config.file_metadata_sync_interval),
415 )
416 },
417 )
418 .await
419 }
420
421 /// Return the shared `MasterInquireClient` (zero-cost Arc clone).
422 pub fn acquire_inquire_client(&self) -> Arc<dyn MasterInquireClient> {
423 self.inquire_client.clone()
424 }
425
426 /// Return the configuration used to build this context.
427 pub fn config(&self) -> &GoosefsConfig {
428 &self.config
429 }
430
431 /// Return the shared `ConfigRefresher` (zero-cost Arc clone).
432 ///
433 /// Use this to query the current transparent acceleration switch values
434 /// or to trigger a config reload check.
435 pub fn acquire_config_refresher(&self) -> Arc<ConfigRefresher> {
436 self.config_refresher.clone()
437 }
438
439 /// Refresh the transparent acceleration switch by reloading config if expired.
440 ///
441 /// Convenience wrapper around `ConfigRefresher::refresh_transparent_acceleration_switch()`.
442 /// Mirrors Java's `AbstractCompatibleFileSystem.refreshTransparentAccelerationSwitch()`.
443 pub fn refresh_transparent_acceleration_switch(&self) -> TransparentAccelerationSwitch {
444 self.config_refresher
445 .refresh_transparent_acceleration_switch()
446 }
447
448 // ── Lifecycle ────────────────────────────────────────────────────────────
449
450 /// Gracefully shut down the context.
451 ///
452 /// Aborts the background refresh task and marks the context as closed.
453 /// Idempotent — safe to call multiple times.
454 pub async fn close(&self) -> Result<()> {
455 if self.closed.swap(true, Ordering::SeqCst) {
456 return Ok(()); // Already closed.
457 }
458
459 // Cancel the background refresh tasks.
460 let worker_handle = self.worker_refresh_task.lock().await.take();
461 if let Some(h) = worker_handle {
462 h.abort();
463 debug!("worker refresh task aborted");
464 }
465 let config_handle = self.config_refresh_task.lock().await.take();
466 if let Some(h) = config_handle {
467 h.abort();
468 debug!("config refresh task aborted");
469 }
470
471 // Gracefully shut down the metrics heartbeat task (performs final flush).
472 if let Some(task) = self.metrics_heartbeat.lock().await.take() {
473 task.shutdown().await;
474 debug!("metrics heartbeat task shut down");
475 }
476
477 // Gracefully shut down the Pushgateway push task (performs final push).
478 #[cfg(feature = "metrics-pushgateway")]
479 if let Some(task) = self.pushgateway_task.lock().await.take() {
480 task.shutdown().await;
481 debug!("pushgateway task shut down");
482 }
483
484 Ok(())
485 }
486
487 /// Return `true` if `close()` has been called.
488 pub fn is_closed(&self) -> bool {
489 self.closed.load(Ordering::SeqCst)
490 }
491
492 // ── Background refresh ────────────────────────────────────────────────────
493
494 /// Start the background worker-list TTL-refresh loop.
495 ///
496 /// The loop wakes every [`REFRESH_CHECK_INTERVAL`] seconds, calls
497 /// [`WorkerRouter::needs_refresh`], and if stale triggers
498 /// [`WorkerRouter::refresh_workers`].
499 async fn start_worker_refresh_task(self: Arc<Self>) {
500 let worker_router = self.worker_router.clone();
501 let worker_manager = self.worker_manager.clone();
502 let closed = self.closed.clone();
503
504 // If no WorkerManager is available, skip the refresh task entirely.
505 let Some(wm) = worker_manager else {
506 debug!("worker refresh task skipped: no WorkerManager available");
507 return;
508 };
509
510 let handle = tokio::spawn(async move {
511 loop {
512 tokio::time::sleep(REFRESH_CHECK_INTERVAL).await;
513
514 // Stop if the context has been closed.
515 if closed.load(Ordering::SeqCst) {
516 debug!("worker refresh task: context closed, exiting");
517 break;
518 }
519
520 // Refresh worker list if stale.
521 if worker_router.needs_refresh().await {
522 if let Err(e) = worker_router.refresh_workers(&wm).await {
523 warn!("worker refresh failed: {}", e);
524 // refresh_workers already resets the TTL clock to avoid
525 // hammering on repeated failures (stale-while-revalidate).
526 } else {
527 debug!("worker list refreshed by background task");
528 }
529 }
530 }
531 });
532
533 *self.worker_refresh_task.lock().await = Some(handle);
534 }
535
536 /// Start the background config refresh loop.
537 ///
538 /// On first invocation the task **immediately** loads the config from
539 /// `goosefs-site.properties` (via `refresh_transparent_acceleration_switch`)
540 /// so that the transparent acceleration switches are up-to-date right after
541 /// `connect()` returns. Subsequent refreshes happen every
542 /// [`CONFIG_REFRESH_INTERVAL`] seconds (default 60s, matching Java's
543 /// `refreshInterval`).
544 ///
545 /// This runs independently from the worker-list refresh task.
546 async fn start_config_refresh_task(self: Arc<Self>) {
547 let config_refresher = self.config_refresher.clone();
548 let closed = self.closed.clone();
549
550 let handle = tokio::spawn(async move {
551 // Eagerly load config on startup so the switches are current
552 // before any file-system operation is issued.
553 let switch = config_refresher.refresh_transparent_acceleration_switch();
554 debug!(
555 transparent_acceleration_enabled = switch.enabled,
556 cosranger_enabled = switch.cosranger_enabled,
557 "config refresh: initial load completed"
558 );
559
560 loop {
561 tokio::time::sleep(CONFIG_REFRESH_INTERVAL).await;
562
563 // Stop if the context has been closed.
564 if closed.load(Ordering::SeqCst) {
565 debug!("config refresh task: context closed, exiting");
566 break;
567 }
568
569 // Refresh transparent acceleration switch (reload config if expired).
570 // Mirrors Java's NamespaceRefreshThread calling
571 // refreshTransparentAccelerationSwitch() each loop iteration.
572 let switch = config_refresher.refresh_transparent_acceleration_switch();
573 debug!(
574 transparent_acceleration_enabled = switch.enabled,
575 cosranger_enabled = switch.cosranger_enabled,
576 "config refresh check completed"
577 );
578 }
579 });
580
581 *self.config_refresh_task.lock().await = Some(handle);
582 }
583}
584
585impl Drop for FileSystemContext {
586 fn drop(&mut self) {
587 // Signal `closed` first so any in-flight background tasks that poll
588 // this flag stop themselves before we try to abort their handles.
589 // This avoids a race where a task wakes up between our abort() call
590 // and the actual cancellation and touches shared state.
591 self.closed.store(true, Ordering::SeqCst);
592
593 // Best-effort abort of the refresh tasks.
594 // `drop` is synchronous, so we use `try_lock`; if we cannot obtain the
595 // lock the task loop will observe `closed == true` on its next iteration
596 // and exit on its own.
597 if let Ok(mut guard) = self.worker_refresh_task.try_lock() {
598 if let Some(h) = guard.take() {
599 h.abort();
600 }
601 }
602 if let Ok(mut guard) = self.config_refresh_task.try_lock() {
603 if let Some(h) = guard.take() {
604 h.abort();
605 }
606 }
607 // Send non-blocking shutdown signal to heartbeat task.
608 // HeartbeatTask::drop() will handle the rest (sets closed + try_send).
609 // We don't await shutdown() here because drop() is synchronous.
610 if let Ok(mut guard) = self.metrics_heartbeat.try_lock() {
611 guard.take(); // dropping the Arc triggers HeartbeatTask::drop
612 }
613 }
614}
615
616impl FileSystemContext {
617 // ── Metrics heartbeat ──────────────────────────────────────────────────
618
619 /// Start the periodic metrics heartbeat background task.
620 ///
621 /// Does nothing when `config.metrics_enabled = false`, so no
622 /// `MetricsMasterClient` is created and no background task is spawned.
623 ///
624 /// The task shares the same [`MasterInquireClient`] as `MasterClient` and
625 /// `WorkerManagerClient` for HA primary discovery.
626 async fn start_metrics_heartbeat_task(self: Arc<Self>) -> Result<()> {
627 if !self.config.metrics_enabled {
628 debug!("metrics disabled — heartbeat task not started");
629 return Ok(());
630 }
631
632 let mm_client =
633 MetricsMasterClient::connect_with_inquire(&self.config, self.inquire_client.clone())
634 .await?;
635
636 let reporter = Arc::new(ClientMetricsReporter::default());
637 let app_id = resolve_app_id(&self.config);
638
639 debug!(
640 app_id = %app_id,
641 interval_ms = self.config.metrics_heartbeat_interval.as_millis(),
642 timeout_ms = self.config.metrics_heartbeat_timeout.as_millis(),
643 "starting metrics heartbeat task"
644 );
645
646 let task = Arc::new(HeartbeatTask::spawn(
647 Arc::new(mm_client) as Arc<dyn MetricsClient>,
648 reporter,
649 app_id,
650 self.config.metrics_heartbeat_interval,
651 self.config.metrics_heartbeat_timeout,
652 self.closed.clone(),
653 ));
654 *self.metrics_heartbeat.lock().await = Some(task);
655 Ok(())
656 }
657
658 // ── Pushgateway ────────────────────────────────────────────────────────
659
660 /// Start the Prometheus Pushgateway background push task.
661 ///
662 /// If the initial config has `pushgateway_enabled = false`, this method
663 /// will attempt to auto-discover pushgateway settings from the properties
664 /// file (via `GoosefsConfig::from_properties_auto()`). This ensures that
665 /// callers using `GoosefsConfig::new(addr)` still get pushgateway reporting
666 /// when the properties file has it enabled.
667 ///
668 /// Does nothing only when **both** the initial config and the properties
669 /// file have pushgateway disabled (or no properties file is found).
670 #[cfg(feature = "metrics-pushgateway")]
671 async fn start_pushgateway_task(self: Arc<Self>) {
672 // Determine the effective pushgateway config: prefer the initial config
673 // if it already has pushgateway enabled; otherwise try auto-discovery.
674 let effective_config = if self.config.pushgateway_enabled {
675 // Caller explicitly enabled pushgateway — use their settings.
676 None
677 } else {
678 // Try loading from properties file to see if pushgateway is enabled there.
679 match GoosefsConfig::from_properties_auto() {
680 Ok(file_cfg) if file_cfg.pushgateway_enabled => {
681 debug!(
682 "pushgateway not enabled in initial config, \
683 but enabled in properties file — using file config"
684 );
685 Some(file_cfg)
686 }
687 _ => {
688 debug!("pushgateway disabled — push task not started");
689 return;
690 }
691 }
692 };
693
694 // Use the effective config (either initial or from properties file).
695 let cfg = effective_config.as_ref().unwrap_or(&self.config);
696
697 let mut pg_config = PushgatewayConfig::new(
698 cfg.pushgateway_endpoint.clone(),
699 cfg.pushgateway_job.clone(),
700 )
701 .with_push_interval(cfg.pushgateway_push_interval);
702
703 if let Some(ref instance) = cfg.pushgateway_instance {
704 pg_config = pg_config.with_instance(instance.clone());
705 } else {
706 // Auto-generate a unique instance identifier using "ip:pid"
707 // so that multiple client processes on the same machine do not
708 // overwrite each other's metrics in Pushgateway.
709 // Using IP is more intuitive and aligns with Prometheus conventions.
710 let pid = std::process::id();
711 let ip = Self::resolve_local_ip();
712 let auto_instance = format!("{}:{}", ip, pid);
713 debug!(auto_instance = %auto_instance, "auto-generated pushgateway instance");
714 pg_config = pg_config.with_instance(auto_instance);
715 }
716
717 debug!(
718 endpoint = %cfg.pushgateway_endpoint,
719 job = %cfg.pushgateway_job,
720 interval_ms = cfg.pushgateway_push_interval.as_millis(),
721 "starting pushgateway push task"
722 );
723
724 let task = PushgatewayTask::spawn(pg_config);
725 *self.pushgateway_task.lock().await = Some(task);
726 }
727
728 // ── Helpers ────────────────────────────────────────────────────────────
729
730 /// Resolve the local outbound IP address.
731 ///
732 /// Uses a UDP socket trick: connect to a public address (without actually
733 /// sending data) and read back the local address the OS chose. This gives
734 /// the correct outbound IP even on multi-homed machines.
735 ///
736 /// Falls back to `"127.0.0.1"` if detection fails.
737 #[cfg(feature = "metrics-pushgateway")]
738 fn resolve_local_ip() -> String {
739 use std::net::UdpSocket;
740 match UdpSocket::bind("0.0.0.0:0") {
741 Ok(socket) => {
742 // Connect to a well-known external address (Google DNS).
743 // No actual traffic is sent — this just triggers route lookup.
744 if socket.connect("8.8.8.8:80").is_ok() {
745 if let Ok(addr) = socket.local_addr() {
746 return addr.ip().to_string();
747 }
748 }
749 "127.0.0.1".to_string()
750 }
751 Err(_) => "127.0.0.1".to_string(),
752 }
753 }
754}
755
756// ---------------------------------------------------------------------------
757// Tests
758// ---------------------------------------------------------------------------
759
760#[cfg(test)]
761mod tests {
762 use super::*;
763 use std::time::Duration;
764
765 /// Verify that the context fields initialise with sane values when connected
766 /// (we can't test actual network here, but we can validate the structure).
767 #[test]
768 fn test_context_closed_starts_false() {
769 let closed = Arc::new(AtomicBool::new(false));
770 assert!(!closed.load(Ordering::SeqCst));
771 }
772
773 #[test]
774 fn test_context_close_is_idempotent() {
775 let closed = Arc::new(AtomicBool::new(false));
776
777 // First close
778 let was_open = !closed.swap(true, Ordering::SeqCst);
779 assert!(was_open);
780
781 // Second close — should be a no-op
782 let was_open2 = !closed.swap(true, Ordering::SeqCst);
783 assert!(!was_open2);
784 }
785
786 /// Verify that the worker refresh check interval constant is 30s.
787 #[test]
788 fn test_refresh_check_interval() {
789 assert_eq!(REFRESH_CHECK_INTERVAL, Duration::from_secs(30));
790 }
791
792 /// Verify that the config refresh interval constant is 60s (matching Java's refreshInterval).
793 #[test]
794 fn test_config_refresh_interval() {
795 assert_eq!(CONFIG_REFRESH_INTERVAL, Duration::from_secs(60));
796 }
797
798 /// Verify that WorkerRouter with_ttls accepts the values used by context.
799 #[test]
800 fn test_worker_router_ttls_accepted() {
801 let router = WorkerRouter::with_ttls(Duration::from_secs(60), Duration::from_secs(30));
802 // Just verifying it constructs without panic — fields are private.
803 drop(router);
804 }
805
806 /// Verify that `resolve_app_id` is accessible from context and returns a
807 /// non-empty string — the full resolution logic is tested in heartbeat tests.
808 #[test]
809 fn test_resolve_app_id_non_empty() {
810 let config = GoosefsConfig::new("127.0.0.1:9200");
811 let id = resolve_app_id(&config);
812 assert!(!id.is_empty());
813 }
814
815 /// Verify that metrics are disabled by default when the builder disables them,
816 /// and that the config field round-trips correctly.
817 #[test]
818 fn test_metrics_enabled_flag_is_accessible() {
819 let config_on = GoosefsConfig::new("127.0.0.1:9200").with_metrics_enabled(true);
820 assert!(config_on.metrics_enabled);
821
822 let config_off = GoosefsConfig::new("127.0.0.1:9200").with_metrics_enabled(false);
823 assert!(!config_off.metrics_enabled);
824 }
825
826 /// Verify that `start_metrics_heartbeat_task` returns `Ok(())` immediately
827 /// without attempting any network connection when `metrics_enabled = false`.
828 ///
829 /// This is the core contract for the `disabled_no_task_spawn` requirement
830 /// from the design spec :
831 /// - `metrics_enabled=false` → no HeartbeatTask spawned, no MetricsMasterClient created.
832 ///
833 /// We test the gate condition directly (the config flag check) rather than
834 /// exercising `FileSystemContext::connect()` which requires a real cluster.
835 #[tokio::test]
836 async fn disabled_no_task_spawn() {
837 // Build a minimal config with metrics disabled.
838 let config = Arc::new(GoosefsConfig::new("127.0.0.1:9200").with_metrics_enabled(false));
839 assert!(!config.metrics_enabled, "metrics_enabled must be false");
840
841 // Simulate the gate condition in `start_metrics_heartbeat_task`:
842 // when `metrics_enabled = false`, the task must not be spawned.
843 let metrics_heartbeat: Mutex<Option<Arc<HeartbeatTask>>> = Mutex::new(None);
844
845 // Replicate the exact guard from start_metrics_heartbeat_task.
846 let task_was_spawned = if config.metrics_enabled {
847 // Would connect to master and spawn task (not reached here).
848 true
849 } else {
850 // early return — no task spawned
851 false
852 };
853
854 assert!(
855 !task_was_spawned,
856 "metrics_enabled=false must prevent task from being spawned"
857 );
858
859 // Verify the Mutex remains None (no task was placed into it).
860 let guard = metrics_heartbeat.lock().await;
861 assert!(
862 guard.is_none(),
863 "metrics_heartbeat field must remain None when metrics are disabled"
864 );
865 }
866
867 /// Verify the `metrics_enabled` default value (true = opt-in enabled by default,
868 /// matching the design spec ).
869 #[test]
870 fn metrics_disabled_by_default() {
871 let config = GoosefsConfig::new("127.0.0.1:9200");
872 // Per config.rs, metrics_enabled defaults to true (align with Java SDK default).
873 // Explicitly disabling sets it to false.
874 let config_off = GoosefsConfig::new("127.0.0.1:9200").with_metrics_enabled(false);
875 assert!(
876 !config_off.metrics_enabled,
877 "with_metrics_enabled(false) must disable metrics"
878 );
879
880 // Verify the default (true).
881 assert!(
882 config.metrics_enabled,
883 "metrics_enabled defaults to true (opt-in enabled by default per Java SDK alignment)"
884 );
885 }
886
887 // ── Metadata cache construction gate ────────────────────────────────
888
889 /// Cache is **off** by default even with `cfg!(feature = "metadata-cache")`
890 /// (Java `USER_METADATA_CACHE_ENABLED`). TTL / size stay Java-aligned.
891 #[cfg(feature = "metadata-cache")]
892 #[test]
893 fn metadata_cache_disabled_by_default() {
894 let cfg = GoosefsConfig::default();
895 assert!(!cfg.metadata_cache_enabled);
896 assert_eq!(
897 cfg.metadata_cache_expiration,
898 Duration::from_secs(600),
899 "Java expiration.time default is 10min"
900 );
901 assert_eq!(
902 cfg.metadata_cache_max_size, 100_000,
903 "Java max.size default is 100000"
904 );
905 let constructed = if cfg.metadata_cache_enabled {
906 crate::metadata_cache::MetadataCache::maybe_new(
907 cfg.metadata_cache_expiration,
908 cfg.metadata_cache_max_size,
909 )
910 } else {
911 None
912 };
913 assert!(
914 constructed.is_none(),
915 "default enabled=false must skip cache construction"
916 );
917 }
918
919 /// Opting out must still bypass construction entirely.
920 #[test]
921 fn metadata_cache_opt_out_constructs_nothing() {
922 let cfg = GoosefsConfig::default().with_metadata_cache_enabled(false);
923 assert!(!cfg.metadata_cache_enabled);
924 let constructed = if cfg.metadata_cache_enabled {
925 crate::metadata_cache::MetadataCache::maybe_new(
926 cfg.metadata_cache_expiration,
927 cfg.metadata_cache_max_size,
928 )
929 } else {
930 None
931 };
932 assert!(constructed.is_none());
933 }
934
935 /// Setting the switch explicitly still keeps Java TTL / size.
936 #[cfg(feature = "metadata-cache")]
937 #[test]
938 fn metadata_cache_opt_in_keeps_java_defaults() {
939 let cfg = GoosefsConfig::new("127.0.0.1:9200").with_metadata_cache_enabled(true);
940 assert!(cfg.metadata_cache_enabled);
941 assert_eq!(cfg.metadata_cache_expiration, Duration::from_secs(600));
942 assert_eq!(cfg.metadata_cache_max_size, 100_000);
943
944 let cache = crate::metadata_cache::MetadataCache::maybe_new(
945 cfg.metadata_cache_expiration,
946 cfg.metadata_cache_max_size,
947 )
948 .expect("enabled + default expiration must produce a live cache");
949 assert_eq!(cache.ttl(), Duration::from_secs(600));
950 }
951
952 /// `with_metadata_cache_max_size(0)` must be clamped to `1`.
953 #[test]
954 fn metadata_cache_max_size_clamped_to_one() {
955 let cfg = GoosefsConfig::new("127.0.0.1:9200").with_metadata_cache_max_size(0);
956 assert_eq!(
957 cfg.metadata_cache_max_size, 1,
958 "with_metadata_cache_max_size(0) must clamp to 1"
959 );
960 }
961}