feat(flags): add metrics to new flag service that mirror existing met…

…rics (#25296)

Co-authored-by: github-actions <41898282+github-actions[bot]@users.noreply.github.com>

rust/common/metrics/src/lib.rs

@@ -5,6 +5,22 @@ use axum::{
     routing::get, Router,
 };
 use metrics_exporter_prometheus::{PrometheusBuilder, PrometheusHandle};
+use std::sync::OnceLock;
+
+type LabelFilterFn =
+    Box<dyn Fn(&[(String, String)]) -> Vec<(String, String)> + Send + Sync + 'static>;
+
+static LABEL_FILTER: OnceLock<LabelFilterFn> = OnceLock::new();
+
+pub fn set_label_filter<F>(filter: F)
+where
+    F: Fn(&[(String, String)]) -> Vec<(String, String)> + Send + Sync + 'static,
+{
+    let boxed_filter: LabelFilterFn = Box::new(filter);
+    if LABEL_FILTER.set(boxed_filter).is_err() {
+        panic!("Label filter already set");
+    }
+}
 
 /// Bind a `TcpListener` on the provided bind address to serve a `Router` on it.
 /// This function is intended to take a Router as returned by `setup_metrics_router`, potentially with more routes added by the caller.
@@ -83,7 +99,16 @@ pub fn get_current_timestamp_seconds() -> f64 {
 
 // Shorthand for common metric types
 pub fn inc(name: &'static str, labels: &[(String, String)], value: u64) {
-    metrics::counter!(name, labels).increment(value);
+    let filtered_labels = apply_label_filter(labels);
+    metrics::counter!(name, &filtered_labels).increment(value);
+}
+
+fn apply_label_filter(labels: &[(String, String)]) -> Vec<(String, String)> {
+    if let Some(filter) = LABEL_FILTER.get() {
+        filter(labels)
+    } else {
+        labels.to_vec()
+    }
 }
 
 pub fn gauge(name: &'static str, lables: &[(String, String)], value: f64) {

rust/feature-flags/src/config.rs

@@ -1,10 +1,72 @@
 use envconfig::Envconfig;
 use once_cell::sync::Lazy;
 use std::net::SocketAddr;
+use std::num::ParseIntError;
 use std::path::{Path, PathBuf};
 use std::str::FromStr;
 
-// TODO rewrite this to follow the AppConfig pattern in other files
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum TeamIdsToTrack {
+    All,
+    TeamIds(Vec<i32>),
+}
+
+#[derive(Debug)]
+pub enum ParseTeamIdsError {
+    InvalidRange(String),
+    InvalidNumber(ParseIntError),
+}
+
+impl std::fmt::Display for ParseTeamIdsError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            ParseTeamIdsError::InvalidRange(r) => write!(f, "Invalid range: {}", r),
+            ParseTeamIdsError::InvalidNumber(e) => write!(f, "Invalid number: {}", e),
+        }
+    }
+}
+
+impl std::error::Error for ParseTeamIdsError {}
+
+impl FromStr for TeamIdsToTrack {
+    type Err = ParseTeamIdsError;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        let s = s.trim();
+        if s.eq_ignore_ascii_case("all") {
+            Ok(TeamIdsToTrack::All)
+        } else {
+            let mut team_ids = Vec::new();
+            for part in s.split(',').map(|p| p.trim()) {
+                if part.contains(':') {
+                    let bounds: Vec<&str> = part.split(':').collect();
+                    if bounds.len() != 2 {
+                        return Err(ParseTeamIdsError::InvalidRange(part.to_string()));
+                    }
+                    let start = bounds[0]
+                        .parse::<i32>()
+                        .map_err(ParseTeamIdsError::InvalidNumber)?;
+                    let end = bounds[1]
+                        .parse::<i32>()
+                        .map_err(ParseTeamIdsError::InvalidNumber)?;
+                    if end < start {
+                        return Err(ParseTeamIdsError::InvalidRange(part.to_string()));
+                    }
+                    for id in start..=end {
+                        team_ids.push(id);
+                    }
+                } else {
+                    let id = part
+                        .parse::<i32>()
+                        .map_err(ParseTeamIdsError::InvalidNumber)?;
+                    team_ids.push(id);
+                }
+            }
+            Ok(TeamIdsToTrack::TeamIds(team_ids))
+        }
+    }
+}
+
 #[derive(Envconfig, Clone, Debug)]
 pub struct Config {
     #[envconfig(default = "127.0.0.1:3001")]
@@ -33,6 +95,9 @@ pub struct Config {
 
     #[envconfig(default = "false")]
     pub enable_metrics: bool,
+
+    #[envconfig(from = "TEAM_IDS_TO_TRACK", default = "all")]
+    pub team_ids_to_track: TeamIdsToTrack,
 }
 
 impl Config {
@@ -48,6 +113,7 @@ impl Config {
             acquire_timeout_secs: 1,
             maxmind_db_path: "".to_string(),
             enable_metrics: false,
+            team_ids_to_track: TeamIdsToTrack::All,
         }
     }
 
@@ -90,6 +156,7 @@ mod tests {
         assert_eq!(config.max_concurrency, 1000);
         assert_eq!(config.max_pg_connections, 10);
         assert_eq!(config.redis_url, "redis://localhost:6379/");
+        assert_eq!(config.team_ids_to_track, TeamIdsToTrack::All);
     }
 
     #[test]
@@ -107,6 +174,7 @@ mod tests {
         assert_eq!(config.max_concurrency, 1000);
         assert_eq!(config.max_pg_connections, 10);
         assert_eq!(config.redis_url, "redis://localhost:6379/");
+        assert_eq!(config.team_ids_to_track, TeamIdsToTrack::All);
     }
 
     #[test]
@@ -124,5 +192,42 @@ mod tests {
         assert_eq!(config.max_concurrency, 1000);
         assert_eq!(config.max_pg_connections, 10);
         assert_eq!(config.redis_url, "redis://localhost:6379/");
+        assert_eq!(config.team_ids_to_track, TeamIdsToTrack::All);
+    }
+
+    #[test]
+    fn test_team_ids_to_track_all() {
+        let team_ids: TeamIdsToTrack = "all".parse().unwrap();
+        assert_eq!(team_ids, TeamIdsToTrack::All);
+    }
+
+    #[test]
+    fn test_team_ids_to_track_single_ids() {
+        let team_ids: TeamIdsToTrack = "1,5,7,13".parse().unwrap();
+        assert_eq!(team_ids, TeamIdsToTrack::TeamIds(vec![1, 5, 7, 13]));
+    }
+
+    #[test]
+    fn test_team_ids_to_track_ranges() {
+        let team_ids: TeamIdsToTrack = "1:3".parse().unwrap();
+        assert_eq!(team_ids, TeamIdsToTrack::TeamIds(vec![1, 2, 3]));
+    }
+
+    #[test]
+    fn test_team_ids_to_track_mixed() {
+        let team_ids: TeamIdsToTrack = "1:3,5,7:9".parse().unwrap();
+        assert_eq!(team_ids, TeamIdsToTrack::TeamIds(vec![1, 2, 3, 5, 7, 8, 9]));
+    }
+
+    #[test]
+    fn test_invalid_range() {
+        let result: Result<TeamIdsToTrack, _> = "5:3".parse();
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn test_invalid_number() {
+        let result: Result<TeamIdsToTrack, _> = "abc".parse();
+        assert!(result.is_err());
     }
 }

rust/feature-flags/src/flag_matching.rs

@@ -3,9 +3,12 @@ use crate::{
     database::Client as DatabaseClient,
     feature_flag_match_reason::FeatureFlagMatchReason,
     flag_definitions::{FeatureFlag, FeatureFlagList, FlagGroupType, PropertyFilter},
+    metrics_consts::{FLAG_EVALUATION_ERROR_COUNTER, FLAG_HASH_KEY_WRITES_COUNTER},
     property_matching::match_property,
+    utils::parse_exception_for_prometheus_label,
 };
 use anyhow::Result;
+use common_metrics::inc;
 use serde_json::Value;
 use sha1::{Digest, Sha1};
 use sqlx::{postgres::PgQueryResult, Acquire, FromRow};
@@ -99,6 +102,7 @@ impl GroupTypeMappingCache {
             Ok(mapping) if !mapping.is_empty() => mapping,
             Ok(_) => {
                 self.failed_to_fetch_flags = true;
+                // TODO add the `"Failed to fetch group"` type of lable.  See posthog/models/feature_flag/flag_matching.py:parse_exception_for_error_message
                 return Err(FlagError::NoGroupTypeMappings);
             }
             Err(e) => {
@@ -126,6 +130,7 @@ impl GroupTypeMappingCache {
             self.group_indexes_to_types.clone_from(&result);
             Ok(result)
         } else {
+            // TODO add the `"Failed to fetch group"` type of lable.  See posthog/models/feature_flag/flag_matching.py:parse_exception_for_error_message
             Err(FlagError::NoGroupTypeMappings)
         }
     }
@@ -154,6 +159,7 @@ impl GroupTypeMappingCache {
             .collect();
 
         if mapping.is_empty() {
+            // TODO add the `"Failed to fetch group"` type of lable.  See posthog/models/feature_flag/flag_matching.py:parse_exception_for_error_message
             Err(FlagError::NoGroupTypeMappings)
         } else {
             Ok(mapping)
@@ -237,6 +243,16 @@ impl FeatureFlagMatcher {
             (None, false)
         };
 
+        // If there was an initial error in processing hash key overrides, increment the error counter
+        if initial_error {
+            let reason = "hash_key_override_error";
+            common_metrics::inc(
+                FLAG_EVALUATION_ERROR_COUNTER,
+                &[("reason".to_string(), reason.to_string())],
+                1,
+            );
+        }
+
         let flags_response = self
             .evaluate_flags_with_overrides(
                 feature_flags,
@@ -272,25 +288,54 @@ impl FeatureFlagMatcher {
                     "Failed to check if hash key override should be written: {:?}",
                     e
                 );
+                let reason = parse_exception_for_prometheus_label(&e);
+                inc(
+                    FLAG_EVALUATION_ERROR_COUNTER,
+                    &[("reason".to_string(), reason.to_string())],
+                    1,
+                );
                 return (None, true);
             }
         };
 
+        let mut writing_hash_key_override = false;
+
         if should_write {
             if let Err(e) = set_feature_flag_hash_key_overrides(
                 // NB: this is the only method that writes to the database, so it's the only one that should use the writer
                 self.postgres_writer.clone(),
                 self.team_id,
                 target_distinct_ids.clone(),
-                hash_key,
+                hash_key.clone(),
             )
             .await
             {
                 error!("Failed to set feature flag hash key overrides: {:?}", e);
+                // Increment the counter for failed write
+                let reason = parse_exception_for_prometheus_label(&e);
+                inc(
+                    FLAG_EVALUATION_ERROR_COUNTER,
+                    &[("reason".to_string(), reason.to_string())],
+                    1,
+                );
                 return (None, true);
             }
+            writing_hash_key_override = true;
         }
 
+        // TODO I'm not sure if this is the right place to increment this counter
+        inc(
+            FLAG_HASH_KEY_WRITES_COUNTER,
+            &[
+                ("team_id".to_string(), self.team_id.to_string()),
+                (
+                    "successful_write".to_string(),
+                    writing_hash_key_override.to_string(),
+                ),
+            ],
+            1,
+        );
+
         match get_feature_flag_hash_key_overrides(
             self.postgres_reader.clone(),
             self.team_id,
@@ -301,6 +346,12 @@ impl FeatureFlagMatcher {
             Ok(overrides) => (Some(overrides), false),
             Err(e) => {
                 error!("Failed to get feature flag hash key overrides: {:?}", e);
+                let reason = parse_exception_for_prometheus_label(&e);
+                common_metrics::inc(
+                    FLAG_EVALUATION_ERROR_COUNTER,
+                    &[("reason".to_string(), reason.to_string())],
+                    1,
+                );
                 (None, true)
             }
         }
@@ -345,6 +396,12 @@ impl FeatureFlagMatcher {
                         "Error evaluating feature flag '{}' with overrides for distinct_id '{}': {:?}",
                         flag.key, self.distinct_id, e
                     );
+                    let reason = parse_exception_for_prometheus_label(&e);
+                    inc(
+                        FLAG_EVALUATION_ERROR_COUNTER,
+                        &[("reason".to_string(), reason.to_string())],
+                        1,
+                    );
                 }
             }
         }
@@ -374,6 +431,12 @@ impl FeatureFlagMatcher {
                     error_while_computing_flags = true;
                     // TODO add sentry exception tracking
                     error!("Error fetching properties: {:?}", e);
+                    let reason = parse_exception_for_prometheus_label(&e);
+                    inc(
+                        FLAG_EVALUATION_ERROR_COUNTER,
+                        &[("reason".to_string(), reason.to_string())],
+                        1,
+                    );
                 }
             }
 
@@ -397,6 +460,12 @@ impl FeatureFlagMatcher {
                             "Error evaluating feature flag '{}' for distinct_id '{}': {:?}",
                             flag.key, self.distinct_id, e
                         );
+                        let reason = parse_exception_for_prometheus_label(&e);
+                        inc(
+                            FLAG_EVALUATION_ERROR_COUNTER,
+                            &[("reason".to_string(), reason.to_string())],
+                            1,
+                        );
                     }
                 }
             }
@@ -1210,7 +1279,7 @@ async fn get_feature_flag_hash_key_overrides(
 }
 
 async fn set_feature_flag_hash_key_overrides(
-    postgres_writer: PostgresReader,
+    postgres_writer: PostgresWriter,
     team_id: TeamId,
     distinct_ids: Vec<String>,
     hash_key_override: String,