neon/libs/posthog_client_lite/src/background_loop.rs

//! A background loop that fetches feature flags from PostHog and updates the feature store.

use std::{sync::Arc, time::Duration};

use arc_swap::ArcSwap;
use tokio_util::sync::CancellationToken;

use crate::{FeatureStore, PostHogClient, PostHogClientConfig};

/// A background loop that fetches feature flags from PostHog and updates the feature store.
pub struct FeatureResolverBackgroundLoop {
    posthog_client: PostHogClient,
    feature_store: ArcSwap<FeatureStore>,
    cancel: CancellationToken,
}

impl FeatureResolverBackgroundLoop {
    pub fn new(config: PostHogClientConfig, shutdown_pageserver: CancellationToken) -> Self {
        Self {
            posthog_client: PostHogClient::new(config),
            feature_store: ArcSwap::new(Arc::new(FeatureStore::new())),
            cancel: shutdown_pageserver,
        }
    }

    pub fn spawn(self: Arc<Self>, handle: &tokio::runtime::Handle, refresh_period: Duration) {
        let this = self.clone();
        let cancel = self.cancel.clone();
        handle.spawn(async move {
            tracing::info!("Starting PostHog feature resolver");
            let mut ticker = tokio::time::interval(refresh_period);
            ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
            loop {
                tokio::select! {
                    _ = ticker.tick() => {}
                    _ = cancel.cancelled() => break
                }
                let resp = match this
                    .posthog_client
                    .get_feature_flags_local_evaluation()
                    .await
                {
                    Ok(resp) => resp,
                    Err(e) => {
                        tracing::warn!("Cannot get feature flags: {}", e);
                        continue;
                    }
                };
                let feature_store = FeatureStore::new_with_flags(resp.flags);
                this.feature_store.store(Arc::new(feature_store));
            }
            tracing::info!("PostHog feature resolver stopped");
        });
    }

    pub fn feature_store(&self) -> Arc<FeatureStore> {
        self.feature_store.load_full()
    }
}