use std::collections::HashMap;
use std::sync::Arc;
use crate::error::{DbError, DbResult};
use crate::txn::participant::{TxLocked, TxParticipant};
use crate::txn::{Action, Op};
use crate::upsert::Upsert;
pub(crate) struct Entry {
pub id: u32,
pub rank: u8,
pub participant: Arc<dyn TxParticipant>,
pub shards: Vec<usize>,
}
pub(crate) struct Write {
pub entry: usize,
pub key: Vec<u8>,
pub value: Option<Vec<u8>>,
pub derived: bool,
}
type ProjectionInput = (usize, Vec<u8>, Option<Vec<u8>>, Option<Vec<u8>>);
type EvalRank0Outcome = (Vec<Write>, Vec<ProjectionInput>);
fn idx_for(entries: &[Entry], participant: &Arc<dyn TxParticipant>) -> Option<usize> {
entries
.iter()
.position(|e| Arc::ptr_eq(&e.participant, participant))
}
fn ensure_entry(
entries: &mut Vec<Entry>,
db: &crate::armour::Db,
participant: Arc<dyn TxParticipant>,
rank: u8,
key: &[u8],
) -> DbResult<usize> {
let sid = participant.shard_of(key)?;
if let Some(pos) = idx_for(entries, &participant) {
if !entries[pos].shards.contains(&sid) {
entries[pos].shards.push(sid);
}
return Ok(pos);
}
let id = db.tx_collection_id(participant.as_ref())?;
entries.push(Entry {
id,
rank,
participant,
shards: vec![sid],
});
Ok(entries.len() - 1)
}
fn ensure_participant(
entries: &mut Vec<Entry>,
db: &crate::armour::Db,
participant: Arc<dyn TxParticipant>,
rank: u8,
) -> DbResult<usize> {
if let Some(pos) = idx_for(entries, &participant) {
return Ok(pos);
}
let id = db.tx_collection_id(participant.as_ref())?;
entries.push(Entry {
id,
rank,
participant,
shards: Vec::new(),
});
Ok(entries.len() - 1)
}
fn sort_entries(entries: &mut [Entry]) {
entries.sort_by_key(|a| (a.rank, a.id));
for e in entries.iter_mut() {
e.shards.sort_unstable();
}
}
fn set_write(
writes: &mut Vec<Write>,
entry: usize,
key: Vec<u8>,
value: Option<Vec<u8>>,
derived: bool,
) {
match writes.iter_mut().find(|w| w.entry == entry && w.key == key) {
Some(existing) => {
existing.value = value;
existing.derived = derived;
}
None => writes.push(Write {
entry,
key,
value,
derived,
}),
}
}
pub(crate) fn abort_locked<'a>(locked: Vec<Box<dyn TxLocked + 'a>>) {
for l in locked.into_iter().rev() {
l.abort();
}
}
pub(crate) struct PlanOutcome {
pub entries: Vec<Entry>,
pub rank0_len: usize,
pub rank0_ops: Vec<Op>,
pub rank1_ops: Vec<Op>,
}
pub(crate) fn plan(db: &crate::armour::Db, ops: Vec<Op>) -> DbResult<PlanOutcome> {
let mut rank0_ops = Vec::new();
let mut rank1_ops = Vec::new();
for op in ops {
if db.tx_rank_of(op.participant.as_ref()) == 0 {
rank0_ops.push(op);
} else {
rank1_ops.push(op);
}
}
let mut entries: Vec<Entry> = Vec::new();
for op in &rank0_ops {
ensure_entry(&mut entries, db, op.participant.clone(), 0, &op.key)?;
}
let sources: Vec<_> = entries.iter().map(|e| e.participant.clone()).collect();
for src in sources {
for target in src.tx_preplan_projection_targets()? {
ensure_participant(&mut entries, db, target, 1)?;
}
}
for op in &rank1_ops {
ensure_entry(&mut entries, db, op.participant.clone(), 1, &op.key)?;
}
sort_entries(&mut entries);
let rank0_len = entries.iter().filter(|e| e.rank == 0).count();
Ok(PlanOutcome {
entries,
rank0_len,
rank0_ops,
rank1_ops,
})
}
pub(crate) fn eval_rank0<'a>(
entries: &'a [Entry],
locked: &mut [Box<dyn TxLocked + 'a>],
ops: Vec<Op>,
) -> DbResult<EvalRank0Outcome> {
let mut writes: Vec<Write> = Vec::new();
let mut projection_inputs = Vec::new();
let mut eval_err = None;
for op in ops {
let Some(entry_pos) = idx_for(entries, &op.participant) else {
eval_err = Some(DbError::Internal("operation lost its participant"));
break;
};
let current: Option<Vec<u8>> = match writes
.iter()
.rev()
.find(|w| w.entry == entry_pos && w.key == op.key)
{
Some(w) => w.value.clone(),
None => match locked[entry_pos].read(&op.key) {
Ok(v) => v,
Err(e) => {
eval_err = Some(e);
break;
}
},
};
match op.action {
Action::Put(bytes) => {
projection_inputs.push((
entry_pos,
op.key.clone(),
current.clone(),
Some(bytes.clone()),
));
set_write(&mut writes, entry_pos, op.key, Some(bytes), false);
}
Action::Delete => {
projection_inputs.push((entry_pos, op.key.clone(), current.clone(), None));
set_write(&mut writes, entry_pos, op.key, None, false);
}
Action::Eval(f) => match f(current.as_deref()) {
Ok(Upsert::Put(bytes)) => {
projection_inputs.push((
entry_pos,
op.key.clone(),
current.clone(),
Some(bytes.clone()),
));
set_write(&mut writes, entry_pos, op.key, Some(bytes), false);
}
Ok(Upsert::Delete) => {
projection_inputs.push((entry_pos, op.key.clone(), current.clone(), None));
set_write(&mut writes, entry_pos, op.key, None, false);
}
Ok(Upsert::Keep) => {}
Err(e) => {
eval_err = Some(e);
break;
}
},
Action::Guard(f) => {
if let Err(e) = f(current.as_deref()) {
eval_err = Some(e);
break;
}
}
}
}
eval_err.map_or(Ok((writes, projection_inputs)), Err)
}
pub(crate) fn project_rank0(
entries: &[Entry],
writes: &mut Vec<Write>,
inputs: Vec<ProjectionInput>,
log_encrypted: bool,
) -> DbResult<HashMap<usize, Vec<usize>>> {
let mut extra_shards: HashMap<usize, Vec<usize>> = HashMap::new();
for (src, key, old, new) in inputs {
let derived = entries[src]
.participant
.tx_project(&key, old.as_deref(), new.as_deref())?;
for row in derived {
if row.participant.encrypted() && !log_encrypted {
return Err(DbError::Config(
"an encrypted collection needs TxConfig::log_encryption_key: \
the transaction log would otherwise hold its values in the clear",
));
}
let Some(target_pos) = idx_for(entries, &row.participant) else {
return Err(DbError::UndeclaredIndexTarget {
index: row.participant.name(),
});
};
let sid = row.participant.shard_of(&row.key)?;
extra_shards.entry(target_pos).or_default().push(sid);
set_write(writes, target_pos, row.key, row.value, true);
}
}
Ok(extra_shards)
}
pub(crate) fn eval_rank1<'a>(
entries: &'a [Entry],
locked: &mut [Box<dyn TxLocked + 'a>],
rank0_len: usize,
ops: Vec<Op>,
writes: &mut Vec<Write>,
) -> DbResult<()> {
for op in ops {
let Some(entry_pos) = idx_for(entries, &op.participant) else {
return Err(DbError::Internal("operation lost its participant"));
};
debug_assert!(entry_pos >= rank0_len);
let locked_pos = entry_pos;
let current: Option<Vec<u8>> = match writes
.iter()
.rev()
.find(|w| w.entry == entry_pos && w.key == op.key)
{
Some(w) => w.value.clone(),
None => locked[locked_pos].read(&op.key)?,
};
match op.action {
Action::Put(bytes) => set_write(writes, entry_pos, op.key, Some(bytes), false),
Action::Delete => set_write(writes, entry_pos, op.key, None, false),
Action::Eval(f) => match f(current.as_deref()) {
Ok(Upsert::Put(bytes)) => set_write(writes, entry_pos, op.key, Some(bytes), false),
Ok(Upsert::Delete) => set_write(writes, entry_pos, op.key, None, false),
Ok(Upsert::Keep) => {}
Err(e) => return Err(e),
},
Action::Guard(f) => f(current.as_deref())?,
}
}
Ok(())
}
pub(crate) fn admit(db: &crate::armour::Db, entries: &[Entry], writes: &[Write]) -> DbResult<()> {
for entry in entries {
db.tx_collection_id(entry.participant.as_ref())?;
}
if writes.len() > db.tx_config().max_ops {
let has_derived = writes.iter().any(|w| w.derived);
return Err(DbError::Client(if has_derived {
"transaction exceeds max_ops (including derived secondary-index rows)"
} else {
"transaction exceeds max_ops"
}));
}
for w in writes {
let value_len = w.value.as_deref().map_or(0, <[u8]>::len);
if let Some(cap) = entries[w.entry].participant.max_entry_bytes()
&& crate::entry::entry_size(w.key.len(), value_len as u32) as usize > cap
{
return Err(DbError::Client(if w.derived {
"derived secondary-index row exceeds write_buffer_size"
} else {
"transaction entry exceeds write_buffer_size"
}));
}
}
Ok(())
}
pub(crate) fn max_bytes_err(writes: &[Write]) -> DbError {
let has_derived = writes.iter().any(|w| w.derived);
DbError::Client(if has_derived {
"transaction exceeds max_bytes (including derived secondary-index rows)"
} else {
"transaction exceeds max_bytes"
})
}