From 706990d0f78439885ce4258b1da41b3514bac60a Mon Sep 17 00:00:00 2001
From: Simon Hornby <liquidwicked64@gmail.com>
Date: Wed, 24 Jan 2024 10:13:04 +0200
Subject: [PATCH] fix: include strategy variant stickiness (#403)

---
 Cargo.lock                                    | 36 ++++++++++--
 server/Cargo.toml                             |  4 +-
 server/src/builder.rs                         | 10 +++-
 server/src/client_api.rs                      | 57 ++++++++++++++++++-
 server/src/http/feature_refresher.rs          | 25 ++++----
 .../client_token_from_frontend_token.rs       |  2 +-
 server/src/offline/offline_hotload.rs         |  7 ++-
 7 files changed, 118 insertions(+), 23 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index e8073130..856b4d31 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1175,6 +1175,17 @@ dependencies = [
  "digest",
 ]
 
+[[package]]
+name = "hostname"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3c731c3e10504cc8ed35cfe2f1db4c9274c3d35fa486e3b31df46f068ef3e867"
+dependencies = [
+ "libc",
+ "match_cfg",
+ "winapi",
+]
+
 [[package]]
 name = "http"
 version = "0.2.11"
@@ -1339,6 +1350,15 @@ version = "2.9.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "8f518f335dce6725a761382244631d86cf0ccb2863413590b31338feb467f9c3"
 
+[[package]]
+name = "ipnetwork"
+version = "0.20.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bf466541e9d546596ee94f9f69590f89473455f88372423e0008fc1a7daf100e"
+dependencies = [
+ "serde",
+]
+
 [[package]]
 name = "is-terminal"
 version = "0.4.9"
@@ -1500,6 +1520,12 @@ version = "1.0.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "3e2e65a1a2e43cfcb47a895c4c8b10d1f4a61097f9f254f183aee60cad9c651d"
 
+[[package]]
+name = "match_cfg"
+version = "0.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ffbee8634e0d45d258acb448e7eaab3fce7a0a467395d4d9f228e3c1f01fb2e4"
+
 [[package]]
 name = "matchers"
 version = "0.1.0"
@@ -3024,9 +3050,9 @@ dependencies = [
 
 [[package]]
 name = "unleash-types"
-version = "0.10.6"
+version = "0.11.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1ff718fdcaecdcc682ef3751b2496670e064f4058408ecb4199f0bf065c27181"
+checksum = "b6db70f5a700ad12ebb8ab08483f81eea23de296d5a094d2164e442a5b701d80"
 dependencies = [
  "base64",
  "chrono",
@@ -3039,14 +3065,16 @@ dependencies = [
 
 [[package]]
 name = "unleash-yggdrasil"
-version = "0.8.0"
+version = "0.9.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "bd9eb99db51abff93ec3e1bff5150add9f42ecd510b6a2e5e59ded8b7f69e008"
+checksum = "37e4c35ce30ecda43c78d1f72ea55fffee3412a5a9c9ddea9de8db4efedaf03e"
 dependencies = [
  "chrono",
  "convert_case 0.6.0",
  "dashmap",
+ "hostname",
  "ipnet",
+ "ipnetwork",
  "lazy_static",
  "murmur3",
  "pest",
diff --git a/server/Cargo.toml b/server/Cargo.toml
index 4369df09..89144e69 100644
--- a/server/Cargo.toml
+++ b/server/Cargo.toml
@@ -79,8 +79,8 @@ tokio = { version = "1.35.1", features = [
 tracing = { version = "0.1.40", features = ["log"] }
 tracing-subscriber = { version = "0.3.18", features = ["json", "env-filter"] }
 ulid = "1.1.0"
-unleash-types = { version = "0.10", features = ["openapi", "hashes"] }
-unleash-yggdrasil = { version = "0.8.0" }
+unleash-types = { version = "0.11", features = ["openapi", "hashes"] }
+unleash-yggdrasil = { version = "0.9.0" }
 utoipa = { version = "4.2.0", features = ["actix_extras", "chrono"] }
 utoipa-swagger-ui = { version = "6", features = ["actix-web"] }
 [dev-dependencies]
diff --git a/server/src/builder.rs b/server/src/builder.rs
index ca1c0688..7ad36846 100644
--- a/server/src/builder.rs
+++ b/server/src/builder.rs
@@ -61,8 +61,14 @@ async fn hydrate_from_persistent_storage(
         tracing::debug!("Hydrating features for {key:?}");
         features_cache.insert(key.clone(), features.clone());
         let mut engine_state = EngineState::default();
-        engine_state.take_state(features);
-        engine_cache.insert(key, engine_state);
+
+        if engine_state.take_state(features).is_err() {
+            tracing::warn!(
+                "Loaded an invalid state from persistent storage, bootstrapping has failed"
+            );
+        } else {
+            engine_cache.insert(key, engine_state);
+        }
     }
 
     for target in refresh_targets {
diff --git a/server/src/client_api.rs b/server/src/client_api.rs
index 1196a18d..9c91107f 100644
--- a/server/src/client_api.rs
+++ b/server/src/client_api.rs
@@ -280,7 +280,9 @@ mod tests {
     use maplit::hashmap;
     use reqwest::StatusCode;
     use ulid::Ulid;
-    use unleash_types::client_features::{ClientFeature, Constraint, Operator, Strategy};
+    use unleash_types::client_features::{
+        ClientFeature, Constraint, Operator, Strategy, StrategyVariant,
+    };
     use unleash_types::client_metrics::{
         ClientMetricsEnv, ConnectViaBuilder, MetricBucket, ToggleStats,
     };
@@ -431,7 +433,12 @@ mod tests {
                     project: Some("default".into()),
                     strategies: Some(vec![
                         Strategy {
-                            variants: None,
+                            variants: Some(vec![StrategyVariant {
+                                name: "test".into(),
+                                payload: None,
+                                weight: 7,
+                                stickiness: Some("sticky-on-something".into()),
+                            }]),
                             name: "standard".into(),
                             sort_order: Some(500),
                             segments: None,
@@ -507,6 +514,52 @@ mod tests {
         }
     }
 
+    #[tokio::test]
+    async fn response_includes_variant_stickiness_for_strategy_variants() {
+        let features_cache: Arc<DashMap<String, ClientFeatures>> = Arc::new(DashMap::default());
+        let token_cache: Arc<DashMap<String, EdgeToken>> = Arc::new(DashMap::default());
+        let app = test::init_service(
+            App::new()
+                .app_data(Data::from(features_cache.clone()))
+                .app_data(Data::from(token_cache.clone()))
+                .service(web::scope("/api/client").service(get_features)),
+        )
+        .await;
+
+        features_cache.insert("production".into(), cached_client_features());
+        let mut production_token = EdgeToken::try_from(
+            "*:production.03fa5f506428fe80ed5640c351c7232e38940814d2923b08f5c05fa7".to_string(),
+        )
+        .unwrap();
+        production_token.token_type = Some(TokenType::Client);
+        production_token.status = TokenValidationStatus::Validated;
+        token_cache.insert(production_token.token.clone(), production_token.clone());
+        let req = make_features_request_with_token(production_token.clone()).await;
+        let res: ClientFeatures = test::call_and_read_body_json(&app, req).await;
+
+        assert_eq!(res.features.len(), cached_client_features().features.len());
+        let strategy_variant_stickiness = res
+            .features
+            .iter()
+            .find(|f| f.name == "feature_one")
+            .unwrap()
+            .strategies
+            .clone()
+            .unwrap()
+            .iter()
+            .find(|s| s.name == "standard")
+            .unwrap()
+            .variants
+            .clone()
+            .unwrap()
+            .iter()
+            .find(|v| v.name == "test")
+            .unwrap()
+            .stickiness
+            .clone();
+        assert!(strategy_variant_stickiness.is_some());
+    }
+
     #[tokio::test]
     async fn register_endpoint_correctly_aggregates_applications() {
         let metrics_cache = Arc::new(MetricsCache::default());
diff --git a/server/src/http/feature_refresher.rs b/server/src/http/feature_refresher.rs
index 738a3967..31a9cd49 100644
--- a/server/src/http/feature_refresher.rs
+++ b/server/src/http/feature_refresher.rs
@@ -5,7 +5,7 @@ use actix_web::http::header::EntityTag;
 use chrono::Utc;
 use dashmap::DashMap;
 use reqwest::StatusCode;
-use tracing::{debug, info};
+use tracing::{debug, error, info};
 use unleash_types::client_features::Segment;
 use unleash_types::client_metrics::ClientApplication;
 use unleash_types::{
@@ -353,13 +353,18 @@ impl FeatureRefresher {
                         .and_modify(|engine| {
                             if let Some(f) = self.features_cache.get(&key) {
                                 let mut new_state = EngineState::default();
-                                new_state.take_state(f.clone());
-                                *engine = new_state;
+                                if new_state.take_state(f.clone()).is_err() {
+                                    error!("Received a response from Unleash that cannot be compiled, this is likely a bug in Edge. Features will not be updated")
+                                } else {
+                                    *engine = new_state;
+                                }
                             }
                         })
                         .or_insert_with(|| {
                             let mut new_state = EngineState::default();
-                            new_state.take_state(features);
+                            if new_state.take_state(features).is_err() {
+                                error!("Received a response from Unleash that cannot be compiled, this is likely a bug in Edge. Features will not be updated")
+                            }
                             new_state
                         });
                 }
@@ -912,7 +917,7 @@ mod tests {
         let example_features = features_from_disk("../examples/features.json");
         let cache_key = cache_key(&token);
         let mut engine_state = EngineState::default();
-        engine_state.take_state(example_features.clone());
+        engine_state.take_state(example_features.clone()).unwrap();
         upstream_features_cache.insert(cache_key.clone(), example_features.clone());
         upstream_engine_cache.insert(cache_key.clone(), engine_state);
         let mut server = client_api_test_server(
@@ -968,7 +973,7 @@ mod tests {
         let example_features = features_from_disk("../examples/features.json");
         let cache_key = cache_key(&valid_token);
         let mut engine_state = EngineState::default();
-        engine_state.take_state(example_features.clone());
+        engine_state.take_state(example_features.clone()).unwrap();
         upstream_features_cache.insert(cache_key.clone(), example_features.clone());
         upstream_engine_cache.insert(cache_key.clone(), engine_state);
         let server = client_api_test_server(
@@ -1013,7 +1018,7 @@ mod tests {
         let example_features = features_from_disk("../examples/hostedexample.json");
         let cache_key = cache_key(&dx_token);
         let mut engine_state = EngineState::default();
-        engine_state.take_state(example_features.clone());
+        engine_state.take_state(example_features.clone()).unwrap();
         upstream_features_cache.insert(cache_key.clone(), example_features.clone());
         upstream_engine_cache.insert(cache_key.clone(), engine_state);
         let server = client_api_test_server(
@@ -1062,7 +1067,7 @@ mod tests {
         let cache_key = cache_key(&dx_token);
         upstream_features_cache.insert(cache_key.clone(), example_features.clone());
         let mut engine_state = EngineState::default();
-        engine_state.take_state(example_features.clone());
+        engine_state.take_state(example_features.clone()).unwrap();
         upstream_engine_cache.insert(cache_key, engine_state);
         let server = client_api_test_server(
             upstream_token_cache,
@@ -1123,7 +1128,7 @@ mod tests {
         let cache_key = cache_key(&dx_token);
         upstream_features_cache.insert(cache_key.clone(), example_features.clone());
         let mut engine_state = EngineState::default();
-        engine_state.take_state(example_features.clone());
+        engine_state.take_state(example_features.clone()).unwrap();
         upstream_engine_cache.insert(cache_key, engine_state);
         let server = client_api_test_server(
             upstream_token_cache,
@@ -1230,7 +1235,7 @@ mod tests {
         let cache_key = cache_key(&eg_token);
         upstream_features_cache.insert(cache_key.clone(), example_features.clone());
         let mut engine_state = EngineState::default();
-        engine_state.take_state(example_features.clone());
+        engine_state.take_state(example_features.clone()).unwrap();
         upstream_engine_cache.insert(cache_key.clone(), engine_state);
         let server = client_api_test_server(
             upstream_token_cache,
diff --git a/server/src/middleware/client_token_from_frontend_token.rs b/server/src/middleware/client_token_from_frontend_token.rs
index bbbfe6d5..9f82d79b 100644
--- a/server/src/middleware/client_token_from_frontend_token.rs
+++ b/server/src/middleware/client_token_from_frontend_token.rs
@@ -195,7 +195,7 @@ mod tests {
         upstream_features_cache.insert(client_token.environment.clone().unwrap(), features.clone());
         let upstream_engine_cache: Arc<DashMap<String, EngineState>> = Arc::new(DashMap::default());
         let mut engine = EngineState::default();
-        engine.take_state(features.clone());
+        engine.take_state(features.clone()).unwrap();
         upstream_engine_cache.insert(client_token.token.clone(), engine);
         let upstream_server = upstream_server(
             upstream_token_cache.clone(),
diff --git a/server/src/offline/offline_hotload.rs b/server/src/offline/offline_hotload.rs
index 2047dfe5..1f707251 100644
--- a/server/src/offline/offline_hotload.rs
+++ b/server/src/offline/offline_hotload.rs
@@ -63,8 +63,11 @@ pub(crate) fn load_offline_engine_cache(
         client_features.clone(),
     );
     let mut engine_state = EngineState::default();
-    engine_state.take_state(client_features);
-    engine_cache.insert(crate::tokens::cache_key(edge_token), engine_state);
+    if engine_state.take_state(client_features).is_err() {
+        tracing::warn!("Loaded an invalid feature set from file, likely due to a manual edit, reverting to previous state");
+    } else {
+        engine_cache.insert(crate::tokens::cache_key(edge_token), engine_state);
+    }
 }
 
 #[derive(Deserialize)]