chore: Limit by connection (#57)

* chore: Limit by connection

* Update error message and proxy config

* Reinstate logger message

* Increase max connections for tier 3

* Fix 429 metrics

Author: gonzalezzfelipe

Diff for: bootstrap/proxy/config.tf

@@ -1,54 +1,42 @@
 locals {
   tiers = [
     {
-      "name" = "0",
+      "name"            = "0",
+      "max_connections" = 2
       "rates" = [
         {
           "interval" = "1m",
-          "limit"    = floor(60 / var.replicas)
-        },
-        {
-          "interval" = "1d",
-          "limit"    = floor(60 * 60 * 24 / var.replicas)
+          "limit"    = 500
         }
       ]
     },
     {
-      "name" = "1",
+      "name"            = "1",
+      "max_connections" = 5
       "rates" = [
         {
           "interval" = "1m",
-          "limit"    = floor(300 / var.replicas)
-        },
-        {
-          "interval" = "1d",
-          "limit"    = floor(300 * 60 * 24 / var.replicas)
+          "limit"    = 500
         }
       ]
     },
     {
-      "name" = "2",
+      "name"            = "2",
+      "max_connections" = 250
       "rates" = [
         {
           "interval" = "1m",
-          "limit"    = floor(2400 / var.replicas)
-        },
-        {
-          "interval" = "1d",
-          "limit"    = floor(2400 * 60 * 24 / var.replicas)
+          "limit"    = 500
         }
       ]
     },
     {
-      "name" = "3",
+      "name"            = "3",
+      "max_connections" = 250
       "rates" = [
         {
           "interval" = "1m",
-          "limit"    = floor(4800 / var.replicas)
-        },
-        {
-          "interval" = "1d",
-          "limit"    = floor(4800 * 60 * 24 / var.replicas)
+          "limit"    = 500
         }
       ]
     }

Diff for: proxy/src/main.rs

@@ -80,6 +80,7 @@ pub struct Consumer {
     key: String,
     network: String,
     version: String,
+    active_connections: usize,
 }
 impl Display for Consumer {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
@@ -102,6 +103,34 @@ impl From<&OgmiosPort> for Consumer {
             key,
             network,
             version,
+            active_connections: 0,
         }
     }
 }
+impl Consumer {
+    pub async fn inc_connections(&self, state: Arc<State>) {
+        state
+            .consumers
+            .write()
+            .await
+            .entry(self.key.clone())
+            .and_modify(|consumer| consumer.active_connections += 1);
+    }
+    pub async fn dec_connections(&self, state: Arc<State>) {
+        state
+            .consumers
+            .write()
+            .await
+            .entry(self.key.clone())
+            .and_modify(|consumer| consumer.active_connections -= 1);
+    }
+    pub async fn get_active_connections(&self, state: Arc<State>) -> usize {
+        state
+            .consumers
+            .read()
+            .await
+            .get(&self.key)
+            .map(|consumer| consumer.active_connections)
+            .unwrap_or_default()
+    }
+}

Diff for: proxy/src/proxy.rs

@@ -110,7 +110,29 @@ async fn handle(
             let proxy_req = proxy_req_result.unwrap();
             let response_result = match proxy_req.protocol {
                 Protocol::Http => handle_http(hyper_req, &proxy_req).await,
-                Protocol::Websocket => handle_websocket(hyper_req, &proxy_req, state.clone()).await,
+                Protocol::Websocket => {
+                    // Before handling the websocket connection, check if consumer has available
+                    // connections.
+                    let tiers = state.tiers.read().await.clone();
+                    match tiers.get(&proxy_req.consumer.tier) {
+                        Some(tier) => {
+                            if proxy_req.consumer.active_connections >= tier.max_connections {
+                                Ok(Response::builder()
+                                    .status(StatusCode::TOO_MANY_REQUESTS)
+                                    .body(full("Connection limit exceeded"))
+                                    .unwrap())
+                            } else {
+                                handle_websocket(hyper_req, &proxy_req, state.clone()).await
+                            }
+                        }
+                        None => Ok(Response::builder()
+                            .status(StatusCode::INTERNAL_SERVER_ERROR)
+                            .body(full(
+                                "Invalid tier value. Contact support team for more information.",
+                            ))
+                            .unwrap()),
+                    }
+                }
             };
 
             match &response_result {
@@ -187,14 +209,24 @@ async fn handle_websocket(
                 let (mut instance_outgoing, instance_incoming) = instance_stream.split();
 
                 state.metrics.inc_ws_total_connection(&proxy_req);
+                proxy_req.consumer.inc_connections(state.clone()).await;
+
+                let active_connections = proxy_req
+                    .consumer
+                    .get_active_connections(state.clone())
+                    .await;
+                info!(
+                    consumer = proxy_req.consumer.to_string(),
+                    active_connections, "client connected"
+                );
 
                 let client_in = async {
                     while let Some(result) = client_incoming.next().await {
                         match result {
                             Ok(data) => {
                                 if let Err(err) = limiter(state.clone(), &proxy_req.consumer).await
                                 {
-                                    error!(error = err.to_string(), "Failed to run limiter");
+                                    error!(error = err.to_string(), "Failed to run limiter.");
                                     break;
                                 };
                                 if let Err(err) = instance_outgoing.send(data).await {
@@ -221,6 +253,16 @@ async fn handle_websocket(
                 select(client_in, instance_in).await;
 
                 state.metrics.dec_ws_total_connection(&proxy_req);
+                proxy_req.consumer.dec_connections(state.clone()).await;
+
+                let active_connections = proxy_req
+                    .consumer
+                    .get_active_connections(state.clone())
+                    .await;
+                info!(
+                    consumer = proxy_req.consumer.to_string(),
+                    active_connections, "client disconnected"
+                );
             }
             Err(err) => {
                 error!(error = err.to_string(), "upgrade error");

Diff for: proxy/src/tiers.rs

@@ -12,6 +12,7 @@ use crate::State;
 pub struct Tier {
     pub name: String,
     pub rates: Vec<TierRate>,
+    pub max_connections: usize,
 }
 #[derive(Debug, Clone, Deserialize)]
 pub struct TierRate {