use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use async_trait::async_trait;
use ciborium::Value as CborValue;
use fake::rand::rngs::StdRng;
use fake::rand::{RngExt as _, SeedableRng as _};
use indexmap::IndexMap;
use tokio::time::Instant;
use vantage_core::{Result, error};
use vantage_types::Record;
use vantage_vista::capabilities::VistaCapabilities;
use vantage_vista::column::Column;
use vantage_vista::reference::Reference;
use vantage_vista::source::TableShell;
use vantage_vista::Vista;
#[derive(Clone, Copy, Debug)]
pub struct Latency {
pub min: Duration,
pub max: Duration,
}
impl Latency {
pub fn fixed(d: Duration) -> Self {
Self { min: d, max: d }
}
pub fn between(min: Duration, max: Duration) -> Self {
Self {
min,
max: max.max(min),
}
}
fn draw(&self, rng: &mut StdRng) -> Duration {
if self.max <= self.min {
return self.min;
}
let spread = (self.max - self.min).as_millis() as u64;
self.min + Duration::from_millis(rng.random_range(0..=spread))
}
}
#[derive(Clone, Debug, Default)]
pub struct LatencyModel {
pub list: Option<Latency>,
pub get: Option<Latency>,
pub window: Option<Latency>,
pub count: Option<Latency>,
pub search_extra: Option<Latency>,
}
#[derive(Clone, Copy, Debug)]
pub struct Offline {
pub down: Duration,
pub period: Duration,
}
#[derive(Clone, Debug, Default)]
pub struct FaultSchedule {
pub error_rate: f64,
pub offline: Option<Offline>,
pub cursor_expiry: Option<Duration>,
pub total_lie: i64,
pub boundary_skew: bool,
}
#[derive(Clone, Copy, Debug)]
pub struct ExtraFields {
pub count: usize,
pub size: usize,
}
#[derive(Clone, Debug)]
pub struct BackendShape {
pub capabilities: VistaCapabilities,
pub page_size: usize,
pub latency: LatencyModel,
pub faults: FaultSchedule,
pub extra_fields: Option<ExtraFields>,
pub weirdness: f64,
pub seed: Option<u64>,
}
impl Default for BackendShape {
fn default() -> Self {
Self {
capabilities: VistaCapabilities {
can_count: true,
can_insert: true,
can_update: true,
can_delete: true,
can_order: true,
can_search: true,
..VistaCapabilities::default()
},
page_size: 25,
latency: LatencyModel::default(),
faults: FaultSchedule::default(),
extra_fields: None,
weirdness: 0.0,
seed: None,
}
}
}
#[derive(Clone, Copy)]
enum OpClass {
List,
Get,
Window,
Count,
}
pub struct ShapedShell {
inner: Box<dyn TableShell>,
shape: Arc<BackendShape>,
rng: Arc<Mutex<StdRng>>,
epoch: Instant,
page_size: usize,
searching: Arc<AtomicBool>,
}
impl ShapedShell {
pub fn new(inner: Box<dyn TableShell>, shape: BackendShape) -> Self {
let rng = match shape.seed {
Some(seed) => StdRng::seed_from_u64(seed ^ 0x5AAD_F00D_u64),
None => crate::value_gen::entropy_rng(),
};
let page_size = shape.page_size.max(1);
Self {
inner,
shape: Arc::new(shape),
rng: Arc::new(Mutex::new(rng)),
epoch: Instant::now(),
page_size,
searching: Arc::new(AtomicBool::new(false)),
}
}
async fn toll(&self, class: OpClass) -> Result<()> {
if let Some(off) = self.shape.faults.offline {
let elapsed = self.epoch.elapsed();
let period = off.period.max(off.down);
let into_period =
Duration::from_nanos((elapsed.as_nanos() % period.as_nanos().max(1)) as u64);
if into_period >= period.saturating_sub(off.down) {
return Err(error!("shaped source is offline (scheduled outage)"));
}
}
let (delay, failed) = {
let rng = &mut *self.rng.lock().unwrap();
let band = match class {
OpClass::List => self.shape.latency.list,
OpClass::Get => self.shape.latency.get,
OpClass::Window => self.shape.latency.window,
OpClass::Count => self.shape.latency.count,
};
let mut delay = band.map(|b| b.draw(rng)).unwrap_or_default();
if matches!(class, OpClass::List | OpClass::Window)
&& self.searching.load(Ordering::Relaxed)
{
if let Some(extra) = self.shape.latency.search_extra {
delay += extra.draw(rng);
}
}
let failed = self.shape.faults.error_rate > 0.0
&& rng.random_range(0.0..1.0) < self.shape.faults.error_rate;
(delay, failed)
};
if !delay.is_zero() {
tokio::time::sleep(delay).await;
}
if failed {
return Err(error!("shaped source request failed (injected fault)"));
}
Ok(())
}
fn skewed_offset(&self, offset: usize) -> usize {
if !self.shape.faults.boundary_skew || offset == 0 {
return offset;
}
let draw: f64 = self.rng.lock().unwrap().random_range(0.0..1.0);
if draw < 0.2 {
offset - 1 } else if draw < 0.4 {
offset + 1 } else {
offset
}
}
async fn lied_total(&self, vista: &Vista) -> Result<i64> {
let truth = self.inner.get_vista_count(vista).await?;
Ok((truth + self.shape.faults.total_lie).max(0))
}
fn cursor_token(&self, offset: usize) -> CborValue {
CborValue::Array(vec![
CborValue::Integer((offset as i64).into()),
CborValue::Integer((self.epoch.elapsed().as_millis() as i64).into()),
])
}
fn decode_cursor(&self, token: &CborValue) -> Result<usize> {
let CborValue::Array(parts) = token else {
return Err(error!("shaped source: malformed cursor token"));
};
let (Some(CborValue::Integer(offset)), Some(CborValue::Integer(issued_ms))) =
(parts.first(), parts.get(1))
else {
return Err(error!("shaped source: malformed cursor token"));
};
if let Some(expiry) = self.shape.faults.cursor_expiry {
let issued = Duration::from_millis(i128::from(*issued_ms).max(0) as u64);
if self.epoch.elapsed().saturating_sub(issued) > expiry {
return Err(error!("shaped source: cursor token expired"));
}
}
Ok(i128::from(*offset).max(0) as usize)
}
fn gate(&self, allowed: bool, method: &str, capability: &str) -> Result<()> {
if allowed {
Ok(())
} else {
Err(self.default_error(method, capability))
}
}
}
#[async_trait]
#[allow(clippy::ptr_arg)]
impl TableShell for ShapedShell {
fn columns(&self) -> &IndexMap<String, Column> {
self.inner.columns()
}
fn references(&self) -> &IndexMap<String, Reference> {
self.inner.references()
}
fn id_column(&self) -> Option<&str> {
self.inner.id_column()
}
fn capabilities(&self) -> &VistaCapabilities {
&self.shape.capabilities
}
fn driver_name(&self) -> &'static str {
"faker-shaped"
}
async fn list_vista_values(
&self,
vista: &Vista,
) -> Result<IndexMap<String, Record<CborValue>>> {
tracing::debug!(target: "vantage_faker::shape", op = "list", "request");
self.toll(OpClass::List).await?;
self.inner.list_vista_values(vista).await
}
async fn get_vista_value(
&self,
vista: &Vista,
id: &String,
) -> Result<Option<Record<CborValue>>> {
tracing::debug!(target: "vantage_faker::shape", op = "get", id = %id, "request");
self.toll(OpClass::Get).await?;
self.inner.get_vista_value(vista, id).await
}
async fn get_vista_some_value(
&self,
vista: &Vista,
) -> Result<Option<(String, Record<CborValue>)>> {
self.toll(OpClass::Get).await?;
self.inner.get_vista_some_value(vista).await
}
async fn fetch_window(
&self,
vista: &Vista,
offset: usize,
limit: usize,
) -> Result<Vec<(String, Record<CborValue>)>> {
self.gate(
self.shape.capabilities.can_fetch_window,
"fetch_window",
"can_fetch_window",
)?;
tracing::debug!(target: "vantage_faker::shape", op = "window", offset, limit, "request");
self.toll(OpClass::Window).await?;
let offset = self.skewed_offset(offset);
self.inner.fetch_window(vista, offset, limit).await
}
async fn fetch_window_counted(
&self,
vista: &Vista,
offset: usize,
limit: usize,
) -> Result<(Vec<(String, Record<CborValue>)>, Option<i64>)> {
let rows = self.fetch_window(vista, offset, limit).await?;
let total = if self.shape.capabilities.can_count {
Some(self.lied_total(vista).await?)
} else {
None
};
Ok((rows, total))
}
async fn fetch_page(
&self,
vista: &Vista,
page: usize,
) -> Result<Vec<(String, Record<CborValue>)>> {
self.gate(
self.shape.capabilities.can_fetch_page,
"fetch_page",
"can_fetch_page",
)?;
self.toll(OpClass::Window).await?;
let offset = self.skewed_offset(page.saturating_sub(1) * self.page_size);
self.inner.fetch_window(vista, offset, self.page_size).await
}
async fn fetch_next(
&self,
vista: &Vista,
token: Option<CborValue>,
) -> Result<(Vec<(String, Record<CborValue>)>, Option<CborValue>)> {
self.gate(
self.shape.capabilities.can_fetch_next,
"fetch_next",
"can_fetch_next",
)?;
self.toll(OpClass::Window).await?;
let offset = match &token {
None => 0,
Some(t) => self.decode_cursor(t)?,
};
let offset = self.skewed_offset(offset);
let rows = self.inner.fetch_window(vista, offset, self.page_size).await?;
let next = (rows.len() == self.page_size).then(|| self.cursor_token(offset + rows.len()));
Ok((rows, next))
}
async fn get_vista_count(&self, vista: &Vista) -> Result<i64> {
self.gate(self.shape.capabilities.can_count, "get_vista_count", "can_count")?;
tracing::debug!(target: "vantage_faker::shape", op = "count", "request");
self.toll(OpClass::Count).await?;
self.lied_total(vista).await
}
async fn insert_vista_value(
&self,
vista: &Vista,
id: &String,
record: &Record<CborValue>,
) -> Result<Record<CborValue>> {
self.gate(self.shape.capabilities.can_insert, "insert_vista_value", "can_insert")?;
self.inner.insert_vista_value(vista, id, record).await
}
async fn replace_vista_value(
&self,
vista: &Vista,
id: &String,
record: &Record<CborValue>,
) -> Result<Record<CborValue>> {
self.gate(self.shape.capabilities.can_update, "replace_vista_value", "can_update")?;
self.inner.replace_vista_value(vista, id, record).await
}
async fn patch_vista_value(
&self,
vista: &Vista,
id: &String,
partial: &Record<CborValue>,
) -> Result<Record<CborValue>> {
self.gate(self.shape.capabilities.can_update, "patch_vista_value", "can_update")?;
self.inner.patch_vista_value(vista, id, partial).await
}
async fn delete_vista_value(&self, vista: &Vista, id: &String) -> Result<()> {
self.gate(self.shape.capabilities.can_delete, "delete_vista_value", "can_delete")?;
self.inner.delete_vista_value(vista, id).await
}
async fn delete_vista_all_values(&self, vista: &Vista) -> Result<()> {
self.gate(
self.shape.capabilities.can_delete,
"delete_vista_all_values",
"can_delete",
)?;
self.inner.delete_vista_all_values(vista).await
}
async fn insert_vista_return_id_value(
&self,
vista: &Vista,
record: &Record<CborValue>,
) -> Result<String> {
self.gate(
self.shape.capabilities.can_insert,
"insert_vista_return_id_value",
"can_insert",
)?;
self.inner.insert_vista_return_id_value(vista, record).await
}
fn add_eq_condition(&mut self, field: &str, value: &CborValue) -> Result<()> {
self.inner.add_eq_condition(field, value)
}
fn add_search(&mut self, text: &str) -> Result<()> {
self.gate(self.shape.capabilities.can_search, "add_search", "can_search")?;
self.inner.add_search(text)?;
self.searching.store(true, Ordering::Relaxed);
Ok(())
}
fn clear_search(&mut self) -> Result<()> {
self.gate(self.shape.capabilities.can_search, "clear_search", "can_search")?;
self.inner.clear_search()?;
self.searching.store(false, Ordering::Relaxed);
Ok(())
}
fn add_order(&mut self, field: &str, dir: vantage_vista::sort::SortDirection) -> Result<()> {
self.gate(self.shape.capabilities.can_order, "add_order", "can_order")?;
self.inner.add_order(field, dir)
}
fn clear_orders(&mut self) -> Result<()> {
self.gate(self.shape.capabilities.can_order, "clear_orders", "can_order")?;
self.inner.clear_orders()
}
fn set_page_size(&mut self, size: usize) -> Result<()> {
self.gate(
self.shape.capabilities.can_set_page_size,
"set_page_size",
"can_set_page_size",
)?;
self.page_size = size.max(1);
Ok(())
}
fn clone_shell(&self) -> Option<Box<dyn TableShell>> {
let inner = self.inner.clone_shell()?;
Some(Box::new(Self {
inner,
shape: self.shape.clone(),
rng: self.rng.clone(),
epoch: self.epoch,
page_size: self.page_size,
searching: Arc::new(AtomicBool::new(false)),
}))
}
}