State, Channels and Backpressure

use std::{collections::HashMap, sync::{Arc, Mutex}};

#[derive(Default)]
struct RateCounters {
    by_key: Mutex<HashMap<String, u64>>,
}

impl RateCounters {
    // Sync method: the guard cannot accidentally live across an .await.
    fn hit(&self, key: &str) -> u64 {
        let mut map = self.by_key.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
        let n = map.entry(key.to_owned()).or_insert(0);
        *n += 1;
        *n
    }
}

async fn handle(counters: Arc<RateCounters>, key: String) -> u64 {
    let n = counters.hit(&key); // lock taken and released inside a sync fn
    tokio::task::yield_now().await;
    n
}

use std::sync::Arc;
use tokio::{io::AsyncWriteExt, net::TcpStream, sync::Mutex};

#[derive(Clone)]
struct AuditSink {
    conn: Arc<Mutex<TcpStream>>,
}

impl AuditSink {
    async fn write_record(&self, line: &str) -> std::io::Result<()> {
        let mut conn = self.conn.lock().await;        // held across the two awaits below: correct
        conn.write_all(line.as_bytes()).await?;
        conn.write_all(b"\n").await?;
        Ok(())
    }
}

use std::time::Duration;
use tokio::sync::watch;

#[derive(Clone, Debug, PartialEq)]
struct Limits { max_batch: usize, upstream_timeout: Duration }

async fn reload_loop(tx: watch::Sender<Limits>) {
    let mut every = tokio::time::interval(Duration::from_secs(30));
    loop {
        every.tick().await;
        let new = Limits { max_batch: 64, upstream_timeout: Duration::from_secs(5) }; // load from source
        tx.send_if_modified(|cur| if *cur != new { *cur = new; true } else { false });
    }
}

async fn worker(mut rx: watch::Receiver<Limits>) {
    loop {
        let limits = rx.borrow_and_update().clone(); // clone out: the Ref is a read guard, !Send
        tracing::debug!(?limits, "using limits");
        if rx.changed().await.is_err() {
            break; // sender dropped
        }
    }
}

use tokio::sync::broadcast::{self, error::RecvError};

async fn subscriber(mut rx: broadcast::Receiver<String>) {
    loop {
        match rx.recv().await {
            Ok(ev) => tracing::debug!(%ev, "event"),
            Err(RecvError::Lagged(skipped)) => tracing::warn!(skipped, "subscriber too slow, events dropped"),
            Err(RecvError::Closed) => break,
        }
    }
}

use axum::{extract::State, http::StatusCode};
use tokio::sync::mpsc::error::TrySendError;

async fn enqueue(State(state): State<AppStateRef>, body: String) -> StatusCode {
    match state.jobs.try_send(Job { body }) {
        Ok(()) => StatusCode::ACCEPTED,
        Err(TrySendError::Full(_)) => StatusCode::SERVICE_UNAVAILABLE, // shed: client retries later
        Err(TrySendError::Closed(_)) => StatusCode::SERVICE_UNAVAILABLE, // shutting down
    }
}

use std::collections::HashMap;
use tokio::sync::{mpsc, oneshot};
use tokio_util::sync::CancellationToken;

enum Msg {
    Get { key: String, reply: oneshot::Sender<Option<String>> },
    Put { key: String, value: String },
}

#[derive(Clone)]
pub struct CacheHandle {
    tx: mpsc::Sender<Msg>,
}

#[derive(Debug, thiserror::Error)]
#[error("cache actor is gone")]
pub struct ActorGone;

impl CacheHandle {
    pub fn spawn(tracker: &tokio_util::task::TaskTracker, token: CancellationToken) -> Self {
        let (tx, rx) = mpsc::channel(512);
        tracker.spawn(run(rx, token));
        Self { tx }
    }

    pub async fn get(&self, key: &str) -> Result<Option<String>, ActorGone> {
        let (reply, rx) = oneshot::channel();
        self.tx.send(Msg::Get { key: key.to_owned(), reply }).await.map_err(|_| ActorGone)?;
        rx.await.map_err(|_| ActorGone)
    }

    pub async fn put(&self, key: String, value: String) -> Result<(), ActorGone> {
        self.tx.send(Msg::Put { key, value }).await.map_err(|_| ActorGone)
    }
}

async fn run(mut rx: mpsc::Receiver<Msg>, token: CancellationToken) {
    let mut map: HashMap<String, String> = HashMap::new();
    loop {
        tokio::select! {
            biased;
            _ = token.cancelled() => break,
            msg = rx.recv() => match msg {
                Some(Msg::Get { key, reply }) => { let _ = reply.send(map.get(&key).cloned()); }
                Some(Msg::Put { key, value }) => { map.insert(key, value); }
                None => break, // every handle dropped
            },
        }
    }
}

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

pub struct Upstream {
    http: reqwest::Client,
    permits: Arc<Semaphore>,
}

impl Upstream {
    pub fn new(http: reqwest::Client, max_in_flight: usize) -> Self {
        Self { http, permits: Arc::new(Semaphore::new(max_in_flight)) }
    }

    pub async fn price(&self, sku: &str) -> anyhow::Result<f64> {
        // wait at most 200 ms for a permit, then fail fast instead of queueing forever
        let _permit = tokio::time::timeout(Duration::from_millis(200), self.permits.acquire())
            .await
            .map_err(|_| anyhow::anyhow!("pricing saturated"))??;
        let price = self
            .http
            .get(format!("https://pricing.internal/v1/{sku}"))
            .timeout(Duration::from_secs(2))
            .send()
            .await?
            .error_for_status()?
            .json::<f64>()
            .await?;
        Ok(price)
    }
}

use std::{future::Future, time::Duration};

pub async fn with_retry<T, E, F, Fut>(mut op: F) -> Result<T, E>
where
    F: FnMut() -> Fut,
    Fut: Future<Output = Result<T, E>>,
    E: std::fmt::Display,
{
    let mut delay = Duration::from_millis(100);
    let max_attempts = 5;
    for attempt in 1.. {
        match op().await {
            Ok(v) => return Ok(v),
            Err(err) if attempt < max_attempts => {
                tracing::warn!(attempt, %err, ?delay, "retrying");
                tokio::time::sleep(delay).await;
                delay = (delay * 2).min(Duration::from_secs(5)); // add random jitter in production
            }
            Err(err) => return Err(err),
        }
    }
    unreachable!("loop returns on the last attempt")
}

pub fn is_transient(err: &reqwest::Error) -> bool {
    err.is_timeout()
        || err.is_connect()
        || err.status().is_some_and(|s| s.as_u16() == 429 || s.is_server_error())
}

use std::time::Duration;

pub fn http_client() -> reqwest::Result<reqwest::Client> {
    reqwest::Client::builder()
        .user_agent(concat!(env!("CARGO_PKG_NAME"), "/", env!("CARGO_PKG_VERSION")))
        .timeout(Duration::from_secs(10))          // total per request, including body
        .connect_timeout(Duration::from_secs(3))
        .pool_idle_timeout(Duration::from_secs(90))
        .pool_max_idle_per_host(32)
        .tcp_keepalive(Duration::from_secs(60))
        .build()
}

use std::time::Duration;
use sqlx::postgres::PgPoolOptions;

pub async fn db_pool(url: &str) -> Result<sqlx::PgPool, sqlx::Error> {
    PgPoolOptions::new()
        .max_connections(20)                      // <= Postgres max_connections / replicas
        .min_connections(2)
        .acquire_timeout(Duration::from_secs(3))  // PoolTimedOut instead of waiting forever
        .idle_timeout(Duration::from_secs(300))
        .max_lifetime(Duration::from_secs(1800))
        .connect(url)
        .await
}

#[derive(Debug, sqlx::FromRow, serde::Serialize)]
pub struct Order {
    pub id: i64,          // Postgres has no unsigned types: use i64/i32
    pub sku: String,
    pub qty: i32,
}

pub async fn find_order(pool: &sqlx::PgPool, id: i64) -> Result<Option<Order>, sqlx::Error> {
    sqlx::query_as::<_, Order>("SELECT id, sku, qty FROM orders WHERE id = $1")
        .bind(id)
        .fetch_optional(pool)
        .await
}

pub async fn transfer(pool: &sqlx::PgPool, from: i64, to: i64, qty: i32) -> Result<(), sqlx::Error> {
    let mut tx = pool.begin().await?;
    sqlx::query("UPDATE stock SET qty = qty - $1 WHERE id = $2").bind(qty).bind(from).execute(&mut *tx).await?;
    sqlx::query("UPDATE stock SET qty = qty + $1 WHERE id = $2").bind(qty).bind(to).execute(&mut *tx).await?;
    tx.commit().await // dropping tx without commit rolls back
}