From 0faa378e1ff9cf9ce0b5b08dac8520a8db49bf2d Mon Sep 17 00:00:00 2001
From: metamuffin <metamuffin@disroot.org>
Date: Sat, 24 Aug 2024 01:06:30 +0200
Subject: load balancing (untested)

---
 src/modules/loadbalance.rs | 58 ++++++++++++++++++++++++++++++++++++++++++++++
 src/modules/mod.rs         |  1 +
 2 files changed, 59 insertions(+)
 create mode 100644 src/modules/loadbalance.rs

diff --git a/src/modules/loadbalance.rs b/src/modules/loadbalance.rs
new file mode 100644
index 0000000..168db33
--- /dev/null
+++ b/src/modules/loadbalance.rs
@@ -0,0 +1,58 @@
+use super::{Node, NodeContext, NodeKind, NodeRequest, NodeResponse};
+use crate::{config::DynNode, error::ServiceError};
+use anyhow::Result;
+use serde::Deserialize;
+use serde_yaml::Value;
+use std::{
+    future::Future,
+    pin::Pin,
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc,
+    },
+};
+
+pub struct LoadBalanceKind;
+
+#[derive(Deserialize)]
+struct LoadBalanceConfig(Vec<DynNode>);
+
+struct LoadBalance {
+    load: Vec<AtomicUsize>,
+    config: LoadBalanceConfig,
+}
+
+impl NodeKind for LoadBalanceKind {
+    fn name(&self) -> &'static str {
+        "loadbalance"
+    }
+    fn instanciate(&self, config: Value) -> Result<Arc<dyn Node>> {
+        let config = serde_yaml::from_value::<LoadBalanceConfig>(config)?;
+        Ok(Arc::new(LoadBalance {
+            load: config.0.iter().map(|_| AtomicUsize::new(0)).collect(),
+            config,
+        }))
+    }
+}
+impl Node for LoadBalance {
+    fn handle<'a>(
+        &'a self,
+        context: &'a mut NodeContext,
+        request: NodeRequest,
+    ) -> Pin<Box<dyn Future<Output = Result<NodeResponse, ServiceError>> + Send + Sync + 'a>> {
+        Box::pin(async move {
+            let index = self
+                .load
+                .iter()
+                .enumerate()
+                .min_by_key(|(_, k)| k.load(Ordering::Relaxed))
+                .map(|(i, _)| i)
+                .ok_or(ServiceError::CustomStatic("zero routes to balance load"))?;
+
+            self.load[index].fetch_add(1, Ordering::Relaxed);
+            let resp = self.config.0[index].handle(context, request).await;
+            self.load[index].fetch_sub(1, Ordering::Relaxed);
+            resp
+        })
+    }
+}
diff --git a/src/modules/mod.rs b/src/modules/mod.rs
index 051299b..e0f74ec 100644
--- a/src/modules/mod.rs
+++ b/src/modules/mod.rs
@@ -20,6 +20,7 @@ pub mod paths;
 pub mod proxy;
 pub mod redirect;
 pub mod switch;
+pub mod loadbalance;
 
 pub type NodeRequest = Request<Incoming>;
 pub type NodeResponse = Response<BoxBody<Bytes, ServiceError>>;
-- 
cgit v1.2.3-70-g09d2