From f163d1d348c78a27bb19a86a6158f4b5d1287f6f Mon Sep 17 00:00:00 2001
From: Alex Auvolat <alex@adnab.me>
Date: Mon, 2 Jan 2023 17:33:05 +0100
Subject: [PATCH] Don't do the hack with same_site and same_node, separate lb
 flags

---
 src/https.rs        |  8 +++----
 src/proxy_config.rs | 56 ++++++++++++++++++++++++++++++---------------
 2 files changed, 42 insertions(+), 22 deletions(-)

diff --git a/src/https.rs b/src/https.rs
index 3718bfd..49df11a 100644
--- a/src/https.rs
+++ b/src/https.rs
@@ -247,8 +247,8 @@ async fn select_target_and_proxy(
 					.as_ref()
 					.map(|x| x.len() as i32)
 					.unwrap_or(0),
-				ent.same_node,
-				ent.same_site,
+				(ent.flags.same_node || ent.flags.site_lb || ent.flags.global_lb),
+				(ent.flags.same_site || ent.flags.global_lb),
 				-ent.calls_in_progress.load(Ordering::SeqCst),
 				-ent.last_call.load(Ordering::SeqCst),
 			)
@@ -260,8 +260,8 @@ async fn select_target_and_proxy(
 			"target_addr",
 			proxy_to.target_addr.to_string(),
 		));
-		tags.push(KeyValue::new("same_node", proxy_to.same_node));
-		tags.push(KeyValue::new("same_site", proxy_to.same_site));
+		tags.push(KeyValue::new("same_node", proxy_to.flags.same_node));
+		tags.push(KeyValue::new("same_site", proxy_to.flags.same_site));
 
 		proxy_to.last_call.fetch_max(
 			(received_time - https_config.time_origin).as_millis() as i64,
diff --git a/src/proxy_config.rs b/src/proxy_config.rs
index af1f576..172bdb8 100644
--- a/src/proxy_config.rs
+++ b/src/proxy_config.rs
@@ -64,12 +64,8 @@ pub struct ProxyEntry {
 	/// Is the target serving HTTPS instead of HTTP?
 	pub https_target: bool,
 
-	/// Is the target the same node as we are running on?
-	/// (if yes priorize it over other matching targets)
-	pub same_node: bool,
-	/// Is the target the same site as this node?
-	/// (if yes priorize it over other matching targets)
-	pub same_site: bool,
+	/// Flags for target selection
+	pub flags: ProxyEntryFlags,
 
 	/// Add the following headers to all responses returned
 	/// when matching this rule
@@ -81,6 +77,21 @@ pub struct ProxyEntry {
 	pub last_call: atomic::AtomicI64,
 }
 
+#[derive(Debug, Clone, Copy)]
+pub struct ProxyEntryFlags {
+	/// Is the target the same node as we are running on?
+	/// (if yes priorize it over other matching targets)
+	pub same_node: bool,
+	/// Is the target the same site as this node?
+	/// (if yes priorize it over other matching targets)
+	pub same_site: bool,
+
+	/// Is site-wide load balancing enabled for this service?
+	pub site_lb: bool,
+	/// Is global load balancing enabled for this service?
+	pub global_lb: bool,
+}
+
 impl std::fmt::Display for ProxyEntry {
 	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
 		if self.https_target {
@@ -94,11 +105,16 @@ impl std::fmt::Display for ProxyEntry {
 			self.path_prefix.as_deref().unwrap_or_default(),
 			self.priority
 		)?;
-		if self.same_node {
+		if self.flags.same_node {
 			write!(f, " OURSELF")?;
-		} else if self.same_site {
+		} else if self.flags.same_site {
 			write!(f, " SAME_SITE")?;
 		}
+		if self.flags.global_lb {
+			write!(f, " GLOBAL-LB")?;
+		} else if self.flags.site_lb {
+			write!(f, " SITE-LB")?;
+		}
 		if !self.add_headers.is_empty() {
 			write!(f, " +Headers: {:?}", self.add_headers)?;
 		}
@@ -126,8 +142,7 @@ fn parse_tricot_tag(
 	tag: &str,
 	target_addr: SocketAddr,
 	add_headers: &[(String, String)],
-	same_node: bool,
-	same_site: bool,
+	flags: ProxyEntryFlags,
 ) -> Option<ProxyEntry> {
 	let splits = tag.split(' ').collect::<Vec<_>>();
 	if (splits.len() != 2 && splits.len() != 3)
@@ -162,8 +177,7 @@ fn parse_tricot_tag(
 		target_addr,
 		https_target: (splits[0] == "tricot-https"),
 		host,
-		same_node,
-		same_site,
+		flags,
 		path_prefix,
 		priority,
 		add_headers: add_headers.to_vec(),
@@ -206,12 +220,19 @@ fn parse_consul_catalog(
 		};
 		let addr = SocketAddr::new(ip_addr, svc.port);
 
-		let (same_node, same_site) = if svc.tags.contains(&"tricot-global-lb".into()) {
-			(false, false)
+		let (site_lb, global_lb) = if svc.tags.contains(&"tricot-global-lb".into()) {
+			(false, true)
 		} else if svc.tags.contains(&"tricot-site-lb".into()) {
-			(false, same_site)
+			(true, false)
 		} else {
-			(same_node, same_site)
+			(false, false)
+		};
+
+		let flags = ProxyEntryFlags {
+			same_node,
+			same_site,
+			site_lb,
+			global_lb,
 		};
 
 		let mut add_headers = vec![];
@@ -227,8 +248,7 @@ fn parse_consul_catalog(
 				tag,
 				addr,
 				&add_headers[..],
-				same_node,
-				same_site,
+				flags,
 			) {
 				entries.push(ent);
 			}