openrtc 2.8.2

use std::collections::HashMap;
use std::hash::Hash;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{LazyLock, Mutex};

static NEXT_TRANSPORT_GENERATION: AtomicU64 = AtomicU64::new(1);
static LIVE_TRANSPORT_GENERATIONS: LazyLock<Mutex<HashMap<usize, u64>>> =
    LazyLock::new(|| Mutex::new(HashMap::new()));

fn next_generation() -> u64 {
    let generation = NEXT_TRANSPORT_GENERATION.fetch_add(1, Ordering::Relaxed);
    assert_ne!(
        generation, 0,
        "OpenRTC transport generation counter exhausted"
    );
    generation
}

/// Translate a runtime connection's address-like `stable_id` into an OpenRTC
/// generation that is never reused during the process lifetime.
///
/// Quinn only guarantees that `stable_id` remains fixed for one connection's
/// lifetime; it is an allocation address and can be reused after retirement.
/// OpenRTC keeps stale callbacks and admission proofs generation-fenced beyond
/// that lifetime, so the raw value is not sufficient as a durable generation.
#[cfg(test)]
pub(crate) fn for_runtime_connection(raw_stable_id: usize) -> u64 {
    let mut generations = LIVE_TRANSPORT_GENERATIONS
        .lock()
        .unwrap_or_else(|poisoned| poisoned.into_inner());
    *generations
        .entry(raw_stable_id)
        .or_insert_with(next_generation)
}

pub(crate) fn for_connection(connection: &iroh::endpoint::Connection) -> u64 {
    let raw_stable_id = connection.stable_id();
    let mut generations = LIVE_TRANSPORT_GENERATIONS
        .lock()
        .unwrap_or_else(|poisoned| poisoned.into_inner());
    if let Some(generation) = generations.get(&raw_stable_id) {
        return *generation;
    }
    let generation = next_generation();
    // A stale callback may retain a closed connection after its physical
    // owner retired the mapping. Give that observation a fail-closed unique
    // value without reintroducing an entry that no connection loop can clean.
    if connection.close_reason().is_none() {
        generations.insert(raw_stable_id, generation);
    }
    generation
}

/// Retire the address-to-generation association after the physical connection
/// loop and its terminal event have both completed. A later allocation that
/// reuses the same address receives a fresh OpenRTC generation.
pub(crate) fn forget_runtime_connection(raw_stable_id: usize) {
    LIVE_TRANSPORT_GENERATIONS
        .lock()
        .unwrap_or_else(|poisoned| poisoned.into_inner())
        .remove(&raw_stable_id);
}

pub(crate) fn forget_connection(connection: &iroh::endpoint::Connection) {
    forget_runtime_connection(connection.stable_id());
}

pub(crate) fn remove_if_generation_matches<K: Eq + Hash, V>(
    values: &mut HashMap<K, V>,
    key: &K,
    expected_generation: u64,
    generation: impl FnOnce(&V) -> u64,
) -> Option<V> {
    if values
        .get(key)
        .is_some_and(|value| generation(value) == expected_generation)
    {
        values.remove(key)
    } else {
        None
    }
}

#[cfg(test)]
mod tests {
    use super::{for_runtime_connection, forget_runtime_connection, remove_if_generation_matches};
    use std::collections::HashMap;

    #[test]
    fn reused_runtime_identifier_receives_a_fresh_openrtc_generation() {
        let raw = usize::MAX - 17;
        let first = for_runtime_connection(raw);
        assert_eq!(for_runtime_connection(raw), first);
        forget_runtime_connection(raw);
        let replacement = for_runtime_connection(raw);
        assert_ne!(replacement, first);
        forget_runtime_connection(raw);
    }

    #[test]
    fn rejected_wasm_generation_cleanup_removes_only_the_exact_generation() {
        let mut connections = HashMap::from([("peer", 41)]);
        assert_eq!(
            remove_if_generation_matches(&mut connections, &"peer", 41, |value| *value),
            Some(41),
        );
        assert!(!connections.contains_key("peer"));

        connections.insert("peer", 42);
        assert_eq!(
            remove_if_generation_matches(&mut connections, &"peer", 41, |value| *value),
            None,
        );
        assert_eq!(connections.get("peer"), Some(&42));
    }
}