garage/src/table/data.rs

use core::borrow::Borrow;
use std::convert::TryInto;
use std::sync::Arc;

use serde_bytes::ByteBuf;
use tokio::sync::Notify;

use garage_db as db;
use garage_db::counted_tree_hack::CountedTree;

use garage_util::data::*;
use garage_util::error::*;
use garage_util::migrate::Migrate;

use garage_rpc::system::System;

use crate::crdt::Crdt;
use crate::gc::GcTodoEntry;
use crate::metrics::*;
use crate::replication::*;
use crate::schema::*;
use crate::util::*;

pub struct TableData<F: TableSchema, R: TableReplication> {
	system: Arc<System>,

	pub instance: F,
	pub replication: R,

	pub store: db::Tree,

	pub(crate) merkle_tree: db::Tree,
	pub(crate) merkle_todo: db::Tree,
	pub(crate) merkle_todo_notify: Notify,

	pub(crate) insert_queue: db::Tree,
	pub(crate) insert_queue_notify: Notify,

	pub(crate) gc_todo: CountedTree,

	pub(crate) metrics: TableMetrics,
}

impl<F: TableSchema, R: TableReplication> TableData<F, R> {
	pub fn new(system: Arc<System>, instance: F, replication: R, db: &db::Db) -> Arc<Self> {
		let store = db
			.open_tree(&format!("{}:table", F::TABLE_NAME))
			.expect("Unable to open DB tree");

		let merkle_tree = db
			.open_tree(&format!("{}:merkle_tree", F::TABLE_NAME))
			.expect("Unable to open DB Merkle tree tree");
		let merkle_todo = db
			.open_tree(&format!("{}:merkle_todo", F::TABLE_NAME))
			.expect("Unable to open DB Merkle TODO tree");

		let insert_queue = db
			.open_tree(&format!("{}:insert_queue", F::TABLE_NAME))
			.expect("Unable to open insert queue DB tree");

		let gc_todo = db
			.open_tree(&format!("{}:gc_todo_v2", F::TABLE_NAME))
			.expect("Unable to open GC DB tree");
		let gc_todo = CountedTree::new(gc_todo).expect("Cannot count gc_todo_v2");

		let metrics = TableMetrics::new(
			F::TABLE_NAME,
			store.clone(),
			merkle_tree.clone(),
			merkle_todo.clone(),
			gc_todo.clone(),
		);

		Arc::new(Self {
			system,
			instance,
			replication,
			store,
			merkle_tree,
			merkle_todo,
			merkle_todo_notify: Notify::new(),
			insert_queue,
			insert_queue_notify: Notify::new(),
			gc_todo,
			metrics,
		})
	}

	// Read functions

	pub fn read_entry(&self, p: &F::P, s: &F::S) -> Result<Option<ByteBuf>, Error> {
		let tree_key = self.tree_key(p, s);
		if let Some(bytes) = self.store.get(&tree_key)? {
			Ok(Some(ByteBuf::from(bytes.to_vec())))
		} else {
			Ok(None)
		}
	}

	pub fn read_range(
		&self,
		partition_key: &F::P,
		start: &Option<F::S>,
		filter: &Option<F::Filter>,
		limit: usize,
		enumeration_order: EnumerationOrder,
	) -> Result<Vec<Arc<ByteBuf>>, Error> {
		let partition_hash = partition_key.hash();
		match enumeration_order {
			EnumerationOrder::Forward => {
				let first_key = match start {
					None => partition_hash.to_vec(),
					Some(sk) => self.tree_key(partition_key, sk),
				};
				let range = self.store.range(first_key..)?;
				self.read_range_aux(partition_hash, range, filter, limit)
			}
			EnumerationOrder::Reverse => match start {
				Some(sk) => {
					let last_key = self.tree_key(partition_key, sk);
					let range = self.store.range_rev(..=last_key)?;
					self.read_range_aux(partition_hash, range, filter, limit)
				}
				None => {
					let mut last_key = partition_hash.to_vec();
					let lower = u128::from_be_bytes(last_key[16..32].try_into().unwrap());
					last_key[16..32].copy_from_slice(&u128::to_be_bytes(lower + 1));
					let range = self.store.range_rev(..last_key)?;
					self.read_range_aux(partition_hash, range, filter, limit)
				}
			},
		}
	}

	fn read_range_aux<'a>(
		&self,
		partition_hash: Hash,
		range: db::ValueIter<'a>,
		filter: &Option<F::Filter>,
		limit: usize,
	) -> Result<Vec<Arc<ByteBuf>>, Error> {
		let mut ret = vec![];
		for item in range {
			let (key, value) = item?;
			if &key[..32] != partition_hash.as_slice() {
				break;
			}
			let keep = match filter {
				None => true,
				Some(f) => {
					let entry = self.decode_entry(value.as_ref())?;
					F::matches_filter(&entry, f)
				}
			};
			if keep {
				ret.push(Arc::new(ByteBuf::from(value)));
			}
			if ret.len() >= limit {
				break;
			}
		}
		Ok(ret)
	}

	// Mutation functions
	// When changing this code, take care of propagating modifications correctly:
	// - When an entry is modified or deleted, call the updated() function
	//   on the table instance
	// - When an entry is modified or deleted, add it to the merkle updater's todo list.
	//   This has to be done atomically with the modification for the merkle updater
	//   to maintain consistency. The merkle updater must then be notified with todo_notify.
	// - When an entry is updated to be a tombstone, add it to the gc_todo tree

	pub(crate) fn update_many<T: Borrow<ByteBuf>>(&self, entries: &[T]) -> Result<(), Error> {
		for update_bytes in entries.iter() {
			self.update_entry(update_bytes.borrow().as_slice())?;
		}
		Ok(())
	}

	pub(crate) fn update_entry(&self, update_bytes: &[u8]) -> Result<(), Error> {
		let update = self.decode_entry(update_bytes)?;

		self.update_entry_with(
			update.partition_key(),
			update.sort_key(),
			|_tx, ent| match ent {
				Some(mut ent) => {
					ent.merge(&update);
					Ok(ent)
				}
				None => Ok(update.clone()),
			},
		)?;
		Ok(())
	}

	pub fn update_entry_with(
		&self,
		partition_key: &F::P,
		sort_key: &F::S,
		update_fn: impl Fn(&mut db::Transaction, Option<F::E>) -> db::TxOpResult<F::E>,
	) -> Result<Option<F::E>, Error> {
		let tree_key = self.tree_key(partition_key, sort_key);

		let changed = self.store.db().transaction(|mut tx| {
			let (old_entry, old_bytes, new_entry) = match tx.get(&self.store, &tree_key)? {
				Some(old_bytes) => {
					let old_entry = self.decode_entry(&old_bytes).map_err(db::TxError::Abort)?;
					let new_entry = update_fn(&mut tx, Some(old_entry.clone()))?;
					(Some(old_entry), Some(old_bytes), new_entry)
				}
				None => (None, None, update_fn(&mut tx, None)?),
			};

			// Changed can be true in two scenarios
			// Scenario 1: the actual represented value changed,
			//   so of course the messagepack encoding changed as well
			// Scenario 2: the value didn't change but due to a migration in the
			//   data format, the messagepack encoding changed. In this case,
			//   we also have to write the migrated value in the table and update
			//   the associated Merkle tree entry.
			let new_bytes = new_entry
				.encode()
				.map_err(Error::RmpEncode)
				.map_err(db::TxError::Abort)?;
			let changed = Some(&new_bytes[..]) != old_bytes.as_deref();
			drop(old_bytes);

			if changed {
				let new_bytes_hash = blake2sum(&new_bytes);
				tx.insert(&self.merkle_todo, &tree_key, new_bytes_hash.as_slice())?;
				tx.insert(&self.store, &tree_key, new_bytes)?;

				self.instance
					.updated(&mut tx, old_entry.as_ref(), Some(&new_entry))?;

				Ok(Some((new_entry, new_bytes_hash)))
			} else {
				Ok(None)
			}
		})?;

		if let Some((new_entry, new_bytes_hash)) = changed {
			self.metrics.internal_update_counter.add(1);

			let is_tombstone = new_entry.is_tombstone();
			self.merkle_todo_notify.notify_one();
			if is_tombstone {
				// We are only responsible for GC'ing this item if we are the
				// "leader" of the partition, i.e. the first node in the
				// set of nodes that replicates this partition.
				// This avoids GC loops and does not change the termination properties
				// of the GC algorithm, as in all cases GC is suspended if
				// any node of the partition is unavailable.
				let pk_hash = Hash::try_from(&tree_key[..32]).unwrap();
				let nodes = self.replication.write_nodes(&pk_hash);
				if nodes.first() == Some(&self.system.id) {
					GcTodoEntry::new(tree_key, new_bytes_hash).save(&self.gc_todo)?;
				}
			}

			Ok(Some(new_entry))
		} else {
			Ok(None)
		}
	}

	pub(crate) fn delete_if_equal(self: &Arc<Self>, k: &[u8], v: &[u8]) -> Result<bool, Error> {
		let removed = self
			.store
			.db()
			.transaction(|mut tx| match tx.get(&self.store, k)? {
				Some(cur_v) if cur_v == v => {
					let old_entry = self.decode_entry(v).map_err(db::TxError::Abort)?;

					tx.remove(&self.store, k)?;
					tx.insert(&self.merkle_todo, k, vec![])?;

					self.instance.updated(&mut tx, Some(&old_entry), None)?;
					Ok(true)
				}
				_ => Ok(false),
			})?;

		if removed {
			self.metrics.internal_delete_counter.add(1);
			self.merkle_todo_notify.notify_one();
		}
		Ok(removed)
	}

	pub(crate) fn delete_if_equal_hash(
		self: &Arc<Self>,
		k: &[u8],
		vhash: Hash,
	) -> Result<bool, Error> {
		let removed = self
			.store
			.db()
			.transaction(|mut tx| match tx.get(&self.store, k)? {
				Some(cur_v) if blake2sum(&cur_v[..]) == vhash => {
					let old_entry = self.decode_entry(&cur_v[..]).map_err(db::TxError::Abort)?;

					tx.remove(&self.store, k)?;
					tx.insert(&self.merkle_todo, k, vec![])?;

					self.instance.updated(&mut tx, Some(&old_entry), None)?;
					Ok(true)
				}
				_ => Ok(false),
			})?;

		if removed {
			self.metrics.internal_delete_counter.add(1);
			self.merkle_todo_notify.notify_one();
		}
		Ok(removed)
	}

	// ---- Insert queue functions ----

	pub(crate) fn queue_insert(
		&self,
		tx: &mut db::Transaction,
		ins: &F::E,
	) -> db::TxResult<(), Error> {
		let tree_key = self.tree_key(ins.partition_key(), ins.sort_key());

		let new_entry = match tx.get(&self.insert_queue, &tree_key)? {
			Some(old_v) => {
				let mut entry = self.decode_entry(&old_v).map_err(db::TxError::Abort)?;
				entry.merge(ins);
				entry.encode()
			}
			None => ins.encode(),
		};
		let new_entry = new_entry
			.map_err(Error::RmpEncode)
			.map_err(db::TxError::Abort)?;
		tx.insert(&self.insert_queue, &tree_key, new_entry)?;
		self.insert_queue_notify.notify_one();

		Ok(())
	}

	// ---- Utility functions ----

	pub fn tree_key(&self, p: &F::P, s: &F::S) -> Vec<u8> {
		let mut ret = p.hash().to_vec();
		ret.extend(s.sort_key());
		ret
	}

	pub fn decode_entry(&self, bytes: &[u8]) -> Result<F::E, Error> {
		match F::E::decode(bytes) {
			Some(x) => Ok(x),
			None => {
				error!("Unable to decode entry of {}", F::TABLE_NAME);
				for line in hexdump::hexdump_iter(bytes) {
					debug!("{}", line);
				}
				Err(Error::Message(format!(
					"Unable to decode entry of {}",
					F::TABLE_NAME
				)))
			}
		}
	}

	pub fn gc_todo_len(&self) -> Result<usize, Error> {
		Ok(self.gc_todo.len())
	}
}