kcode-k1-txn-ordering 0.4.0

use kcode_k1_canonical_chain::{CanonicalChain, CommitOutcome};
pub use kcode_k1_canonical_chain::{SubmitError, TxId};
use kcode_k1_transaction::Transaction;
pub use kcode_k1_transaction::{GENESIS_PARENT, REGISTER_AT_TIP, SubsystemId};
use std::collections::HashMap;
use std::panic::{AssertUnwindSafe, catch_unwind};
use std::path::Path;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Condvar, Mutex, MutexGuard};
use std::thread;
use std::time::{Duration, Instant};

pub trait Subsystem: Send + Sync + 'static {
    fn submit_txn(&self, id: TxId, payload: &[u8]) -> Result<(), String>;
    fn reorg(&self) -> Result<(), String>;
}

pub struct K1TxnOrdering {
    writer: Mutex<WriterState>,
    chain: Mutex<CanonicalChain>,
}

struct WriterState {
    registrations: HashMap<SubsystemId, Registration>,
}

struct Registration {
    lane: Arc<Lane>,
    latest: Option<TxId>,
    next_ticket: u64,
    mode: RegistrationMode,
}

#[derive(Clone, Copy, Eq, PartialEq)]
enum RegistrationMode {
    Replaying,
    Active,
    OutOfCommission,
}

struct Lane {
    handler: Arc<dyn Subsystem>,
    progress: Mutex<LaneProgress>,
    changed: Condvar,
    running: AtomicBool,
    faulted: AtomicBool,
}

struct LaneProgress {
    serving: u64,
    fault: Option<String>,
}

struct Reservation {
    lane: Arc<Lane>,
    ticket: u64,
}

enum CommittedWork {
    Duplicate,
    Extension {
        id: TxId,
        delivery: Option<Reservation>,
    },
    Reorganization {
        id: TxId,
        reorgs: Vec<Reservation>,
        replacement: Option<Reservation>,
    },
}

impl K1TxnOrdering {
    pub fn open(root: &Path) -> Result<Self, String> {
        let started = Instant::now();
        let result = CanonicalChain::open(root);
        let elapsed = started.elapsed();
        if elapsed > Duration::from_millis(100) {
            eprintln!(
                "{{\"module\":\"kcode-k1-txn-ordering\",\"operation\":\"open\",\"elapsed_microseconds\":{},\"outcome\":\"{}\"}}",
                elapsed.as_micros(),
                if result.is_ok() { "ready" } else { "error" }
            );
        }
        result.map(|chain| Self {
            writer: Mutex::new(WriterState {
                registrations: HashMap::new(),
            }),
            chain: Mutex::new(chain),
        })
    }

    pub fn register_subsystem(
        &self,
        subsystem: SubsystemId,
        after: Option<TxId>,
        handler: Arc<dyn Subsystem>,
    ) -> Result<(), String> {
        let replay_after = match after {
            None | Some(GENESIS_PARENT) => None,
            Some(id) => Some(id),
        };
        let (mut cursor, lane) = {
            let mut writer = self.lock_writer();
            if let Some(registration) = writer.registrations.get(&subsystem) {
                ensure_replaceable(registration)?;
            }
            let chain = self.lock_chain();
            let lane = Arc::new(Lane::new(handler));
            if replay_after == Some(REGISTER_AT_TIP) {
                writer.registrations.insert(
                    subsystem,
                    Registration {
                        lane,
                        latest: None,
                        next_ticket: 0,
                        mode: RegistrationMode::Active,
                    },
                );
                return Ok(());
            }
            let cursor = chain.replay_cursor(subsystem, replay_after)?;
            writer.registrations.insert(
                subsystem,
                Registration {
                    lane: lane.clone(),
                    latest: replay_after,
                    next_ticket: 0,
                    mode: RegistrationMode::Replaying,
                },
            );
            (cursor, lane)
        };
        let mut latest = replay_after;
        loop {
            let mut next = match self.lock_chain().replay_next(&mut cursor) {
                Ok(next) => next,
                Err(message) => return Err(fault_cursor(&lane, message)),
            };
            if next.is_none() {
                let final_result = {
                    let mut writer = self.lock_writer();
                    let chain = self.lock_chain();
                    match chain.replay_next(&mut cursor) {
                        Ok(None) => {
                            let registration = writer
                                .registrations
                                .get_mut(&subsystem)
                                .expect("replaying registration exists");
                            if !Arc::ptr_eq(&registration.lane, &lane) {
                                return Err("registration changed during replay".to_owned());
                            }
                            registration.latest = latest;
                            registration.mode = RegistrationMode::Active;
                            return Ok(());
                        }
                        result => result,
                    }
                };
                next = match final_result {
                    Ok(next) => next,
                    Err(message) => return Err(fault_cursor(&lane, message)),
                };
            }
            let replayed = next.expect("replay result contains a transaction");
            let transaction = match Transaction::parse(&replayed.bytes) {
                Ok(transaction) => transaction,
                Err(message) => {
                    let failure = committed_error(
                        replayed.id,
                        vec![format!(
                            "canonical replay transaction was invalid: {message}"
                        )],
                    );
                    lane.record_fault(failure.clone());
                    return Err(failure);
                }
            };
            if let Err(message) = lane.run_replay(replayed.id, transaction.payload()) {
                return Err(committed_error(replayed.id, vec![message]));
            }
            latest = Some(replayed.id);
        }
    }

    pub fn submit_txn(&self, transaction: &[u8]) -> Result<(), SubmitError> {
        let parsed = Transaction::parse(transaction)
            .map_err(|message| SubmitError::Other(format!("invalid transaction: {message}")))?;
        let payload = parsed.payload();
        let work = {
            let mut writer = self.lock_writer();
            let mut chain = self.lock_chain();
            match chain.submit_validated(transaction)? {
                CommitOutcome::Duplicate => CommittedWork::Duplicate,
                CommitOutcome::Extension { id, subsystem } => CommittedWork::Extension {
                    id,
                    delivery: reserve_active_delivery(&mut writer, subsystem, id),
                },
                CommitOutcome::Reorganization { id, subsystem } => {
                    let affected: Vec<SubsystemId> = writer
                        .registrations
                        .iter()
                        .filter_map(|(&registered, registration)| {
                            (registration.is_active()
                                && registration
                                    .latest
                                    .is_some_and(|latest| !chain.contains(latest)))
                            .then_some(registered)
                        })
                        .collect();
                    let mut reorgs = Vec::with_capacity(affected.len());
                    for affected_subsystem in affected {
                        let registration = writer
                            .registrations
                            .get_mut(&affected_subsystem)
                            .expect("affected registration exists");
                        reorgs.push(registration.reserve_reorg());
                        registration.mode = RegistrationMode::OutOfCommission;
                    }
                    let replacement = reserve_active_delivery(&mut writer, subsystem, id);
                    CommittedWork::Reorganization {
                        id,
                        reorgs,
                        replacement,
                    }
                }
            }
        };
        match work {
            CommittedWork::Duplicate => Ok(()),
            CommittedWork::Extension { id, delivery } => {
                let Some(delivery) = delivery else {
                    return Ok(());
                };
                delivery
                    .run_delivery(id, payload)
                    .map_err(|message| SubmitError::Other(committed_error(id, vec![message])))
            }
            CommittedWork::Reorganization {
                id,
                reorgs,
                replacement,
            } => {
                let failures = run_reorganization(id, payload, reorgs, replacement);
                if failures.is_empty() {
                    Ok(())
                } else {
                    Err(SubmitError::Other(committed_error(id, failures)))
                }
            }
        }
    }

    pub fn submit_local_txn<F, Q>(
        &self,
        timestamp: u64,
        creator: [u8; 32],
        subsystem: SubsystemId,
        payload: &[u8],
        signer: F,
        queue_propagation: Q,
    ) -> Result<(TxId, Vec<u8>), String>
    where
        F: FnOnce(&[u8]) -> Result<[u8; 64], String>,
        Q: FnOnce(&[u8]) -> Result<(), String>,
    {
        let (bytes, id, delivery) = {
            let mut writer = self.lock_writer();
            let lane = writer
                .registrations
                .get(&subsystem)
                .filter(|registration| registration.is_active())
                .map(|registration| registration.lane.clone())
                .ok_or_else(|| "target subsystem is not registered and active".to_owned())?;
            let mut chain = self.lock_chain();
            let bytes =
                chain.submit_local(timestamp, creator, subsystem, payload, move |prefix| {
                    match catch_unwind(AssertUnwindSafe(|| signer(prefix))) {
                        Ok(result) => result,
                        Err(_) => Err("signer panicked".to_owned()),
                    }
                })?;
            let id = TxId::for_transaction(&bytes);
            let registration = writer
                .registrations
                .get_mut(&subsystem)
                .expect("prechecked registration exists");
            if !Arc::ptr_eq(&registration.lane, &lane) {
                return Err(committed_error(
                    id,
                    vec!["target registration changed during local commit".to_owned()],
                ));
            }
            let delivery = registration.reserve_delivery(id);
            (bytes, id, delivery)
        };
        let mut failures = Vec::new();
        match catch_unwind(AssertUnwindSafe(|| queue_propagation(&bytes))) {
            Ok(Ok(())) => {}
            Ok(Err(message)) => failures.push(format!("queue propagation failed: {message}")),
            Err(_) => failures.push("queue propagation panicked".to_owned()),
        }
        if let Err(message) = delivery.run_delivery(id, payload) {
            failures.push(message);
        }
        if failures.is_empty() {
            Ok((id, bytes))
        } else {
            Err(committed_error(id, failures))
        }
    }

    pub fn contains(&self, id: TxId) -> bool {
        self.lock_chain().contains(id)
    }

    pub fn tip(&self) -> Option<TxId> {
        self.lock_chain().tip()
    }

    pub fn between_txids(&self, older: TxId, newer: TxId) -> Result<Vec<TxId>, String> {
        self.lock_chain().between_txids(older, newer)
    }

    pub fn get_txn(&self, id: TxId) -> Result<Option<Vec<u8>>, String> {
        self.lock_chain().get_txn(id)
    }

    fn lock_writer(&self) -> MutexGuard<'_, WriterState> {
        self.writer.lock().expect("KTO writer mutex poisoned")
    }

    fn lock_chain(&self) -> MutexGuard<'_, CanonicalChain> {
        self.chain.lock().expect("KTO chain mutex poisoned")
    }
}

impl Registration {
    fn is_active(&self) -> bool {
        self.mode == RegistrationMode::Active && !self.lane.faulted.load(Ordering::Acquire)
    }

    fn reserve_delivery(&mut self, id: TxId) -> Reservation {
        let reservation = Reservation {
            lane: self.lane.clone(),
            ticket: self.next_ticket,
        };
        self.next_ticket += 1;
        self.latest = Some(id);
        reservation
    }

    fn reserve_reorg(&mut self) -> Reservation {
        let reservation = Reservation {
            lane: self.lane.clone(),
            ticket: self.next_ticket,
        };
        self.next_ticket += 1;
        reservation
    }
}

impl Lane {
    fn new(handler: Arc<dyn Subsystem>) -> Self {
        Self {
            handler,
            progress: Mutex::new(LaneProgress {
                serving: 0,
                fault: None,
            }),
            changed: Condvar::new(),
            running: AtomicBool::new(false),
            faulted: AtomicBool::new(false),
        }
    }

    fn run_replay(&self, id: TxId, payload: &[u8]) -> Result<(), String> {
        self.running.store(true, Ordering::Release);
        let result = invoke_callback("replay submit", || self.handler.submit_txn(id, payload));
        if let Err(message) = &result {
            self.record_fault(message.clone());
        }
        self.running.store(false, Ordering::Release);
        result
    }

    fn run_ticket<F>(&self, ticket: u64, operation: &str, callback: F) -> Result<(), String>
    where
        F: FnOnce() -> Result<(), String>,
    {
        let mut progress = self.progress.lock().expect("subsystem lane mutex poisoned");
        while progress.serving < ticket {
            progress = self
                .changed
                .wait(progress)
                .expect("subsystem lane mutex poisoned");
        }
        if progress.serving > ticket {
            return Err(format!("{operation} ticket was already completed"));
        }
        if let Some(reason) = progress.fault.clone() {
            progress.serving += 1;
            self.changed.notify_all();
            return Err(format!(
                "{operation} callback was not run because the registration faulted: {reason}"
            ));
        }
        self.running.store(true, Ordering::Release);
        drop(progress);
        let result = invoke_callback(operation, callback);
        let mut progress = self.progress.lock().expect("subsystem lane mutex poisoned");
        if let Err(message) = &result
            && progress.fault.is_none()
        {
            progress.fault = Some(message.clone());
            self.faulted.store(true, Ordering::Release);
        }
        progress.serving += 1;
        self.running.store(false, Ordering::Release);
        self.changed.notify_all();
        result
    }

    fn record_fault(&self, message: String) {
        let mut progress = self.progress.lock().expect("subsystem lane mutex poisoned");
        if progress.fault.is_none() {
            progress.fault = Some(message);
            self.faulted.store(true, Ordering::Release);
        }
    }

    fn is_quiescent(&self, next_ticket: u64) -> bool {
        let progress = self.progress.lock().expect("subsystem lane mutex poisoned");
        progress.serving == next_ticket && !self.running.load(Ordering::Acquire)
    }
}

impl Reservation {
    fn run_delivery(self, id: TxId, payload: &[u8]) -> Result<(), String> {
        self.lane.run_ticket(self.ticket, "submit", || {
            self.lane.handler.submit_txn(id, payload)
        })
    }

    fn run_reorg(self) -> Result<(), String> {
        self.lane
            .run_ticket(self.ticket, "reorg", || self.lane.handler.reorg())
    }
}

fn invoke_callback<F>(operation: &str, callback: F) -> Result<(), String>
where
    F: FnOnce() -> Result<(), String>,
{
    match catch_unwind(AssertUnwindSafe(callback)) {
        Ok(Ok(())) => Ok(()),
        Ok(Err(message)) => Err(format!("{operation} callback failed: {message}")),
        Err(_) => Err(format!("{operation} callback panicked")),
    }
}

fn ensure_replaceable(registration: &Registration) -> Result<(), String> {
    if registration.mode == RegistrationMode::Replaying
        && !registration.lane.faulted.load(Ordering::Acquire)
    {
        return Err("subsystem registration is replaying".to_owned());
    }
    if registration.mode == RegistrationMode::Active
        && !registration.lane.faulted.load(Ordering::Acquire)
    {
        return Err("subsystem is already registered and active".to_owned());
    }
    if !registration.lane.is_quiescent(registration.next_ticket) {
        return Err("previous subsystem lane work has not quiesced".to_owned());
    }
    Ok(())
}

fn reserve_active_delivery(
    writer: &mut WriterState,
    subsystem: SubsystemId,
    id: TxId,
) -> Option<Reservation> {
    writer
        .registrations
        .get_mut(&subsystem)
        .filter(|registration| registration.is_active())
        .map(|registration| registration.reserve_delivery(id))
}

fn run_reorganization(
    id: TxId,
    payload: &[u8],
    reorgs: Vec<Reservation>,
    replacement: Option<Reservation>,
) -> Vec<String> {
    thread::scope(|scope| {
        let mut jobs = Vec::with_capacity(reorgs.len() + usize::from(replacement.is_some()));
        for reservation in reorgs {
            jobs.push(scope.spawn(move || reservation.run_reorg()));
        }
        if let Some(reservation) = replacement {
            jobs.push(scope.spawn(move || reservation.run_delivery(id, payload)));
        }
        let mut failures = Vec::new();
        for job in jobs {
            match job.join() {
                Ok(Ok(())) => {}
                Ok(Err(message)) => failures.push(message),
                Err(_) => failures.push("callback lane panicked".to_owned()),
            }
        }
        failures
    })
}

fn fault_cursor(lane: &Lane, message: String) -> String {
    let failure = format!("subsystem replay cursor failed: {message}");
    lane.record_fault(failure.clone());
    failure
}

fn committed_error(id: TxId, failures: Vec<String>) -> String {
    format!("TxId {id:?} was committed; {}", failures.join("; "))
}

#[cfg(test)]
mod tests {
    use super::*;
    use kcode_k1_txn_ordering_live_testkit::{
        OrderingCandidate, OrderingHarness, TestQueuePropagation, TestSigner, TestSubsystem,
        verify_independent_subsystem_lanes, verify_queue_before_callback,
        verify_registration_sentinels, verify_signing_commit_exclusion,
        verify_unrelated_callback_reentry,
    };
    use kcode_k1_txn_ordering_recovery_testkit::{
        verify_callback_failure_isolation, verify_reorganization_isolation,
        verify_replay_live_handoff, verify_restart_duplicates_and_queries,
    };
    use std::path::Path;

    struct Harness;

    struct TestSubsystemAdapter {
        inner: Arc<dyn TestSubsystem>,
    }

    impl Subsystem for TestSubsystemAdapter {
        fn submit_txn(&self, id: TxId, payload: &[u8]) -> Result<(), String> {
            self.inner.submit_txn(id, payload)
        }

        fn reorg(&self) -> Result<(), String> {
            self.inner.reorg()
        }
    }

    impl OrderingCandidate for K1TxnOrdering {
        fn register_subsystem(
            &self,
            subsystem: SubsystemId,
            after: Option<TxId>,
            handler: Arc<dyn TestSubsystem>,
        ) -> Result<(), String> {
            K1TxnOrdering::register_subsystem(
                self,
                subsystem,
                after,
                Arc::new(TestSubsystemAdapter { inner: handler }),
            )
        }

        fn submit_txn(&self, transaction: &[u8]) -> Result<(), SubmitError> {
            K1TxnOrdering::submit_txn(self, transaction)
        }

        fn submit_local_txn(
            &self,
            timestamp: u64,
            creator: [u8; 32],
            subsystem: SubsystemId,
            payload: &[u8],
            signer: TestSigner,
            queue_propagation: TestQueuePropagation,
        ) -> Result<(TxId, Vec<u8>), String> {
            K1TxnOrdering::submit_local_txn(
                self,
                timestamp,
                creator,
                subsystem,
                payload,
                signer,
                queue_propagation,
            )
        }

        fn contains(&self, id: TxId) -> bool {
            K1TxnOrdering::contains(self, id)
        }

        fn tip(&self) -> Option<TxId> {
            K1TxnOrdering::tip(self)
        }

        fn between_txids(&self, older: TxId, newer: TxId) -> Result<Vec<TxId>, String> {
            K1TxnOrdering::between_txids(self, older, newer)
        }

        fn get_txn(&self, id: TxId) -> Result<Option<Vec<u8>>, String> {
            K1TxnOrdering::get_txn(self, id)
        }
    }

    impl OrderingHarness for Harness {
        fn open(&self, root: &Path) -> Result<Arc<dyn OrderingCandidate>, String> {
            Ok(Arc::new(K1TxnOrdering::open(root)?))
        }
    }

    #[test]
    fn registration_sentinels() {
        verify_registration_sentinels(&Harness).unwrap();
    }

    #[test]
    fn queue_before_callback() {
        verify_queue_before_callback(&Harness).unwrap();
    }

    #[test]
    fn independent_subsystem_lanes() {
        verify_independent_subsystem_lanes(&Harness).unwrap();
    }

    #[test]
    fn signing_commit_exclusion() {
        verify_signing_commit_exclusion(&Harness).unwrap();
    }

    #[test]
    fn replay_live_handoff() {
        verify_replay_live_handoff(&Harness).unwrap();
    }

    #[test]
    fn callback_failure_isolation() {
        verify_callback_failure_isolation(&Harness).unwrap();
    }

    #[test]
    fn reorganization_isolation() {
        verify_reorganization_isolation(&Harness).unwrap();
    }

    #[test]
    fn unrelated_callback_reentry() {
        verify_unrelated_callback_reentry(&Harness).unwrap();
    }

    #[test]
    fn restart_duplicates_and_queries() {
        verify_restart_duplicates_and_queries(&Harness).unwrap();
    }
}