fix: Update limiters on port modifications (#50)

* fix: Refill limiter with limit

* fix: Update limiters on port modifications

* Clear limiter map

* Remove consumer from limiter on port update

Author: gonzalezzfelipe

Diff for: proxy/src/auth.rs

@@ -1,18 +1,14 @@
 use futures_util::TryStreamExt;
 use operator::{
     kube::{
-        api::ListParams,
-        runtime::{
-            watcher::{self, Config},
-            WatchStreamExt,
-        },
+        runtime::watcher::{self, Config, Event},
         Api, Client, ResourceExt,
     },
     OgmiosPort,
 };
 use std::{collections::HashMap, sync::Arc};
 use tokio::pin;
-use tracing::{error, instrument};
+use tracing::{error, info, instrument};
 
 use crate::{Consumer, State};
 
@@ -24,49 +20,67 @@ pub fn start(state: Arc<State>) {
             .expect("failed to create kube client");
 
         let api = Api::<OgmiosPort>::all(client.clone());
-        update_auth(state.clone(), api.clone()).await;
-
-        let stream = watcher::watcher(api.clone(), Config::default()).touched_objects();
+        let stream = watcher::watcher(api.clone(), Config::default());
         pin!(stream);
 
         loop {
             let result = stream.try_next().await;
-            if let Err(err) = result {
-                error!(error = err.to_string(), "fail crd auth watcher");
-                continue;
-            }
+            match result {
+                // Stream restart, also run on startup.
+                Ok(Some(Event::Restarted(crds))) => {
+                    info!("auth: Watcher restarted, reseting consumers");
+                    let consumers: HashMap<String, Consumer> = crds
+                        .iter()
+                        .map(|crd| {
+                            let consumer = Consumer::from(crd);
+                            (consumer.hash_key.clone(), consumer)
+                        })
+                        .collect();
+                    *state.consumers.write().await = consumers;
 
-            update_auth(state.clone(), api.clone()).await;
+                    // When the watcher is restarted, we reset the limiter because a user
+                    // could have changed the tier on the watcher restart.
+                    state.limiter.write().await.clear();
+                }
+                // New port created or updated.
+                Ok(Some(Event::Applied(crd))) => match crd.status {
+                    Some(_) => {
+                        info!("auth: Adding new consumer: {}", crd.name_any());
+                        let consumer = Consumer::from(&crd);
+                        state.limiter.write().await.remove(&consumer.key);
+                        state
+                            .consumers
+                            .write()
+                            .await
+                            .insert(consumer.hash_key.clone(), consumer);
+                    }
+                    None => {
+                        // New ports are created without status. When the status is added, a new
+                        // Applied event is triggered.
+                        info!("auth: New port created: {}", crd.name_any());
+                    }
+                },
+                // Port deleted.
+                Ok(Some(Event::Deleted(crd))) => {
+                    info!(
+                        "auth: Port deleted, removing from state: {}",
+                        crd.name_any()
+                    );
+                    let consumer = Consumer::from(&crd);
+                    state.consumers.write().await.remove(&consumer.hash_key);
+                    state.limiter.write().await.remove(&consumer.key);
+                }
+                // Empty response from stream. Should never happen.
+                Ok(None) => {
+                    error!("auth: Empty response from watcher.");
+                    continue;
+                }
+                // Unexpected error when streaming CRDs.
+                Err(err) => {
+                    error!(error = err.to_string(), "auth: Failed to update crds.");
+                    std::process::exit(1);
+                }
+            }
         }
     });
 }
-
-async fn update_auth(state: Arc<State>, api: Api<OgmiosPort>) {
-    let result = api.list(&ListParams::default()).await;
-    if let Err(err) = result {
-        error!(
-            error = err.to_string(),
-            "error to get crds while updating auth keys"
-        );
-        return;
-    }
-
-    let mut consumers = HashMap::new();
-    for crd in result.unwrap().items.iter() {
-        if crd.status.is_some() {
-            let network = crd.spec.network.to_string();
-            let version = crd.spec.version;
-            let tier = crd.spec.throughput_tier.to_string();
-            let key = crd.status.as_ref().unwrap().auth_token.clone();
-            let namespace = crd.metadata.namespace.as_ref().unwrap().clone();
-            let port_name = crd.name_any();
-
-            let hash_key = format!("{}.{}.{}", network, version, key);
-            let consumer = Consumer::new(namespace, port_name, tier, key);
-
-            consumers.insert(hash_key, consumer);
-        }
-    }
-
-    *state.consumers.write().await = consumers;
-}

Diff for: proxy/src/main.rs

@@ -2,6 +2,7 @@ use config::Config;
 use dotenv::dotenv;
 use leaky_bucket::RateLimiter;
 use metrics::Metrics;
+use operator::{kube::ResourceExt, OgmiosPort};
 use prometheus::Registry;
 use regex::Regex;
 use std::collections::HashMap;
@@ -79,19 +80,29 @@ pub struct Consumer {
     port_name: String,
     tier: String,
     key: String,
+    hash_key: String,
 }
-impl Consumer {
-    pub fn new(namespace: String, port_name: String, tier: String, key: String) -> Self {
+impl Display for Consumer {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}.{}", self.namespace, self.port_name)
+    }
+}
+impl From<&OgmiosPort> for Consumer {
+    fn from(value: &OgmiosPort) -> Self {
+        let network = value.spec.network.to_string();
+        let version = value.spec.version;
+        let tier = value.spec.throughput_tier.to_string();
+        let key = value.status.as_ref().unwrap().auth_token.clone();
+        let namespace = value.metadata.namespace.as_ref().unwrap().clone();
+        let port_name = value.name_any();
+
+        let hash_key = format!("{}.{}.{}", network, version, key);
         Self {
             namespace,
             port_name,
             tier,
             key,
+            hash_key,
         }
     }
 }
-impl Display for Consumer {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{}.{}", self.namespace, self.port_name)
-    }
-}