Skip to content

Commit

Permalink
Make storage transactions specific to a client_id (#67)
Browse files Browse the repository at this point in the history 
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.
  • Loading branch information
@djmitche
djmitche authored Nov 27, 2024
1 parent 4029c03 commit 1828a31
Show file tree
Hide file tree
Showing 9 changed files with 244 additions and 317 deletions.
134 changes: 62 additions & 72 deletions core/src/inmemory.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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>,
Expand All @@ -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(())
Expand Down Expand Up @@ -182,29 +179,29 @@ 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(())
}

#[test]
fn test_client_storage() -> 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 latest_version_id = Uuid::new_v4();
txn.new_client(client_id, latest_version_id)?;
txn.new_client(latest_version_id)?;

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());

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());

Expand All @@ -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);

Expand All @@ -226,42 +223,36 @@ 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(())
}

#[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,
parent_version_id,
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()?;
Expand All @@ -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(())
Expand Down
Loading

0 comments on commit 1828a31

Please sign in to comment.