load balancing (untested)

2 files changed, 59 insertions, 0 deletions

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>>;