| @@ -0,0 +1,143 @@ | |||
| 1 | + | use futures::{StreamExt, TryStreamExt}; | |
| 2 | + | use k8s_openapi::api::core::v1::Service; | |
| 3 | + | use kube::{ | |
| 4 | + | Client, | |
| 5 | + | api::{Api, ListParams, ResourceExt, WatchEvent, WatchParams}, | |
| 6 | + | }; | |
| 7 | + | use std::collections::HashMap; | |
| 8 | + | use std::sync::Arc; | |
| 9 | + | ||
| 10 | + | const LOCAL_DOMAINS_LABEL: &'static str = "dns.tensor.ru/local"; | |
| 11 | + | const LOCAL_DOMAINS_ANNOTATION: &'static str = "dns.tensor.ru/local-domains"; | |
| 12 | + | const HOSTS_DIR: &'static str = "/data"; | |
| 13 | + | ||
| 14 | + | use tokio::fs::OpenOptions; | |
| 15 | + | use tokio::io::AsyncWriteExt; | |
| 16 | + | use tokio::sync::Mutex; | |
| 17 | + | ||
| 18 | + | struct Informer { | |
| 19 | + | domains: Mutex<HashMap<String, Vec<String>>>, | |
| 20 | + | client: Arc<Client>, | |
| 21 | + | } | |
| 22 | + | ||
| 23 | + | impl Informer { | |
| 24 | + | async fn add_service(&mut self, service: Service) { | |
| 25 | + | let annotations = service.annotations(); | |
| 26 | + | if let Some(local_domains) = annotations.get(LOCAL_DOMAINS_ANNOTATION) { | |
| 27 | + | let domains_list: Vec<String> = local_domains | |
| 28 | + | .split_whitespace() | |
| 29 | + | .map(|a| a.to_string()) | |
| 30 | + | .collect(); | |
| 31 | + | if let Some(spec) = service.spec { | |
| 32 | + | if let Some(cluster_ip) = spec.cluster_ip | |
| 33 | + | && domains_list.len() > 0 | |
| 34 | + | { | |
| 35 | + | println!("Service added {} -> {:?}", cluster_ip, domains_list); | |
| 36 | + | let mut lock = self.domains.lock().await; | |
| 37 | + | lock.insert(cluster_ip, domains_list); | |
| 38 | + | } | |
| 39 | + | } | |
| 40 | + | } | |
| 41 | + | } | |
| 42 | + | ||
| 43 | + | async fn delete_service(&mut self, service: Service) { | |
| 44 | + | if let Some(spec) = service.spec { | |
| 45 | + | if let Some(cluster_ip) = spec.cluster_ip { | |
| 46 | + | let mut lock = self.domains.lock().await; | |
| 47 | + | match lock.remove(&cluster_ip) { | |
| 48 | + | Some(domains) => { | |
| 49 | + | println!("Service deleted {} -> {:?}", cluster_ip, domains); | |
| 50 | + | } | |
| 51 | + | None => {} | |
| 52 | + | } | |
| 53 | + | } | |
| 54 | + | } | |
| 55 | + | } | |
| 56 | + | ||
| 57 | + | pub async fn fill(&mut self) { | |
| 58 | + | let client: Client = (*self.client).clone(); | |
| 59 | + | let services: Api<Service> = Api::all(client); | |
| 60 | + | let lp = ListParams { | |
| 61 | + | label_selector: Some(LOCAL_DOMAINS_LABEL.to_string()), | |
| 62 | + | ..Default::default() | |
| 63 | + | }; | |
| 64 | + | while let Ok(services_list) = services.list(&lp).await { | |
| 65 | + | for service in services_list { | |
| 66 | + | self.add_service(service).await; | |
| 67 | + | } | |
| 68 | + | } | |
| 69 | + | let _ = self.write_hosts_file().await; | |
| 70 | + | } | |
| 71 | + | ||
| 72 | + | pub async fn get_cache(&self) -> HashMap<String, Vec<String>> { | |
| 73 | + | let lock = self.domains.lock().await; | |
| 74 | + | lock.clone() | |
| 75 | + | } | |
| 76 | + | ||
| 77 | + | pub async fn watch(mut self) -> Result<(), Box<dyn std::error::Error>> { | |
| 78 | + | let client: Client = (*self.client).clone(); | |
| 79 | + | let services: Api<Service> = Api::all(client); | |
| 80 | + | let wp = WatchParams { | |
| 81 | + | label_selector: Some("dns.tensor.ru/local".to_string()), | |
| 82 | + | ..Default::default() | |
| 83 | + | }; | |
| 84 | + | let mut stream = services.watch(&wp, "0").await?.boxed(); | |
| 85 | + | ||
| 86 | + | while let Some(status) = stream.try_next().await? { | |
| 87 | + | match status { | |
| 88 | + | WatchEvent::Added(s) | WatchEvent::Modified(s) => { | |
| 89 | + | self.add_service(s).await; | |
| 90 | + | self.write_hosts_file().await?; | |
| 91 | + | } | |
| 92 | + | WatchEvent::Deleted(s) => { | |
| 93 | + | self.delete_service(s).await; | |
| 94 | + | self.write_hosts_file().await?; | |
| 95 | + | } | |
| 96 | + | WatchEvent::Bookmark(s) => {} | |
| 97 | + | WatchEvent::Error(s) => println!("{}", s), | |
| 98 | + | } | |
| 99 | + | } | |
| 100 | + | ||
| 101 | + | Ok(()) | |
| 102 | + | } | |
| 103 | + | ||
| 104 | + | async fn write_hosts_file(&self) -> Result<(), Box<dyn std::error::Error>> { | |
| 105 | + | let mut file = OpenOptions::new() | |
| 106 | + | .write(true) | |
| 107 | + | .create(true) | |
| 108 | + | .truncate(true) | |
| 109 | + | .open(format!("{}/hosts", HOSTS_DIR)) | |
| 110 | + | .await?; | |
| 111 | + | ||
| 112 | + | // Собираем все строки в один буфер | |
| 113 | + | let mut content = String::new(); | |
| 114 | + | let lock = self.domains.lock().await; | |
| 115 | + | for (k, v) in lock.iter() { | |
| 116 | + | content.push_str(&format!("{} {}\n", k, v.join(" "))); | |
| 117 | + | } | |
| 118 | + | file.write_all(content.as_bytes()).await?; | |
| 119 | + | file.flush().await?; | |
| 120 | + | Ok(()) | |
| 121 | + | } | |
| 122 | + | } | |
| 123 | + | ||
| 124 | + | #[tokio::main] | |
| 125 | + | async fn main() -> Result<(), Box<dyn std::error::Error>> { | |
| 126 | + | let mut informer = Informer { | |
| 127 | + | domains: Mutex::new(HashMap::new()), | |
| 128 | + | client: Arc::new(Client::try_default().await?), | |
| 129 | + | }; | |
| 130 | + | informer.fill().await; | |
| 131 | + | println!("{:?}", informer.get_cache().await); | |
| 132 | + | ||
| 133 | + | let handle = tokio::spawn(async move { | |
| 134 | + | if let Err(e) = informer.watch().await { | |
| 135 | + | eprintln!("Watcher error: {}", e); | |
| 136 | + | } | |
| 137 | + | }); | |
| 138 | + | ||
| 139 | + | // Ждем выполнения фоновой задачи (в этот момент main не завершится) | |
| 140 | + | handle.await?; | |
| 141 | + | ||
| 142 | + | Ok(()) | |
| 143 | + | } | |
aquafire / coredns_configurator.rs
Last active 6 months ago
aquafire revised this gist 11 months ago · 48045a6
1 file changed, 143 insertions