Make storage transactions specific to a client_id (#67)

Transactions for different client_ids cannot interfere with one another,
so this provides an opportunity for the sort of concurrency that a
mult-client hosting solution might need. For example, a postgres backend
could lock the client row in each transaction.
This commit is contained in:
Dustin J. Mitchell
2024-11-27 00:09:03 -05:00
committed by GitHub
parent 4029c03479
commit 1828a31a24
9 changed files with 246 additions and 319 deletions

View File

@ -38,14 +38,16 @@ impl InMemoryStorage {
}
struct InnerTxn<'a> {
client_id: Uuid,
guard: MutexGuard<'a, Inner>,
written: bool,
committed: bool,
}
impl Storage for InMemoryStorage {
fn txn(&self) -> anyhow::Result<Box<dyn StorageTxn + '_>> {
fn txn(&self, client_id: Uuid) -> anyhow::Result<Box<dyn StorageTxn + '_>> {
Ok(Box::new(InnerTxn {
client_id,
guard: self.0.lock().expect("poisoned lock"),
written: false,
committed: false,
@ -54,16 +56,16 @@ impl Storage for InMemoryStorage {
}
impl<'a> StorageTxn for InnerTxn<'a> {
fn get_client(&mut self, client_id: Uuid) -> anyhow::Result<Option<Client>> {
Ok(self.guard.clients.get(&client_id).cloned())
fn get_client(&mut self) -> anyhow::Result<Option<Client>> {
Ok(self.guard.clients.get(&self.client_id).cloned())
}
fn new_client(&mut self, client_id: Uuid, latest_version_id: Uuid) -> anyhow::Result<()> {
if self.guard.clients.contains_key(&client_id) {
return Err(anyhow::anyhow!("Client {} already exists", client_id));
fn new_client(&mut self, latest_version_id: Uuid) -> anyhow::Result<()> {
if self.guard.clients.contains_key(&self.client_id) {
return Err(anyhow::anyhow!("Client {} already exists", self.client_id));
}
self.guard.clients.insert(
client_id,
self.client_id,
Client {
latest_version_id,
snapshot: None,
@ -73,64 +75,57 @@ impl<'a> StorageTxn for InnerTxn<'a> {
Ok(())
}
fn set_snapshot(
&mut self,
client_id: Uuid,
snapshot: Snapshot,
data: Vec<u8>,
) -> anyhow::Result<()> {
fn set_snapshot(&mut self, snapshot: Snapshot, data: Vec<u8>) -> anyhow::Result<()> {
let client = self
.guard
.clients
.get_mut(&client_id)
.get_mut(&self.client_id)
.ok_or_else(|| anyhow::anyhow!("no such client"))?;
client.snapshot = Some(snapshot);
self.guard.snapshots.insert(client_id, data);
self.guard.snapshots.insert(self.client_id, data);
self.written = true;
Ok(())
}
fn get_snapshot_data(
&mut self,
client_id: Uuid,
version_id: Uuid,
) -> anyhow::Result<Option<Vec<u8>>> {
fn get_snapshot_data(&mut self, version_id: Uuid) -> anyhow::Result<Option<Vec<u8>>> {
// sanity check
let client = self.guard.clients.get(&client_id);
let client = self.guard.clients.get(&self.client_id);
let client = client.ok_or_else(|| anyhow::anyhow!("no such client"))?;
if Some(&version_id) != client.snapshot.as_ref().map(|snap| &snap.version_id) {
return Err(anyhow::anyhow!("unexpected snapshot_version_id"));
}
Ok(self.guard.snapshots.get(&client_id).cloned())
Ok(self.guard.snapshots.get(&self.client_id).cloned())
}
fn get_version_by_parent(
&mut self,
client_id: Uuid,
parent_version_id: Uuid,
) -> anyhow::Result<Option<Version>> {
if let Some(parent_version_id) = self.guard.children.get(&(client_id, parent_version_id)) {
if let Some(parent_version_id) = self
.guard
.children
.get(&(self.client_id, parent_version_id))
{
Ok(self
.guard
.versions
.get(&(client_id, *parent_version_id))
.get(&(self.client_id, *parent_version_id))
.cloned())
} else {
Ok(None)
}
}
fn get_version(
&mut self,
client_id: Uuid,
version_id: Uuid,
) -> anyhow::Result<Option<Version>> {
Ok(self.guard.versions.get(&(client_id, version_id)).cloned())
fn get_version(&mut self, version_id: Uuid) -> anyhow::Result<Option<Version>> {
Ok(self
.guard
.versions
.get(&(self.client_id, version_id))
.cloned())
}
fn add_version(
&mut self,
client_id: Uuid,
version_id: Uuid,
parent_version_id: Uuid,
history_segment: Vec<u8>,
@ -142,19 +137,21 @@ impl<'a> StorageTxn for InnerTxn<'a> {
history_segment,
};
if let Some(client) = self.guard.clients.get_mut(&client_id) {
if let Some(client) = self.guard.clients.get_mut(&self.client_id) {
client.latest_version_id = version_id;
if let Some(ref mut snap) = client.snapshot {
snap.versions_since += 1;
}
} else {
return Err(anyhow::anyhow!("Client {} does not exist", client_id));
return Err(anyhow::anyhow!("Client {} does not exist", self.client_id));
}
self.guard
.children
.insert((client_id, parent_version_id), version_id);
self.guard.versions.insert((client_id, version_id), version);
.insert((self.client_id, parent_version_id), version_id);
self.guard
.versions
.insert((self.client_id, version_id), version);
self.written = true;
Ok(())
@ -182,8 +179,8 @@ mod test {
#[test]
fn test_get_client_empty() -> anyhow::Result<()> {
let storage = InMemoryStorage::new();
let mut txn = storage.txn()?;
let maybe_client = txn.get_client(Uuid::new_v4())?;
let mut txn = storage.txn(Uuid::new_v4())?;
let maybe_client = txn.get_client()?;
assert!(maybe_client.is_none());
Ok(())
}
@ -191,20 +188,20 @@ mod test {
#[test]
fn test_client_storage() -> anyhow::Result<()> {
let storage = InMemoryStorage::new();
let mut txn = storage.txn()?;
let client_id = Uuid::new_v4();
let latest_version_id = Uuid::new_v4();
txn.new_client(client_id, latest_version_id)?;
let mut txn = storage.txn(client_id)?;
let client = txn.get_client(client_id)?.unwrap();
let latest_version_id = Uuid::new_v4();
txn.new_client(latest_version_id)?;
let client = txn.get_client()?.unwrap();
assert_eq!(client.latest_version_id, latest_version_id);
assert!(client.snapshot.is_none());
let latest_version_id = Uuid::new_v4();
txn.add_version(client_id, latest_version_id, Uuid::new_v4(), vec![1, 1])?;
txn.add_version(latest_version_id, Uuid::new_v4(), vec![1, 1])?;
let client = txn.get_client(client_id)?.unwrap();
let client = txn.get_client()?.unwrap();
assert_eq!(client.latest_version_id, latest_version_id);
assert!(client.snapshot.is_none());
@ -213,9 +210,9 @@ mod test {
timestamp: Utc::now(),
versions_since: 4,
};
txn.set_snapshot(client_id, snap.clone(), vec![1, 2, 3])?;
txn.set_snapshot(snap.clone(), vec![1, 2, 3])?;
let client = txn.get_client(client_id)?.unwrap();
let client = txn.get_client()?.unwrap();
assert_eq!(client.latest_version_id, latest_version_id);
assert_eq!(client.snapshot.unwrap(), snap);
@ -226,8 +223,9 @@ mod test {
#[test]
fn test_gvbp_empty() -> anyhow::Result<()> {
let storage = InMemoryStorage::new();
let mut txn = storage.txn()?;
let maybe_version = txn.get_version_by_parent(Uuid::new_v4(), Uuid::new_v4())?;
let client_id = Uuid::new_v4();
let mut txn = storage.txn(client_id)?;
let maybe_version = txn.get_version_by_parent(Uuid::new_v4())?;
assert!(maybe_version.is_none());
Ok(())
}
@ -235,20 +233,15 @@ mod test {
#[test]
fn test_add_version_and_get_version() -> anyhow::Result<()> {
let storage = InMemoryStorage::new();
let mut txn = storage.txn()?;
let client_id = Uuid::new_v4();
let mut txn = storage.txn(client_id)?;
let version_id = Uuid::new_v4();
let parent_version_id = Uuid::new_v4();
let history_segment = b"abc".to_vec();
txn.new_client(client_id, parent_version_id)?;
txn.add_version(
client_id,
version_id,
parent_version_id,
history_segment.clone(),
)?;
txn.new_client(parent_version_id)?;
txn.add_version(version_id, parent_version_id, history_segment.clone())?;
let expected = Version {
version_id,
@ -256,12 +249,10 @@ mod test {
history_segment,
};
let version = txn
.get_version_by_parent(client_id, parent_version_id)?
.unwrap();
let version = txn.get_version_by_parent(parent_version_id)?.unwrap();
assert_eq!(version, expected);
let version = txn.get_version(client_id, version_id)?.unwrap();
let version = txn.get_version(version_id)?.unwrap();
assert_eq!(version, expected);
txn.commit()?;
@ -271,41 +262,40 @@ mod test {
#[test]
fn test_snapshots() -> anyhow::Result<()> {
let storage = InMemoryStorage::new();
let mut txn = storage.txn()?;
let client_id = Uuid::new_v4();
let mut txn = storage.txn(client_id)?;
txn.new_client(client_id, Uuid::new_v4())?;
assert!(txn.get_client(client_id)?.unwrap().snapshot.is_none());
txn.new_client(Uuid::new_v4())?;
assert!(txn.get_client()?.unwrap().snapshot.is_none());
let snap = Snapshot {
version_id: Uuid::new_v4(),
timestamp: Utc::now(),
versions_since: 3,
};
txn.set_snapshot(client_id, snap.clone(), vec![9, 8, 9])?;
txn.set_snapshot(snap.clone(), vec![9, 8, 9])?;
assert_eq!(
txn.get_snapshot_data(client_id, snap.version_id)?.unwrap(),
txn.get_snapshot_data(snap.version_id)?.unwrap(),
vec![9, 8, 9]
);
assert_eq!(txn.get_client(client_id)?.unwrap().snapshot, Some(snap));
assert_eq!(txn.get_client()?.unwrap().snapshot, Some(snap));
let snap2 = Snapshot {
version_id: Uuid::new_v4(),
timestamp: Utc::now(),
versions_since: 10,
};
txn.set_snapshot(client_id, snap2.clone(), vec![0, 2, 4, 6])?;
txn.set_snapshot(snap2.clone(), vec![0, 2, 4, 6])?;
assert_eq!(
txn.get_snapshot_data(client_id, snap2.version_id)?.unwrap(),
txn.get_snapshot_data(snap2.version_id)?.unwrap(),
vec![0, 2, 4, 6]
);
assert_eq!(txn.get_client(client_id)?.unwrap().snapshot, Some(snap2));
assert_eq!(txn.get_client()?.unwrap().snapshot, Some(snap2));
// check that mismatched version is detected
assert!(txn.get_snapshot_data(client_id, Uuid::new_v4()).is_err());
assert!(txn.get_snapshot_data(Uuid::new_v4()).is_err());
txn.commit()?;
Ok(())