use std::collections::HashSet;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::{
AtomicU64,
AtomicUsize,
Ordering,
};
use std::time::{
Duration,
Instant,
};
use dashmap::DashMap;
use kftray_commons::models::config_model::Config;
use log::warn;
use once_cell::sync::Lazy;
use tokio::sync::Mutex;
use crate::kube::client::{
KubeConnection,
create_client_with_specific_context,
get_kubeconfig_paths_from_option,
merge_kubeconfigs,
};
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct ServiceClientKey {
pub context_name: Option<String>,
pub kubeconfig_path: Option<String>,
}
impl ServiceClientKey {
pub fn new(context_name: Option<String>, kubeconfig_path: Option<String>) -> Self {
Self {
context_name: context_name.filter(|name| name != Config::DEFAULT_CONTEXT),
kubeconfig_path,
}
}
}
struct CachedClient {
connection: Arc<KubeConnection>,
created_at: Instant,
signature: KubeconfigSignature,
}
impl CachedClient {
fn new(connection: KubeConnection, signature: KubeconfigSignature) -> Self {
Self {
connection: Arc::new(connection),
created_at: Instant::now(),
signature,
}
}
fn is_expired(&self, ttl: Duration) -> bool {
self.created_at.elapsed() > ttl
}
fn is_fresh(&self, ttl: Duration, signature: &KubeconfigSignature) -> bool {
!self.is_expired(ttl) && self.signature == *signature
}
}
type CreationLock = Arc<Creation>;
struct Creation {
failures: Mutex<Vec<SharedFailure>>,
interest: AtomicUsize,
sequence: AtomicU64,
}
fn new_creation_lock() -> CreationLock {
Arc::new(Creation {
failures: Mutex::new(Vec::new()),
interest: AtomicUsize::new(0),
sequence: AtomicU64::new(0),
})
}
struct SharedFailure {
sequence: u64,
signature: KubeconfigSignature,
error: anyhow::Error,
}
fn clone_error_text(error: &anyhow::Error) -> anyhow::Error {
let mut causes = error
.chain()
.rev()
.map(|cause| anyhow::Error::msg(cause.to_string()));
let root = causes
.next()
.unwrap_or_else(|| anyhow::Error::msg(error.to_string()));
causes.fold(root, |copy, cause| copy.context(cause.to_string()))
}
struct CreationLockInterest<'a> {
manager: &'a SharedClientManager,
key: ServiceClientKey,
lock: CreationLock,
}
impl<'a> CreationLockInterest<'a> {
fn already_registered(
manager: &'a SharedClientManager, key: ServiceClientKey, lock: CreationLock,
) -> Self {
Self { manager, key, lock }
}
}
impl Drop for CreationLockInterest<'_> {
fn drop(&mut self) {
self.manager.release_creation_lock(&self.key, &self.lock);
}
}
type KubeconfigSignature = Vec<(PathBuf, Option<std::time::SystemTime>)>;
fn kubeconfig_signature(paths: &[PathBuf]) -> KubeconfigSignature {
paths
.iter()
.map(|path| {
let mtime = std::fs::metadata(path)
.and_then(|meta| meta.modified())
.ok();
(path.clone(), mtime)
})
.collect()
}
pub struct SharedClientManager {
clients: DashMap<ServiceClientKey, CachedClient>,
client_ttl: Duration,
creation_locks: DashMap<ServiceClientKey, CreationLock>,
resolved_contexts: DashMap<Option<String>, (String, Instant, KubeconfigSignature)>,
}
impl SharedClientManager {
pub fn new() -> Self {
Self::with_ttl(Duration::from_secs(3600))
}
fn with_ttl(client_ttl: Duration) -> Self {
Self {
clients: DashMap::new(),
client_ttl,
creation_locks: DashMap::new(),
resolved_contexts: DashMap::new(),
}
}
pub async fn get_connection(
&self, key: ServiceClientKey,
) -> anyhow::Result<Arc<KubeConnection>> {
let (key, signature) = self.resolve_key(key).await?;
let context_name = key
.context_name
.as_deref()
.expect("resolve_key always returns a key carrying a context");
if let Some(cached) = self.clients.get(&key)
&& cached.is_fresh(self.client_ttl, &signature)
{
return Ok(cached.connection.clone());
}
let (lock, failures_seen) = {
let entry = self
.creation_locks
.entry(key.clone())
.or_insert_with(new_creation_lock);
let lock = entry.clone();
let failures_seen = lock.sequence.load(Ordering::SeqCst);
lock.interest.fetch_add(1, Ordering::SeqCst);
(lock, failures_seen)
};
let _interest = CreationLockInterest::already_registered(self, key.clone(), lock.clone());
let mut guard = lock.failures.lock().await;
if let Some(cached) = self.clients.get(&key)
&& cached.is_fresh(self.client_ttl, &signature)
{
let client_arc = cached.connection.clone();
drop(cached);
drop(guard);
return Ok(client_arc);
}
if let Some(failure) = guard
.iter()
.find(|failure| failure.sequence > failures_seen && failure.signature == signature)
{
let error = clone_error_text(&failure.error);
drop(guard);
return Err(error);
}
self.clients.remove(&key);
match create_client_with_specific_context(key.kubeconfig_path.clone(), context_name).await {
Ok(connection) => {
let cached_client = CachedClient::new(connection, signature);
let client_arc = cached_client.connection.clone();
self.clients.insert(key.clone(), cached_client);
drop(guard);
Ok(client_arc)
}
Err(error) => {
guard.push(SharedFailure {
sequence: lock.sequence.fetch_add(1, Ordering::SeqCst) + 1,
signature,
error: clone_error_text(&error),
});
drop(guard);
Err(error)
}
}
}
async fn resolve_key(
&self, key: ServiceClientKey,
) -> anyhow::Result<(ServiceClientKey, KubeconfigSignature)> {
let path_option = key.kubeconfig_path.clone();
let (current_paths, current_signature) = tokio::task::spawn_blocking(move || {
let current_paths = get_kubeconfig_paths_from_option(path_option)?;
let current_signature = kubeconfig_signature(¤t_paths);
Ok::<_, anyhow::Error>((current_paths, current_signature))
})
.await
.map_err(|e| anyhow::anyhow!("kubeconfig path resolution task panicked: {e}"))??;
if key.context_name.is_some() {
return Ok((key, current_signature));
}
if let Some(resolved) = self.resolved_contexts.get(&key.kubeconfig_path) {
let (context_name, resolved_at, signature) = resolved.value();
if resolved_at.elapsed() <= self.client_ttl && *signature == current_signature {
return Ok((
ServiceClientKey {
context_name: Some(context_name.clone()),
kubeconfig_path: key.kubeconfig_path,
},
current_signature,
));
}
}
let kubeconfig_path = key.kubeconfig_path.clone();
let current_context = tokio::task::spawn_blocking(move || {
let (kubeconfig, errors) = merge_kubeconfigs(¤t_paths)?;
match kubeconfig.current_context {
Some(current_context) => {
if !errors.is_empty() {
warn!(
"Resolved current-context for kubeconfig {} despite {} file(s) \
failing to read or merge: {}",
kubeconfig_path.as_deref().unwrap_or("default"),
errors.len(),
errors.join("; ")
);
}
Ok(current_context)
}
None if errors.is_empty() => Err(anyhow::anyhow!(
"Kubernetes context is required (kubeconfig: {}) and it has no \
current-context",
kubeconfig_path.as_deref().unwrap_or("default")
)),
None => Err(anyhow::anyhow!(
"Kubernetes context is required (kubeconfig: {}) and it has no \
current-context; it also failed to read: {}",
kubeconfig_path.as_deref().unwrap_or("default"),
errors.join("; ")
)),
}
})
.await
.map_err(|e| anyhow::anyhow!("kubeconfig resolution task panicked: {e}"))??;
self.resolved_contexts.insert(
key.kubeconfig_path.clone(),
(
current_context.clone(),
Instant::now(),
current_signature.clone(),
),
);
Ok((
ServiceClientKey {
context_name: Some(current_context),
kubeconfig_path: key.kubeconfig_path,
},
current_signature,
))
}
fn release_creation_lock(&self, key: &ServiceClientKey, lock: &CreationLock) {
if lock.interest.fetch_sub(1, Ordering::SeqCst) == 1 {
self.creation_locks.remove_if(key, |_, entry| {
Arc::ptr_eq(entry, lock) && entry.interest.load(Ordering::SeqCst) == 0
});
}
}
pub fn invalidate_client(&self, key: &ServiceClientKey) {
if key.context_name.is_none() {
if let Some((_, (resolved, _, _))) = self.resolved_contexts.remove(&key.kubeconfig_path)
{
let resolved_key = ServiceClientKey {
context_name: Some(resolved),
kubeconfig_path: key.kubeconfig_path.clone(),
};
self.clients.remove(&resolved_key);
return;
}
self.clients
.retain(|cached_key, _| cached_key.kubeconfig_path != key.kubeconfig_path);
return;
}
if let Some(ctx) = key.context_name.as_deref() {
self.resolved_contexts
.remove_if(&key.kubeconfig_path, |_, (resolved, _, _)| resolved == ctx);
}
self.clients.remove(key);
}
pub fn release_unused(&self, in_use: &HashSet<ServiceClientKey>) {
self.clients.retain(|key, _| {
in_use.contains(key)
|| in_use.contains(&ServiceClientKey::new(None, key.kubeconfig_path.clone()))
});
}
pub fn cleanup_expired(&self) {
self.clients
.retain(|_, cached| !cached.is_expired(self.client_ttl));
self.creation_locks.retain(|key, lock| {
self.clients.contains_key(key) || lock.interest.load(Ordering::SeqCst) > 0
});
self.resolved_contexts
.retain(|kubeconfig_path, (context_name, resolved_at, _)| {
if resolved_at.elapsed() > self.client_ttl {
return false;
}
let resolved_key = ServiceClientKey {
context_name: Some(context_name.clone()),
kubeconfig_path: kubeconfig_path.clone(),
};
self.clients.contains_key(&resolved_key)
});
}
}
impl Default for SharedClientManager {
fn default() -> Self {
Self::new()
}
}
pub static SHARED_CLIENT_MANAGER: Lazy<SharedClientManager> = Lazy::new(SharedClientManager::new);
#[cfg(test)]
mod tests {
use super::*;
async fn wait_for_interest_count(lock: &CreationLock, expected: usize) {
for _ in 0..500 {
if lock.interest.load(Ordering::SeqCst) == expected {
return;
}
tokio::time::sleep(Duration::from_millis(1)).await;
}
panic!(
"creation lock interest count did not reach {expected} in time (last seen {})",
lock.interest.load(Ordering::SeqCst)
);
}
#[tokio::test]
async fn get_connection_errors_when_kubeconfig_has_no_current_context() {
let manager = SharedClientManager::new();
let key = ServiceClientKey::new(
None,
Some("/nonexistent/kftray-test-kubeconfig-no-current-context".to_string()),
);
let err = manager.get_connection(key).await.expect_err(
"a key without a context and without a resolvable current-context must not resolve",
);
assert!(
err.to_string().contains("current-context"),
"error should name the missing current-context: {err}"
);
assert!(
err.to_string().contains("failed to read"),
"error should surface the underlying kubeconfig read failure instead of \
discarding it: {err}"
);
}
#[tokio::test]
async fn get_connection_resolves_none_context_to_kubeconfig_current_context() {
use std::io::Write;
crate::ssl::ensure_crypto_provider_installed();
let mut kubeconfig_file = tempfile::NamedTempFile::new().expect("create temp kubeconfig");
write!(
kubeconfig_file,
"apiVersion: v1\n\
kind: Config\n\
current-context: kftray-test-current-context\n\
contexts:\n\
- name: kftray-test-current-context\n \
context:\n \
cluster: kftray-test-cluster\n \
user: kftray-test-user\n\
clusters:\n\
- name: kftray-test-cluster\n \
cluster:\n \
server: https://127.0.0.1:1\n\
users:\n\
- name: kftray-test-user\n \
user: {{}}\n"
)
.expect("write temp kubeconfig");
let manager = SharedClientManager::new();
let key = ServiceClientKey::new(
None,
Some(kubeconfig_file.path().to_string_lossy().to_string()),
);
let err = manager
.get_connection(key)
.await
.expect_err("an unreachable resolved cluster must not produce a client");
assert!(
err.to_string().contains("kftray-test-current-context"),
"resolving a None context must use the kubeconfig's current-context name so the \
connection attempt (and its error) names it: {err}"
);
assert!(
!err.to_string().contains("Kubernetes context is required"),
"resolving a None context via the kubeconfig's current-context must not fail with \
a missing-context error: {err}"
);
}
#[tokio::test]
async fn a_config_imported_without_a_context_connects_with_the_current_context() {
use kftray_commons::utils::config::{
import_configs_with_mode,
read_configs_with_mode,
};
use kftray_commons::utils::db_mode::DatabaseMode;
crate::ssl::ensure_crypto_provider_installed();
let _db = kftray_commons::test_utils::test_db().await;
let kubeconfig_file = tempfile::NamedTempFile::new().unwrap();
std::fs::write(
kubeconfig_file.path(),
kubeconfig_yaml("kftray-test-current-context"),
)
.unwrap();
import_configs_with_mode(
r#"[{"service":"web","namespace":"default","workload_type":"service",
"protocol":"tcp","remote_port":80,"local_port":18080}]"#
.to_string(),
DatabaseMode::Memory,
)
.await
.unwrap();
let config = read_configs_with_mode(DatabaseMode::Memory)
.await
.unwrap()
.remove(0);
let key = ServiceClientKey::new(
config.context,
Some(kubeconfig_file.path().to_string_lossy().to_string()),
);
let err = SharedClientManager::new()
.get_connection(key)
.await
.expect_err("an unreachable cluster must not produce a client");
assert!(
err.to_string().contains("kftray-test-current-context"),
"{err}"
);
}
fn kubeconfig_yaml(context_name: &str) -> String {
format!(
"apiVersion: v1\n\
kind: Config\n\
current-context: {context_name}\n\
contexts:\n\
- name: {context_name}\n \
context:\n \
cluster: kftray-test-cluster\n \
user: kftray-test-user\n\
clusters:\n\
- name: kftray-test-cluster\n \
cluster:\n \
server: https://127.0.0.1:1\n\
users:\n\
- name: kftray-test-user\n \
user: {{}}\n"
)
}
fn rewrite_kubeconfig_preserving_mtime(
path: &std::path::Path, contents: &str, mtime: std::time::SystemTime,
) {
std::fs::write(path, contents).expect("rewrite temp kubeconfig");
std::fs::OpenOptions::new()
.write(true)
.open(path)
.and_then(|file| file.set_times(std::fs::FileTimes::new().set_modified(mtime)))
.expect("restore the kubeconfig's mtime");
}
#[tokio::test]
async fn resolve_key_with_none_context_is_cached_on_a_hit_within_ttl() {
use std::io::Write;
let mut kubeconfig_file = tempfile::NamedTempFile::new().expect("create temp kubeconfig");
write!(
kubeconfig_file,
"{}",
kubeconfig_yaml("kftray-test-cache-hit-context")
)
.expect("write temp kubeconfig");
let manager = SharedClientManager::new();
let key = ServiceClientKey::new(
None,
Some(kubeconfig_file.path().to_string_lossy().to_string()),
);
let (first_resolved, _) = manager
.resolve_key(key.clone())
.await
.expect("a kubeconfig with a current-context must resolve");
assert_eq!(
first_resolved.context_name.as_deref(),
Some("kftray-test-cache-hit-context")
);
let original_mtime = std::fs::metadata(kubeconfig_file.path())
.and_then(|meta| meta.modified())
.expect("read the kubeconfig's mtime");
rewrite_kubeconfig_preserving_mtime(
kubeconfig_file.path(),
&kubeconfig_yaml("kftray-test-changed-context"),
original_mtime,
);
let (second_resolved, _) = manager
.resolve_key(key)
.await
.expect("a cache hit must resolve without needing the rewritten kubeconfig to parse");
assert_eq!(
second_resolved.context_name.as_deref(),
Some("kftray-test-cache-hit-context"),
"a cache hit with an unchanged mtime must reuse the recorded resolved context"
);
}
#[tokio::test]
async fn resolve_key_re_resolves_when_the_kubeconfig_mtime_changes() {
use std::io::Write;
let mut kubeconfig_file = tempfile::NamedTempFile::new().expect("create temp kubeconfig");
write!(
kubeconfig_file,
"{}",
kubeconfig_yaml("kftray-test-original-context")
)
.expect("write temp kubeconfig");
let manager = SharedClientManager::new();
let key = ServiceClientKey::new(
None,
Some(kubeconfig_file.path().to_string_lossy().to_string()),
);
let (first_resolved, _) = manager
.resolve_key(key.clone())
.await
.expect("a kubeconfig with a current-context must resolve");
assert_eq!(
first_resolved.context_name.as_deref(),
Some("kftray-test-original-context")
);
let advanced_mtime = std::time::SystemTime::now() + Duration::from_secs(60);
rewrite_kubeconfig_preserving_mtime(
kubeconfig_file.path(),
&kubeconfig_yaml("kftray-test-switched-context"),
advanced_mtime,
);
let (second_resolved, _) = manager
.resolve_key(key)
.await
.expect("the rewritten kubeconfig must still resolve");
assert_eq!(
second_resolved.context_name.as_deref(),
Some("kftray-test-switched-context"),
"a changed mtime must force a re-resolution even though the TTL has not expired"
);
}
#[tokio::test]
async fn invalidate_client_with_none_context_uses_recorded_resolved_context_only() {
use http::{
Request,
Response,
};
use kube::client::Body;
use tower_test::mock;
let manager = SharedClientManager::new();
let kubeconfig_path = Some("kftray-test-recorded-context-kubeconfig".to_string());
let make_cached_client = || {
let (mock_service, _handle) = mock::pair::<Request<Body>, Response<Body>>();
CachedClient::new(
KubeConnection {
client: kube::Client::new(mock_service, "default"),
cluster_url: "https://example.invalid".parse().unwrap(),
},
Vec::new(),
)
};
let recorded_key =
ServiceClientKey::new(Some("recorded-ctx".to_string()), kubeconfig_path.clone());
let other_key =
ServiceClientKey::new(Some("other-ctx".to_string()), kubeconfig_path.clone());
manager
.clients
.insert(recorded_key.clone(), make_cached_client());
manager
.clients
.insert(other_key.clone(), make_cached_client());
manager.resolved_contexts.insert(
kubeconfig_path.clone(),
("recorded-ctx".to_string(), Instant::now(), Vec::new()),
);
manager.invalidate_client(&ServiceClientKey::new(None, kubeconfig_path.clone()));
assert!(
!manager.clients.contains_key(&recorded_key),
"the recorded resolved-context entry must be evicted"
);
assert!(
manager.clients.contains_key(&other_key),
"invalidation with a recorded resolved context must hit only that exact key, not \
sweep every entry for the kubeconfig path"
);
assert!(
!manager.resolved_contexts.contains_key(&kubeconfig_path),
"invalidation must drop the recorded resolved context so a later resolve re-reads \
the current context from disk"
);
}
#[tokio::test]
async fn get_connection_releases_lock_after_failed_creation() {
let manager = SharedClientManager::new();
let key = ServiceClientKey::new(
Some("kftray-test-nonexistent-context".to_string()),
Some("/nonexistent/kftray-test-kubeconfig".to_string()),
);
manager.get_connection(key.clone()).await.expect_err(
"a context absent from the (empty) kubeconfig must fail to create a client",
);
assert!(
!manager.creation_locks.contains_key(&key),
"a failed creation with no other waiter must not leak its lock entry"
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 8)]
async fn concurrent_get_connection_calls_never_leak_a_fresh_creation_lock_entry() {
let manager = Arc::new(SharedClientManager::new());
let key = ServiceClientKey::new(
Some("kftray-test-race-context".to_string()),
Some("/nonexistent/kftray-test-race-kubeconfig".to_string()),
);
for _ in 0..25 {
let mut handles = Vec::with_capacity(32);
for _ in 0..32 {
let manager = Arc::clone(&manager);
let key = key.clone();
handles.push(tokio::spawn(
async move { manager.get_connection(key).await },
));
}
for handle in handles {
let _ = handle
.await
.expect("a racing get_connection call must not panic");
}
assert!(
!manager.creation_locks.contains_key(&key),
"every caller failed and released; the creation-lock entry must not be leaked"
);
}
}
#[tokio::test]
async fn invalidate_client_with_none_context_removes_resolved_entry() {
use std::io::Write;
use http::{
Request,
Response,
};
use kube::client::Body;
use tower_test::mock;
let mut kubeconfig_file = tempfile::NamedTempFile::new().expect("create temp kubeconfig");
write!(
kubeconfig_file,
"apiVersion: v1\n\
kind: Config\n\
current-context: kftray-test-invalidate-context\n\
contexts:\n\
- name: kftray-test-invalidate-context\n \
context:\n \
cluster: kftray-test-cluster\n \
user: kftray-test-user\n\
clusters:\n\
- name: kftray-test-cluster\n \
cluster:\n \
server: https://127.0.0.1:1\n\
users:\n\
- name: kftray-test-user\n \
user: {{}}\n"
)
.expect("write temp kubeconfig");
let manager = SharedClientManager::new();
let kubeconfig_path = kubeconfig_file.path().to_string_lossy().to_string();
let raw_key = ServiceClientKey::new(None, Some(kubeconfig_path.clone()));
let resolved_key = ServiceClientKey::new(
Some("kftray-test-invalidate-context".to_string()),
Some(kubeconfig_path.clone()),
);
let (mock_service, _handle) = mock::pair::<Request<Body>, Response<Body>>();
let connection = KubeConnection {
client: kube::Client::new(mock_service, "default"),
cluster_url: "https://example.invalid".parse().unwrap(),
};
manager.clients.insert(
resolved_key.clone(),
CachedClient::new(connection, Vec::new()),
);
manager.resolved_contexts.insert(
Some(kubeconfig_path.clone()),
(
"kftray-test-invalidate-context".to_string(),
Instant::now(),
Vec::new(),
),
);
manager.invalidate_client(&raw_key);
assert!(
!manager.clients.contains_key(&resolved_key),
"invalidate_client with a None-context key must remove the entry cached under \
the resolved current-context key, not just the never-cached raw key"
);
assert!(
!manager
.resolved_contexts
.contains_key(&Some(kubeconfig_path)),
"invalidation must drop the recorded resolved context so a later resolve re-reads \
the current context from disk"
);
}
#[tokio::test]
async fn invalidate_client_with_none_context_clears_every_entry_for_kubeconfig_path() {
use http::{
Request,
Response,
};
use kube::client::Body;
use tower_test::mock;
let manager = SharedClientManager::new();
let kubeconfig_path = Some("kftray-test-shared-kubeconfig".to_string());
let make_cached_client = || {
let (mock_service, _handle) = mock::pair::<Request<Body>, Response<Body>>();
CachedClient::new(
KubeConnection {
client: kube::Client::new(mock_service, "default"),
cluster_url: "https://example.invalid".parse().unwrap(),
},
Vec::new(),
)
};
let key_a = ServiceClientKey::new(Some("context-a".to_string()), kubeconfig_path.clone());
let key_b = ServiceClientKey::new(Some("context-b".to_string()), kubeconfig_path.clone());
let other_kubeconfig_key = ServiceClientKey::new(
Some("context-a".to_string()),
Some("kftray-test-other-kubeconfig".to_string()),
);
manager.clients.insert(key_a.clone(), make_cached_client());
manager.clients.insert(key_b.clone(), make_cached_client());
manager
.clients
.insert(other_kubeconfig_key.clone(), make_cached_client());
manager.invalidate_client(&ServiceClientKey::new(None, kubeconfig_path));
assert!(
!manager.clients.contains_key(&key_a),
"every cached context for the kubeconfig path must be dropped"
);
assert!(
!manager.clients.contains_key(&key_b),
"every cached context for the kubeconfig path must be dropped, not just one \
re-resolved current-context entry"
);
assert!(
manager.clients.contains_key(&other_kubeconfig_key),
"entries for a different kubeconfig path must be left untouched"
);
}
#[tokio::test]
async fn release_unused_keeps_clients_a_forward_still_needs() {
use http::{
Request,
Response,
};
use kube::client::Body;
use tower_test::mock;
let manager = SharedClientManager::new();
let kubeconfig_path = Some("kftray-test-shared-kubeconfig".to_string());
let make_cached_client = || {
let (mock_service, _handle) = mock::pair::<Request<Body>, Response<Body>>();
CachedClient::new(
KubeConnection {
client: kube::Client::new(mock_service, "default"),
cluster_url: "https://example.invalid".parse().unwrap(),
},
Vec::new(),
)
};
let key_a = ServiceClientKey::new(Some("context-a".to_string()), kubeconfig_path.clone());
let key_b = ServiceClientKey::new(Some("context-b".to_string()), kubeconfig_path.clone());
let other_kubeconfig_key = ServiceClientKey::new(
Some("context-a".to_string()),
Some("kftray-test-other-kubeconfig".to_string()),
);
manager.clients.insert(key_a.clone(), make_cached_client());
manager.clients.insert(key_b.clone(), make_cached_client());
manager
.clients
.insert(other_kubeconfig_key.clone(), make_cached_client());
manager.release_unused(&HashSet::from([
key_b.clone(),
ServiceClientKey::new(None, other_kubeconfig_key.kubeconfig_path.clone()),
]));
assert!(
!manager.clients.contains_key(&key_a),
"a client no forward uses must be dropped"
);
assert!(
manager.clients.contains_key(&key_b),
"a client a forward names by context must stay"
);
assert!(
manager.clients.contains_key(&other_kubeconfig_key),
"a forward without a context keeps every client for its kubeconfig"
);
manager.release_unused(&HashSet::new());
assert!(
manager.clients.is_empty(),
"with nothing forwarding, no client is kept"
);
}
#[tokio::test]
async fn cancelling_get_connection_releases_creation_lock_interest() {
let manager = Arc::new(SharedClientManager::new());
let key = ServiceClientKey::new(
Some("ctx".to_string()),
Some("kftray-test-cancel-kubeconfig".to_string()),
);
let lock = manager
.creation_locks
.entry(key.clone())
.or_insert_with(new_creation_lock)
.clone();
let guard = lock.failures.lock().await;
let waiter_manager = manager.clone();
let waiter_key = key.clone();
let handle = tokio::spawn(async move { waiter_manager.get_connection(waiter_key).await });
wait_for_interest_count(&lock, 1).await;
assert_eq!(
lock.interest.load(Ordering::SeqCst),
1,
"waiter must have registered interest before being cancelled"
);
handle.abort();
let _ = handle.await;
assert_eq!(
lock.interest.load(Ordering::SeqCst),
0,
"a cancelled get_connection must release its creation-lock interest"
);
drop(guard);
}
#[test]
fn release_creation_lock_keeps_entry_when_another_waiter_holds_it() {
let manager = SharedClientManager::new();
let key = ServiceClientKey::new(Some("ctx".to_string()), None);
let lock = manager
.creation_locks
.entry(key.clone())
.or_insert_with(new_creation_lock)
.clone();
lock.interest.fetch_add(2, Ordering::SeqCst);
manager.release_creation_lock(&key, &lock);
assert!(
manager.creation_locks.contains_key(&key),
"a lock another waiter still references must not be dropped"
);
}
#[test]
fn cleanup_expired_retains_lock_entries_still_in_use() {
let manager = SharedClientManager::new();
let key = ServiceClientKey::new(Some("ctx".to_string()), None);
let lock = manager
.creation_locks
.entry(key.clone())
.or_insert_with(new_creation_lock)
.clone();
lock.interest.fetch_add(1, Ordering::SeqCst);
manager.cleanup_expired();
assert!(
manager.creation_locks.contains_key(&key),
"cleanup_expired must not drop a lock entry an in-flight creation still holds"
);
lock.interest.fetch_sub(1, Ordering::SeqCst);
manager.cleanup_expired();
assert!(
!manager.creation_locks.contains_key(&key),
"cleanup_expired should drop an unheld lock entry once no client is cached for it"
);
}
#[tokio::test]
async fn get_connection_releases_lock_on_post_lock_cache_hit() {
use http::{
Request,
Response,
};
use kube::client::Body;
use tower_test::mock;
let manager = Arc::new(SharedClientManager::new());
let kubeconfig_path = "kftray-test-post-lock-kubeconfig";
let key = ServiceClientKey::new(Some("ctx".to_string()), Some(kubeconfig_path.to_string()));
let lock = manager
.creation_locks
.entry(key.clone())
.or_insert_with(new_creation_lock)
.clone();
lock.interest.fetch_add(1, Ordering::SeqCst);
let guard = lock.failures.lock().await;
let waiter_manager = manager.clone();
let waiter_key = key.clone();
let waiter = tokio::spawn(async move { waiter_manager.get_connection(waiter_key).await });
wait_for_interest_count(&lock, 2).await;
let (mock_service, _handle) = mock::pair::<Request<Body>, Response<Body>>();
let connection = KubeConnection {
client: kube::Client::new(mock_service, "default"),
cluster_url: "https://example.invalid".parse().unwrap(),
};
manager.clients.insert(
key.clone(),
CachedClient::new(
connection,
kubeconfig_signature(&[PathBuf::from(kubeconfig_path)]),
),
);
drop(guard);
manager.release_creation_lock(&key, &lock);
let result = tokio::time::timeout(Duration::from_secs(2), waiter)
.await
.expect("post-lock cache hit must resolve, not hang")
.expect("waiter task must not panic");
assert!(
result.is_ok(),
"post-lock cache hit must resolve to the cached connection"
);
assert!(
!manager.creation_locks.contains_key(&key),
"a post-lock cache hit must release its clone of the creation lock"
);
}
const FAILING_KUBECONFIG: &str = "/nonexistent/kftray-test-shared-failure-kubeconfig";
fn failing_signature() -> KubeconfigSignature {
kubeconfig_signature(&[PathBuf::from(FAILING_KUBECONFIG)])
}
fn stale_signature() -> KubeconfigSignature {
vec![(
PathBuf::from(FAILING_KUBECONFIG),
Some(std::time::SystemTime::UNIX_EPOCH),
)]
}
fn record_failure(
lock: &CreationLock, failures: &mut Vec<SharedFailure>, signature: KubeconfigSignature,
message: &str,
) {
failures.push(SharedFailure {
sequence: lock.sequence.fetch_add(1, Ordering::SeqCst) + 1,
signature,
error: anyhow::anyhow!("{message}").context("Failed to create client"),
});
}
async fn waiter_behind_failed_creations(
stored: Vec<(KubeconfigSignature, &str)>,
) -> (ServiceClientKey, Arc<SharedClientManager>, anyhow::Error) {
let manager = Arc::new(SharedClientManager::new());
let key = ServiceClientKey::new(
Some("ctx".to_string()),
Some(FAILING_KUBECONFIG.to_string()),
);
let lock = manager
.creation_locks
.entry(key.clone())
.or_insert_with(new_creation_lock)
.clone();
lock.interest.fetch_add(1, Ordering::SeqCst);
let mut guard = lock.failures.lock().await;
let waiter_manager = manager.clone();
let waiter_key = key.clone();
let waiter = tokio::spawn(async move { waiter_manager.get_connection(waiter_key).await });
wait_for_interest_count(&lock, 2).await;
for (signature, message) in stored {
record_failure(&lock, &mut guard, signature, message);
}
drop(guard);
manager.release_creation_lock(&key, &lock);
let error = tokio::time::timeout(Duration::from_secs(5), waiter)
.await
.expect("a waiter behind a failed creation must not hang")
.expect("waiter task must not panic")
.expect_err("no client can be created from a nonexistent kubeconfig");
(key, manager, error)
}
#[tokio::test]
async fn a_waiter_queued_behind_a_failed_creation_gets_that_failure() {
let (key, manager, error) =
waiter_behind_failed_creations(vec![(failing_signature(), "cluster unreachable")])
.await;
assert_eq!(
format!("{error:#}"),
"Failed to create client: cluster unreachable"
);
assert!(
!manager.creation_locks.contains_key(&key),
"the shared failure must leave with the last queued caller"
);
}
#[tokio::test]
async fn a_failure_of_another_kubeconfig_state_is_not_shared() {
let (_, _, error) =
waiter_behind_failed_creations(vec![(stale_signature(), "cluster unreachable")]).await;
assert!(
!error.to_string().contains("cluster unreachable"),
"a waiter whose kubeconfig differs from the failed attempt must try again: {error}"
);
}
#[tokio::test]
async fn a_later_failure_of_another_kubeconfig_state_keeps_the_earlier_one() {
let (_, _, error) = waiter_behind_failed_creations(vec![
(failing_signature(), "cluster unreachable"),
(stale_signature(), "edited kubeconfig rejected"),
])
.await;
assert_eq!(
format!("{error:#}"),
"Failed to create client: cluster unreachable"
);
}
#[tokio::test]
async fn a_failure_shared_past_a_later_success_of_another_state_leaves_that_success_cached() {
crate::ssl::ensure_crypto_provider_installed();
let port = fake_api_server().await;
let kubeconfig = tempfile::NamedTempFile::new().expect("create temp kubeconfig");
std::fs::write(kubeconfig.path(), kubeconfig_yaml("ctx")).expect("write temp kubeconfig");
let kubeconfig_path = kubeconfig.path().to_string_lossy().to_string();
let signature = || kubeconfig_signature(&[kubeconfig.path().to_path_buf()]);
let manager = Arc::new(SharedClientManager::new());
let key = ServiceClientKey::new(Some("ctx".to_string()), Some(kubeconfig_path.clone()));
let lock = manager
.creation_locks
.entry(key.clone())
.or_insert_with(new_creation_lock)
.clone();
lock.interest.fetch_add(1, Ordering::SeqCst);
let mut guard = lock.failures.lock().await;
let waiter_manager = manager.clone();
let waiter_key = key.clone();
let waiter = tokio::spawn(async move { waiter_manager.get_connection(waiter_key).await });
wait_for_interest_count(&lock, 2).await;
record_failure(&lock, &mut guard, signature(), "cluster unreachable");
rewrite_kubeconfig_preserving_mtime(
kubeconfig.path(),
&kubeconfig_for_server(port),
std::time::SystemTime::now() + Duration::from_secs(60),
);
let created = create_client_with_specific_context(Some(kubeconfig_path), "ctx")
.await
.expect("the edited kubeconfig reaches the fake API server");
let cached = CachedClient::new(created, signature());
let newer_connection = cached.connection.clone();
manager.clients.insert(key.clone(), cached);
drop(guard);
manager.release_creation_lock(&key, &lock);
let error = tokio::time::timeout(Duration::from_secs(5), waiter)
.await
.expect("a waiter given a shared failure must not hang")
.expect("waiter task must not panic")
.expect_err("the older caller is given the failure of its own kubeconfig state");
assert_eq!(
format!("{error:#}"),
"Failed to create client: cluster unreachable"
);
let retained = manager
.clients
.get(&key)
.expect("the newer client must stay cached")
.connection
.clone();
assert!(
Arc::ptr_eq(&retained, &newer_connection),
"the older caller's failure must not evict the newer client"
);
let reused = manager
.get_connection(key)
.await
.expect("a caller on the newer kubeconfig state reuses the cached client");
assert!(
Arc::ptr_eq(&reused, &newer_connection),
"a caller on the newer kubeconfig state must not recreate the client"
);
let version = reused
.client
.apiserver_version()
.await
.expect("the retained client still reaches the cluster");
assert_eq!(version.git_version, "v1.30.0");
}
#[tokio::test]
async fn a_caller_still_waiting_for_the_lock_leaves_a_stale_client_cached() {
use http::{
Request,
Response,
};
use kube::client::Body;
use tower_test::mock;
let manager = Arc::new(SharedClientManager::new());
let key = ServiceClientKey::new(
Some("ctx".to_string()),
Some(FAILING_KUBECONFIG.to_string()),
);
let (mock_service, _handle) = mock::pair::<Request<Body>, Response<Body>>();
let connection = KubeConnection {
client: kube::Client::new(mock_service, "default"),
cluster_url: "https://example.invalid".parse().unwrap(),
};
manager.clients.insert(
key.clone(),
CachedClient::new(connection, stale_signature()),
);
let lock = manager
.creation_locks
.entry(key.clone())
.or_insert_with(new_creation_lock)
.clone();
lock.interest.fetch_add(1, Ordering::SeqCst);
let guard = lock.failures.lock().await;
let waiter_manager = manager.clone();
let waiter_key = key.clone();
let waiter = tokio::spawn(async move { waiter_manager.get_connection(waiter_key).await });
wait_for_interest_count(&lock, 2).await;
assert!(
manager.clients.contains_key(&key),
"a caller that does not hold the creation lock must not evict the cached client"
);
drop(guard);
manager.release_creation_lock(&key, &lock);
tokio::time::timeout(Duration::from_secs(5), waiter)
.await
.expect("the waiter must not hang")
.expect("waiter task must not panic")
.expect_err("no client can be created from a nonexistent kubeconfig");
assert!(
!manager.clients.contains_key(&key),
"the caller that replaces a client built for another kubeconfig state must evict it"
);
}
#[tokio::test]
async fn a_caller_arriving_after_the_failure_tries_again() {
let manager = SharedClientManager::new();
let key = ServiceClientKey::new(
Some("ctx".to_string()),
Some(FAILING_KUBECONFIG.to_string()),
);
let lock = manager
.creation_locks
.entry(key.clone())
.or_insert_with(new_creation_lock)
.clone();
lock.interest.fetch_add(1, Ordering::SeqCst);
let mut failures = lock.failures.lock().await;
record_failure(
&lock,
&mut failures,
failing_signature(),
"cluster unreachable",
);
drop(failures);
let error = manager
.get_connection(key)
.await
.expect_err("no client can be created from a nonexistent kubeconfig");
assert!(
!error.to_string().contains("cluster unreachable"),
"a caller that registered after the failure must try again: {error}"
);
}
#[test]
fn clone_error_text_keeps_the_message_and_its_causes() {
let error = anyhow::anyhow!("root cause")
.context("middle")
.context("outer");
let copy = clone_error_text(&error);
assert_eq!(format!("{copy:#}"), format!("{error:#}"));
assert_eq!(copy.to_string(), "outer");
}
#[tokio::test]
async fn invalidate_client_with_explicit_context_drops_matching_resolved_context() {
use http::{
Request,
Response,
};
use kube::client::Body;
use tower_test::mock;
let manager = SharedClientManager::new();
let kubeconfig_path = Some("kftray-test-explicit-invalidate-kubeconfig".to_string());
let (mock_service, _handle) = mock::pair::<Request<Body>, Response<Body>>();
let connection = KubeConnection {
client: kube::Client::new(mock_service, "default"),
cluster_url: "https://example.invalid".parse().unwrap(),
};
let key = ServiceClientKey::new(Some("explicit-ctx".to_string()), kubeconfig_path.clone());
manager
.clients
.insert(key.clone(), CachedClient::new(connection, Vec::new()));
manager.resolved_contexts.insert(
kubeconfig_path.clone(),
("explicit-ctx".to_string(), Instant::now(), Vec::new()),
);
manager.invalidate_client(&key);
assert!(
!manager.clients.contains_key(&key),
"the invalidated client entry must be removed"
);
assert!(
!manager.resolved_contexts.contains_key(&kubeconfig_path),
"invalidating the context currently recorded as resolved must drop the recorded \
resolution too, so a later None-context call re-reads the kubeconfig instead of \
replaying the now-invalid context"
);
}
#[tokio::test]
async fn invalidate_client_with_explicit_context_leaves_unrelated_resolved_context() {
let manager = SharedClientManager::new();
let kubeconfig_path = Some("kftray-test-unrelated-kubeconfig".to_string());
manager.resolved_contexts.insert(
kubeconfig_path.clone(),
("current-ctx".to_string(), Instant::now(), Vec::new()),
);
manager.invalidate_client(&ServiceClientKey::new(
Some("some-other-ctx".to_string()),
kubeconfig_path.clone(),
));
assert!(
manager.resolved_contexts.contains_key(&kubeconfig_path),
"invalidating a context that is not the currently recorded resolution must not \
touch the recorded resolution"
);
}
#[tokio::test]
async fn resolve_key_re_resolves_after_ttl_expires() {
use std::io::Write;
let mut kubeconfig_file = tempfile::NamedTempFile::new().expect("create temp kubeconfig");
write!(
kubeconfig_file,
"apiVersion: v1\n\
kind: Config\n\
current-context: kftray-test-ttl-context\n\
contexts:\n\
- name: kftray-test-ttl-context\n \
context:\n \
cluster: kftray-test-cluster\n \
user: kftray-test-user\n\
clusters:\n\
- name: kftray-test-cluster\n \
cluster:\n \
server: https://127.0.0.1:1\n\
users:\n\
- name: kftray-test-user\n \
user: {{}}\n"
)
.expect("write temp kubeconfig");
let manager = SharedClientManager::with_ttl(Duration::from_millis(20));
let key = ServiceClientKey::new(
None,
Some(kubeconfig_file.path().to_string_lossy().to_string()),
);
let (first, _) = manager
.resolve_key(key.clone())
.await
.expect("initial resolution must succeed");
assert_eq!(
first.context_name.as_deref(),
Some("kftray-test-ttl-context")
);
manager.resolved_contexts.insert(
key.kubeconfig_path.clone(),
(
"stale-context".to_string(),
Instant::now() - Duration::from_millis(21),
Vec::new(),
),
);
let (second, _) = manager
.resolve_key(key)
.await
.expect("expired cache entry must fall back to a fresh kubeconfig read");
assert_eq!(
second.context_name.as_deref(),
Some("kftray-test-ttl-context"),
"an expired resolved-context cache entry must not be reused; resolve_key must \
re-read the kubeconfig instead of returning the stale recorded context"
);
}
async fn fake_api_server() -> u16 {
use tokio::io::{
AsyncReadExt,
AsyncWriteExt,
};
use tokio::net::TcpListener;
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let port = listener.local_addr().unwrap().port();
tokio::spawn(async move {
loop {
let Ok((mut socket, _)) = listener.accept().await else {
return;
};
tokio::spawn(async move {
let mut buf = [0u8; 4096];
let _ = socket.read(&mut buf).await;
let body = r#"{"major":"1","minor":"30","gitVersion":"v1.30.0","gitCommit":"x","gitTreeState":"clean","buildDate":"x","goVersion":"x","compiler":"gc","platform":"linux/amd64"}"#;
let response = format!(
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: \
{}\r\nconnection: close\r\n\r\n{}",
body.len(),
body
);
let _ = socket.write_all(response.as_bytes()).await;
});
}
});
port
}
fn kubeconfig_for_server(port: u16) -> String {
format!(
"apiVersion: v1\nkind: Config\ncurrent-context: ctx\ncontexts:\n- name: ctx\n \
context:\n cluster: c\n user: u\nclusters:\n- name: c\n cluster:\n \
server: http://127.0.0.1:{port}\nusers:\n- name: u\n user: {{}}\n"
)
}
#[tokio::test]
async fn explicit_context_client_follows_kubeconfig_edit() {
crate::ssl::ensure_crypto_provider_installed();
let a = fake_api_server().await;
let b = fake_api_server().await;
let file = tempfile::NamedTempFile::new().unwrap();
std::fs::write(file.path(), kubeconfig_for_server(a)).unwrap();
let manager = SharedClientManager::new();
let key = ServiceClientKey::new(
Some("ctx".to_string()),
Some(file.path().to_string_lossy().to_string()),
);
let first = manager.get_connection(key.clone()).await.unwrap();
assert_eq!(first.cluster_url.port_u16(), Some(a));
tokio::time::sleep(Duration::from_millis(1100)).await;
std::fs::write(file.path(), kubeconfig_for_server(b)).unwrap();
let second = manager.get_connection(key).await.unwrap();
assert_eq!(
second.cluster_url.port_u16(),
Some(b),
"served the cached cluster {} after the kubeconfig moved to port {b}",
second.cluster_url
);
}
}