130 lines
3.4 KiB
Rust
130 lines
3.4 KiB
Rust
use reqwest::Client;
|
|
use serde::Serialize;
|
|
use std::collections::HashMap;
|
|
use std::env;
|
|
use std::fs;
|
|
use std::net::ToSocketAddrs;
|
|
use std::path::Path;
|
|
use uuid::Uuid;
|
|
use warp::Filter;
|
|
|
|
const VERSION: &'static str = env!("CARGO_PKG_VERSION");
|
|
|
|
#[derive(Serialize)]
|
|
struct ConsulRegistration {
|
|
id: String,
|
|
name: String,
|
|
address: String,
|
|
port: u16,
|
|
tags: Vec<String>,
|
|
meta: HashMap<String, String>,
|
|
check: ConsulCheck,
|
|
}
|
|
|
|
#[derive(Serialize)]
|
|
struct ConsulCheck {
|
|
http: String,
|
|
interval: String,
|
|
}
|
|
|
|
pub fn generate_service_id() -> String {
|
|
Uuid::new_v4().to_string()
|
|
}
|
|
pub fn get_or_generate_service_id(package_name: &str) -> String {
|
|
// let package_name = env!("CARGO_PKG_NAME");
|
|
let file_name = format!("/services/{}_service_id.txt", package_name);
|
|
let path = Path::new(&file_name);
|
|
let _ = fs::create_dir_all("/services"); // make sure the folders exist
|
|
|
|
if path.exists() {
|
|
// Read the service ID from the file
|
|
if let Ok(service_id) = fs::read_to_string(path) {
|
|
return service_id.trim().to_string();
|
|
}
|
|
}
|
|
|
|
// Generate a new service ID and save it to disk
|
|
let service_id = generate_service_id();
|
|
if let Err(err) = fs::write(path, &service_id) {
|
|
eprintln!("Failed to write service ID to disk: {}", err);
|
|
}
|
|
|
|
service_id
|
|
}
|
|
|
|
pub async fn register_service(
|
|
consul_url: &str,
|
|
service_id: &str,
|
|
service_name: &str,
|
|
service_address: &str,
|
|
service_port: u16,
|
|
tags: Vec<String>,
|
|
mut meta: HashMap<String, String>,
|
|
health_check_url: &str,
|
|
) -> Result<(), Box<dyn std::error::Error>> {
|
|
meta.insert("version".to_string(), VERSION.to_string());
|
|
let registration = ConsulRegistration {
|
|
id: service_id.to_string(),
|
|
name: service_name.to_string(),
|
|
address: service_address.to_string(),
|
|
port: service_port,
|
|
tags,
|
|
meta,
|
|
check: ConsulCheck {
|
|
http: health_check_url.to_string(),
|
|
interval: "10s".to_string(), // Health check interval
|
|
},
|
|
};
|
|
|
|
let client = Client::new();
|
|
let consul_register_url = format!("{}/v1/agent/service/register", consul_url);
|
|
|
|
client
|
|
.put(&consul_register_url)
|
|
.json(®istration)
|
|
.send()
|
|
.await?
|
|
.error_for_status()?; // Ensure response is successful
|
|
|
|
Ok(())
|
|
}
|
|
|
|
pub async fn deregister_service(
|
|
consul_url: &str,
|
|
service_id: &str,
|
|
) -> Result<(), Box<dyn std::error::Error>> {
|
|
let client = Client::new();
|
|
let consul_deregister_url =
|
|
format!("{}/v1/agent/service/deregister/{}", consul_url, service_id);
|
|
|
|
client
|
|
.put(&consul_deregister_url)
|
|
.send()
|
|
.await?
|
|
.error_for_status()?; // Ensure response is successful
|
|
|
|
Ok(())
|
|
}
|
|
|
|
pub async fn start_health_check(service_address: &str) -> Result<(), Box<dyn std::error::Error>> {
|
|
let health_port = env::var("HEALTH_CHECK_PORT").unwrap_or_else(|_| "8082".to_string());
|
|
let health_check_endpoint_addr = format!("{}:{}", service_address, health_port);
|
|
|
|
// Start health-check endpoint
|
|
let log = warp::log("health_check");
|
|
let health_route = warp::path!("health")
|
|
.map(|| warp::reply::with_status("OK", warp::http::StatusCode::OK))
|
|
.with(log);
|
|
|
|
tokio::spawn(
|
|
warp::serve(health_route).run(
|
|
health_check_endpoint_addr
|
|
.to_socket_addrs()?
|
|
.next()
|
|
.unwrap(),
|
|
),
|
|
);
|
|
|
|
Ok(())
|
|
}
|