286 lines
9.3 KiB
Rust
286 lines
9.3 KiB
Rust
use futures::{StreamExt, TryStreamExt};
|
|
use k8s_openapi::api::core::v1::Pod;
|
|
use kube::{
|
|
api::Api,
|
|
runtime::watcher::{watcher, Config},
|
|
runtime::WatchStreamExt,
|
|
Client,
|
|
};
|
|
use std::collections::{HashMap, HashSet};
|
|
use std::sync::{Arc, Mutex, RwLock};
|
|
use std::time::Duration;
|
|
use tokio::task;
|
|
use tokio::time;
|
|
use tracing::{error, info, warn};
|
|
|
|
/// Represents the service discovery configuration
|
|
#[derive(Debug, Clone)]
|
|
pub struct ServiceDiscoveryConfig {
|
|
pub enabled: bool,
|
|
pub selector: HashMap<String, String>,
|
|
pub check_interval: Duration,
|
|
pub port: u16,
|
|
pub namespace: Option<String>,
|
|
}
|
|
|
|
impl Default for ServiceDiscoveryConfig {
|
|
fn default() -> Self {
|
|
ServiceDiscoveryConfig {
|
|
enabled: false,
|
|
selector: HashMap::new(),
|
|
check_interval: Duration::from_secs(60),
|
|
port: 80, // Default port to connect to pods
|
|
namespace: None, // None means watch all namespaces
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Represents a Kubernetes pod's information used for worker management
|
|
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
|
|
pub struct PodInfo {
|
|
pub name: String,
|
|
pub ip: String,
|
|
pub status: String,
|
|
pub is_ready: bool,
|
|
}
|
|
|
|
impl PodInfo {
|
|
pub fn from_pod(pod: &Pod) -> Option<Self> {
|
|
let name = pod.metadata.name.clone()?;
|
|
let status = pod.status.clone()?;
|
|
let pod_ip = status.pod_ip?;
|
|
|
|
let is_ready = if let Some(conditions) = &status.conditions {
|
|
conditions
|
|
.iter()
|
|
.any(|condition| condition.type_ == "Ready" && condition.status == "True")
|
|
} else {
|
|
false
|
|
};
|
|
|
|
let pod_status = status.phase.unwrap_or_else(|| "Unknown".to_string());
|
|
|
|
Some(PodInfo {
|
|
name,
|
|
ip: pod_ip,
|
|
status: pod_status,
|
|
is_ready,
|
|
})
|
|
}
|
|
|
|
/// Returns true if the pod is in a state where it can accept traffic
|
|
pub fn is_healthy(&self) -> bool {
|
|
self.is_ready && self.status == "Running"
|
|
}
|
|
|
|
/// Generates a worker URL for this pod
|
|
pub fn worker_url(&self, port: u16) -> String {
|
|
format!("http://{}:{}", self.ip, port)
|
|
}
|
|
}
|
|
|
|
pub async fn start_service_discovery(
|
|
config: ServiceDiscoveryConfig,
|
|
worker_urls: Arc<RwLock<Vec<String>>>,
|
|
) -> Result<task::JoinHandle<()>, kube::Error> {
|
|
// Don't initialize anything if service discovery is disabled
|
|
if !config.enabled {
|
|
// Return a generic error when service discovery is disabled
|
|
return Err(kube::Error::Api(kube::error::ErrorResponse {
|
|
status: "Disabled".to_string(),
|
|
message: "Service discovery is disabled".to_string(),
|
|
reason: "ConfigurationError".to_string(),
|
|
code: 400,
|
|
}));
|
|
}
|
|
|
|
// Initialize Kubernetes client
|
|
let client = Client::try_default().await?;
|
|
|
|
// Construct label selector string from map
|
|
let label_selector = config
|
|
.selector
|
|
.iter()
|
|
.map(|(k, v)| format!("{}={}", k, v))
|
|
.collect::<Vec<_>>()
|
|
.join(",");
|
|
|
|
info!(
|
|
"Starting Kubernetes service discovery with selector: {}",
|
|
label_selector
|
|
);
|
|
|
|
// Create the task that will run in the background
|
|
let handle = task::spawn(async move {
|
|
// We'll track pods we've already added to avoid duplicates
|
|
let tracked_pods = Arc::new(Mutex::new(HashSet::new()));
|
|
|
|
// Create a watcher for pods
|
|
let pods: Api<Pod> = if let Some(namespace) = &config.namespace {
|
|
Api::namespaced(client, namespace)
|
|
} else {
|
|
Api::all(client)
|
|
};
|
|
|
|
info!("Kubernetes service discovery initialized successfully");
|
|
|
|
// Create an Arc for the selector map
|
|
let selector = Arc::new(config.selector);
|
|
let port = config.port;
|
|
|
|
loop {
|
|
// Create a watcher with the proper parameters according to the kube-rs API
|
|
let watcher_config = Config::default();
|
|
let watcher_stream = watcher(pods.clone(), watcher_config).applied_objects();
|
|
|
|
// Clone Arcs for the closures
|
|
let selector_clone = Arc::clone(&selector);
|
|
let tracked_pods_clone = Arc::clone(&tracked_pods);
|
|
let worker_urls_clone = Arc::clone(&worker_urls);
|
|
|
|
// Apply label selector filter separately since we can't do it directly with the watcher anymore
|
|
let filtered_stream = watcher_stream.filter_map(move |obj_res| {
|
|
let selector_inner = Arc::clone(&selector_clone);
|
|
|
|
async move {
|
|
match obj_res {
|
|
Ok(pod) => {
|
|
// Only process pods matching our label selector
|
|
if pod.metadata.labels.as_ref().map_or(false, |labels| {
|
|
// Check if the pod has all the labels from our selector
|
|
selector_inner.iter().all(|(k, v)| {
|
|
labels.get(k).map_or(false, |label_value| label_value == v)
|
|
})
|
|
}) {
|
|
Some(Ok(pod))
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
Err(e) => Some(Err(e)),
|
|
}
|
|
}
|
|
});
|
|
|
|
// Clone again for the next closure
|
|
let tracked_pods_clone2 = Arc::clone(&tracked_pods_clone);
|
|
let worker_urls_clone2 = Arc::clone(&worker_urls_clone);
|
|
|
|
match filtered_stream
|
|
.try_for_each(move |pod| {
|
|
let tracked_pods_inner = Arc::clone(&tracked_pods_clone2);
|
|
let worker_urls_inner = Arc::clone(&worker_urls_clone2);
|
|
|
|
async move {
|
|
if let Some(pod_info) = PodInfo::from_pod(&pod) {
|
|
if pod.metadata.deletion_timestamp.is_some() {
|
|
handle_pod_deletion(
|
|
&pod_info,
|
|
tracked_pods_inner,
|
|
worker_urls_inner,
|
|
port,
|
|
)
|
|
.await;
|
|
} else {
|
|
handle_pod_event(
|
|
&pod_info,
|
|
tracked_pods_inner,
|
|
worker_urls_inner,
|
|
port,
|
|
)
|
|
.await;
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
})
|
|
.await
|
|
{
|
|
Ok(_) => {}
|
|
Err(err) => {
|
|
error!("Error in Kubernetes watcher: {}", err);
|
|
// Wait a bit before retrying
|
|
time::sleep(Duration::from_secs(5)).await;
|
|
}
|
|
}
|
|
|
|
// If the watcher exits for some reason, wait a bit before restarting
|
|
warn!(
|
|
"Kubernetes watcher exited, restarting in {} seconds",
|
|
config.check_interval.as_secs()
|
|
);
|
|
time::sleep(config.check_interval).await;
|
|
}
|
|
});
|
|
|
|
Ok(handle)
|
|
}
|
|
|
|
async fn handle_pod_event(
|
|
pod_info: &PodInfo,
|
|
tracked_pods: Arc<Mutex<HashSet<PodInfo>>>,
|
|
worker_urls: Arc<RwLock<Vec<String>>>,
|
|
port: u16,
|
|
) {
|
|
let worker_url = pod_info.worker_url(port);
|
|
|
|
// Check if pod is already tracked
|
|
let already_tracked = {
|
|
let tracker = tracked_pods.lock().unwrap();
|
|
tracker.contains(pod_info)
|
|
};
|
|
|
|
// If pod is healthy and not already tracked, add it
|
|
if pod_info.is_healthy() {
|
|
if !already_tracked {
|
|
info!(
|
|
"Adding healthy pod {} ({}) as worker",
|
|
pod_info.name, pod_info.ip
|
|
);
|
|
|
|
// Add URL to worker list
|
|
let mut urls = worker_urls.write().unwrap();
|
|
if !urls.contains(&worker_url) {
|
|
urls.push(worker_url.clone());
|
|
info!("Added new worker URL: {}", worker_url);
|
|
}
|
|
|
|
// Track this pod
|
|
let mut tracker = tracked_pods.lock().unwrap();
|
|
tracker.insert(pod_info.clone());
|
|
}
|
|
} else if already_tracked {
|
|
// If pod was healthy before but not anymore, remove it
|
|
handle_pod_deletion(pod_info, tracked_pods, worker_urls, port).await;
|
|
}
|
|
}
|
|
|
|
async fn handle_pod_deletion(
|
|
pod_info: &PodInfo,
|
|
tracked_pods: Arc<Mutex<HashSet<PodInfo>>>,
|
|
worker_urls: Arc<RwLock<Vec<String>>>,
|
|
port: u16,
|
|
) {
|
|
let worker_url = pod_info.worker_url(port);
|
|
|
|
// Remove the pod from our tracking
|
|
let was_tracked = {
|
|
let mut tracker = tracked_pods.lock().unwrap();
|
|
tracker.remove(pod_info)
|
|
};
|
|
|
|
if was_tracked {
|
|
info!(
|
|
"Removing pod {} ({}) from workers",
|
|
pod_info.name, pod_info.ip
|
|
);
|
|
|
|
// Remove URL from worker list
|
|
let mut urls = worker_urls.write().unwrap();
|
|
if let Some(idx) = urls.iter().position(|url| url == &worker_url) {
|
|
urls.remove(idx);
|
|
info!("Removed worker URL: {}", worker_url);
|
|
}
|
|
}
|
|
}
|