fix: Only replace projects covered by token (#466)

* fix: Only replace projects covered by token

Previously, if the update was empty, we assumed we needed to replace all
stored features. This is not the case. We only need to remove the
projects that are connected to the token we're using for the merge.

In addition, debug output on `/internal-backstage/tokens` now includes how many features was received in
the update

server/src/auth/token_validator.rs

@@ -172,7 +172,7 @@ mod tests {
     use actix_http::HttpService;
     use actix_http_test::{test_server, TestServer};
     use actix_service::map_config;
-    use actix_web::{App, dev::AppConfig, HttpResponse, web};
+    use actix_web::{dev::AppConfig, web, App, HttpResponse};
     use dashmap::DashMap;
     use serde::{Deserialize, Serialize};
 

server/src/builder.rs

@@ -46,10 +46,7 @@ fn build_caches() -> CacheContainer {
     )
 }
 
-async fn hydrate_from_persistent_storage(
-    cache: CacheContainer,
-    storage: Arc<dyn EdgePersistence>,
-) {
+async fn hydrate_from_persistent_storage(cache: CacheContainer, storage: Arc<dyn EdgePersistence>) {
     let (token_cache, features_cache, engine_cache) = cache;
     let tokens = storage.load_tokens().await.unwrap_or_else(|error| {
         warn!("Failed to load tokens from cache {error:?}");

server/src/edge_api.rs

@@ -1,7 +1,7 @@
 use actix_web::{
-    HttpRequest,
     post,
     web::{self, Data, Json},
+    HttpRequest,
 };
 use dashmap::DashMap;
 use utoipa;
@@ -58,9 +58,9 @@ pub fn configure_edge_api(cfg: &mut web::ServiceConfig) {
 mod tests {
     use std::sync::Arc;
 
-    use actix_web::{App, test, web};
     use actix_web::http::header::ContentType;
     use actix_web::web::Json;
+    use actix_web::{test, web, App};
     use dashmap::DashMap;
 
     use crate::auth::token_validator::TokenValidator;

server/src/http/background_send_metrics.rs

@@ -5,14 +5,14 @@ use actix_web::http::StatusCode;
 use chrono::Duration;
 use dashmap::DashMap;
 use lazy_static::lazy_static;
-use prometheus::{IntGauge, IntGaugeVec, Opts, register_int_gauge, register_int_gauge_vec};
+use prometheus::{register_int_gauge, register_int_gauge_vec, IntGauge, IntGaugeVec, Opts};
 use tracing::{error, info, trace, warn};
 
+use crate::types::TokenRefresh;
 use crate::{
     error::EdgeError,
-    metrics::client_metrics::{MetricsCache, size_of_batch},
+    metrics::client_metrics::{size_of_batch, MetricsCache},
 };
-use crate::types::TokenRefresh;
 
 use super::feature_refresher::FeatureRefresher;
 

server/src/http/feature_refresher.rs

@@ -1,27 +1,27 @@
-use std::{sync::Arc, time::Duration};
 use std::collections::HashSet;
+use std::{sync::Arc, time::Duration};
 
 use actix_web::http::header::EntityTag;
 use chrono::Utc;
 use dashmap::DashMap;
 use reqwest::StatusCode;
 use tracing::{debug, info, warn};
+use unleash_types::client_features::Segment;
+use unleash_types::client_metrics::ClientApplication;
 use unleash_types::{
     client_features::{ClientFeature, ClientFeatures},
     Deduplicate,
 };
-use unleash_types::client_features::Segment;
-use unleash_types::client_metrics::ClientApplication;
 use unleash_yggdrasil::EngineState;
 
+use crate::error::{EdgeError, FeatureError};
+use crate::filters::{filter_client_features, FeatureFilterSet};
+use crate::types::{build, EdgeResult, TokenType, TokenValidationStatus};
 use crate::{
     persistence::EdgePersistence,
     tokens::{cache_key, simplify},
     types::{ClientFeaturesRequest, ClientFeaturesResponse, EdgeToken, TokenRefresh},
 };
-use crate::error::{EdgeError, FeatureError};
-use crate::filters::{FeatureFilterSet, filter_client_features};
-use crate::types::{build, EdgeResult, TokenType, TokenValidationStatus};
 
 use super::unleash_client::UnleashClient;
 
@@ -70,29 +70,25 @@ fn merge_segments_update(
         (None, None) => None,
     }
 }
-fn update_projects_from_feature_update(
+pub(crate) fn update_projects_from_feature_update(
     token: &EdgeToken,
     original: &[ClientFeature],
     updated: &[ClientFeature],
 ) -> Vec<ClientFeature> {
-    if updated.is_empty() {
-        vec![]
+    let projects_to_update = &token.projects;
+    if projects_to_update.contains(&"*".into()) {
+        updated.into()
     } else {
-        let projects_to_update = &token.projects;
-        if projects_to_update.contains(&"*".into()) {
-            updated.into()
-        } else {
-            let mut to_keep: Vec<ClientFeature> = original
-                .iter()
-                .filter(|toggle| {
-                    let p = toggle.project.clone().unwrap_or_else(|| "default".into());
-                    !projects_to_update.contains(&p)
-                })
-                .cloned()
-                .collect();
-            to_keep.extend(updated.iter().cloned());
-            to_keep
-        }
+        let mut to_keep: Vec<ClientFeature> = original
+            .iter()
+            .filter(|toggle| {
+                let p = toggle.project.clone().unwrap_or_else(|| "default".into());
+                !projects_to_update.contains(&p)
+            })
+            .cloned()
+            .collect();
+        to_keep.extend(updated.iter().cloned());
+        to_keep
     }
 }
 
@@ -258,8 +254,7 @@ impl FeatureRefresher {
     /// Registers a token for refresh, the token will be discarded if it can be subsumed by another previously registered token
     pub async fn register_token_for_refresh(&self, token: EdgeToken, etag: Option<EntityTag>) {
         if !self.tokens_to_refresh.contains_key(&token.token) {
-            self
-                .unleash_client
+            self.unleash_client
                 .register_as_client(
                     token.token.clone(),
                     client_application_from_token(
@@ -324,7 +319,7 @@ impl FeatureRefresher {
                 ClientFeaturesResponse::Updated(features, etag) => {
                     debug!("Got updated client features. Updating features with {etag:?}");
                     let key = cache_key(&refresh.token);
-                    self.update_last_refresh(&refresh.token, etag);
+                    self.update_last_refresh(&refresh.token, etag, features.features.len());
                     self.features_cache
                         .entry(key.clone())
                         .and_modify(|existing_data| {
@@ -416,10 +411,15 @@ impl FeatureRefresher {
             });
     }
 
-    pub fn update_last_refresh(&self, token: &EdgeToken, etag: Option<EntityTag>) {
+    pub fn update_last_refresh(
+        &self,
+        token: &EdgeToken,
+        etag: Option<EntityTag>,
+        feature_count: usize,
+    ) {
         self.tokens_to_refresh
             .alter(&token.token, |_k, old_refresh| {
-                old_refresh.successful_refresh(&self.refresh_interval, etag)
+                old_refresh.successful_refresh(&self.refresh_interval, etag, feature_count)
             });
     }
 }
@@ -432,26 +432,26 @@ mod tests {
     use actix_http::HttpService;
     use actix_http_test::{test_server, TestServer};
     use actix_service::map_config;
-    use actix_web::{App, web};
     use actix_web::dev::AppConfig;
     use actix_web::http::header::EntityTag;
+    use actix_web::{web, App};
     use chrono::{Duration, Utc};
     use dashmap::DashMap;
     use reqwest::Url;
     use unleash_types::client_features::{ClientFeature, ClientFeatures};
     use unleash_yggdrasil::EngineState;
 
+    use crate::filters::{project_filter, FeatureFilterSet};
+    use crate::tests::features_from_disk;
+    use crate::tokens::cache_key;
+    use crate::types::TokenValidationStatus::Validated;
+    use crate::types::{TokenType, TokenValidationStatus};
     use crate::{
         http::unleash_client::UnleashClient,
         types::{EdgeToken, TokenRefresh},
     };
-    use crate::filters::{FeatureFilterSet, project_filter};
-    use crate::tests::features_from_disk;
-    use crate::tokens::cache_key;
-    use crate::types::{TokenType, TokenValidationStatus};
-    use crate::types::TokenValidationStatus::Validated;
 
-    use super::{FeatureRefresher, frontend_token_is_covered_by_tokens};
+    use super::{frontend_token_is_covered_by_tokens, FeatureRefresher};
 
     impl PartialEq for TokenRefresh {
         fn eq(&self, other: &Self) -> bool {
@@ -790,6 +790,7 @@ mod tests {
             last_refreshed: None,
             last_check: None,
             failure_count: 0,
+            last_feature_count: None,
         };
         let etag_and_last_refreshed_token =
             EdgeToken::try_from("projectb:development.etag_and_last_refreshed_token".to_string())
@@ -801,6 +802,7 @@ mod tests {
             last_refreshed: Some(Utc::now()),
             last_check: Some(Utc::now()),
             failure_count: 0,
+            last_feature_count: None,
         };
         let etag_but_old_token =
             EdgeToken::try_from("projectb:development.etag_but_old_token".to_string()).unwrap();
@@ -813,6 +815,7 @@ mod tests {
             last_refreshed: Some(ten_seconds_ago),
             last_check: Some(ten_seconds_ago),
             failure_count: 0,
+            last_feature_count: None,
         };
         feature_refresher.tokens_to_refresh.insert(
             etag_but_last_refreshed_ten_seconds_ago.token.token.clone(),
@@ -1199,6 +1202,7 @@ mod tests {
             last_refreshed: None,
             last_check: None,
             failure_count: 0,
+            last_feature_count: None,
         };
 
         current_tokens.insert(wildcard_token.token, token_refresh);
@@ -1420,4 +1424,103 @@ mod tests {
         assert_eq!(updated.len(), 1);
         assert!(updated.iter().all(|f| f.project == Some("default".into())))
     }
+
+    #[test]
+    pub fn token_with_access_to_different_project_than_exists_in_cache_should_never_delete_features_from_other_projects(
+    ) {
+        // Added after customer issue in May '24 when tokens unrelated to projects in cache with no actual features connected to them removed existing features in cache for unrelated projects
+        let features = vec![
+            ClientFeature {
+                name: "my.first.toggle.in.default".to_string(),
+                feature_type: Some("release".into()),
+                description: None,
+                created_at: None,
+                last_seen_at: None,
+                enabled: true,
+                stale: None,
+                impression_data: None,
+                project: Some("testproject".into()),
+                strategies: None,
+                variants: None,
+                dependencies: None,
+            },
+            ClientFeature {
+                name: "my.second.toggle.in.testproject".to_string(),
+                feature_type: Some("release".into()),
+                description: None,
+                created_at: None,
+                last_seen_at: None,
+                enabled: false,
+                stale: None,
+                impression_data: None,
+                project: Some("testproject".into()),
+                strategies: None,
+                variants: None,
+                dependencies: None,
+            },
+        ];
+        let empty_features = vec![];
+        let unrelated_token_to_existing_features = EdgeToken {
+            token: "someotherproject:dev.myextralongsecretstringwithfeatures".to_string(),
+            token_type: Some(TokenType::Client),
+            environment: Some("dev".into()),
+            projects: vec![String::from("someother")],
+            status: TokenValidationStatus::Validated,
+        };
+        let updated = super::update_projects_from_feature_update(
+            &unrelated_token_to_existing_features,
+            &features,
+            &empty_features,
+        );
+        assert_eq!(updated.len(), 2);
+    }
+    #[test]
+    pub fn token_with_access_to_both_a_different_project_than_exists_in_cache_and_the_cached_project_should_delete_features_from_both_projects(
+    ) {
+        // Added after customer issue in May '24 when tokens unrelated to projects in cache with no actual features connected to them removed existing features in cache for unrelated projects
+        let features = vec![
+            ClientFeature {
+                name: "my.first.toggle.in.default".to_string(),
+                feature_type: Some("release".into()),
+                description: None,
+                created_at: None,
+                last_seen_at: None,
+                enabled: true,
+                stale: None,
+                impression_data: None,
+                project: Some("testproject".into()),
+                strategies: None,
+                variants: None,
+                dependencies: None,
+            },
+            ClientFeature {
+                name: "my.second.toggle.in.testproject".to_string(),
+                feature_type: Some("release".into()),
+                description: None,
+                created_at: None,
+                last_seen_at: None,
+                enabled: false,
+                stale: None,
+                impression_data: None,
+                project: Some("testproject".into()),
+                strategies: None,
+                variants: None,
+                dependencies: None,
+            },
+        ];
+        let empty_features = vec![];
+        let token_with_access_to_both_empty_and_full_project = EdgeToken {
+            token: "[]:dev.myextralongsecretstringwithfeatures".to_string(),
+            token_type: Some(TokenType::Client),
+            environment: Some("dev".into()),
+            projects: vec![String::from("testproject"), String::from("someother")],
+            status: TokenValidationStatus::Validated,
+        };
+        let updated = super::update_projects_from_feature_update(
+            &token_with_access_to_both_empty_and_full_project,
+            &features,
+            &empty_features,
+        );
+        assert_eq!(updated.len(), 0);
+    }
 }