apollo-router 2.16.0

// tonic does not derive `Eq` for the gRPC message types, which causes a warning from Clippy. The
// current suggestion is to explicitly allow the lint in the module that imports the protos.
// Read more: https://github.com/hyperium/tonic/issues/1056
#![allow(clippy::derive_partial_eq_without_eq)]

use std::collections::HashSet;
use std::pin::Pin;
use std::str::FromStr;
use std::sync::Arc;
use std::task::Context;
use std::task::Poll;
use std::time::Duration;
use std::time::SystemTime;

use futures::Stream;
use futures::StreamExt;
use futures::future::Ready;
use futures::stream::FilterMap;
use futures::stream::Fuse;
use futures::stream::Repeat;
use futures::stream::Zip;
use graphql_client::GraphQLQuery;
use pin_project_lite::pin_project;
use strum::IntoEnumIterator;
use tokio::time::Instant;
use tokio_util::time::DelayQueue;

use super::license_enforcement::LicenseLimits;
use super::license_enforcement::TpsLimit;
use crate::AllowedFeature;
use crate::router::Event;
use crate::uplink::UplinkRequest;
use crate::uplink::UplinkResponse;
use crate::uplink::license_enforcement::Audience;
use crate::uplink::license_enforcement::Claims;
use crate::uplink::license_enforcement::License;
use crate::uplink::license_enforcement::LicenseState;
use crate::uplink::license_enforcement::OneOrMany;
use crate::uplink::license_stream::license_query::FetchErrorCode;
use crate::uplink::license_stream::license_query::LicenseQueryRouterEntitlements;

const APOLLO_ROUTER_LICENSE_OFFLINE_UNSUPPORTED: &str = "APOLLO_ROUTER_LICENSE_OFFLINE_UNSUPPORTED";

/// The documented grace window between `warn_at` and `halt_at` (days).
/// Used both to define the maximum documented halt offset and to restore
/// the grace window when both deadlines are clamped.
const HALT_GRACE_PERIOD_DAYS: u64 = 28;

/// Maximum span from "now" to `halt_at` for offline licenses in
/// `docs/source/routing/license.mdx`: at most one year of validity plus a
/// 28-day grace period (~393 days).
const DOCUMENTED_MAX_HALT_OFFSET_DAYS: u64 = 365 + HALT_GRACE_PERIOD_DAYS;

/// Slack above the documented span so legitimate `warn_at` / `halt_at` values are
/// not truncated if validity or grace windows change.
const LICENSE_SCHEDULE_HEADROOM_DAYS: u64 = 90;

const SECS_PER_DAY: u64 = 24 * 60 * 60;

/// Tokio's timer wheel rejects deadlines beyond `(1 << 36) - 1` ms (~795 days).
const TOKIO_TIMER_WHEEL_MAX_MILLIS: u64 = (1 << 36) - 1;

const MAX_TIMER_DURATION_SECS: u64 =
    (DOCUMENTED_MAX_HALT_OFFSET_DAYS + LICENSE_SCHEDULE_HEADROOM_DAYS) * SECS_PER_DAY;

const _: () = assert!(MAX_TIMER_DURATION_SECS * 1000 < TOKIO_TIMER_WHEEL_MAX_MILLIS);

/// Upper bound for scheduling `warn_at` / `halt_at` in `DelayQueue`.
///
/// Derived from the documented license timeline plus headroom, and kept below
/// Tokio's timer wheel maximum so `insert_at` cannot panic. When both deadlines
/// exceed this cap, `warn_at` is anchored `WARN_BEFORE_HALT_GRACE` before the
/// clamped `halt_at` so the soft-then-hard grace ordering is preserved.
const MAX_TIMER_DURATION: Duration = Duration::from_secs(MAX_TIMER_DURATION_SECS);

/// How far before the clamped `halt_at` to schedule `warn_at` when both deadlines
/// exceed [`MAX_TIMER_DURATION`]. Matches the documented [`HALT_GRACE_PERIOD_DAYS`]-day grace window.
const WARN_BEFORE_HALT_GRACE: Duration = Duration::from_secs(HALT_GRACE_PERIOD_DAYS * SECS_PER_DAY);

#[derive(GraphQLQuery)]
#[graphql(
    query_path = "src/uplink/license_query.graphql",
    schema_path = "src/uplink/uplink.graphql",
    request_derives = "Debug",
    response_derives = "PartialEq, Debug, Deserialize",
    deprecated = "warn"
)]
pub(crate) struct LicenseQuery {}

impl From<UplinkRequest> for license_query::Variables {
    fn from(req: UplinkRequest) -> Self {
        license_query::Variables {
            api_key: req.api_key,
            graph_ref: req.graph_ref,
            if_after_id: req.id,
        }
    }
}

impl From<license_query::ResponseData> for UplinkResponse<License> {
    fn from(response: license_query::ResponseData) -> Self {
        match response.router_entitlements {
            LicenseQueryRouterEntitlements::RouterEntitlementsResult(result) => {
                if let Some(license) = result.entitlement {
                    match License::from_str(&license.jwt) {
                        Ok(jwt) => UplinkResponse::New {
                            response: jwt,
                            id: result.id,
                            // this will truncate the number of seconds to under u64::MAX, which should be
                            // a large enough delay anyway
                            delay: result.min_delay_seconds as u64,
                        },
                        Err(error) => UplinkResponse::Error {
                            retry_later: true,
                            code: "INVALID_LICENSE".to_string(),
                            message: error.to_string(),
                        },
                    }
                } else {
                    UplinkResponse::New {
                        response: License::default(),
                        id: result.id,
                        // this will truncate the number of seconds to under u64::MAX, which should be
                        // a large enough delay anyway
                        delay: result.min_delay_seconds as u64,
                    }
                }
            }
            LicenseQueryRouterEntitlements::Unchanged(response) => UplinkResponse::Unchanged {
                id: Some(response.id),
                delay: Some(response.min_delay_seconds as u64),
            },
            LicenseQueryRouterEntitlements::FetchError(error) => UplinkResponse::Error {
                retry_later: error.code == FetchErrorCode::RETRY_LATER,
                code: match error.code {
                    FetchErrorCode::AUTHENTICATION_FAILED => "AUTHENTICATION_FAILED".to_string(),
                    FetchErrorCode::ACCESS_DENIED => "ACCESS_DENIED".to_string(),
                    FetchErrorCode::UNKNOWN_REF => "UNKNOWN_REF".to_string(),
                    FetchErrorCode::RETRY_LATER => "RETRY_LATER".to_string(),
                    FetchErrorCode::NOT_IMPLEMENTED_ON_THIS_INSTANCE => {
                        "NOT_IMPLEMENTED_ON_THIS_INSTANCE".to_string()
                    }
                    FetchErrorCode::Other(other) => other,
                },
                message: error.message,
            },
        }
    }
}

pin_project! {
    /// This stream wrapper will cause check the current license at the point of warn_at or halt_at.
    /// This means that the state machine can be kept clean, and not have to deal with setting it's own timers and also avoids lots of racy scenarios as license checks are guaranteed to happen after a license update even if they were in the past.
    #[must_use = "streams do nothing unless polled"]
    #[project = LicenseExpanderProj]
    pub(crate) struct LicenseExpander<Upstream>
    where
        Upstream: Stream<Item = License>,
    {
        #[pin]
        checks: DelayQueue<Event>,
        #[pin]
        upstream: Fuse<Upstream>,
    }
}

impl<Upstream> Stream for LicenseExpander<Upstream>
where
    Upstream: Stream<Item = License>,
{
    type Item = Event;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let mut this = self.project();
        let checks = this.checks.poll_expired(cx);
        // Only check downstream if checks was not Some
        let next = if matches!(checks, Poll::Ready(Some(_))) {
            None
        } else {
            // Poll upstream. Note that it is OK for this to be called again after it has finished as the stream is fused and if it is exhausted it will return Poll::Ready(None).
            Some(this.upstream.poll_next(cx))
        };

        match (checks, next) {
            // Checks has an expired claim that needs checking.
            // This is the ONLY arm where upstream.poll_next has not been called, and this is OK because we are not returning pending.
            (Poll::Ready(Some(item)), _) => Poll::Ready(Some(item.into_inner())),
            // Upstream has a new license with a claim
            (_, Some(Poll::Ready(Some(license)))) if license.claims.is_some() => {
                // If we got a new license then we need to reset the stream of events and return the new license event.
                reset_checks_for_licenses(&mut this.checks, license)
            }
            // Upstream has a new license with no claim.
            (_, Some(Poll::Ready(Some(_)))) => {
                // We don't clear the checks if there is a license with no claim.
                Poll::Ready(Some(Event::UpdateLicense(Arc::new(
                    LicenseState::Unlicensed,
                ))))
            }
            // If either checks or upstream returned pending then we need to return pending.
            // It is the responsibility of upstream and checks to schedule wakeup.
            // If we have got to this line then checks.poll_expired and upstream.poll_next *will* have been called.
            (Poll::Pending, _) | (_, Some(Poll::Pending)) => Poll::Pending,
            // If both stream are exhausted then return none.
            (Poll::Ready(None), Some(Poll::Ready(None))) => Poll::Ready(None),
            (Poll::Ready(None), None) => {
                unreachable!("upstream will have been called as checks did not have a value")
            }
        }
    }
}

/// This function takes a license and returns the appropriate event for that license.
/// If warn at or halt at are in the future it will register appropriate checks to trigger at such times.
fn reset_checks_for_licenses(
    checks: &mut DelayQueue<Event>,
    license: License,
) -> Poll<Option<Event>> {
    // We got a new claim, so clear the previous checks.
    checks.clear();
    let claims = license.claims.as_ref().expect("claims is gated, qed");

    // Router limitations based on claims
    let limits = match (claims.tps, &claims.allowed_features) {
        (None, None) => None,
        (Some(tps_limit), Some(features)) => Some(
            LicenseLimits::builder()
                .tps(
                    TpsLimit::builder()
                        .capacity(tps_limit.capacity)
                        .interval(tps_limit.interval)
                        .build(),
                )
                .allowed_features(HashSet::from_iter(features.clone()))
                .build(),
        ),
        (Some(tps_limit), None) => Some(LicenseLimits {
            tps: Some(TpsLimit {
                capacity: tps_limit.capacity,
                interval: tps_limit.interval,
            }),
            allowed_features: HashSet::from_iter(AllowedFeature::iter()),
        }),
        (None, Some(features)) => Some(
            LicenseLimits::builder()
                .allowed_features(HashSet::from_iter(features.clone()))
                .build(),
        ),
    };

    let halt_at = to_positive_instant(claims.halt_at);
    let now_system = SystemTime::now();

    // If halt_at was clamped, ensure warn_at is anchored at least WARN_BEFORE_HALT_GRACE
    // before it. This covers two cases: (1) both deadlines exceed the cap (warn_at_raw and
    // halt_at collapse to the same clamped instant), and (2) only halt_at is clamped but
    // warn_at is close enough to MAX_TIMER_DURATION that the remaining gap is shorter than
    // the documented grace window.
    let halt_exceeds_cap = claims
        .halt_at
        .duration_since(now_system)
        .map(|d| d > MAX_TIMER_DURATION)
        .unwrap_or(false);
    let warn_at_raw = to_positive_instant(claims.warn_at);
    let warn_at = if halt_exceeds_cap && warn_at_raw + WARN_BEFORE_HALT_GRACE > halt_at {
        halt_at - WARN_BEFORE_HALT_GRACE
    } else {
        warn_at_raw
    };

    // Capture `now` after all to_positive_instant calls so that a past warn_at
    // (returned as Instant::now() inside that function) correctly compares as <= now.
    let now = Instant::now();
    // Insert the new checks. If any of the boundaries are in the past then just return the immediate result
    if halt_at > now {
        // Only add halt if it isn't immediately going to be triggered.
        checks.insert_at(
            Event::UpdateLicense(Arc::new(LicenseState::LicensedHalt {
                limits: limits.clone(),
            })),
            halt_at,
        );
    } else {
        return Poll::Ready(Some(Event::UpdateLicense(Arc::new(
            LicenseState::LicensedHalt {
                limits: limits.clone(),
            },
        ))));
    }
    if warn_at > now {
        // Only add warn if it isn't immediately going to be triggered and halt is not already set.
        // Something that is halted is by definition also warn.
        checks.insert_at(
            Event::UpdateLicense(Arc::new(LicenseState::LicensedWarn {
                limits: limits.clone(),
            })),
            warn_at,
        );
    } else {
        return Poll::Ready(Some(Event::UpdateLicense(Arc::new(
            LicenseState::LicensedWarn {
                limits: limits.clone(),
            },
        ))));
    }

    Poll::Ready(Some(Event::UpdateLicense(Arc::new(
        LicenseState::Licensed {
            limits: limits.clone(),
        },
    ))))
}

/// Converts an externally generated `SystemTime` (e.g. JWT `warn_at` / `halt_at`) into a
/// `tokio::time::Instant` for scheduling license state transitions in a `DelayQueue`.
///
/// The result is approximate: there is no exact conversion between `SystemTime` and
/// `Instant`. If `system_time` is in the past, returns `Instant::now()` — we only need to
/// know the event already happened, not how far in the past. (Subtracting from instants is
/// not supported on all platforms, which motivates treating past times as "now".)
///
/// Uses `tokio::time::Instant` (not `std::time::Instant`) so scheduling respects
/// `tokio::time::pause()` / `tokio::time::advance()` in tests and stays consistent with the
/// queue's clock.
///
/// Future times are clamped to [`MAX_TIMER_DURATION`] so `DelayQueue::insert_at` cannot
/// overflow Tokio's timer wheel. When both `warn_at` and `halt_at` exceed the cap,
/// `reset_checks_for_licenses` anchors `warn_at` to [`WARN_BEFORE_HALT_GRACE`] before the
/// clamped `halt_at` so the soft-then-hard grace ordering is preserved.
fn to_positive_instant(system_time: SystemTime) -> Instant {
    let now_instant = Instant::now();
    let now_system_time = SystemTime::now();

    match system_time.duration_since(now_system_time) {
        Ok(duration) => now_instant + duration.min(MAX_TIMER_DURATION),
        Err(_) => now_instant,
    }
}

type ValidateAudience<T> = FilterMap<
    Zip<T, Repeat<Arc<HashSet<Audience>>>>,
    Ready<Option<License>>,
    fn((License, Arc<HashSet<Audience>>)) -> Ready<Option<License>>,
>;

pub(crate) trait LicenseStreamExt: Stream<Item = License> {
    fn expand_licenses(self) -> LicenseExpander<Self>
    where
        Self: Sized,
    {
        LicenseExpander {
            checks: Default::default(),
            upstream: self.fuse(),
        }
    }

    fn validate_audience(self, audiences: impl Into<HashSet<Audience>>) -> ValidateAudience<Self>
    where
        Self: Sized,
    {
        // Zip is used to inject the data into the stream, and then filter_map can be used to actually deal with the data.
        // There's no way to do this with a closure without hitting compiler issues.
        // In the past we have implemented our own steps where we have needed to inject state, but this is the recommended way to do it.
        let audiences: Arc<HashSet<Audience>> = Arc::new(audiences.into());
        self.zip(futures::stream::repeat(audiences))
            .filter_map(|(license, audiences)| {
                let matches = match &license {
                    License {
                        claims:
                            Some(Claims {
                                aud: OneOrMany::Many(aud),
                                ..
                            }),
                    } => aud.iter().any(|aud| audiences.contains(aud)),
                    License {
                        claims:
                            Some(Claims {
                                aud: OneOrMany::One(aud),
                                ..
                            }),
                    } => audiences.contains(aud),
                    // A license with no claims is always valid. We will check later if any commercial features are in use.
                    License { claims: None } => true,
                };

                if !matches {
                    tracing::error!(
                        code = APOLLO_ROUTER_LICENSE_OFFLINE_UNSUPPORTED,
                        "the license file was valid, but was not enabled offline use",
                    );
                }
                futures::future::ready(if matches { Some(license) } else { None })
            })
    }
}

impl<T: Stream<Item = License>> LicenseStreamExt for T {}

#[cfg(test)]
mod test {
    use std::future::ready;
    use std::time::Duration;
    use std::time::SystemTime;

    use futures::StreamExt;
    use futures_test::stream::StreamTestExt;
    use tokio::time::Instant;
    use tracing::instrument::WithSubscriber;

    use crate::assert_snapshot_subscriber;
    use crate::router::Event;
    use crate::uplink::UplinkConfig;
    use crate::uplink::license_enforcement::Audience;
    use crate::uplink::license_enforcement::Claims;
    use crate::uplink::license_enforcement::License;
    use crate::uplink::license_enforcement::LicenseState;
    use crate::uplink::license_enforcement::OneOrMany;
    use crate::uplink::license_stream::LicenseQuery;
    use crate::uplink::license_stream::LicenseStreamExt;
    use crate::uplink::license_stream::to_positive_instant;
    use crate::uplink::stream_from_uplink;

    #[tokio::test]
    async fn integration_test() {
        if let (Ok(apollo_key), Ok(apollo_graph_ref)) = (
            std::env::var("TEST_APOLLO_KEY"),
            std::env::var("TEST_APOLLO_GRAPH_REF"),
        ) {
            let results = stream_from_uplink::<LicenseQuery, License>(UplinkConfig {
                apollo_key,
                apollo_graph_ref,
                endpoints: None,
                poll_interval: Duration::from_secs(1),
                timeout: Duration::from_secs(5),
            })
            .take(1)
            .collect::<Vec<_>>()
            .await;

            assert!(
                results
                    .first()
                    .expect("expected one result")
                    .as_ref()
                    .expect("license should be OK")
                    .claims
                    .is_some()
            )
        }
    }

    #[test]
    fn test_to_instant() {
        let now_system_time = SystemTime::now();
        let now_instant = Instant::now();
        let future_system_time = now_system_time + Duration::from_secs(1024);
        let future_instant = to_positive_instant(future_system_time);
        assert!(future_instant < now_instant + Duration::from_secs(1025));
        assert!(future_instant > now_instant + Duration::from_secs(1023));

        // One day below the scheduling cap: a realistic license halt window must not be clamped.
        let one_day_before_cap =
            super::MAX_TIMER_DURATION - Duration::from_secs(super::SECS_PER_DAY);
        let near_cap_system_time = now_system_time + one_day_before_cap;
        let near_cap_instant = to_positive_instant(near_cap_system_time);
        assert!(
            near_cap_instant > now_instant + one_day_before_cap - Duration::from_secs(1),
            "deadline one day below MAX_TIMER_DURATION must not be clamped"
        );
        assert!(
            near_cap_instant < now_instant + one_day_before_cap + Duration::from_secs(1),
            "deadline one day below MAX_TIMER_DURATION must not be clamped"
        );

        // An instant in the past will return something greater than the original
        // now_instant, and at most equal to a new instant. The upper bound is
        // inclusive because on low-resolution monotonic clocks (Windows ticks at
        // ~16ms) the `Instant::now()` here can read the same value
        // `to_positive_instant` did during the same tick.
        let past_system_time = now_system_time - Duration::from_secs(1024);
        let past_instant = to_positive_instant(past_system_time);
        assert!(past_instant > now_instant);
        assert!(past_instant <= Instant::now());
    }

    #[tokio::test(start_paused = true)]
    async fn test_to_instant_far_future_is_clamped() {
        let now_instant = Instant::now();
        let three_years = Duration::from_secs(3 * 365 * 24 * 3600);
        let far_future = SystemTime::now() + three_years;

        let result = to_positive_instant(far_future);

        assert_eq!(
            result,
            now_instant + super::MAX_TIMER_DURATION,
            "far-future SystemTime must clamp to MAX_TIMER_DURATION from the current Instant"
        );
    }

    #[tokio::test(start_paused = true)]
    async fn license_expander_far_future_does_not_panic() {
        let three_years_ms: u64 = 3 * 365 * 24 * 3600 * 1000;
        let license = license_with_claim(three_years_ms - 1_000_000, three_years_ms);

        // Before the clamp fix this would panic with "invalid deadline; err=Invalid"
        // because the deadline exceeded tokio's timer wheel maximum of ~2^36 ms.
        let events_stream = futures::stream::iter(vec![license])
            .expand_licenses()
            .map(SimpleEvent::from);

        let events = events_stream.collect::<Vec<_>>().await;
        // When both deadlines exceed MAX_TIMER_DURATION, warn_at is anchored
        // WARN_BEFORE_HALT_GRACE before the clamped halt_at, so staged ordering
        // is preserved even for far-future licenses.
        assert_eq!(
            events,
            &[
                SimpleEvent::UpdateLicense,
                SimpleEvent::WarnLicense,
                SimpleEvent::HaltLicense,
            ]
        );
    }

    #[tokio::test(start_paused = true)]
    async fn license_expander_halt_clamped_warn_close_to_cap() {
        // warn_at is 14 days before MAX_TIMER_DURATION — under the cap on its own, but
        // within WARN_BEFORE_HALT_GRACE (28 days) of the clamped halt_at. The old code
        // only adjusted warn_at when both deadlines exceeded the cap, leaving a < 28-day
        // gap here. The new code anchors warn_at to halt_at - WARN_BEFORE_HALT_GRACE.
        let half_grace_ms = (super::HALT_GRACE_PERIOD_DAYS / 2) * super::SECS_PER_DAY * 1000;
        let warn_delta_ms = super::MAX_TIMER_DURATION_SECS * 1000 - half_grace_ms;
        let three_years_ms: u64 = 3 * 365 * 24 * 3600 * 1000;

        let license = license_with_claim(warn_delta_ms, three_years_ms);

        let events_stream = futures::stream::iter(vec![license])
            .expand_licenses()
            .map(SimpleEvent::from);

        let events = events_stream.collect::<Vec<_>>().await;
        assert_eq!(
            events,
            &[
                SimpleEvent::UpdateLicense,
                SimpleEvent::WarnLicense,
                SimpleEvent::HaltLicense,
            ]
        );
    }

    #[tokio::test(start_paused = true)]
    async fn license_expander_far_future_halt_only_clamped() {
        let three_years_ms: u64 = 3 * 365 * 24 * 3600 * 1000;
        // warn_at stays under the cap; only halt_at is clamped. This is the more likely
        // real-world shape and must preserve warn-before-halt ordering.
        let license = license_with_claim(15_000, three_years_ms);

        let events_stream = futures::stream::iter(vec![license])
            .expand_licenses()
            .map(SimpleEvent::from);

        let events = events_stream.collect::<Vec<_>>().await;
        assert_eq!(
            events,
            &[
                SimpleEvent::UpdateLicense,
                SimpleEvent::WarnLicense,
                SimpleEvent::HaltLicense,
            ]
        );
    }

    #[tokio::test]
    async fn license_expander() {
        let events_stream = futures::stream::iter(vec![license_with_claim(15, 30)])
            .expand_licenses()
            .map(SimpleEvent::from);

        let events = events_stream.collect::<Vec<_>>().await;
        assert_eq!(
            events,
            &[
                SimpleEvent::UpdateLicense,
                SimpleEvent::WarnLicense,
                SimpleEvent::HaltLicense
            ]
        );
    }

    #[tokio::test]
    async fn license_expander_warn_now() {
        let events_stream = futures::stream::iter(vec![license_with_claim(0, 15)])
            .interleave_pending()
            .expand_licenses()
            .map(SimpleEvent::from);

        let events = events_stream.collect::<Vec<_>>().await;
        assert_eq!(
            events,
            &[SimpleEvent::WarnLicense, SimpleEvent::HaltLicense]
        );
    }

    #[tokio::test]
    async fn license_expander_halt_now() {
        let events_stream = futures::stream::iter(vec![license_with_claim(0, 0)])
            .interleave_pending()
            .expand_licenses()
            .map(SimpleEvent::from);

        let events = events_stream.collect::<Vec<_>>().await;
        assert_eq!(events, &[SimpleEvent::HaltLicense]);
    }

    #[tokio::test]
    async fn license_expander_no_claim() {
        let events_stream = futures::stream::iter(vec![license_with_no_claim()])
            .interleave_pending()
            .expand_licenses()
            .map(SimpleEvent::from);

        let events = events_stream.collect::<Vec<_>>().await;
        assert_eq!(events, &[SimpleEvent::UpdateLicense]);
    }

    #[tokio::test(start_paused = true)]
    async fn license_expander_claim_no_claim() {
        // Licenses with no claim do not clear checks as they are ignored if we move from entitled to unentitled, this is handled at the state machine level.
        //
        // Use paused virtual time so that the warn/halt boundaries (10ms in the future at the
        // moment of claim arrival) cannot fire prematurely between the time we enqueue them and
        // the time the consumer polls `checks.poll_expired`. Under real time, scheduler jitter
        // could push the boundaries into the past before the upstream `no_claim()` was polled,
        // taking the early-return branches in `reset_checks_for_licenses` and producing a
        // different event ordering than the snapshot. Anchoring `to_positive_instant` to
        // `tokio::time::Instant::now()` (see refactor in this file) means the queue's deadlines
        // and `Instant::now()` share the same paused clock here, so deadlines only advance when
        // the test driver decides — `collect::<Vec<_>>().await` auto-advances paused time when
        // the runtime is otherwise idle, which deterministically fires the warn/halt entries.
        let events_stream =
            futures::stream::iter(vec![license_with_claim(10, 10), license_with_no_claim()])
                .interleave_pending()
                .expand_licenses()
                .map(SimpleEvent::from);

        let events = events_stream.collect::<Vec<_>>().await;
        assert_eq!(
            events,
            &[
                SimpleEvent::UpdateLicense,
                SimpleEvent::UpdateLicense,
                SimpleEvent::WarnLicense,
                SimpleEvent::HaltLicense
            ]
        );
    }

    #[tokio::test]
    async fn license_expander_no_claim_claim() {
        let events_stream =
            futures::stream::iter(vec![license_with_no_claim(), license_with_claim(15, 30)])
                .interleave_pending()
                .expand_licenses()
                .map(SimpleEvent::from);

        let events = events_stream.collect::<Vec<_>>().await;
        assert_eq!(
            events,
            &[
                SimpleEvent::UpdateLicense,
                SimpleEvent::UpdateLicense,
                SimpleEvent::WarnLicense,
                SimpleEvent::HaltLicense
            ]
        );
    }

    #[tokio::test(start_paused = true)]
    async fn license_expander_claim_pause_claim() {
        // Use paused virtual time so the schedule of claim arrivals vs. warn/halt
        // expirations is deterministic. The previous version of this test used real
        // `tokio::time::sleep(200ms)` and asserted on event ordering, which raced with
        // the producer task on slow / loaded systems.
        //
        // `to_positive_instant` returns a `tokio::time::Instant` anchored to the same
        // virtual clock that `DelayQueue` reads, so advancing time here precisely fires
        // the inserted warn/halt entries.
        let (tx, rx) = tokio::sync::mpsc::channel(10);
        let rx_stream = tokio_stream::wrappers::ReceiverStream::new(rx);
        let events_stream = rx_stream.expand_licenses().map(SimpleEvent::from);

        tokio::task::spawn(async move {
            // Simulate a new claim coming in between the warning and halt of the first.
            // First claim: warn_at = now + 100ms, halt_at = now + 300ms.
            let _ = tx.send(license_with_claim(100, 300)).await;
            // Advance past the warn boundary (100ms) but before the halt boundary (300ms).
            tokio::time::sleep(Duration::from_millis(200)).await;
            // Second claim resets the schedule from "now" (200ms after start).
            let _ = tx.send(license_with_claim(100, 300)).await;
        });
        let events = events_stream.collect::<Vec<_>>().await;
        assert_eq!(
            events,
            &[
                SimpleEvent::UpdateLicense,
                SimpleEvent::WarnLicense,
                SimpleEvent::UpdateLicense,
                SimpleEvent::WarnLicense,
                SimpleEvent::HaltLicense
            ]
        );
    }

    fn license_with_claim(warn_delta: u64, halt_delta: u64) -> License {
        let now = SystemTime::now();
        License {
            claims: Some(Claims {
                iss: "".to_string(),
                sub: "".to_string(),
                aud: OneOrMany::One(Audience::SelfHosted),
                warn_at: now + Duration::from_millis(warn_delta),
                halt_at: now + Duration::from_millis(halt_delta),
                tps: Default::default(),
                allowed_features: Default::default(),
            }),
        }
    }

    fn license_with_no_claim() -> License {
        License { claims: None }
    }

    #[derive(Eq, PartialEq, Debug)]
    enum SimpleEvent {
        UpdateConfiguration,
        NoMoreConfiguration,
        UpdateSchema,
        NoMoreSchema,
        UpdateLicense,
        HaltLicense,
        WarnLicense,
        NoMoreLicense,
        ForcedHotReload,
        Shutdown,
    }

    impl From<Event> for SimpleEvent {
        fn from(value: Event) -> Self {
            match value {
                Event::UpdateConfiguration(_) => SimpleEvent::UpdateConfiguration,
                Event::NoMoreConfiguration => SimpleEvent::NoMoreConfiguration,
                Event::UpdateSchema(_) => SimpleEvent::UpdateSchema,
                Event::NoMoreSchema => SimpleEvent::NoMoreSchema,
                Event::UpdateLicense(license) => match *license {
                    LicenseState::LicensedHalt { limits: _ } => SimpleEvent::HaltLicense,
                    LicenseState::LicensedWarn { limits: _ } => SimpleEvent::WarnLicense,
                    _ => SimpleEvent::UpdateLicense,
                },
                Event::NoMoreLicense => SimpleEvent::NoMoreLicense,
                Event::Reload | Event::RhaiReload => SimpleEvent::ForcedHotReload,
                Event::Shutdown => SimpleEvent::Shutdown,
            }
        }
    }

    #[tokio::test]
    async fn test_validate_audience_single() {
        assert_eq!(
            futures::stream::once(ready(License {
                claims: Some(Claims {
                    iss: "".to_string(),
                    sub: "".to_string(),
                    aud: OneOrMany::One(Audience::Offline),
                    warn_at: SystemTime::now(),
                    halt_at: SystemTime::now(),
                    tps: Default::default(),
                    allowed_features: Default::default(),
                }),
            }))
            .validate_audience([Audience::Offline, Audience::Cloud])
            .count()
            .with_subscriber(assert_snapshot_subscriber!())
            .await,
            1
        );
    }

    #[tokio::test]
    async fn test_validate_audience_single_filtered() {
        assert_eq!(
            futures::stream::once(ready(License {
                claims: Some(Claims {
                    iss: "".to_string(),
                    sub: "".to_string(),
                    aud: OneOrMany::One(Audience::SelfHosted),
                    warn_at: SystemTime::now(),
                    halt_at: SystemTime::now(),
                    tps: Default::default(),
                    allowed_features: Default::default(),
                }),
            }))
            .validate_audience([Audience::Offline, Audience::Cloud])
            .count()
            .with_subscriber(assert_snapshot_subscriber!())
            .await,
            0
        );
    }

    #[tokio::test]
    async fn test_validate_audience_multiple() {
        assert_eq!(
            futures::stream::once(ready(License {
                claims: Some(Claims {
                    iss: "".to_string(),
                    sub: "".to_string(),
                    aud: OneOrMany::Many(vec![Audience::SelfHosted, Audience::Offline]),
                    warn_at: SystemTime::now(),
                    halt_at: SystemTime::now(),
                    tps: Default::default(),
                    allowed_features: Default::default(),
                }),
            }))
            .validate_audience([Audience::Offline, Audience::Cloud])
            .count()
            .with_subscriber(assert_snapshot_subscriber!())
            .await,
            1
        );
    }

    #[tokio::test]
    async fn test_validate_audience_multiple_filtered() {
        assert_eq!(
            futures::stream::once(ready(License {
                claims: Some(Claims {
                    iss: "".to_string(),
                    sub: "".to_string(),
                    aud: OneOrMany::Many(vec![Audience::SelfHosted, Audience::SelfHosted]),
                    warn_at: SystemTime::now(),
                    halt_at: SystemTime::now(),
                    tps: Default::default(),
                    allowed_features: Default::default(),
                }),
            }))
            .validate_audience([Audience::Offline, Audience::Cloud])
            .count()
            .with_subscriber(assert_snapshot_subscriber!())
            .await,
            0
        );
    }

    #[tokio::test]
    async fn test_validate_no_claim() {
        assert_eq!(
            futures::stream::once(ready(License::default()))
                .validate_audience([Audience::Offline, Audience::Cloud])
                .count()
                .with_subscriber(assert_snapshot_subscriber!())
                .await,
            1
        );
    }
}