use std::collections::BTreeMap;
use std::panic::{AssertUnwindSafe, catch_unwind};
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use rustev_contract::ids::ArtifactId;
use rustev_contract::output::RawOutput;
use rustev_contract::remote::{Attribution, ExchangeMember, ExchangeRecord, ServedIdentity};
use rustev_contract::run::Charge;
use rustev_contract::schema;
use crate::wire::digest_bytes;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RetainedBytes {
pub request: Vec<u8>,
pub response: Option<Vec<u8>>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ReconciliationOffer {
pub attempt_id: String,
pub observed_units: u64,
pub binding: ArtifactId,
}
pub trait ExchangeSink: Send + Sync {
fn deliver(
&self,
record: &ExchangeRecord,
retained: Option<&RetainedBytes>,
) -> Result<(), String>;
fn offer(&self, offer: &ReconciliationOffer) -> Result<(), String> {
let _ = offer;
Ok(())
}
fn deliver_located(
&self,
record: &ExchangeRecord,
retained: Option<&RetainedBytes>,
) -> Result<Option<String>, String> {
self.deliver(record, retained).map(|()| None)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DeliveryNote {
pub digest: Option<String>,
pub located: Result<Option<String>, ()>,
}
pub struct DiscardExchanges;
impl ExchangeSink for DiscardExchanges {
fn deliver(&self, _: &ExchangeRecord, _: Option<&RetainedBytes>) -> Result<(), String> {
Ok(())
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct ExchangeStats {
pub delivered: u64,
pub failed: u64,
pub offers: u64,
pub offers_failed: u64,
}
pub struct Exchanges {
sink: Arc<dyn ExchangeSink>,
retain_bytes: bool,
delivered: AtomicU64,
failed: AtomicU64,
offers: AtomicU64,
offers_failed: AtomicU64,
}
impl Exchanges {
pub fn new(sink: Arc<dyn ExchangeSink>, retain_bytes: bool) -> Self {
Exchanges {
sink,
retain_bytes,
delivered: AtomicU64::new(0),
failed: AtomicU64::new(0),
offers: AtomicU64::new(0),
offers_failed: AtomicU64::new(0),
}
}
pub fn retains_bytes(&self) -> bool {
self.retain_bytes
}
pub fn deliver(&self, record: ExchangeRecord, bytes: Option<RetainedBytes>) -> DeliveryNote {
use rustev_contract::Document;
let record = record.bounded();
let digest = record.record_digest().ok().map(|d| d.as_str().to_string());
let retained = if self.retain_bytes { bytes } else { None };
let located = match catch_unwind(AssertUnwindSafe(|| {
self.sink.deliver_located(&record, retained.as_ref())
})) {
Ok(Ok(locator)) => Ok(locator),
_ => Err(()),
};
let counter = if located.is_ok() {
&self.delivered
} else {
&self.failed
};
counter.fetch_add(1, Ordering::SeqCst);
DeliveryNote { digest, located }
}
pub fn offer(&self, offer: ReconciliationOffer) {
let ok = matches!(
catch_unwind(AssertUnwindSafe(|| self.sink.offer(&offer))),
Ok(Ok(()))
);
let counter = if ok {
&self.offers
} else {
&self.offers_failed
};
counter.fetch_add(1, Ordering::SeqCst);
}
pub fn stats(&self) -> ExchangeStats {
ExchangeStats {
delivered: self.delivered.load(Ordering::SeqCst),
failed: self.failed.load(Ordering::SeqCst),
offers: self.offers.load(Ordering::SeqCst),
offers_failed: self.offers_failed.load(Ordering::SeqCst),
}
}
}
pub fn output_digest(output: &RawOutput) -> Option<String> {
rustev_contract::canonical::record_canonical_bytes(output)
.ok()
.map(|b| digest_bytes(&b))
}
pub fn blank_record(
protocol: &str,
binding: ArtifactId,
members: Vec<ExchangeMember>,
) -> ExchangeRecord {
ExchangeRecord {
schema: schema::REMOTE_EXCHANGE.to_string(),
protocol: protocol.to_string(),
binding,
members,
request_digest: None,
response_digest: None,
bytes_sent: false,
status: None,
retry_after: None,
requested_model: None,
served: ServedIdentity::unknown(),
route: vec![],
generation_id: None,
usage: BTreeMap::new(),
reported_cost: None,
exchange_charge: Charge::Unknown,
attribution: Attribution::Single,
privacy_requested: BTreeMap::new(),
provider_extras: BTreeMap::new(),
late: false,
truncated: false,
phase_expiry: None,
principal: None,
connection: None,
}
}
#[cfg(test)]
mod tests {
use super::*;
use rustev_contract::remote::{MemberOutcome, RemoteCode};
use std::sync::Mutex;
const A: &str = "sha256:00000000000000000000000000000000000000000000000000000000000000aa";
#[derive(Default)]
struct Keep {
records: Mutex<Vec<(ExchangeRecord, Option<RetainedBytes>)>>,
fail: bool,
panic: bool,
}
impl ExchangeSink for Keep {
fn deliver(&self, r: &ExchangeRecord, b: Option<&RetainedBytes>) -> Result<(), String> {
if self.panic {
panic!("SYNTHETIC sink panic");
}
self.records.lock().unwrap().push((r.clone(), b.cloned()));
if self.fail {
Err("SYNTHETIC refusal".into())
} else {
Ok(())
}
}
}
fn record() -> ExchangeRecord {
blank_record(
"rustev.remote/1",
ArtifactId::parse(A).unwrap(),
vec![ExchangeMember {
attempt_id: "d/s//1".into(),
outcome: MemberOutcome::Failure {
code: RemoteCode::RemoteError,
detail: "x".repeat(1000),
},
charge: Charge::Unknown,
}],
)
}
fn bytes() -> Option<RetainedBytes> {
Some(RetainedBytes {
request: b"{}".to_vec(),
response: None,
})
}
#[test]
fn records_are_bounded_and_bytes_kept_only_under_retention() {
let keep = Arc::new(Keep::default());
let x = Exchanges::new(keep.clone(), false);
x.deliver(record(), bytes());
let (r, b) = keep.records.lock().unwrap()[0].clone();
assert!(r.is_bounded() && r.truncated);
assert!(b.is_none(), "digest-only by default");
let x = Exchanges::new(keep.clone(), true);
x.deliver(record(), bytes());
assert!(keep.records.lock().unwrap()[1].1.is_some());
}
#[test]
fn sink_failures_and_panics_are_counted() {
let x = Exchanges::new(
Arc::new(Keep {
fail: true,
..Keep::default()
}),
false,
);
x.deliver(record(), None);
let p = Exchanges::new(
Arc::new(Keep {
panic: true,
..Keep::default()
}),
false,
);
p.deliver(record(), None);
assert_eq!(x.stats().failed, 1);
assert_eq!(p.stats().failed, 1);
assert_eq!(p.stats().delivered, 0);
p.offer(ReconciliationOffer {
attempt_id: "a".into(),
observed_units: 1,
binding: ArtifactId::parse(A).unwrap(),
});
assert_eq!(p.stats().offers, 1);
}
}