From 5c61552561d62430c4268ee7ee362e04c7e8f122 Mon Sep 17 00:00:00 2001
From: Carson McManus <dyc3@users.noreply.github.com>
Date: Fri, 3 Nov 2023 16:02:17 -0400
Subject: [PATCH] add `region` field to balancer protocol (#1124)

---
 crates/harness/src/monolith.rs               |  5 ++++-
 crates/ott-balancer-protocol/src/monolith.rs |  1 +
 server/balancer.ts                           |  2 ++
 server/generated.ts                          |  1 +
 server/ott-config.ts                         | 11 +++++++++++
 5 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/crates/harness/src/monolith.rs b/crates/harness/src/monolith.rs
index 1f7564ef0..adf8225d0 100644
--- a/crates/harness/src/monolith.rs
+++ b/crates/harness/src/monolith.rs
@@ -85,7 +85,10 @@ impl Monolith {
                 loop {
                     let (stream, _) = _listener.accept().await.unwrap();
                     let mut ws = tokio_tungstenite::accept_async(stream).await.unwrap();
-                    let init = M2BInit { port: http_port };
+                    let init = M2BInit {
+                        port: http_port,
+                        region: "unknown".into(),
+                    };
                     let msg = serde_json::to_string(&MsgM2B::Init(init)).unwrap();
                     ws.send(Message::Text(msg)).await.unwrap();
                     state.lock().unwrap().connected = true;
diff --git a/crates/ott-balancer-protocol/src/monolith.rs b/crates/ott-balancer-protocol/src/monolith.rs
index 1bea847fc..7c6f3b6ed 100644
--- a/crates/ott-balancer-protocol/src/monolith.rs
+++ b/crates/ott-balancer-protocol/src/monolith.rs
@@ -93,6 +93,7 @@ pub enum MsgM2B {
 pub struct M2BInit {
     /// The port that the monolith is listening for HTTP requests on.
     pub port: u16,
+    pub region: String,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
diff --git a/server/balancer.ts b/server/balancer.ts
index 8d0a44e39..b64cc38d5 100644
--- a/server/balancer.ts
+++ b/server/balancer.ts
@@ -75,6 +75,7 @@ class BalancerManager {
 			type: "init",
 			payload: {
 				port: conf.get("port"),
+				region: conf.get("balancing.region"),
 			},
 		};
 		conn.send(init);
@@ -170,6 +171,7 @@ export class BalancerConnection {
 			type: "init",
 			payload: {
 				port: conf.get("port"),
+				region: conf.get("balancing.region"),
 			},
 		};
 		this.send(init);
diff --git a/server/generated.ts b/server/generated.ts
index 714c06fa5..c5017f846 100644
--- a/server/generated.ts
+++ b/server/generated.ts
@@ -36,6 +36,7 @@ export interface B2MClientMsg<T = unknown> {
 export interface M2BInit {
 	/** The port that the monolith is listening for HTTP requests on. */
 	port: number;
+	region: string;
 }
 
 export enum Visibility {
diff --git a/server/ott-config.ts b/server/ott-config.ts
index 750ee2070..a7f6315b8 100644
--- a/server/ott-config.ts
+++ b/server/ott-config.ts
@@ -360,6 +360,12 @@ export const conf = convict({
 			default: 3002,
 			env: "BALANCING_PORT",
 		},
+		region: {
+			doc: "The region that this server is in.",
+			format: String,
+			default: "unknown",
+			env: "BALANCING_REGION",
+		},
 	},
 	mail: {
 		enabled: {
@@ -513,6 +519,11 @@ function postProcessConfig(): void {
 	if (conf.get("mail.enabled")) {
 		validateMail();
 	}
+
+	if (process.env.FLY_REGION) {
+		log.info("Found FLY_REGION. Using it for balancing.region.");
+		conf.set("balancing.region", process.env.FLY_REGION);
+	}
 }
 
 function validateMail() {