feat(featureflags): get some metrics wired up (#25021)

Co-authored-by: Dylan Martin <[email protected]>
PostHog · Sep 18, 2024 · f7eb6e3 · f7eb6e3
1 parent 8ba6871
commit f7eb6e3
Show file tree

Hide file tree

Showing 15 changed files with 703 additions and 960 deletions.
diff --git a/rust/Cargo.lock b/rust/Cargo.lock
diff --git a/rust/capture/src/exceptions.rs b/rust/capture/src/exceptions.rs
@@ -0,0 +1,11 @@
+pub trait Exception {
+    fn fingerprint(&self) -> String;
+}
+
+pub trait Stacked {
+    type Frame;
+
+    fn raw_ident(&self) -> String;
+    fn stack(&self) -> Vec<Self::Frame>;
+    fn lang_hint(&self) -> String;
+}
diff --git a/rust/feature-flags/Cargo.toml b/rust/feature-flags/Cargo.toml
@@ -34,6 +34,9 @@ uuid = { workspace = true }
 base64.workspace = true
 flate2.workspace = true
 common-alloc = { path = "../common/alloc" }
+health = { path = "../common/health" }
+common-metrics = { path = "../common/metrics" }
+tower = { workspace = true }
 
 [lints]
 workspace = true

diff --git a/rust/feature-flags/src/api.rs b/rust/feature-flags/src/api.rs
@@ -201,7 +201,6 @@ impl From<CustomRedisError> for FlagError {
 impl From<CustomDatabaseError> for FlagError {
     fn from(e: CustomDatabaseError) -> Self {
         match e {
-            CustomDatabaseError::NotFound => FlagError::TokenValidationError,
             CustomDatabaseError::Other(_) => {
                 tracing::error!("failed to get connection: {}", e);
                 FlagError::DatabaseUnavailable

diff --git a/rust/feature-flags/src/config.rs b/rust/feature-flags/src/config.rs
@@ -16,10 +16,10 @@ pub struct Config {
     #[envconfig(default = "postgres://posthog:posthog@localhost:5432/posthog")]
     pub read_database_url: String,
 
-    #[envconfig(default = "1024")]
-    pub max_concurrent_jobs: usize,
+    #[envconfig(default = "1000")]
+    pub max_concurrency: usize,
 
-    #[envconfig(default = "100")]
+    #[envconfig(default = "10")]
     pub max_pg_connections: u32,
 
     #[envconfig(default = "redis://localhost:6379/")]
@@ -30,6 +30,9 @@ pub struct Config {
 
     #[envconfig(from = "MAXMIND_DB_PATH", default = "")]
     pub maxmind_db_path: String,
+
+    #[envconfig(default = "false")]
+    pub enable_metrics: bool,
 }
 
 impl Config {
@@ -40,10 +43,11 @@ impl Config {
             write_database_url: "postgres://posthog:posthog@localhost:5432/test_posthog"
                 .to_string(),
             read_database_url: "postgres://posthog:posthog@localhost:5432/test_posthog".to_string(),
-            max_concurrent_jobs: 1024,
-            max_pg_connections: 100,
+            max_concurrency: 1000,
+            max_pg_connections: 10,
             acquire_timeout_secs: 1,
             maxmind_db_path: "".to_string(),
+            enable_metrics: false,
         }
     }
 
@@ -83,8 +87,8 @@ mod tests {
             config.read_database_url,
             "postgres://posthog:posthog@localhost:5432/posthog"
         );
-        assert_eq!(config.max_concurrent_jobs, 1024);
-        assert_eq!(config.max_pg_connections, 100);
+        assert_eq!(config.max_concurrency, 1000);
+        assert_eq!(config.max_pg_connections, 10);
         assert_eq!(config.redis_url, "redis://localhost:6379/");
     }
 
@@ -100,8 +104,8 @@ mod tests {
             config.read_database_url,
             "postgres://posthog:posthog@localhost:5432/test_posthog"
         );
-        assert_eq!(config.max_concurrent_jobs, 1024);
-        assert_eq!(config.max_pg_connections, 100);
+        assert_eq!(config.max_concurrency, 1000);
+        assert_eq!(config.max_pg_connections, 10);
         assert_eq!(config.redis_url, "redis://localhost:6379/");
     }
 
@@ -117,8 +121,8 @@ mod tests {
             config.read_database_url,
             "postgres://posthog:posthog@localhost:5432/test_posthog"
         );
-        assert_eq!(config.max_concurrent_jobs, 1024);
-        assert_eq!(config.max_pg_connections, 100);
+        assert_eq!(config.max_concurrency, 1000);
+        assert_eq!(config.max_pg_connections, 10);
         assert_eq!(config.redis_url, "redis://localhost:6379/");
     }
 }
diff --git a/rust/feature-flags/src/database.rs b/rust/feature-flags/src/database.rs
@@ -5,20 +5,15 @@ use async_trait::async_trait;
 use sqlx::{
     pool::PoolConnection,
     postgres::{PgPoolOptions, PgRow},
-    Postgres,
+    PgPool, Postgres,
 };
 use thiserror::Error;
 use tokio::time::timeout;
 
-use crate::config::Config;
-
 const DATABASE_TIMEOUT_MILLISECS: u64 = 1000;
 
 #[derive(Error, Debug)]
 pub enum CustomDatabaseError {
-    #[error("Not found in database")]
-    NotFound,
-
     #[error("Pg error: {0}")]
     Other(#[from] sqlx::Error),
 
@@ -40,36 +35,17 @@ pub trait Client {
     ) -> Result<Vec<PgRow>, CustomDatabaseError>;
 }
 
-pub struct PgClient {
-    pool: sqlx::PgPool,
-}
-
-impl PgClient {
-    pub async fn new_read_client(config: &Config) -> Result<PgClient, CustomDatabaseError> {
-        let pool = PgPoolOptions::new()
-            .max_connections(config.max_pg_connections)
-            .acquire_timeout(Duration::from_secs(1))
-            .test_before_acquire(true)
-            .connect(&config.read_database_url)
-            .await?;
-
-        Ok(PgClient { pool })
-    }
-
-    pub async fn new_write_client(config: &Config) -> Result<PgClient, CustomDatabaseError> {
-        let pool = PgPoolOptions::new()
-            .max_connections(config.max_pg_connections)
-            .acquire_timeout(Duration::from_secs(1))
-            .test_before_acquire(true)
-            .connect(&config.write_database_url)
-            .await?;
-
-        Ok(PgClient { pool })
-    }
+pub async fn get_pool(url: &str, max_connections: u32) -> Result<PgPool, sqlx::Error> {
+    PgPoolOptions::new()
+        .max_connections(max_connections)
+        .acquire_timeout(Duration::from_secs(1))
+        .test_before_acquire(true)
+        .connect(url)
+        .await
 }
 
 #[async_trait]
-impl Client for PgClient {
+impl Client for PgPool {
     async fn run_query(
         &self,
         query: String,
@@ -80,7 +56,7 @@ impl Client for PgClient {
         let built_query = parameters
             .iter()
             .fold(built_query, |acc, param| acc.bind(param));
-        let query_results = built_query.fetch_all(&self.pool);
+        let query_results = built_query.fetch_all(self);
 
         let timeout_ms = match timeout_ms {
             Some(ms) => ms,
@@ -93,6 +69,6 @@ impl Client for PgClient {
     }
 
     async fn get_connection(&self) -> Result<PoolConnection<Postgres>, CustomDatabaseError> {
-        Ok(self.pool.acquire().await?)
+        Ok(self.acquire().await?)
     }
 }
diff --git a/rust/feature-flags/src/flag_definitions.rs b/rust/feature-flags/src/flag_definitions.rs
@@ -237,7 +237,10 @@ mod tests {
             .await
             .expect("Failed to fetch flags from redis");
         assert_eq!(flags_from_redis.flags.len(), 1);
-        let flag = flags_from_redis.flags.get(0).expect("Empty flags in redis");
+        let flag = flags_from_redis
+            .flags
+            .first()
+            .expect("Empty flags in redis");
         assert_eq!(flag.key, "flag1");
         assert_eq!(flag.team_id, team.id);
         assert_eq!(flag.filters.groups.len(), 1);
@@ -288,7 +291,7 @@ mod tests {
             .expect("Failed to fetch flags from pg");
 
         assert_eq!(flags_from_pg.flags.len(), 1);
-        let flag = flags_from_pg.flags.get(0).expect("Flags should be in pg");
+        let flag = flags_from_pg.flags.first().expect("Flags should be in pg");
 
         assert_eq!(flag.key, "flag1");
         assert_eq!(flag.team_id, team.id);
@@ -414,13 +417,11 @@ mod tests {
     async fn test_fetch_empty_team_from_pg() {
         let client = setup_pg_client(None).await;
 
-        match FeatureFlagList::from_pg(client.clone(), 1234)
+        let FeatureFlagList { flags } = FeatureFlagList::from_pg(client.clone(), 1234)
             .await
-            .expect("Failed to fetch flags from pg")
+            .expect("Failed to fetch flags from pg");
         {
-            FeatureFlagList { flags } => {
-                assert_eq!(flags.len(), 0);
-            }
+            assert_eq!(flags.len(), 0);
         }
     }
 

diff --git a/rust/feature-flags/src/flag_matching.rs b/rust/feature-flags/src/flag_matching.rs
@@ -91,7 +91,7 @@ impl GroupTypeMappingCache {
                 return Err(e);
             }
         };
-        self.group_types_to_indexes = mapping.clone();
+        self.group_types_to_indexes.clone_from(&mapping);
 
         Ok(mapping)
     }
@@ -108,7 +108,7 @@ impl GroupTypeMappingCache {
             types_to_indexes.into_iter().map(|(k, v)| (v, k)).collect();
 
         if !result.is_empty() {
-            self.group_indexes_to_types = result.clone();
+            self.group_indexes_to_types.clone_from(&result);
             Ok(result)
         } else {
             Err(FlagError::NoGroupTypeMappings)
@@ -836,7 +836,7 @@ mod tests {
             None,
         );
         let match_result = matcher.get_match(&flag, None).await.unwrap();
-        assert_eq!(match_result.matches, true);
+        assert!(match_result.matches);
         assert_eq!(match_result.variant, None);
 
         let mut matcher = FeatureFlagMatcher::new(
@@ -848,7 +848,7 @@ mod tests {
             None,
         );
         let match_result = matcher.get_match(&flag, None).await.unwrap();
-        assert_eq!(match_result.matches, false);
+        assert!(!match_result.matches);
         assert_eq!(match_result.variant, None);
 
         let mut matcher = FeatureFlagMatcher::new(
@@ -860,7 +860,7 @@ mod tests {
             None,
         );
         let match_result = matcher.get_match(&flag, None).await.unwrap();
-        assert_eq!(match_result.matches, false);
+        assert!(!match_result.matches);
         assert_eq!(match_result.variant, None);
     }
 
@@ -1041,7 +1041,7 @@ mod tests {
             .is_condition_match(&flag, &condition, None)
             .await
             .unwrap();
-        assert_eq!(is_match, true);
+        assert!(is_match);
         assert_eq!(reason, "CONDITION_MATCH");
     }
 

diff --git a/rust/feature-flags/src/flag_request.rs b/rust/feature-flags/src/flag_request.rs
@@ -210,7 +210,7 @@ mod tests {
     #[test]
     fn too_large_distinct_id_is_truncated() {
         let json = json!({
-            "distinct_id": std::iter::repeat("a").take(210).collect::<String>(),
+            "distinct_id": "a".repeat(210),
             "token": "my_token1",
         });
         let bytes = Bytes::from(json.to_string());

diff --git a/rust/feature-flags/src/metrics_consts.rs b/rust/feature-flags/src/metrics_consts.rs