coredns_configurator.rs

coredns_configurator.rs · 4.7 KiB · Rust Raw

use futures::{StreamExt, TryStreamExt};
use k8s_openapi::api::core::v1::Service;
use kube::{
    Client,
    api::{Api, ListParams, ResourceExt, WatchEvent, WatchParams},
};
use std::collections::HashMap;
use std::sync::Arc;

const LOCAL_DOMAINS_LABEL: &'static str = "dns.tensor.ru/local";
const LOCAL_DOMAINS_ANNOTATION: &'static str = "dns.tensor.ru/local-domains";
const HOSTS_DIR: &'static str = "/data";

use tokio::fs::OpenOptions;
use tokio::io::AsyncWriteExt;
use tokio::sync::Mutex;

struct Informer {
    domains: Mutex<HashMap<String, Vec<String>>>,
    client: Arc<Client>,
}

impl Informer {
    async fn add_service(&mut self, service: Service) {
        let annotations = service.annotations();
        if let Some(local_domains) = annotations.get(LOCAL_DOMAINS_ANNOTATION) {
            let domains_list: Vec<String> = local_domains
                .split_whitespace()
                .map(|a| a.to_string())
                .collect();
            if let Some(spec) = service.spec {
                if let Some(cluster_ip) = spec.cluster_ip
                    && domains_list.len() > 0
                {
                    println!("Service added {} -> {:?}", cluster_ip, domains_list);
                    let mut lock = self.domains.lock().await;
                    lock.insert(cluster_ip, domains_list);
                }
            }
        }
    }

async fn delete_service(&mut self, service: Service) {
        if let Some(spec) = service.spec {
            if let Some(cluster_ip) = spec.cluster_ip {
                let mut lock = self.domains.lock().await;
                match lock.remove(&cluster_ip) {
                    Some(domains) => {
                        println!("Service deleted {} -> {:?}", cluster_ip, domains);
                    }
                    None => {}
                }
            }
        }
    }

pub async fn fill(&mut self) {
        let client: Client = (*self.client).clone();
        let services: Api<Service> = Api::all(client);
        let lp = ListParams {
            label_selector: Some(LOCAL_DOMAINS_LABEL.to_string()),
            ..Default::default()
        };
        while let Ok(services_list) = services.list(&lp).await {
            for service in services_list {
                self.add_service(service).await;
            }
        }
        let _ = self.write_hosts_file().await;
    }

pub async fn get_cache(&self) -> HashMap<String, Vec<String>> {
        let lock = self.domains.lock().await;
        lock.clone()
    }

pub async fn watch(mut self) -> Result<(), Box<dyn std::error::Error>> {
        let client: Client = (*self.client).clone();
        let services: Api<Service> = Api::all(client);
        let wp = WatchParams {
            label_selector: Some("dns.tensor.ru/local".to_string()),
            ..Default::default()
        };
        let mut stream = services.watch(&wp, "0").await?.boxed();

while let Some(status) = stream.try_next().await? {
            match status {
                WatchEvent::Added(s) | WatchEvent::Modified(s) => {
                    self.add_service(s).await;
                    self.write_hosts_file().await?;
                }
                WatchEvent::Deleted(s) => {
                    self.delete_service(s).await;
                    self.write_hosts_file().await?;
                }
                WatchEvent::Bookmark(s) => {}
                WatchEvent::Error(s) => println!("{}", s),
            }
        }

Ok(())
    }

async fn write_hosts_file(&self) -> Result<(), Box<dyn std::error::Error>> {
        let mut file = OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .open(format!("{}/hosts", HOSTS_DIR))
            .await?;

// Собираем все строки в один буфер
        let mut content = String::new();
        let lock = self.domains.lock().await;
        for (k, v) in lock.iter() {
            content.push_str(&format!("{} {}\n", k, v.join(" ")));
        }
        file.write_all(content.as_bytes()).await?;
        file.flush().await?;
        Ok(())
    }
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let mut informer = Informer {
        domains: Mutex::new(HashMap::new()),
        client: Arc::new(Client::try_default().await?),
    };
    informer.fill().await;
    println!("{:?}", informer.get_cache().await);

let handle = tokio::spawn(async move {
        if let Err(e) = informer.watch().await {
            eprintln!("Watcher error: {}", e);
        }
    });

// Ждем выполнения фоновой задачи (в этот момент main не завершится)
    handle.await?;

Ok(())
}

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	}