AkurAI Build
Menu

Public repository / 6d386d9876a747c093481f05e926b1d43ccfb711

AkurAI-IDP / src/lib/telemetry.rs

Latest change e1794828e5fae9a8545a66e77b2e8200cf8c090b - feat: add organization lifecycle telemetry by Ólafur Búi Ólafsson

use akurai_telemetry::{
    EmitOutcome, PrivacyClassification, TelemetryContext, TelemetryEnvelope, TelemetryReceiver,
    TelemetrySender,
};
use rusqlite::Connection;
use sha2::{Digest, Sha256};
use std::sync::{
    Arc,
    atomic::{AtomicU64, Ordering},
};

use crate::db;

pub const APP_ID: &str = "akurai-idp";
pub const EVENT_VERSION: &str = "1";

pub const ORGANIZATION_CREATED: &str = "organization.created";
pub const INVITATION_DRAFT_CREATED: &str = "organization.invitation_draft_created";
pub const INVITATION_ACCEPTED: &str = "organization.invitation_accepted";
pub const FIRST_ADMIN_LOGIN: &str = "organization.first_admin_login";
pub const ACTIVE_MEMBER_COUNT: &str = "organization.active_member_count";

#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct TelemetryEmitMetrics {
    pub queued: u64,
    pub duplicates: u64,
    pub buffer_full: u64,
    pub closed: u64,
    pub rejected: u64,
}

/// The request-path handle. It only validates and uses the shared contract's
/// non-blocking `try_send` queue; receiver failures are isolated by the worker.
/// Event IDs contain only organization and operation identifiers and remain
/// deterministic for retries. Receivers must durably deduplicate `event_id`
/// across process restarts.
#[derive(Clone)]
pub struct OrganizationTelemetry {
    sender: TelemetrySender,
    metrics: Arc<OutcomeCounters>,
}

#[derive(Default)]
struct OutcomeCounters {
    queued: AtomicU64,
    duplicates: AtomicU64,
    buffer_full: AtomicU64,
    closed: AtomicU64,
    rejected: AtomicU64,
}

impl OrganizationTelemetry {
    pub fn from_sender(sender: TelemetrySender) -> Self {
        Self {
            sender,
            metrics: Arc::new(OutcomeCounters::default()),
        }
    }

    pub fn metrics(&self) -> TelemetryEmitMetrics {
        TelemetryEmitMetrics {
            queued: self.metrics.queued.load(Ordering::Relaxed),
            duplicates: self.metrics.duplicates.load(Ordering::Relaxed),
            buffer_full: self.metrics.buffer_full.load(Ordering::Relaxed),
            closed: self.metrics.closed.load(Ordering::Relaxed),
            rejected: self.metrics.rejected.load(Ordering::Relaxed),
        }
    }

    fn observe(&self, event_name: &str, outcome: EmitOutcome) -> EmitOutcome {
        let counter = match outcome {
            EmitOutcome::Queued => &self.metrics.queued,
            EmitOutcome::Duplicate => &self.metrics.duplicates,
            EmitOutcome::BufferFull => &self.metrics.buffer_full,
            EmitOutcome::Closed => &self.metrics.closed,
            EmitOutcome::Rejected => &self.metrics.rejected,
        };
        let count = counter.fetch_add(1, Ordering::Relaxed) + 1;
        if matches!(
            outcome,
            EmitOutcome::BufferFull | EmitOutcome::Closed | EmitOutcome::Rejected
        ) && (count == 1 || count.is_power_of_two())
        {
            tracing::warn!(
                event_name,
                outcome = ?outcome,
                count,
                "organization lifecycle telemetry enqueue degraded"
            );
        }
        outcome
    }

    fn emit_organization_event(
        &self,
        organization_id: &str,
        event_name: &str,
        source_event_id: &str,
        occurred_at: i64,
        attributes: Vec<(String, String)>,
    ) -> EmitOutcome {
        let (workspace_id, installation_id, app_id) =
            db::with_db(|conn| organization_context(conn, organization_id));
        self.emit_with_context(
            organization_id,
            &workspace_id,
            &installation_id,
            &app_id,
            event_name,
            source_event_id,
            occurred_at,
            attributes,
        )
    }

    fn emit_with_context(
        &self,
        organization_id: &str,
        workspace_id: &str,
        installation_id: &str,
        app_id: &str,
        event_name: &str,
        source_event_id: &str,
        occurred_at: i64,
        mut attributes: Vec<(String, String)>,
    ) -> EmitOutcome {
        attributes.extend([
            ("event_version".to_owned(), EVENT_VERSION.to_owned()),
            ("source_event_id".to_owned(), source_event_id.to_owned()),
        ]);
        let context = TelemetryContext::new(
            stable_id(organization_id, "unknown-organization"),
            stable_id(
                workspace_id,
                &bounded_source_id("organization.workspace", organization_id, None, None),
            ),
            stable_id(
                installation_id,
                &bounded_source_id("organization.installation", organization_id, None, None),
            ),
            stable_id(app_id, APP_ID),
        );
        let event_id = format!("{event_name}:{source_event_id}");
        let Ok(event) = TelemetryEnvelope::new(
            event_id,
            &context,
            event_name,
            occurred_at,
            PrivacyClassification::Operational,
            attributes,
        ) else {
            return self.observe(event_name, EmitOutcome::Rejected);
        };
        self.observe(event_name, self.sender.emit(event))
    }

    pub fn organization_created(&self, organization_id: &str, occurred_at: i64) -> EmitOutcome {
        let source_event_id = bounded_source_id(ORGANIZATION_CREATED, organization_id, None, None);
        self.emit_organization_event(
            organization_id,
            ORGANIZATION_CREATED,
            &source_event_id,
            occurred_at,
            Vec::new(),
        )
    }

    pub fn invitation_draft_created(
        &self,
        organization_id: &str,
        invitation_id: &str,
        role: &str,
        occurred_at: i64,
    ) -> EmitOutcome {
        let source_event_id = bounded_source_id(
            INVITATION_DRAFT_CREATED,
            organization_id,
            Some(invitation_id),
            None,
        );
        self.emit_organization_event(
            organization_id,
            INVITATION_DRAFT_CREATED,
            &source_event_id,
            occurred_at,
            vec![("role".to_owned(), role.to_owned())],
        )
    }

    pub fn invitation_accepted(
        &self,
        organization_id: &str,
        invitation_id: &str,
        role: &str,
        occurred_at: i64,
    ) -> EmitOutcome {
        let source_event_id = bounded_source_id(
            INVITATION_ACCEPTED,
            organization_id,
            Some(invitation_id),
            None,
        );
        self.emit_organization_event(
            organization_id,
            INVITATION_ACCEPTED,
            &source_event_id,
            occurred_at,
            vec![("role".to_owned(), role.to_owned())],
        )
    }

    pub fn first_admin_login(&self, organization_id: &str, occurred_at: i64) -> EmitOutcome {
        let source_event_id = bounded_source_id(FIRST_ADMIN_LOGIN, organization_id, None, None);
        self.emit_organization_event(
            organization_id,
            FIRST_ADMIN_LOGIN,
            &source_event_id,
            occurred_at,
            Vec::new(),
        )
    }

    fn active_member_count(
        &self,
        organization_id: &str,
        source_event_id: &str,
        count: i64,
        occurred_at: i64,
    ) -> EmitOutcome {
        self.emit_organization_event(
            organization_id,
            ACTIVE_MEMBER_COUNT,
            source_event_id,
            occurred_at,
            vec![("active_member_count".to_owned(), count.to_string())],
        )
    }

    pub fn active_member_count_after_organization_created(
        &self,
        organization_id: &str,
        occurred_at: i64,
    ) -> Option<EmitOutcome> {
        let source_event_id = bounded_source_id(ORGANIZATION_CREATED, organization_id, None, None);
        self.emit_active_member_count(organization_id, &source_event_id, occurred_at)
    }

    pub fn active_member_count_after_invitation(
        &self,
        organization_id: &str,
        invitation_id: &str,
        count: i64,
        occurred_at: i64,
    ) -> EmitOutcome {
        let source_event_id = bounded_source_id(
            INVITATION_ACCEPTED,
            organization_id,
            Some(invitation_id),
            None,
        );
        self.active_member_count(organization_id, &source_event_id, count, occurred_at)
    }

    pub fn active_member_count_after_membership_revoked(
        &self,
        organization_id: &str,
        occurred_at: i64,
    ) -> Option<EmitOutcome> {
        let source_event_id = bounded_source_id(
            "organization.membership-revoked",
            organization_id,
            None,
            Some(occurred_at),
        );
        self.emit_active_member_count(organization_id, &source_event_id, occurred_at)
    }

    pub fn active_member_count_after_membership_left(
        &self,
        organization_id: &str,
        occurred_at: i64,
    ) -> Option<EmitOutcome> {
        let source_event_id = bounded_source_id(
            "organization.membership-left",
            organization_id,
            None,
            Some(occurred_at),
        );
        self.emit_active_member_count(organization_id, &source_event_id, occurred_at)
    }

    /// Count only authoritative active memberships. A query failure means the
    /// fact is unknown and is deliberately not coerced to zero.
    fn emit_active_member_count(
        &self,
        organization_id: &str,
        source_event_id: &str,
        occurred_at: i64,
    ) -> Option<EmitOutcome> {
        let count = match db::with_db(|conn| {
            conn.query_row(
                "SELECT COUNT(*) FROM organization_memberships
                 WHERE organization_id = ?1 AND status = 'active'",
                [organization_id],
                |row| row.get::<_, i64>(0),
            )
        }) {
            Ok(count) => count,
            Err(_) => {
                tracing::warn!("unable to determine active member count for telemetry");
                return None;
            }
        };
        Some(self.active_member_count(organization_id, source_event_id, count, occurred_at))
    }

    pub fn sender(&self) -> &TelemetrySender {
        &self.sender
    }
}

/// Default sink for the IDP process. A downstream collector can replace this
/// receiver without changing request handlers or the shared envelope.
pub struct TracingTelemetryReceiver;

impl TelemetryReceiver for TracingTelemetryReceiver {
    fn receive(&self, event: TelemetryEnvelope) -> Result<(), akurai_telemetry::ReceiverError> {
        tracing::info!(
            event_id = %event.event_id,
            organization_id = %event.organization_id,
            workspace_id = %event.workspace_id,
            event_name = %event.event_name,
            schema_version = event.schema_version,
            attributes = ?event.attributes,
            occurred_at = event.occurred_at,
            "organization lifecycle telemetry"
        );
        Ok(())
    }
}

/// Hash source dimensions instead of embedding caller-controlled identifiers in
/// telemetry IDs. Length-prefixing each component keeps the digest stable even
/// when values contain separator characters, while the fixed SHA-256 hex output
/// always fits the shared 128-byte identifier contract.
fn bounded_source_id(
    operation: &str,
    organization_id: &str,
    invitation_id: Option<&str>,
    occurred_at: Option<i64>,
) -> String {
    let mut hasher = Sha256::new();
    hash_part(&mut hasher, operation);
    hash_part(&mut hasher, organization_id);
    if let Some(invitation_id) = invitation_id {
        hash_part(&mut hasher, invitation_id);
    }
    if let Some(occurred_at) = occurred_at {
        hasher.update(occurred_at.to_be_bytes());
    }
    format!("sha256:{:x}", hasher.finalize())
}

fn hash_part(hasher: &mut Sha256, value: &str) {
    hasher.update((value.len() as u64).to_be_bytes());
    hasher.update(value.as_bytes());
}

fn organization_context(conn: &Connection, organization_id: &str) -> (String, String, String) {
    // Organization-only facts have no workspace yet; use a deterministic
    // organization scope rather than inventing a zero or nullable ID.
    conn.query_row(
        "SELECT workspace_id,
                COALESCE(installation_id, ''),
                COALESCE(app_id, '')
         FROM organization_workspaces
         WHERE organization_id = ?1
         ORDER BY created_at, id LIMIT 1",
        [organization_id],
        |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)),
    )
    .unwrap_or_else(|_| {
        (
            bounded_source_id("organization.workspace", organization_id, None, None),
            bounded_source_id("organization.installation", organization_id, None, None),
            APP_ID.to_owned(),
        )
    })
}

fn stable_id(value: &str, fallback: &str) -> String {
    if value.is_empty()
        || value.len() > 128
        || !value.is_ascii()
        || value.chars().any(char::is_whitespace)
        || value.chars().any(char::is_control)
    {
        fallback.to_owned()
    } else {
        value.to_owned()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use akurai_telemetry::{TelemetryDelivery, TelemetryDeliveryConfig};
    use std::sync::{
        Arc, Mutex,
        mpsc::{Receiver, Sender, channel},
    };
    use std::time::Duration;

    struct RecordingReceiver(Arc<Mutex<Vec<TelemetryEnvelope>>>);

    impl TelemetryReceiver for RecordingReceiver {
        fn receive(&self, event: TelemetryEnvelope) -> Result<(), akurai_telemetry::ReceiverError> {
            if let Ok(mut events) = self.0.lock() {
                events.push(event);
            }
            Ok(())
        }
    }

    #[test]
    fn lifecycle_events_are_versioned_operation_scoped_and_policy_neutral() {
        let events = Arc::new(Mutex::new(Vec::new()));
        let delivery = TelemetryDelivery::new_with_config(
            RecordingReceiver(Arc::clone(&events)),
            TelemetryDeliveryConfig {
                queue_capacity: 8,
                dedupe_capacity: 8,
                shutdown_timeout: Duration::from_secs(1),
            },
        )
        .unwrap();
        let telemetry = OrganizationTelemetry::from_sender(delivery.sender());
        assert_eq!(
            telemetry.invitation_draft_created("org-1", "invite-1", "admin", 1),
            EmitOutcome::Queued
        );
        assert_eq!(
            telemetry.invitation_draft_created("org-1", "invite-1", "admin", 1),
            EmitOutcome::Duplicate
        );
        assert_eq!(telemetry.first_admin_login("org-1", 1), EmitOutcome::Queued);
        delivery.shutdown();

        let events = events
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        assert_eq!(events.len(), 2);
        assert_eq!(
            events[0].event_id,
            format!(
                "{INVITATION_DRAFT_CREATED}:{}",
                bounded_source_id(INVITATION_DRAFT_CREATED, "org-1", Some("invite-1"), None)
            )
        );
        assert_eq!(
            events[1].event_id,
            format!(
                "{FIRST_ADMIN_LOGIN}:{}",
                bounded_source_id(FIRST_ADMIN_LOGIN, "org-1", None, None)
            )
        );
        for event in events.iter() {
            assert_eq!(event.schema_version, akurai_telemetry::SCHEMA_VERSION);
            assert_eq!(event.organization_id, "org-1");
            assert!(!event.attributes.iter().any(|(key, _)| {
                matches!(
                    key.as_str(),
                    "retention_class" | "consent_basis" | "consent_state"
                )
            }));
            for (key, value) in &event.attributes {
                let lower = format!("{key}:{value}").to_ascii_lowercase();
                for forbidden in [
                    "password",
                    "token",
                    "email",
                    "full_name",
                    "user_id",
                    "member_id",
                    "profile",
                ] {
                    assert!(
                        !lower.contains(forbidden),
                        "forbidden telemetry field: {lower}"
                    );
                }
            }
            assert!(!event.event_id.contains("user-"));
            assert!(!event.event_id.contains("member-"));
        }
    }

    #[test]
    fn organization_route_emissions_use_only_operation_dimensions() {
        let events = Arc::new(Mutex::new(Vec::new()));
        let delivery = TelemetryDelivery::new_with_config(
            RecordingReceiver(Arc::clone(&events)),
            TelemetryDeliveryConfig {
                queue_capacity: 8,
                dedupe_capacity: 8,
                shutdown_timeout: Duration::from_secs(1),
            },
        )
        .unwrap();
        let telemetry = OrganizationTelemetry::from_sender(delivery.sender());
        assert_eq!(
            telemetry.organization_created("org-1", 1),
            EmitOutcome::Queued
        );
        assert_eq!(
            telemetry.invitation_draft_created("org-1", "invite-1", "member", 2),
            EmitOutcome::Queued
        );
        assert_eq!(
            telemetry.invitation_accepted("org-1", "invite-1", "member", 3),
            EmitOutcome::Queued
        );
        assert_eq!(
            telemetry.active_member_count_after_invitation("org-1", "invite-1", 2, 3),
            EmitOutcome::Queued
        );
        delivery.shutdown();

        let events = events
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        assert_eq!(events.len(), 4);
        for event in events.iter() {
            assert!(event.event_id.len() <= 128);
            assert!(!event.event_id.contains("org-1"));
            assert!(!event.event_id.contains("invite-1"));
            assert!(event.attributes.iter().all(|(key, value)| {
                !matches!(
                    key.as_str(),
                    "user_id" | "member_id" | "email" | "full_name" | "profile"
                ) && !value.contains("user-")
                    && !value.contains("member-")
                    && !value.contains("org-1")
                    && !value.contains("invite-1")
            }));
        }
    }

    #[test]
    fn maximum_length_and_uuid_ids_queue_without_rejection() {
        let events = Arc::new(Mutex::new(Vec::new()));
        let delivery = TelemetryDelivery::new_with_config(
            RecordingReceiver(Arc::clone(&events)),
            TelemetryDeliveryConfig {
                queue_capacity: 32,
                dedupe_capacity: 32,
                shutdown_timeout: Duration::from_secs(1),
            },
        )
        .unwrap();
        let telemetry = OrganizationTelemetry::from_sender(delivery.sender());
        let cases = [
            (
                "550e8400-e29b-41d4-a716-446655440000".to_owned(),
                "6ba7b810-9dad-11d1-80b4-00c04fd430c8".to_owned(),
            ),
            ("o".repeat(128), "i".repeat(128)),
        ];

        for (case_index, (organization_id, invitation_id)) in cases.iter().enumerate() {
            let occurred_at = 10 + case_index as i64;
            let outcomes = [
                telemetry.organization_created(organization_id, occurred_at),
                telemetry.invitation_draft_created(
                    organization_id,
                    invitation_id,
                    "admin",
                    occurred_at,
                ),
                telemetry.invitation_accepted(
                    organization_id,
                    invitation_id,
                    "member",
                    occurred_at,
                ),
                telemetry.first_admin_login(organization_id, occurred_at),
                telemetry
                    .active_member_count_after_organization_created(organization_id, occurred_at)
                    .expect("active member count query"),
                telemetry.active_member_count_after_invitation(
                    organization_id,
                    invitation_id,
                    1,
                    occurred_at,
                ),
                telemetry
                    .active_member_count_after_membership_revoked(organization_id, occurred_at)
                    .expect("active member count query"),
                telemetry
                    .active_member_count_after_membership_left(organization_id, occurred_at)
                    .expect("active member count query"),
            ];
            assert!(
                outcomes
                    .iter()
                    .all(|outcome| *outcome == EmitOutcome::Queued),
                "unexpected outcomes for case {case_index}: {outcomes:?}"
            );
            assert_eq!(
                telemetry.invitation_draft_created(
                    organization_id,
                    invitation_id,
                    "different-role",
                    occurred_at,
                ),
                EmitOutcome::Duplicate
            );
        }

        assert_eq!(telemetry.metrics().rejected, 0);
        delivery.shutdown();

        let events = events
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        assert_eq!(events.len(), cases.len() * 8);
        for event in events.iter() {
            assert!(event.event_id.len() <= 128);
            assert!(event.event_id.starts_with("organization."));
            assert!(
                event
                    .attributes
                    .iter()
                    .filter(|(key, _)| key == "source_event_id")
                    .all(|(_, value)| value.starts_with("sha256:") && value.len() == 71)
            );
            assert!(
                !event
                    .event_id
                    .contains("550e8400-e29b-41d4-a716-446655440000")
            );
            assert!(
                !event
                    .event_id
                    .contains("6ba7b810-9dad-11d1-80b4-00c04fd430c8")
            );
            assert!(!event.event_id.contains(&"o".repeat(128)));
            assert!(!event.event_id.contains(&"i".repeat(128)));
        }
    }
    struct BlockingReceiver {
        started: Sender<()>,
        release: Mutex<Receiver<()>>,
        block_once: Mutex<bool>,
    }

    impl TelemetryReceiver for BlockingReceiver {
        fn receive(
            &self,
            _event: TelemetryEnvelope,
        ) -> Result<(), akurai_telemetry::ReceiverError> {
            let should_block = self
                .block_once
                .lock()
                .map(|mut block_once| {
                    let should_block = *block_once;
                    *block_once = false;
                    should_block
                })
                .unwrap_or(false);
            if should_block {
                let _ = self.started.send(());
                let _ = self
                    .release
                    .lock()
                    .ok()
                    .and_then(|release| release.recv().ok());
            }
            Ok(())
        }
    }

    #[test]
    fn bounded_queue_drop_is_observed_without_blocking_emitters() {
        let (started_tx, started_rx) = channel();
        let (release_tx, release_rx) = channel();
        let delivery = TelemetryDelivery::new_with_config(
            BlockingReceiver {
                started: started_tx,
                release: Mutex::new(release_rx),
                block_once: Mutex::new(true),
            },
            TelemetryDeliveryConfig {
                queue_capacity: 1,
                dedupe_capacity: 8,
                shutdown_timeout: Duration::from_secs(1),
            },
        )
        .unwrap();
        let telemetry = OrganizationTelemetry::from_sender(delivery.sender());
        assert_eq!(
            telemetry.organization_created("org-1", 1),
            EmitOutcome::Queued
        );
        started_rx
            .recv_timeout(Duration::from_secs(1))
            .expect("receiver started");
        assert_eq!(telemetry.first_admin_login("org-1", 2), EmitOutcome::Queued);
        assert_eq!(
            telemetry.first_admin_login("org-2", 3),
            EmitOutcome::BufferFull
        );
        assert_eq!(telemetry.metrics().buffer_full, 1);
        release_tx.send(()).expect("release receiver");
        assert_eq!(telemetry.first_admin_login("org-2", 4), EmitOutcome::Queued);
        delivery.shutdown();
    }

    #[test]
    fn receiver_outage_is_visible_in_delivery_metrics_not_request_outcome() {
        struct FailingReceiver;
        impl TelemetryReceiver for FailingReceiver {
            fn receive(
                &self,
                _event: TelemetryEnvelope,
            ) -> Result<(), akurai_telemetry::ReceiverError> {
                Err(akurai_telemetry::ReceiverError::new("offline"))
            }
        }
        let delivery = TelemetryDelivery::new(FailingReceiver, 2).unwrap();
        let sender = delivery.sender();
        let telemetry = OrganizationTelemetry::from_sender(sender.clone());
        assert_eq!(
            telemetry.organization_created("org-1", 1),
            EmitOutcome::Queued
        );
        delivery.shutdown();
        let metrics = sender.metrics();
        assert_eq!(metrics.receiver_failures, 1);
        assert_eq!(metrics.receiver_drops, 1);
        assert_eq!(telemetry.metrics().queued, 1);
    }
}