libdd_telemetry/worker/

mod.rs

// Copyright 2021-Present Datadog, Inc. https://www.datadoghq.com/
// SPDX-License-Identifier: Apache-2.0

pub mod http_client;
mod scheduler;
pub mod store;

use crate::{
    config::Config,
    data::{self, Application, Dependency, Endpoint, Host, Integration, Log, Payload, Telemetry},
    metrics::{ContextKey, MetricBuckets, MetricContexts},
};

use async_trait::async_trait;
use libdd_common::{http_common, tag::Tag};
use libdd_shared_runtime::Worker;

use std::iter::Sum;
use std::ops::Add;
use std::{
    collections::hash_map::DefaultHasher,
    hash::{Hash, Hasher},
    ops::ControlFlow,
    sync::{
        atomic::{AtomicU64, Ordering},
        Arc, Condvar, Mutex,
    },
    time,
};
use std::{collections::HashSet, fmt::Debug, time::Duration};

use crate::metrics::MetricBucketStats;
use futures::{
    channel::oneshot,
    future::{self},
};
use http::{header, HeaderValue};
use serde::{Deserialize, Serialize};
use tokio::{
    runtime::{self, Handle},
    sync::mpsc,
    task::JoinHandle,
};
use tokio_util::sync::CancellationToken;
use tracing::debug;

const CONTINUE: ControlFlow<()> = ControlFlow::Continue(());
const BREAK: ControlFlow<()> = ControlFlow::Break(());

fn time_now() -> f64 {
    #[allow(clippy::unwrap_used)]
    std::time::SystemTime::UNIX_EPOCH
        .elapsed()
        .unwrap_or_default()
        .as_secs_f64()
}

macro_rules! telemetry_worker_log {
    ($worker:expr , ERROR , $fmt_str:tt, $($arg:tt)*) => {
        {
            debug!(
                worker.runtime_id = %$worker.runtime_id,
                worker.debug_logging = $worker.config.telemetry_debug_logging_enabled,
                $fmt_str,
                $($arg)*
            );
            if $worker.config.telemetry_debug_logging_enabled {
                eprintln!(concat!("{}: Telemetry worker ERROR: ", $fmt_str), time_now(), $($arg)*);
            }
        }
    };
    ($worker:expr , DEBUG , $fmt_str:tt, $($arg:tt)*) => {
        {
            debug!(
                worker.runtime_id = %$worker.runtime_id,
                worker.debug_logging = $worker.config.telemetry_debug_logging_enabled,
                $fmt_str,
                $($arg)*
            );
            if $worker.config.telemetry_debug_logging_enabled {
                println!(concat!("{}: Telemetry worker DEBUG: ", $fmt_str), time_now(), $($arg)*);
            }
        }
    };
}

#[derive(Debug, Serialize, Deserialize)]
pub enum TelemetryActions {
    AddPoint((f64, ContextKey, Vec<Tag>)),
    AddConfig(data::Configuration),
    AddDependency(Dependency),
    AddIntegration(Integration),
    AddLog((LogIdentifier, Log)),
    AddEndpoint(Endpoint),
    Lifecycle(LifecycleAction),
    #[serde(skip)]
    CollectStats(oneshot::Sender<TelemetryWorkerStats>),
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum LifecycleAction {
    Start,
    Stop,
    FlushMetricAggr,
    FlushData,
    ExtendedHeartbeat,
}

/// Identifies a logging location uniquely
///
/// The identifier is a single 64 bit integer to save space an memory
/// and to be able to generic on the way different languages handle
#[derive(Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct LogIdentifier {
    // Collisions? Never heard of them
    pub identifier: u64,
}

// Holds the current state of the telemetry worker
#[derive(Debug)]
struct TelemetryWorkerData {
    started: bool,
    dependencies: store::Store<Dependency>,
    configurations: store::Store<data::Configuration>,
    integrations: store::Store<data::Integration>,
    endpoints: HashSet<data::Endpoint>,
    logs: store::QueueHashMap<LogIdentifier, Log>,
    metric_contexts: MetricContexts,
    metric_buckets: MetricBuckets,
    host: Host,
    app: Application,
}

pub struct TelemetryWorker {
    flavor: TelemetryWorkerFlavor,
    config: Config,
    mailbox: mpsc::Receiver<TelemetryActions>,
    cancellation_token: CancellationToken,
    seq_id: AtomicU64,
    runtime_id: String,
    client: Box<dyn http_client::HttpClient + Sync + Send>,
    metrics_flush_interval: Duration,
    deadlines: scheduler::Scheduler<LifecycleAction>,
    data: TelemetryWorkerData,
    next_action: Option<TelemetryActions>,
    stopped: bool,
}
impl Debug for TelemetryWorker {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("TelemetryWorker")
            .field("flavor", &self.flavor)
            .field("config", &self.config)
            .field("mailbox", &self.mailbox)
            .field("cancellation_token", &self.cancellation_token)
            .field("seq_id", &self.seq_id)
            .field("runtime_id", &self.runtime_id)
            .field("metrics_flush_interval", &self.metrics_flush_interval)
            .field("deadlines", &self.deadlines)
            .field("data", &self.data)
            .finish()
    }
}

#[async_trait]
impl Worker for TelemetryWorker {
    async fn trigger(&mut self) {
        if self.next_action.is_some() {
            // An action is already available and hasn't been executed
            return;
        }
        if self.stopped {
            // Channel is closed and Stop has already been dispatched. Park forever to avoid
            // a hot loop re-emitting Lifecycle::Stop on every iteration; the runtime will
            // tear the worker down via the handle.
            debug!(
                worker.runtime_id = %self.runtime_id,
                "Telemetry worker mailbox closed; parking until shutdown"
            );
            std::future::pending::<()>().await;
        }
        // Wait for the next action and store it
        let action = self.recv_next_action().await;
        self.next_action = Some(action);
    }

    // Processes a single action from the state machine
    async fn run(&mut self) {
        // Take the action that was stored by trigger()
        if let Some(action) = self.next_action.take() {
            debug!(
                worker.runtime_id = %self.runtime_id,
                action = ?action,
                "Received telemetry action"
            );

            // When running as a [libdd_shared_runtime::Worker] Shutdown is handled by stopping the
            // Worker from the handle and not by sending stop action
            let _action_result = match self.flavor {
                TelemetryWorkerFlavor::Full => self.dispatch_action(action).await,
                TelemetryWorkerFlavor::MetricsLogs => {
                    self.dispatch_metrics_logs_action(action).await
                }
            };
        }
    }

    /// Reset the worker state in the child process after a fork.
    ///
    /// Discards inherited pending telemetry state and dedupe history without sending anything, and
    /// drains the mailbox so that actions queued before the fork are not processed by the
    /// child.
    fn reset(&mut self) {
        // Drain all actions queued in the mailbox before the fork.
        while self.mailbox.try_recv().is_ok() {}

        // Discard any action that was staged by the last trigger() call.
        self.next_action = None;

        // Clear all unbuffered telemetry data; the child must not send pre-fork data.
        self.data.logs = store::QueueHashMap::default();
        self.data.metric_buckets = MetricBuckets::default();
        self.data.dependencies.clear();
        self.data.integrations.clear();
        self.data.configurations.clear();
        self.data.endpoints.clear();
    }

    async fn shutdown(&mut self) {
        let stop_action = TelemetryActions::Lifecycle(LifecycleAction::Stop);
        let _action_result = match self.flavor {
            TelemetryWorkerFlavor::Full => self.dispatch_action(stop_action).await,
            TelemetryWorkerFlavor::MetricsLogs => {
                self.dispatch_metrics_logs_action(stop_action).await
            }
        };
    }
}

#[derive(Debug, Default, Serialize, Deserialize)]
pub struct TelemetryWorkerStats {
    pub dependencies_stored: u32,
    pub dependencies_unflushed: u32,
    pub configurations_stored: u32,
    pub configurations_unflushed: u32,
    pub integrations_stored: u32,
    pub integrations_unflushed: u32,
    pub logs: u32,
    pub metric_contexts: u32,
    pub metric_buckets: MetricBucketStats,
}

impl Add for TelemetryWorkerStats {
    type Output = Self;

    fn add(self, rhs: Self) -> Self::Output {
        TelemetryWorkerStats {
            dependencies_stored: self.dependencies_stored + rhs.dependencies_stored,
            dependencies_unflushed: self.dependencies_unflushed + rhs.dependencies_unflushed,
            configurations_stored: self.configurations_stored + rhs.configurations_stored,
            configurations_unflushed: self.configurations_unflushed + rhs.configurations_unflushed,
            integrations_stored: self.integrations_stored + rhs.integrations_stored,
            integrations_unflushed: self.integrations_unflushed + rhs.integrations_unflushed,
            logs: self.logs + rhs.logs,
            metric_contexts: self.metric_contexts + rhs.metric_contexts,
            metric_buckets: MetricBucketStats {
                buckets: self.metric_buckets.buckets + rhs.metric_buckets.buckets,
                series: self.metric_buckets.series + rhs.metric_buckets.series,
                series_points: self.metric_buckets.series_points + rhs.metric_buckets.series_points,
                distributions: self.metric_buckets.distributions + rhs.metric_buckets.distributions,
                distributions_points: self.metric_buckets.distributions_points
                    + rhs.metric_buckets.distributions_points,
            },
        }
    }
}

impl Sum for TelemetryWorkerStats {
    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
        iter.fold(Self::default(), |a, b| a + b)
    }
}

mod serialize {
    use crate::data;
    use http::HeaderValue;
    #[allow(clippy::declare_interior_mutable_const)]
    pub const CONTENT_TYPE_VALUE: HeaderValue = libdd_common::header::APPLICATION_JSON;
    pub fn serialize(telemetry: &data::Telemetry) -> anyhow::Result<Vec<u8>> {
        Ok(serde_json::to_vec(telemetry)?)
    }
}

impl TelemetryWorker {
    fn log_err(&self, err: &anyhow::Error) {
        telemetry_worker_log!(self, ERROR, "{}", err);
    }

    async fn recv_next_action(&mut self) -> TelemetryActions {
        let action = if let Some((deadline, deadline_action)) = self.deadlines.next_deadline() {
            // If deadline passed, directly return associated action
            if deadline
                .checked_duration_since(time::Instant::now())
                .is_none()
            {
                return TelemetryActions::Lifecycle(*deadline_action);
            };

            // Otherwise run it in a timeout against the mailbox
            match tokio::time::timeout_at(deadline.into(), self.mailbox.recv()).await {
                Ok(mailbox_action) => mailbox_action,
                Err(_) => Some(TelemetryActions::Lifecycle(*deadline_action)),
            }
        } else {
            self.mailbox.recv().await
        };

        // if no action is received, then it means the channel is stopped
        action.unwrap_or_else(|| {
            // the worker handle no longer lives - we must remove restartable here to avoid leaks
            self.config.restartable = false;
            self.stopped = true;
            TelemetryActions::Lifecycle(LifecycleAction::Stop)
        })
    }

    async fn dispatch_metrics_logs_action(&mut self, action: TelemetryActions) -> ControlFlow<()> {
        telemetry_worker_log!(self, DEBUG, "Handling metric action {:?}", action);
        use LifecycleAction::*;
        use TelemetryActions::*;
        match action {
            Lifecycle(Start) => {
                if !self.data.started {
                    #[allow(clippy::unwrap_used)]
                    self.deadlines
                        .schedule_event(LifecycleAction::FlushMetricAggr)
                        .unwrap();

                    #[allow(clippy::unwrap_used)]
                    self.deadlines
                        .schedule_event(LifecycleAction::FlushData)
                        .unwrap();
                    self.data.started = true;
                }
            }
            AddLog((identifier, log)) => {
                let (l, new) = self.data.logs.get_mut_or_insert(identifier, log);
                if !new {
                    l.count += 1;
                }
            }
            AddPoint((point, key, extra_tags)) => {
                self.data.metric_buckets.add_point(key, point, extra_tags)
            }
            Lifecycle(FlushMetricAggr) => {
                self.data.metric_buckets.flush_aggregates();

                #[allow(clippy::unwrap_used)]
                self.deadlines
                    .schedule_event(LifecycleAction::FlushMetricAggr)
                    .unwrap();
            }
            Lifecycle(FlushData) => {
                if !(self.data.started || self.config.restartable) {
                    return CONTINUE;
                }

                #[allow(clippy::unwrap_used)]
                self.deadlines
                    .schedule_event(LifecycleAction::FlushData)
                    .unwrap();

                let batch = self.build_observability_batch();
                if !batch.is_empty() {
                    let payload = data::Payload::MessageBatch(batch);
                    match self.send_payload(&payload).await {
                        Ok(()) => self.payload_sent_success(&payload),
                        Err(e) => self.log_err(&e),
                    }
                }
            }
            AddConfig(_)
            | AddDependency(_)
            | AddIntegration(_)
            | AddEndpoint(_)
            | Lifecycle(ExtendedHeartbeat) => {}
            Lifecycle(Stop) => {
                if !self.data.started {
                    return BREAK;
                }
                self.data.metric_buckets.flush_aggregates();

                let batch = self.build_observability_batch();
                if !batch.is_empty() {
                    let payload = data::Payload::MessageBatch(batch);
                    match self.send_payload(&payload).await {
                        Ok(()) => {
                            if self.config.restartable {
                                self.payload_sent_success(&payload)
                            }
                        }
                        Err(e) => self.log_err(&e),
                    }
                }

                self.data.started = false;
                if !self.config.restartable {
                    self.deadlines.clear_pending();
                }
                return BREAK;
            }
            CollectStats(stats_sender) => {
                stats_sender.send(self.stats()).ok();
            }
        };
        CONTINUE
    }

    async fn dispatch_action(&mut self, action: TelemetryActions) -> ControlFlow<()> {
        telemetry_worker_log!(self, DEBUG, "Handling action {:?}", action);

        use LifecycleAction::*;
        use TelemetryActions::*;
        match action {
            Lifecycle(Start) => {
                if !self.data.started {
                    let app_started = data::Payload::AppStarted(self.build_app_started());
                    match self.send_payload(&app_started).await {
                        Ok(()) => self.payload_sent_success(&app_started),
                        Err(err) => self.log_err(&err),
                    }

                    #[allow(clippy::unwrap_used)]
                    self.deadlines
                        .schedule_event(LifecycleAction::FlushMetricAggr)
                        .unwrap();

                    #[allow(clippy::unwrap_used)]
                    // flush data should be last to previously flushed metrics are sent
                    self.deadlines
                        .schedule_event(LifecycleAction::FlushData)
                        .unwrap();

                    #[allow(clippy::unwrap_used)]
                    self.deadlines
                        .schedule_event(LifecycleAction::ExtendedHeartbeat)
                        .unwrap();
                    self.data.started = true;
                }
            }
            AddDependency(dep) => self.data.dependencies.insert(dep),
            AddIntegration(integration) => self.data.integrations.insert(integration),
            AddConfig(cfg) => self.data.configurations.insert(cfg),
            AddEndpoint(endpoint) => {
                self.data.endpoints.insert(endpoint);
            }
            AddLog((identifier, log)) => {
                let (l, new) = self.data.logs.get_mut_or_insert(identifier, log);
                if !new {
                    l.count += 1;
                }
            }
            AddPoint((point, key, extra_tags)) => {
                self.data.metric_buckets.add_point(key, point, extra_tags)
            }
            Lifecycle(FlushMetricAggr) => {
                self.data.metric_buckets.flush_aggregates();

                #[allow(clippy::unwrap_used)]
                self.deadlines
                    .schedule_event(LifecycleAction::FlushMetricAggr)
                    .unwrap();
            }
            Lifecycle(FlushData) => {
                if !(self.data.started || self.config.restartable) {
                    return CONTINUE;
                }

                #[allow(clippy::unwrap_used)]
                self.deadlines
                    .schedule_event(LifecycleAction::FlushData)
                    .unwrap();

                let mut batch = self.build_app_events_batch();
                let payload = if batch.is_empty() {
                    data::Payload::AppHeartbeat(())
                } else {
                    batch.push(data::Payload::AppHeartbeat(()));
                    data::Payload::MessageBatch(batch)
                };
                match self.send_payload(&payload).await {
                    Ok(()) => self.payload_sent_success(&payload),
                    Err(err) => self.log_err(&err),
                }

                let batch = self.build_observability_batch();
                if !batch.is_empty() {
                    let payload = data::Payload::MessageBatch(batch);
                    match self.send_payload(&payload).await {
                        Ok(()) => self.payload_sent_success(&payload),
                        Err(err) => self.log_err(&err),
                    }
                }
            }
            Lifecycle(ExtendedHeartbeat) => {
                self.data.dependencies.unflush_stored();
                self.data.integrations.unflush_stored();
                self.data.configurations.unflush_stored();

                let extended_hb = data::Payload::AppExtendedHeartbeat(self.build_app_started());
                match self.send_payload(&extended_hb).await {
                    Ok(()) => self.payload_sent_success(&extended_hb),
                    Err(err) => self.log_err(&err),
                }
                // Only re-schedule self. Resetting `FlushData` here would replace its
                // existing deadline with `now + heartbeat_interval`, starving FlushData
                // when `extended_heartbeat_interval < heartbeat_interval` because each
                // ExtendedHeartbeat firing pushes FlushData out before it can fire.
                #[allow(clippy::unwrap_used)]
                self.deadlines
                    .schedule_event(LifecycleAction::ExtendedHeartbeat)
                    .unwrap();
            }
            Lifecycle(Stop) => {
                if !self.data.started {
                    return BREAK;
                }
                self.data.metric_buckets.flush_aggregates();

                let mut app_events = self.build_app_events_batch();
                app_events.push(data::Payload::AppClosing(()));

                let observability_events = self.build_observability_batch();

                let mut payloads = vec![data::Payload::MessageBatch(app_events)];
                if !observability_events.is_empty() {
                    payloads.push(data::Payload::MessageBatch(observability_events));
                }

                let self_arc = Arc::new(tokio::sync::RwLock::new(&mut *self));
                let futures = payloads.into_iter().map(|payload| {
                    let self_arc = self_arc.clone();
                    async move {
                        // This is different from the non-functional:
                        // match self_arc.read().await.send_payload(&payload).await { ... }
                        // presumably because the temp read guard would live till end of match
                        let res = {
                            let self_rguard = self_arc.read().await;
                            self_rguard.send_payload(&payload).await
                        };
                        match res {
                            Ok(()) => self_arc.write().await.payload_sent_success(&payload),
                            Err(err) => self_arc.read().await.log_err(&err),
                        }
                    }
                });
                future::join_all(futures).await;

                self.data.started = false;
                if !self.config.restartable {
                    self.deadlines.clear_pending();
                }

                return BREAK;
            }
            CollectStats(stats_sender) => {
                stats_sender.send(self.stats()).ok();
            }
        }

        CONTINUE
    }

    // Builds telemetry payloads containing lifecycle events
    fn build_app_events_batch(&mut self) -> Vec<Payload> {
        let mut payloads = Vec::new();

        if self.data.dependencies.flush_not_empty() {
            payloads.push(data::Payload::AppDependenciesLoaded(
                data::AppDependenciesLoaded {
                    dependencies: self.data.dependencies.unflushed().cloned().collect(),
                },
            ))
        }
        if self.data.integrations.flush_not_empty() {
            payloads.push(data::Payload::AppIntegrationsChange(
                data::AppIntegrationsChange {
                    integrations: self.data.integrations.unflushed().cloned().collect(),
                },
            ))
        }
        if self.data.configurations.flush_not_empty() {
            payloads.push(data::Payload::AppClientConfigurationChange(
                data::AppClientConfigurationChange {
                    configuration: self.data.configurations.unflushed().cloned().collect(),
                },
            ))
        }
        if !self.data.endpoints.is_empty() {
            payloads.push(data::Payload::AppEndpoints(data::AppEndpoints {
                is_first: true,
                endpoints: self
                    .data
                    .endpoints
                    .iter()
                    .map(|e| e.to_json_value().unwrap_or_default())
                    .filter(|e| e.is_object())
                    .collect(),
            }));
        }
        payloads
    }

    // Builds telemetry payloads containing logs, metrics and distributions
    fn build_observability_batch(&mut self) -> Vec<Payload> {
        let mut payloads = Vec::new();

        let logs = self.build_logs();
        if !logs.logs.is_empty() {
            payloads.push(data::Payload::Logs(logs));
        }
        let metrics = self.build_metrics_series();
        if !metrics.series.is_empty() {
            payloads.push(data::Payload::GenerateMetrics(metrics))
        }
        let distributions = self.build_metrics_distributions();
        if !distributions.series.is_empty() {
            payloads.push(data::Payload::Sketches(distributions))
        }
        payloads
    }

    fn build_metrics_distributions(&mut self) -> data::Distributions {
        let mut series = Vec::new();
        let context_guard = self.data.metric_contexts.lock();
        for (context_key, extra_tags, points) in self.data.metric_buckets.flush_distributions() {
            let Some(context) = context_guard.read(context_key) else {
                telemetry_worker_log!(self, ERROR, "Context not found for key {:?}", context_key);
                continue;
            };
            let mut tags = extra_tags;
            tags.extend(context.tags.iter().cloned());
            series.push(data::metrics::Distribution {
                namespace: context.namespace,
                metric: context.name.clone(),
                tags,
                sketch: data::metrics::SerializedSketch::B64 {
                    sketch_b64: base64::Engine::encode(
                        &base64::engine::general_purpose::STANDARD,
                        points.encode_to_vec(),
                    ),
                },
                common: context.common,
                _type: context.metric_type,
                interval: self.metrics_flush_interval.as_secs(),
            });
        }
        data::Distributions { series }
    }

    fn build_metrics_series(&mut self) -> data::GenerateMetrics {
        let mut series = Vec::new();
        let context_guard = self.data.metric_contexts.lock();
        for (context_key, extra_tags, points) in self.data.metric_buckets.flush_series() {
            let Some(context) = context_guard.read(context_key) else {
                telemetry_worker_log!(self, ERROR, "Context not found for key {:?}", context_key);
                continue;
            };

            let mut tags = extra_tags;
            tags.extend(context.tags.iter().cloned());
            series.push(data::metrics::Serie {
                namespace: context.namespace,
                metric: context.name.clone(),
                tags,
                points,
                common: context.common,
                _type: context.metric_type,
                interval: self.metrics_flush_interval.as_secs(),
            });
        }

        data::GenerateMetrics { series }
    }

    fn build_app_started(&mut self) -> data::AppStarted {
        data::AppStarted {
            configuration: self.data.configurations.unflushed().cloned().collect(),
            dependencies: self.data.dependencies.unflushed().cloned().collect(),
            integrations: self.data.integrations.unflushed().cloned().collect(),
        }
    }

    fn app_started_sent_success(&mut self, p: &data::AppStarted) {
        self.data
            .configurations
            .removed_flushed(p.configuration.len());
        self.data.dependencies.removed_flushed(p.dependencies.len());
        self.data.integrations.removed_flushed(p.integrations.len());
    }

    fn payload_sent_success(&mut self, payload: &data::Payload) {
        use data::Payload::*;
        match payload {
            AppStarted(p) => self.app_started_sent_success(p),
            AppExtendedHeartbeat(p) => self.app_started_sent_success(p),
            AppDependenciesLoaded(p) => {
                self.data.dependencies.removed_flushed(p.dependencies.len())
            }
            AppIntegrationsChange(p) => {
                self.data.integrations.removed_flushed(p.integrations.len())
            }
            AppClientConfigurationChange(p) => self
                .data
                .configurations
                .removed_flushed(p.configuration.len()),
            AppEndpoints(_) => self.data.endpoints.clear(),
            MessageBatch(batch) => {
                for p in batch {
                    self.payload_sent_success(p);
                }
            }
            Logs(p) => {
                for _ in &p.logs {
                    self.data.logs.pop_front();
                }
            }
            AppHeartbeat(()) | AppClosing(()) => {}
            GenerateMetrics(_) | Sketches(_) => {}
        }
    }

    fn build_logs(&self) -> data::Logs {
        // TODO: change the data model to take a &[Log] so don't have to clone data here
        let logs = self.data.logs.iter().map(|(_, l)| l.clone()).collect();
        data::Logs { logs }
    }

    fn next_seq_id(&self) -> u64 {
        self.seq_id.fetch_add(1, Ordering::Release)
    }

    async fn send_payload(&self, payload: &data::Payload) -> anyhow::Result<()> {
        debug!(
            worker.runtime_id = %self.runtime_id,
            payload.type = payload.request_type(),
            seq_id = self.seq_id.load(Ordering::Acquire),
            "Sending telemetry payload"
        );
        let req = self.build_request(payload)?;
        let result = self.send_request(req).await;
        match &result {
            Ok(resp) => debug!(
                worker.runtime_id = %self.runtime_id,
                payload.type = payload.request_type(),
                response.status = resp.status().as_u16(),
                "Successfully sent telemetry payload"
            ),
            Err(e) => debug!(
                worker.runtime_id = %self.runtime_id,
                payload.type = payload.request_type(),
                error = ?e,
                "Failed to send telemetry payload"
            ),
        }
        Ok(())
    }

    fn build_request(&self, payload: &data::Payload) -> anyhow::Result<http_common::HttpRequest> {
        let seq_id = self.next_seq_id();
        let tel = Telemetry {
            api_version: data::ApiVersion::V2,
            tracer_time: time::SystemTime::UNIX_EPOCH
                .elapsed()
                .map_or(0, |d| d.as_secs()),
            runtime_id: &self.runtime_id,
            seq_id,
            host: &self.data.host,
            origin: None,
            application: &self.data.app,
            payload,
        };

        telemetry_worker_log!(self, DEBUG, "Prepared payload: {:?}", tel);

        let req = http_client::request_builder(&self.config)?
            .method(http::Method::POST)
            .header(header::CONTENT_TYPE, serialize::CONTENT_TYPE_VALUE)
            .header(
                http_client::header::REQUEST_TYPE,
                HeaderValue::from_static(payload.request_type()),
            )
            .header(
                http_client::header::API_VERSION,
                HeaderValue::from_static(data::ApiVersion::V2.to_str()),
            )
            .header(
                http_client::header::LIBRARY_LANGUAGE,
                tel.application.language_name.clone(),
            )
            .header(
                http_client::header::LIBRARY_VERSION,
                tel.application.tracer_version.clone(),
            );
        let req = http_client::add_instrumentation_session_headers(
            req,
            self.config.session_id.as_deref(),
            self.config.parent_session_id.as_deref(),
            self.config.root_session_id.as_deref(),
        );

        let body = http_common::Body::from(serialize::serialize(&tel)?);
        Ok(req.body(body)?)
    }

    async fn send_request(
        &self,
        req: http_common::HttpRequest,
    ) -> Result<http_common::HttpResponse, http_common::Error> {
        let timeout_ms = if let Some(endpoint) = self.config.endpoint.as_ref() {
            endpoint.timeout_ms
        } else {
            libdd_common::Endpoint::DEFAULT_TIMEOUT
        };

        debug!(
            worker.runtime_id = %self.runtime_id,
            http.timeout_ms = timeout_ms,
            "Sending HTTP request"
        );

        tokio::select! {
            _ = self.cancellation_token.cancelled() => {
                debug!(
                    worker.runtime_id = %self.runtime_id,
                    "Telemetry request cancelled"
                );
                Err(http_common::Error::Other(anyhow::anyhow!("Request cancelled")))
            },
            _ = tokio::time::sleep(time::Duration::from_millis(timeout_ms)) => {
                debug!(
                    worker.runtime_id = %self.runtime_id,
                    http.timeout_ms = timeout_ms,
                    "Telemetry request timed out"
                );
                Err(http_common::Error::Other(anyhow::anyhow!("Request timed out")))
            },
            r = self.client.request(req) => {
                match r {
                    Ok(resp) => {
                        Ok(resp)
                    }
                    Err(e) => {
                        Err(e)
                    },
                }
            }
        }
    }

    fn stats(&self) -> TelemetryWorkerStats {
        TelemetryWorkerStats {
            dependencies_stored: self.data.dependencies.len_stored() as u32,
            dependencies_unflushed: self.data.dependencies.len_unflushed() as u32,
            configurations_stored: self.data.configurations.len_stored() as u32,
            configurations_unflushed: self.data.configurations.len_unflushed() as u32,
            integrations_stored: self.data.integrations.len_stored() as u32,
            integrations_unflushed: self.data.integrations.len_unflushed() as u32,
            logs: self.data.logs.len() as u32,
            metric_contexts: self.data.metric_contexts.lock().len() as u32,
            metric_buckets: self.data.metric_buckets.stats(),
        }
    }

    // Runs a state machine that waits for actions, either from the worker's
    // mailbox, or scheduled actions from the worker's deadline object.
    async fn run_loop(mut self) {
        debug!(
            worker.flavor = ?self.flavor,
            worker.runtime_id = %self.runtime_id,
            "Starting telemetry worker"
        );

        loop {
            if self.cancellation_token.is_cancelled() {
                debug!(
                    worker.runtime_id = %self.runtime_id,
                    "Telemetry worker cancelled, shutting down"
                );
                return;
            }

            let action = self.recv_next_action().await;
            debug!(
                worker.runtime_id = %self.runtime_id,
                action = ?action,
                "Received telemetry action"
            );

            let action_result = match self.flavor {
                TelemetryWorkerFlavor::Full => self.dispatch_action(action).await,
                TelemetryWorkerFlavor::MetricsLogs => {
                    self.dispatch_metrics_logs_action(action).await
                }
            };

            match action_result {
                ControlFlow::Continue(()) => {}
                ControlFlow::Break(()) => {
                    debug!(
                        worker.runtime_id = %self.runtime_id,
                        worker.restartable = self.config.restartable,
                        "Telemetry worker received break signal"
                    );
                    if !self.config.restartable {
                        break;
                    }
                }
            };
        }

        debug!(
            worker.runtime_id = %self.runtime_id,
            "Telemetry worker stopped"
        );
    }
}

#[derive(Debug)]
struct InnerTelemetryShutdown {
    is_shutdown: Mutex<bool>,
    condvar: Condvar,
}

impl InnerTelemetryShutdown {
    fn wait_for_shutdown(&self) {
        drop(
            #[allow(clippy::unwrap_used)]
            self.condvar
                .wait_while(self.is_shutdown.lock().unwrap(), |is_shutdown| {
                    !*is_shutdown
                })
                .unwrap(),
        )
    }

    #[allow(clippy::unwrap_used)]
    fn shutdown_finished(&self) {
        *self.is_shutdown.lock().unwrap() = true;
        self.condvar.notify_all();
    }
}

#[derive(Clone, Debug)]
/// TelemetryWorkerHandle is a handle which allows interactions with the telemetry worker.
/// The handle is safe to use across threads.
///
/// The worker won't send data to the agent until you call `TelemetryWorkerHandle::send_start`
///
/// To stop the worker, call `TelemetryWorkerHandle::send_stop` which trigger flush asynchronously
/// then `TelemetryWorkerHandle::wait_for_shutdown`
pub struct TelemetryWorkerHandle {
    sender: mpsc::Sender<TelemetryActions>,
    shutdown: Arc<InnerTelemetryShutdown>,
    cancellation_token: CancellationToken,
    // Used to spawn cancellation tasks. Should be None when running as a SharedRuntime worker,
    // since the runtime is not guaranteed to exist for the lifetime of the worker.
    runtime: Option<runtime::Handle>,

    contexts: MetricContexts,
}

impl TelemetryWorkerHandle {
    pub fn register_metric_context(
        &self,
        name: String,
        tags: Vec<Tag>,
        metric_type: data::metrics::MetricType,
        common: bool,
        namespace: data::metrics::MetricNamespace,
    ) -> ContextKey {
        self.contexts
            .register_metric_context(name, tags, metric_type, common, namespace)
    }

    pub fn try_send_msg(&self, msg: TelemetryActions) -> anyhow::Result<()> {
        Ok(self.sender.try_send(msg)?)
    }

    pub async fn send_msg(&self, msg: TelemetryActions) -> anyhow::Result<()> {
        Ok(self.sender.send(msg).await?)
    }

    pub async fn send_msgs<T>(&self, msgs: T) -> anyhow::Result<()>
    where
        T: IntoIterator<Item = TelemetryActions>,
    {
        for msg in msgs {
            self.sender.send(msg).await?;
        }

        Ok(())
    }

    pub async fn send_msg_timeout(
        &self,
        msg: TelemetryActions,
        timeout: time::Duration,
    ) -> anyhow::Result<()> {
        Ok(self.sender.send_timeout(msg, timeout).await?)
    }

    pub fn send_start(&self) -> anyhow::Result<()> {
        Ok(self
            .sender
            .try_send(TelemetryActions::Lifecycle(LifecycleAction::Start))?)
    }

    pub fn send_stop(&self) -> anyhow::Result<()> {
        Ok(self
            .sender
            .try_send(TelemetryActions::Lifecycle(LifecycleAction::Stop))?)
    }

    fn cancel_requests_with_deadline(&self, deadline: time::Instant) {
        let Some(runtime) = &self.runtime else {
            tracing::error!("Cannot schedule cancellation deadline: no runtime handle available");
            return;
        };
        let token = self.cancellation_token.clone();
        let f = async move {
            tokio::time::sleep_until(deadline.into()).await;
            token.cancel()
        };
        runtime.spawn(f);
    }

    pub fn wait_for_shutdown_deadline(&self, deadline: time::Instant) {
        self.cancel_requests_with_deadline(deadline);
        self.wait_for_shutdown()
    }

    pub fn add_dependency(&self, name: String, version: Option<String>) -> anyhow::Result<()> {
        self.sender
            .try_send(TelemetryActions::AddDependency(Dependency {
                name,
                version,
            }))?;
        Ok(())
    }

    pub fn add_integration(
        &self,
        name: String,
        enabled: bool,
        version: Option<String>,
        compatible: Option<bool>,
        auto_enabled: Option<bool>,
    ) -> anyhow::Result<()> {
        self.sender
            .try_send(TelemetryActions::AddIntegration(Integration {
                name,
                version,
                compatible,
                enabled,
                auto_enabled,
            }))?;
        Ok(())
    }

    pub fn add_log<T: Hash>(
        &self,
        identifier: T,
        message: String,
        level: data::LogLevel,
        stack_trace: Option<String>,
    ) -> anyhow::Result<()> {
        let mut hasher = DefaultHasher::new();
        identifier.hash(&mut hasher);
        self.sender.try_send(TelemetryActions::AddLog((
            LogIdentifier {
                identifier: hasher.finish(),
            },
            data::Log {
                message,
                level,
                stack_trace,
                count: 1,
                tags: String::new(),
                is_sensitive: false,
                is_crash: false,
            },
        )))?;
        Ok(())
    }

    pub fn add_point(
        &self,
        value: f64,
        context: &ContextKey,
        extra_tags: Vec<Tag>,
    ) -> anyhow::Result<()> {
        self.sender
            .try_send(TelemetryActions::AddPoint((value, *context, extra_tags)))?;
        Ok(())
    }

    pub fn wait_for_shutdown(&self) {
        self.shutdown.wait_for_shutdown();
    }

    pub fn stats(&self) -> anyhow::Result<oneshot::Receiver<TelemetryWorkerStats>> {
        let (sender, receiver) = oneshot::channel();
        self.sender
            .try_send(TelemetryActions::CollectStats(sender))?;
        Ok(receiver)
    }
}

/// How many dependencies/integrations/configs we keep in memory at most
pub const MAX_ITEMS: usize = 5000;

#[derive(Debug, Default, Clone, Copy)]
pub enum TelemetryWorkerFlavor {
    /// Send all telemetry messages including lifecycle events like app-started, heartbeats,
    /// dependencies and configurations
    #[default]
    Full,
    /// Only send telemetry data not tied to the lifecycle of the app like logs and metrics
    MetricsLogs,
}

pub struct TelemetryWorkerBuilder {
    pub host: Host,
    pub application: Application,
    pub runtime_id: Option<String>,
    pub dependencies: store::Store<data::Dependency>,
    pub integrations: store::Store<data::Integration>,
    pub configurations: store::Store<data::Configuration>,
    pub endpoints: HashSet<data::Endpoint>,
    pub native_deps: bool,
    pub rust_shared_lib_deps: bool,
    pub config: Config,
    pub flavor: TelemetryWorkerFlavor,
}

impl TelemetryWorkerBuilder {
    /// Creates a new telemetry worker builder and infer host information automatically
    pub fn new_fetch_host(
        service_name: String,
        language_name: String,
        language_version: String,
        tracer_version: String,
    ) -> Self {
        Self {
            host: crate::build_host(),
            ..Self::new(
                String::new(),
                service_name,
                language_name,
                language_version,
                tracer_version,
            )
        }
    }

    /// Creates a new telemetry worker builder with the given hostname
    pub fn new(
        hostname: String,
        service_name: String,
        language_name: String,
        language_version: String,
        tracer_version: String,
    ) -> Self {
        Self {
            host: Host {
                hostname,
                ..Default::default()
            },
            application: Application {
                service_name,
                language_name,
                language_version,
                tracer_version,
                ..Default::default()
            },
            runtime_id: None,
            dependencies: store::Store::new(MAX_ITEMS),
            integrations: store::Store::new(MAX_ITEMS),
            configurations: store::Store::new(MAX_ITEMS),
            endpoints: HashSet::new(),
            native_deps: true,
            rust_shared_lib_deps: false,
            config: Config::default(),
            flavor: TelemetryWorkerFlavor::default(),
        }
    }

    /// Build the corresponding worker and its handle.
    ///
    /// The optional runtime handle is stored in the worker handle and should be the one used to run
    /// the worker task cancellation deadlines. Pass `None` when the worker will be run via a
    /// [`SharedRuntime`](libdd_shared_runtime::SharedRuntime).
    pub fn build_worker(
        self,
        tokio_runtime: Option<Handle>,
    ) -> (TelemetryWorkerHandle, TelemetryWorker) {
        let (tx, mailbox) = mpsc::channel(5000);
        let shutdown = Arc::new(InnerTelemetryShutdown {
            is_shutdown: Mutex::new(false),
            condvar: Condvar::new(),
        });
        let contexts = MetricContexts::default();
        let token = CancellationToken::new();
        let config = self.config;
        let telemetry_heartbeat_interval = config.telemetry_heartbeat_interval;
        let telemetry_extended_heartbeat_interval = config.telemetry_extended_heartbeat_interval;
        let client = http_client::from_config(&config);

        let metrics_flush_interval =
            telemetry_heartbeat_interval.min(MetricBuckets::METRICS_FLUSH_INTERVAL);

        #[allow(clippy::unwrap_used)]
        let worker = TelemetryWorker {
            flavor: self.flavor,
            data: TelemetryWorkerData {
                started: false,
                dependencies: self.dependencies,
                integrations: self.integrations,
                configurations: self.configurations,
                endpoints: self.endpoints,
                logs: store::QueueHashMap::default(),
                metric_contexts: contexts.clone(),
                metric_buckets: MetricBuckets::default(),
                host: self.host,
                app: self.application,
            },
            config,
            mailbox,
            seq_id: AtomicU64::new(1),
            runtime_id: self
                .runtime_id
                .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()),
            client,
            metrics_flush_interval,
            deadlines: scheduler::Scheduler::new(vec![
                (metrics_flush_interval, LifecycleAction::FlushMetricAggr),
                (telemetry_heartbeat_interval, LifecycleAction::FlushData),
                (
                    telemetry_extended_heartbeat_interval,
                    LifecycleAction::ExtendedHeartbeat,
                ),
            ]),
            cancellation_token: token.clone(),
            next_action: None,
            stopped: false,
        };

        (
            TelemetryWorkerHandle {
                sender: tx,
                shutdown,
                cancellation_token: token,
                runtime: tokio_runtime,

                contexts,
            },
            worker,
        )
    }

    /// Spawns a telemetry worker task in the current tokio runtime
    /// The worker will capture a reference to the runtime and use it to run its tasks
    pub fn spawn(self) -> (TelemetryWorkerHandle, JoinHandle<()>) {
        let tokio_runtime = tokio::runtime::Handle::current();

        let (worker_handle, worker) = self.build_worker(Some(tokio_runtime.clone()));

        let join_handle = tokio_runtime.spawn(async move { worker.run_loop().await });

        (worker_handle, join_handle)
    }

    /// Spawns a telemetry worker in a new thread and returns a handle to interact with it
    pub fn run(self) -> anyhow::Result<TelemetryWorkerHandle> {
        let runtime = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()?;
        let (handle, worker) = self.build_worker(Some(runtime.handle().clone()));
        let notify_shutdown = handle.shutdown.clone();
        std::thread::spawn(move || {
            runtime.block_on(worker.run_loop());
            runtime.shutdown_background();
            notify_shutdown.shutdown_finished();
        });

        Ok(handle)
    }
}

#[cfg(test)]
mod tests {
    use crate::data::Payload;
    use crate::worker::http_client::header::{
        DD_PARENT_SESSION_ID, DD_ROOT_SESSION_ID, DD_SESSION_ID,
    };
    use crate::worker::{
        LifecycleAction, TelemetryActions, TelemetryWorker, TelemetryWorkerBuilder,
        TelemetryWorkerFlavor, TelemetryWorkerHandle,
    };
    use libdd_common::{http_common, Endpoint};
    use tokio::runtime::Runtime;

    fn is_send<T: Send>(_: T) {}
    fn is_sync<T: Sync>(_: T) {}

    #[test]
    fn test_handle_sync_send() {
        #[allow(clippy::redundant_closure)]
        let _ = |h: TelemetryWorkerHandle| is_send(h);
        #[allow(clippy::redundant_closure)]
        let _ = |h: TelemetryWorkerHandle| is_sync(h);
    }

    fn test_worker(
        session_id: Option<String>,
        root_session_id: Option<String>,
        parent_session_id: Option<String>,
    ) -> TelemetryWorker {
        let mut b = TelemetryWorkerBuilder::new(
            "h".into(),
            "svc".into(),
            "lang".into(),
            "1".into(),
            "tv".into(),
        );
        b.config
            .set_endpoint(Endpoint::from_slice("http://127.0.0.1:1"))
            .unwrap();
        b.runtime_id = Some("rid".into());
        b.config.session_id = session_id;
        b.config.parent_session_id = parent_session_id;
        b.config.root_session_id = root_session_id;
        let rt = Runtime::new().unwrap();
        b.build_worker(Some(rt.handle().clone())).1
    }

    #[test]
    fn telemetry_http_includes_dd_session_id() {
        let req = test_worker(Some("sess".into()), None, None)
            .build_request(&Payload::AppHeartbeat(()))
            .unwrap();
        assert_eq!(
            req.headers().get(DD_SESSION_ID).unwrap().to_str().unwrap(),
            "sess"
        );
        assert!(req.headers().get(DD_ROOT_SESSION_ID).is_none());
        assert!(req.headers().get(DD_PARENT_SESSION_ID).is_none());
    }

    #[test]
    fn telemetry_http_omits_root_session_id_when_same_as_session_id() {
        let req = test_worker(
            Some("sess-id".into()),
            Some("sess-id".into()),
            Some("parent".into()),
        )
        .build_request(&Payload::AppHeartbeat(()))
        .unwrap();
        assert_eq!(
            req.headers().get(DD_SESSION_ID).unwrap().to_str().unwrap(),
            "sess-id"
        );
        assert!(req.headers().get(DD_ROOT_SESSION_ID).is_none());
        assert_eq!(
            req.headers()
                .get(DD_PARENT_SESSION_ID)
                .unwrap()
                .to_str()
                .unwrap(),
            "parent"
        );
    }

    #[test]
    fn telemetry_http_omits_parent_session_id_when_same_as_session_id() {
        let req = test_worker(
            Some("sess-id".into()),
            Some("root".into()),
            Some("sess-id".into()),
        )
        .build_request(&Payload::AppHeartbeat(()))
        .unwrap();
        assert_eq!(
            req.headers().get(DD_SESSION_ID).unwrap().to_str().unwrap(),
            "sess-id"
        );
        assert_eq!(
            req.headers()
                .get(DD_ROOT_SESSION_ID)
                .unwrap()
                .to_str()
                .unwrap(),
            "root"
        );
        assert!(req.headers().get(DD_PARENT_SESSION_ID).is_none());
    }

    #[test]
    fn telemetry_http_omits_session_family_without_valid_session_id() {
        let assert_no_session_headers = |req: &http_common::HttpRequest| {
            assert!(req.headers().get(DD_SESSION_ID).is_none());
            assert!(req.headers().get(DD_ROOT_SESSION_ID).is_none());
            assert!(req.headers().get(DD_PARENT_SESSION_ID).is_none());
        };

        let req = test_worker(None, Some("root".into()), Some("parent".into()))
            .build_request(&Payload::AppHeartbeat(()))
            .unwrap();
        assert_no_session_headers(&req);

        let req = test_worker(
            Some(String::new()),
            Some("root".into()),
            Some("parent".into()),
        )
        .build_request(&Payload::AppHeartbeat(()))
        .unwrap();
        assert_no_session_headers(&req);
    }

    #[test]
    fn telemetry_http_includes_dd_session_root_and_parent_session_ids() {
        let req = test_worker(
            Some("sess".into()),
            Some("root".into()),
            Some("parent".into()),
        )
        .build_request(&Payload::AppHeartbeat(()))
        .unwrap();
        assert_eq!(
            req.headers().get(DD_SESSION_ID).unwrap().to_str().unwrap(),
            "sess"
        );
        assert_eq!(
            req.headers()
                .get(DD_ROOT_SESSION_ID)
                .unwrap()
                .to_str()
                .unwrap(),
            "root"
        );
        assert_eq!(
            req.headers()
                .get(DD_PARENT_SESSION_ID)
                .unwrap()
                .to_str()
                .unwrap(),
            "parent"
        );
    }

    fn build_test_worker_with_flavor(flavor: TelemetryWorkerFlavor) -> TelemetryWorker {
        let mut b = TelemetryWorkerBuilder::new(
            "h".into(),
            "svc".into(),
            "lang".into(),
            "1".into(),
            "tv".into(),
        );
        b.config
            .set_endpoint(Endpoint::from_slice("http://127.0.0.1:1"))
            .unwrap();
        b.runtime_id = Some("rid".into());
        b.flavor = flavor;
        b.build_worker(Some(tokio::runtime::Handle::current())).1
    }

    /// Every event with a delay must be scheduled on Start; otherwise it sits in
    /// `delays` forever and its handler never fires. Walking `delays` (rather than
    /// enumerating variants) guards against future periodic actions regressing.
    #[tokio::test]
    #[cfg_attr(miri, ignore)] // reqwest in dispatch_action
    async fn full_flavor_start_schedules_every_periodic_action() {
        let mut worker = build_test_worker_with_flavor(TelemetryWorkerFlavor::Full);

        let _ = worker
            .dispatch_action(TelemetryActions::Lifecycle(LifecycleAction::Start))
            .await;

        let delays: Vec<LifecycleAction> =
            worker.deadlines.delays.iter().map(|(_, k)| *k).collect();
        let scheduled: Vec<LifecycleAction> =
            worker.deadlines.deadlines.iter().map(|(_, k)| *k).collect();

        assert!(!delays.is_empty(), "scheduler should have periodic actions");
        for ev in &delays {
            assert!(
                scheduled.contains(ev),
                "{ev:?} has a delay but was not scheduled on Start; scheduled={scheduled:?}",
            );
        }
    }

    /// `MetricsLogs` flavor intentionally excludes lifecycle events. Negative guard
    /// so any future change emitting them from this flavor has to update the test.
    #[tokio::test]
    #[cfg_attr(miri, ignore)] // reqwest in build_worker
    async fn metrics_logs_flavor_start_does_not_schedule_extended_heartbeat() {
        let mut worker = build_test_worker_with_flavor(TelemetryWorkerFlavor::MetricsLogs);

        let _ = worker
            .dispatch_metrics_logs_action(TelemetryActions::Lifecycle(LifecycleAction::Start))
            .await;

        let scheduled: Vec<LifecycleAction> =
            worker.deadlines.deadlines.iter().map(|(_, k)| *k).collect();

        assert!(scheduled.contains(&LifecycleAction::FlushMetricAggr));
        assert!(scheduled.contains(&LifecycleAction::FlushData));
        assert!(
            !scheduled.contains(&LifecycleAction::ExtendedHeartbeat),
            "MetricsLogs should not schedule ExtendedHeartbeat; scheduled={scheduled:?}",
        );
    }

    /// Regression: when `extended_heartbeat_interval < heartbeat_interval`, the
    /// ExtendedHeartbeat handler must not reset FlushData's deadline. If it did, each
    /// firing would push FlushData to `now + heartbeat_interval` and the next
    /// (sooner) ExtendedHeartbeat would push it again — starving FlushData forever.
    #[tokio::test]
    #[cfg_attr(miri, ignore)] // reqwest in dispatch_action
    async fn extended_heartbeat_does_not_reset_flush_data() {
        let mut worker = build_test_worker_with_flavor(TelemetryWorkerFlavor::Full);

        let _ = worker
            .dispatch_action(TelemetryActions::Lifecycle(LifecycleAction::Start))
            .await;

        let flush_data_before = worker
            .deadlines
            .deadlines
            .iter()
            .find(|(_, k)| *k == LifecycleAction::FlushData)
            .map(|(d, _)| *d)
            .expect("FlushData scheduled on Start");

        let _ = worker
            .dispatch_action(TelemetryActions::Lifecycle(
                LifecycleAction::ExtendedHeartbeat,
            ))
            .await;

        let flush_data_after = worker
            .deadlines
            .deadlines
            .iter()
            .find(|(_, k)| *k == LifecycleAction::FlushData)
            .map(|(d, _)| *d)
            .expect("FlushData should still be scheduled after ExtendedHeartbeat fires");

        assert_eq!(
            flush_data_before, flush_data_after,
            "ExtendedHeartbeat must not reset FlushData's deadline",
        );
    }

    mod reset {
        use super::super::*;
        use crate::data::{
            metrics::{MetricNamespace, MetricType},
            Configuration, ConfigurationOrigin, Dependency, Endpoint, Integration, Log, LogLevel,
        };
        use libdd_shared_runtime::Worker;

        fn build_test_worker() -> (TelemetryWorkerHandle, TelemetryWorker) {
            let builder = TelemetryWorkerBuilder::new(
                "hostname".to_string(),
                "test-service".to_string(),
                "rust".to_string(),
                "1.0.0".to_string(),
                "1.0.0".to_string(),
            );
            // build_worker requires a tokio Handle; tests using this must be #[tokio::test]
            builder.build_worker(Some(tokio::runtime::Handle::current()))
        }

        fn make_log(id: u64, message: &str) -> (LogIdentifier, Log) {
            (
                LogIdentifier { identifier: id },
                Log {
                    message: message.to_string(),
                    level: LogLevel::Warn,
                    stack_trace: None,
                    count: 1,
                    tags: String::new(),
                    is_sensitive: false,
                    is_crash: false,
                },
            )
        }

        /// After reset(), pending buffered telemetry and dedupe history is cleared.
        #[tokio::test]
        async fn test_reset_clears_buffered_data() {
            let (handle, mut worker) = build_test_worker();

            // Populate every data field that reset() should clear.
            worker.data.dependencies.insert(Dependency {
                name: "dep".to_string(),
                version: None,
            });
            worker.data.integrations.insert(Integration {
                name: "integration".to_string(),
                version: None,
                enabled: true,
                compatible: None,
                auto_enabled: None,
            });
            worker.data.configurations.insert(Configuration {
                name: "cfg".to_string(),
                value: "true".to_string(),
                origin: ConfigurationOrigin::Code,
                config_id: None,
                seq_id: None,
            });
            worker.data.endpoints.insert(Endpoint {
                operation_name: "GET /health".to_string(),
                resource_name: "/health".to_string(),
                ..Default::default()
            });
            let (id, log) = make_log(42, "msg");
            worker.data.logs.get_mut_or_insert(id, log);

            // Register a metric context and add a data point.
            let key = handle.register_metric_context(
                "test.metric".to_string(),
                vec![],
                MetricType::Count,
                false,
                MetricNamespace::Tracers,
            );
            worker.data.metric_buckets.add_point(key, 1.0, vec![]);

            worker.reset();

            let stats = worker.stats();
            assert_eq!(
                stats.dependencies_stored, 0,
                "dependency dedupe history should be cleared"
            );
            assert_eq!(
                stats.dependencies_unflushed, 0,
                "dependency pending queue should be cleared"
            );
            assert_eq!(
                stats.integrations_stored, 0,
                "integration dedupe history should be cleared"
            );
            assert_eq!(
                stats.integrations_unflushed, 0,
                "integration pending queue should be cleared"
            );
            assert_eq!(
                stats.configurations_stored, 0,
                "configuration dedupe history should be cleared"
            );
            assert_eq!(
                stats.configurations_unflushed, 0,
                "configuration pending queue should be cleared"
            );
            assert_eq!(stats.logs, 0, "logs should be cleared");
            assert_eq!(
                stats.metric_buckets.buckets, 0,
                "metric buckets should be cleared"
            );
            assert_eq!(
                stats.metric_buckets.series, 0,
                "metric series should be cleared"
            );
            assert!(
                worker.data.endpoints.is_empty(),
                "endpoints should be cleared"
            );
            assert!(worker.next_action.is_none(), "next_action should be None");
        }

        /// After reset(), actions queued in the mailbox before the fork are discarded.
        #[tokio::test]
        async fn test_reset_drains_mailbox() {
            let (handle, mut worker) = build_test_worker();

            // Enqueue several actions that should be discarded.
            handle
                .try_send_msg(TelemetryActions::AddDependency(Dependency {
                    name: "dep".to_string(),
                    version: None,
                }))
                .unwrap();
            let (id, log) = make_log(1, "pre-fork log");
            handle
                .try_send_msg(TelemetryActions::AddLog((id, log)))
                .unwrap();

            // Stage one action as if trigger() had already stored it.
            worker.next_action = Some(TelemetryActions::Lifecycle(LifecycleAction::Start));

            worker.reset();

            // The mailbox must be empty and next_action cleared.
            assert!(
                worker.mailbox.try_recv().is_err(),
                "mailbox should be empty"
            );
            assert!(worker.next_action.is_none(), "next_action should be None");
            // None of the queued actions should have been applied to pending state.
            let stats = worker.stats();
            assert_eq!(
                stats.dependencies_stored, 0,
                "queued AddDependency must not be applied"
            );
            assert_eq!(
                stats.dependencies_unflushed, 0,
                "queued AddDependency must not be pending"
            );
            assert_eq!(stats.logs, 0, "queued AddLog must be discarded");
        }

        /// After reset(), the worker accepts new telemetry and processes it normally.
        #[tokio::test]
        async fn test_worker_accepts_new_data_after_reset() {
            let (handle, mut worker) = build_test_worker();
            worker.flavor = TelemetryWorkerFlavor::MetricsLogs;

            // Populate state before reset – this data must not survive.
            let (id, log) = make_log(99, "pre-fork");
            worker.data.logs.get_mut_or_insert(id, log);

            worker.reset();

            // Send a new log from the child side.
            let (id2, log2) = make_log(1, "post-fork");
            handle
                .try_send_msg(TelemetryActions::AddLog((id2, log2)))
                .unwrap();

            // Simulate one trigger() + run() cycle.
            worker.trigger().await;
            worker.run().await;

            let stats = worker.stats();
            // Only the new post-fork log should be buffered.
            assert_eq!(stats.logs, 1, "only post-fork log should be present");
        }

        /// After reset(), lifecycle state needed to keep periodic flushing alive is preserved.
        #[tokio::test]
        async fn test_reset_preserves_started_and_deadlines() {
            let (_handle, mut worker) = build_test_worker();

            worker.data.started = true;
            worker
                .deadlines
                .schedule_event(LifecycleAction::FlushMetricAggr)
                .unwrap();
            worker
                .deadlines
                .schedule_event(LifecycleAction::FlushData)
                .unwrap();

            let deadlines_before = worker.deadlines.deadlines.clone();

            worker.reset();

            assert!(worker.data.started, "started flag should be preserved");
            assert_eq!(
                worker.deadlines.deadlines.len(),
                deadlines_before.len(),
                "scheduled deadlines should be preserved"
            );
            for ((_, actual), (_, expected)) in worker
                .deadlines
                .deadlines
                .iter()
                .zip(deadlines_before.iter())
            {
                assert_eq!(
                    actual, expected,
                    "deadline kinds should be preserved across reset"
                );
            }
        }
    }

    #[cfg_attr(miri, ignore)]
    #[test]
    fn test_channel_close_flushes_and_parks_via_shared_runtime() {
        use httpmock::prelude::*;
        use libdd_common::Endpoint;
        use libdd_shared_runtime::SharedRuntime;
        use std::time::Duration;

        const TELEMETRY_PATH: &str = "/telemetry/proxy/api/v2/apmtelemetry";

        let server = MockServer::start();
        let mock = server.mock(|when, then| {
            when.method(POST).path(TELEMETRY_PATH);
            then.status(202).body("");
        });

        let mut builder = TelemetryWorkerBuilder::new(
            "host".into(),
            "svc".into(),
            "lang".into(),
            "1".into(),
            "tv".into(),
        );
        builder
            .config
            .set_endpoint(Endpoint::from_slice(&server.url("/")))
            .unwrap();
        builder.runtime_id = Some("rid".into());

        let shared_runtime = SharedRuntime::new().expect("SharedRuntime::new");
        let runtime_handle = shared_runtime
            .block_on(async { tokio::runtime::Handle::current() })
            .expect("runtime handle");
        let (telemetry_handle, worker) = builder.build_worker(Some(runtime_handle));

        let _worker_handle = shared_runtime
            .spawn_worker(worker, false)
            .expect("spawn_worker");

        // Drive the worker into the started state so Stop has work to flush.
        telemetry_handle.send_start().expect("send_start");

        // Wait for the AppStarted batch so we know the worker reached started == true
        // before we close the channel.
        for _ in 0..50 {
            if mock.calls() >= 1 {
                break;
            }
            std::thread::sleep(Duration::from_millis(20));
        }
        assert!(
            mock.calls() >= 1,
            "worker should POST at least once after Start"
        );

        // Close the mailbox by dropping the handle.
        let hits_before_close = mock.calls();
        drop(telemetry_handle);

        // The worker must dispatch Lifecycle::Stop, which flushes a final batch, then park.
        for _ in 0..50 {
            if mock.calls() > hits_before_close {
                break;
            }
            std::thread::sleep(Duration::from_millis(20));
        }
        assert!(
            mock.calls() > hits_before_close,
            "worker should flush a final batch after the channel is closed"
        );

        // Once parked, the worker must stop POSTing. Sample for a while and require the
        // hit count to stabilise (proves no Stop-emit loop).
        let stable_hits = mock.calls();
        std::thread::sleep(Duration::from_millis(300));
        assert_eq!(
            mock.calls(),
            stable_hits,
            "worker must stop POSTing after parking; observed {} extra hits",
            mock.calls().saturating_sub(stable_hits),
        );
    }
}