harmony/harmony-rs/harmony/src/domain/topology/load_balancer.rs

use std::{net::SocketAddr, str::FromStr};

use async_trait::async_trait;
use log::debug;

use super::{IpAddress, LogicalHost};
use crate::executors::ExecutorError;

#[async_trait]
pub trait LoadBalancer: Send + Sync {
    fn get_ip(&self) -> IpAddress;
    fn get_host(&self) -> LogicalHost;
    async fn add_service(&self, service: &LoadBalancerService) -> Result<(), ExecutorError>;
    async fn remove_service(&self, service: &LoadBalancerService) -> Result<(), ExecutorError>;
    async fn list_services(&self) -> Vec<LoadBalancerService>;
    async fn ensure_initialized(&self) -> Result<(), ExecutorError>;
    async fn commit_config(&self) -> Result<(), ExecutorError>;
    async fn reload_restart(&self) -> Result<(), ExecutorError>;
    async fn ensure_service_exists(
        &self,
        service: &LoadBalancerService,
    ) -> Result<(), ExecutorError> {
        debug!(
            "Listing LoadBalancer services {:?}",
            self.list_services().await
        );
        if !self.list_services().await.contains(service) {
            self.add_service(service).await?;
        }
        Ok(())
    }
}

impl std::fmt::Debug for dyn LoadBalancer {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_fmt(format_args!("LoadBalancer {}", self.get_ip()))
    }
}
#[derive(Debug, PartialEq, Clone)]
pub struct LoadBalancerService {
    pub backend_servers: Vec<BackendServer>,
    pub listening_port: SocketAddr,
    pub health_check: Option<HealthCheck>,
}

#[derive(Debug, PartialEq, Clone)]
pub struct BackendServer {
    pub address: String,
    pub port: u16,
}

#[derive(Debug, Clone, PartialEq)]
pub enum HttpMethod {
    GET,
    POST,
    PUT,
    PATCH,
    DELETE,
}

impl From<String> for HttpMethod {
    fn from(value: String) -> Self {
        Self::from_str(&value).unwrap()
    }
}

impl FromStr for HttpMethod {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_uppercase().as_str() {
            "GET" => Ok(HttpMethod::GET),
            "POST" => Ok(HttpMethod::POST),
            "PUT" => Ok(HttpMethod::PUT),
            "PATCH" => Ok(HttpMethod::PATCH),
            "DELETE" => Ok(HttpMethod::DELETE),
            _ => Err(format!("Invalid HTTP method: {}", s)),
        }
    }
}

impl std::fmt::Display for HttpMethod {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            HttpMethod::GET => write!(f, "GET"),
            HttpMethod::POST => write!(f, "POST"),
            HttpMethod::PUT => write!(f, "PUT"),
            HttpMethod::PATCH => write!(f, "PATCH"),
            HttpMethod::DELETE => write!(f, "DELETE"),
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
pub enum HttpStatusCode {
    Success2xx,
    UserError4xx,
    ServerError5xx,
}

#[derive(Debug, Clone, PartialEq)]
pub enum HealthCheck {
    HTTP(String, HttpMethod, HttpStatusCode),
    TCP(Option<u16>),
}