add region field to balancer protocol (#1124)

dyc3 · Nov 3, 2023 · 5c61552 · 5c61552
1 parent f5a8664
commit 5c61552
Show file tree

Hide file tree

Showing 5 changed files with 19 additions and 1 deletion.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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() {