use std::sync::Arc;
use std::time::{
Duration,
Instant,
};
use anyhow::{
anyhow,
Result,
};
use futures::StreamExt;
use k8s_openapi::api::core::v1::{
Pod,
Service,
};
use kube::{
Api,
Client,
ResourceExt,
};
use kube_runtime::{
reflector::{
self,
ReflectHandle,
Store,
},
watcher::{
self,
Config as WatcherConfig,
},
WatchStreamExt,
};
use tokio::sync::RwLock;
use tokio::task::JoinHandle;
use tracing::{
debug,
error,
info,
};
use crate::kube::models::{
Port,
Target,
TargetPod,
TargetSelector,
};
pub struct PodWatcher {
store: Store<Pod>,
subscriber: ReflectHandle<Pod>,
latest_ready_pod: Arc<RwLock<Option<TargetPod>>>,
_reflector_task: JoinHandle<()>,
_subscriber_task: JoinHandle<()>,
}
impl Drop for PodWatcher {
fn drop(&mut self) {
self._reflector_task.abort();
self._subscriber_task.abort();
debug!("PodWatcher background tasks aborted");
}
}
impl PodWatcher {
pub async fn new(client: Client, target: Target) -> Result<Self> {
let namespace = target.namespace.name_any();
let label_selector = Self::resolve_label_selector(&client, &namespace, &target).await?;
info!(
"Starting pod watcher for namespace '{}' with labels: {}",
namespace, label_selector
);
let (store, writer) = reflector::store_shared(256);
let subscriber = writer
.subscribe()
.ok_or_else(|| anyhow!("Failed to create subscriber"))?;
let latest_ready_pod = Arc::new(RwLock::new(None));
let latest_pod_clone = latest_ready_pod.clone();
let target_clone = target.clone();
let namespace_clone = namespace.clone();
let pods_api: Api<Pod> = Api::namespaced(client, &namespace);
let watcher_config = WatcherConfig::default().labels(&label_selector);
let reflector_task = tokio::spawn(async move {
let stream = watcher::watcher(pods_api, watcher_config)
.default_backoff()
.modify(|pod| {
pod.managed_fields_mut().clear();
pod.annotations_mut().clear();
if let Some(status) = &mut pod.status {
status.container_statuses = None;
status.init_container_statuses = None;
status.ephemeral_container_statuses = None;
}
})
.reflect_shared(writer)
.applied_objects();
let mut stream = std::pin::pin!(stream);
while let Some(result) = stream.next().await {
match result {
Ok(pod) => {
debug!("Pod reflected: {}", pod.name_any());
}
Err(e) => {
error!("Reflector error: {}", e);
tokio::time::sleep(Duration::from_secs(1)).await;
}
}
}
});
let subscriber_clone = subscriber.clone();
let subscriber_task = tokio::spawn(async move {
let filtered_stream = subscriber_clone.filter_map(move |pod| {
let namespace = namespace_clone.clone();
async move {
if pod.namespace().as_deref() != Some(&namespace) {
return None;
}
Some(pod)
}
});
let mut stream = std::pin::pin!(filtered_stream);
while let Some(pod) = stream.next().await {
Self::update_latest_pod(&latest_pod_clone, &pod, &target_clone).await;
}
});
Ok(Self {
store,
subscriber,
latest_ready_pod,
_reflector_task: reflector_task,
_subscriber_task: subscriber_task,
})
}
async fn update_latest_pod(
latest_ready_pod: &Arc<RwLock<Option<TargetPod>>>, pod: &Pod, target: &Target,
) {
if Self::is_pod_ready(pod) {
if let Ok(port_number) = Self::extract_port_from_pod(pod, &target.port) {
let target_pod = TargetPod {
pod_name: pod.name_any(),
port_number,
};
let mut latest = latest_ready_pod.write().await;
*latest = Some(target_pod);
debug!("Updated latest ready pod: {}", pod.name_any());
}
} else {
let mut latest = latest_ready_pod.write().await;
if let Some(current) = latest.as_ref() {
if current.pod_name == pod.name_any() {
*latest = None;
debug!("Cleared no-longer-ready pod: {}", pod.name_any());
}
}
}
}
pub async fn get_ready_pod(&self) -> Option<TargetPod> {
self.latest_ready_pod.read().await.clone()
}
pub async fn wait_for_ready_pod(&self, timeout: Duration) -> Option<TargetPod> {
let start = Instant::now();
while start.elapsed() < timeout {
if let Some(pod) = self.get_ready_pod().await {
return Some(pod);
}
tokio::time::sleep(Duration::from_millis(100)).await;
}
None
}
pub fn get_store(&self) -> Store<Pod> {
self.store.clone()
}
pub fn create_subscriber(&self) -> ReflectHandle<Pod> {
self.subscriber.clone()
}
async fn resolve_label_selector(
client: &Client, namespace: &str, target: &Target,
) -> Result<String> {
match &target.selector {
TargetSelector::PodLabel(label_selector) => Ok(label_selector.clone()),
TargetSelector::ServiceName(service_name) => {
let service_api: Api<Service> = Api::namespaced(client.clone(), namespace);
let service = service_api
.get(service_name)
.await
.map_err(|e| anyhow!("Service '{}' not found: {}", service_name, e))?;
let selector = service
.spec
.as_ref()
.and_then(|spec| spec.selector.as_ref())
.ok_or_else(|| anyhow!("Service '{}' has no selector", service_name))?;
let mut label_selector = String::with_capacity(selector.len() * 20);
let mut first = true;
for (k, v) in selector {
if !first {
label_selector.push(',');
}
label_selector.push_str(k);
label_selector.push('=');
label_selector.push_str(v);
first = false;
}
Ok(label_selector)
}
}
}
fn is_pod_ready(pod: &Pod) -> bool {
pod.status
.as_ref()
.map(|status| {
let is_running = status.phase.as_deref() == Some("Running");
if let Some(name) = &pod.metadata.name {
if name.starts_with("kftray-forward-") && is_running {
return true;
}
}
let is_ready = status
.conditions
.as_ref()
.map(|conditions| {
conditions
.iter()
.any(|c| c.type_ == "Ready" && c.status == "True")
})
.unwrap_or(false);
is_running && is_ready
})
.unwrap_or(false)
}
fn extract_port_from_pod(pod: &Pod, target_port: &Port) -> Result<u16> {
match target_port {
Port::Number(num) => (*num)
.try_into()
.map_err(|_| anyhow!("Invalid port number: {}", num)),
Port::Name(port_name) => pod
.spec
.as_ref()
.and_then(|spec| {
spec.containers
.iter()
.filter_map(|container| container.ports.as_ref())
.flatten()
.find(|p| p.name.as_deref() == Some(port_name))
.map(|p| p.container_port as u16)
})
.ok_or_else(|| anyhow!("Port '{}' not found in pod", port_name)),
}
}
}