From ccb4e87658f622edbd57cc2b5a058c969643bfe2 Mon Sep 17 00:00:00 2001
From: Alex Auvolat <alex@adnab.me>
Date: Tue, 7 Dec 2021 16:37:22 +0100
Subject: [PATCH] Round robin backends

---
 src/https.rs         | 5 ++++-
 src/proxy_config.rs  | 9 ++++++++-
 src/reverse_proxy.rs | 6 ++----
 3 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/src/https.rs b/src/https.rs
index c80d51c..a62ebea 100644
--- a/src/https.rs
+++ b/src/https.rs
@@ -1,5 +1,5 @@
 use std::net::SocketAddr;
-use std::sync::Arc;
+use std::sync::{atomic::Ordering, Arc};
 
 use anyhow::Result;
 use log::*;
@@ -108,10 +108,13 @@ async fn handle(
 					.as_ref()
 					.map(|x| x.len() as i32)
 					.unwrap_or(0)),
+				ent.calls.load(Ordering::SeqCst),
 			)
 		});
 
 	if let Some(proxy_to) = ent {
+		proxy_to.calls.fetch_add(1, Ordering::SeqCst);
+
 		let to_addr = format!("http://{}", proxy_to.target_addr);
 		info!("Proxying {} {} -> {}", host, path, to_addr);
 
diff --git a/src/proxy_config.rs b/src/proxy_config.rs
index 891fdef..8c2444a 100644
--- a/src/proxy_config.rs
+++ b/src/proxy_config.rs
@@ -1,5 +1,5 @@
 use std::net::SocketAddr;
-use std::sync::Arc;
+use std::sync::{atomic, Arc};
 use std::{cmp, time::Duration};
 
 use log::*;
@@ -12,9 +12,15 @@ use crate::consul::*;
 #[derive(Debug)]
 pub struct ProxyEntry {
 	pub target_addr: SocketAddr,
+
 	pub host: String,
 	pub path_prefix: Option<String>,
 	pub priority: u32,
+
+	// Counts the number of times this proxy server has been called to
+	// This implements a round-robin load balancer if there are multiple
+	// entries for the same host and same path prefix.
+	pub calls: atomic::AtomicU64,
 }
 
 #[derive(Debug)]
@@ -53,6 +59,7 @@ fn parse_tricot_tag(target_addr: SocketAddr, tag: &str) -> Option<ProxyEntry> {
 		host: host.to_string(),
 		path_prefix,
 		priority,
+		calls: atomic::AtomicU64::from(0),
 	})
 }
 
diff --git a/src/reverse_proxy.rs b/src/reverse_proxy.rs
index 82533d8..046808f 100644
--- a/src/reverse_proxy.rs
+++ b/src/reverse_proxy.rs
@@ -72,10 +72,8 @@ fn create_proxied_request<B>(
 
 	// If request does not have host header, add it from original URI authority
 	if let Some(authority) = request.uri().authority() {
-		if let hyper::header::Entry::Vacant(entry) = builder
-			.headers_mut()
-			.unwrap()
-			.entry(host_header_name)
+		if let hyper::header::Entry::Vacant(entry) =
+			builder.headers_mut().unwrap().entry(host_header_name)
 		{
 			entry.insert(authority.as_str().parse()?);
 		}