Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Switch to serializable isolation

Browse source 
Co-Authored-By: Nathan Sobo <nathan@zed.dev>
This commit is contained in:
Antonio Scandurra 2022-11-14 17:16:50 +01:00
parent d7369ace6a
commit b9af2ae66e
4 changed files with 298 additions and 160 deletions
Download patch file Download diff file Expand all files Collapse all files

1
crates/collab/migrations.sqlite/20221109000000_test_schema.sql
View file

@ -56,6 +56,7 @@ CREATE TABLE "project_collaborators" (
"is_host" BOOLEAN NOT NULL "is_host" BOOLEAN NOT NULL
); );
CREATE INDEX "index_project_collaborators_on_project_id" ON "project_collaborators" ("project_id"); CREATE INDEX "index_project_collaborators_on_project_id" ON "project_collaborators" ("project_id");
CREATE UNIQUE INDEX "index_project_collaborators_on_project_id" ON "project_collaborators" ("project_id", "replica_id");
CREATE TABLE "worktrees" ( CREATE TABLE "worktrees" (
"id" INTEGER NOT NULL, "id" INTEGER NOT NULL,

404
crates/collab/src/db.rs
View file

@ -2,7 +2,7 @@ use crate::{Error, Result};
use anyhow::anyhow; use anyhow::anyhow;
use axum::http::StatusCode; use axum::http::StatusCode;
use collections::HashMap; use collections::HashMap;
use futures::StreamExt; use futures::{future::BoxFuture, FutureExt, StreamExt};
use rpc::{proto, ConnectionId}; use rpc::{proto, ConnectionId};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use sqlx::{ use sqlx::{
@ -10,7 +10,7 @@ use sqlx::{
types::Uuid, types::Uuid,
FromRow, FromRow,
}; };
use std::{path::Path, time::Duration}; use std::{future::Future, path::Path, time::Duration};
use time::{OffsetDateTime, PrimitiveDateTime}; use time::{OffsetDateTime, PrimitiveDateTime};
#[cfg(test)] #[cfg(test)]
@ -27,27 +27,34 @@ pub struct Db<D: sqlx::Database> {
runtime: Option<tokio::runtime::Runtime>, runtime: Option<tokio::runtime::Runtime>,
} }
macro_rules! test_support { pub trait BeginTransaction: Send + Sync {
($self:ident, { $($token:tt)* }) => {{ type Database: sqlx::Database;
let body = async {
$($token)*
};
if cfg!(test) { fn begin_transaction(&self) -> BoxFuture<Result<sqlx::Transaction<'static, Self::Database>>>;
#[cfg(not(test))] }
unreachable!();
#[cfg(test)] // In Postgres, serializable transactions are opt-in
if let Some(background) = $self.background.as_ref() { impl BeginTransaction for Db<sqlx::Postgres> {
background.simulate_random_delay().await; type Database = sqlx::Postgres;
fn begin_transaction(&self) -> BoxFuture<Result<sqlx::Transaction<'static, sqlx::Postgres>>> {
async move {
let mut tx = self.pool.begin().await?;
sqlx::Executor::execute(&mut tx, "SET TRANSACTION ISOLATION LEVEL SERIALIZABLE;")
.await?;
Ok(tx)
} }
.boxed()
#[cfg(test)] }
$self.runtime.as_ref().unwrap().block_on(body) }
} else {
body.await // In Sqlite, transactions are inherently serializable.
impl BeginTransaction for Db<sqlx::Sqlite> {
type Database = sqlx::Sqlite;
fn begin_transaction(&self) -> BoxFuture<Result<sqlx::Transaction<'static, sqlx::Sqlite>>> {
async move { Ok(self.pool.begin().await?) }.boxed()
} }
}};
} }
pub trait RowsAffected { pub trait RowsAffected {
@ -88,7 +95,8 @@ impl Db<sqlx::Sqlite> {
} }
pub async fn get_users_by_ids(&self, ids: Vec<UserId>) -> Result<Vec<User>> { pub async fn get_users_by_ids(&self, ids: Vec<UserId>) -> Result<Vec<User>> {
test_support!(self, { self.transact(|tx| async {
let mut tx = tx;
let query = " let query = "
SELECT users.* SELECT users.*
FROM users FROM users
@ -96,13 +104,14 @@ impl Db<sqlx::Sqlite> {
"; ";
Ok(sqlx::query_as(query) Ok(sqlx::query_as(query)
.bind(&serde_json::json!(ids)) .bind(&serde_json::json!(ids))
.fetch_all(&self.pool) .fetch_all(&mut tx)
.await?) .await?)
}) })
.await
} }
pub async fn get_user_metrics_id(&self, id: UserId) -> Result<String> { pub async fn get_user_metrics_id(&self, id: UserId) -> Result<String> {
test_support!(self, { self.transact(|mut tx| async move {
let query = " let query = "
SELECT metrics_id SELECT metrics_id
FROM users FROM users
@ -110,9 +119,10 @@ impl Db<sqlx::Sqlite> {
"; ";
Ok(sqlx::query_scalar(query) Ok(sqlx::query_scalar(query)
.bind(id) .bind(id)
.fetch_one(&self.pool) .fetch_one(&mut tx)
.await?) .await?)
}) })
.await
} }
pub async fn create_user( pub async fn create_user(
@ -121,7 +131,7 @@ impl Db<sqlx::Sqlite> {
admin: bool, admin: bool,
params: NewUserParams, params: NewUserParams,
) -> Result<NewUserResult> { ) -> Result<NewUserResult> {
test_support!(self, { self.transact(|mut tx| async {
let query = " let query = "
INSERT INTO users (email_address, github_login, github_user_id, admin, metrics_id) INSERT INTO users (email_address, github_login, github_user_id, admin, metrics_id)
VALUES ($1, $2, $3, $4, $5) VALUES ($1, $2, $3, $4, $5)
@ -131,12 +141,13 @@ impl Db<sqlx::Sqlite> {
let (user_id, metrics_id): (UserId, String) = sqlx::query_as(query) let (user_id, metrics_id): (UserId, String) = sqlx::query_as(query)
.bind(email_address) .bind(email_address)
.bind(params.github_login) .bind(&params.github_login)
.bind(params.github_user_id) .bind(&params.github_user_id)
.bind(admin) .bind(admin)
.bind(Uuid::new_v4().to_string()) .bind(Uuid::new_v4().to_string())
.fetch_one(&self.pool) .fetch_one(&mut tx)
.await?; .await?;
tx.commit().await?;
Ok(NewUserResult { Ok(NewUserResult {
user_id, user_id,
metrics_id, metrics_id,
@ -144,6 +155,7 @@ impl Db<sqlx::Sqlite> {
inviting_user_id: None, inviting_user_id: None,
}) })
}) })
.await
} }
pub async fn fuzzy_search_users(&self, _name_query: &str, _limit: u32) -> Result<Vec<User>> { pub async fn fuzzy_search_users(&self, _name_query: &str, _limit: u32) -> Result<Vec<User>> {
@ -209,7 +221,8 @@ impl Db<sqlx::Postgres> {
} }
pub async fn fuzzy_search_users(&self, name_query: &str, limit: u32) -> Result<Vec<User>> { pub async fn fuzzy_search_users(&self, name_query: &str, limit: u32) -> Result<Vec<User>> {
test_support!(self, { self.transact(|tx| async {
let mut tx = tx;
let like_string = Self::fuzzy_like_string(name_query); let like_string = Self::fuzzy_like_string(name_query);
let query = " let query = "
SELECT users.* SELECT users.*
@ -222,27 +235,28 @@ impl Db<sqlx::Postgres> {
.bind(like_string) .bind(like_string)
.bind(name_query) .bind(name_query)
.bind(limit as i32) .bind(limit as i32)
.fetch_all(&self.pool) .fetch_all(&mut tx)
.await?) .await?)
}) })
.await
} }
pub async fn get_users_by_ids(&self, ids: Vec<UserId>) -> Result<Vec<User>> { pub async fn get_users_by_ids(&self, ids: Vec<UserId>) -> Result<Vec<User>> {
test_support!(self, { let ids = ids.iter().map(|id| id.0).collect::<Vec<_>>();
self.transact(|tx| async {
let mut tx = tx;
let query = " let query = "
SELECT users.* SELECT users.*
FROM users FROM users
WHERE users.id = ANY ($1) WHERE users.id = ANY ($1)
"; ";
Ok(sqlx::query_as(query) Ok(sqlx::query_as(query).bind(&ids).fetch_all(&mut tx).await?)
.bind(&ids.into_iter().map(|id| id.0).collect::<Vec<_>>())
.fetch_all(&self.pool)
.await?)
}) })
.await
} }
pub async fn get_user_metrics_id(&self, id: UserId) -> Result<String> { pub async fn get_user_metrics_id(&self, id: UserId) -> Result<String> {
test_support!(self, { self.transact(|mut tx| async move {
let query = " let query = "
SELECT metrics_id::text SELECT metrics_id::text
FROM users FROM users
@ -250,9 +264,10 @@ impl Db<sqlx::Postgres> {
"; ";
Ok(sqlx::query_scalar(query) Ok(sqlx::query_scalar(query)
.bind(id) .bind(id)
.fetch_one(&self.pool) .fetch_one(&mut tx)
.await?) .await?)
}) })
.await
} }
pub async fn create_user( pub async fn create_user(
@ -261,7 +276,7 @@ impl Db<sqlx::Postgres> {
admin: bool, admin: bool,
params: NewUserParams, params: NewUserParams,
) -> Result<NewUserResult> { ) -> Result<NewUserResult> {
test_support!(self, { self.transact(|mut tx| async {
let query = " let query = "
INSERT INTO users (email_address, github_login, github_user_id, admin) INSERT INTO users (email_address, github_login, github_user_id, admin)
VALUES ($1, $2, $3, $4) VALUES ($1, $2, $3, $4)
@ -271,11 +286,13 @@ impl Db<sqlx::Postgres> {
let (user_id, metrics_id): (UserId, String) = sqlx::query_as(query) let (user_id, metrics_id): (UserId, String) = sqlx::query_as(query)
.bind(email_address) .bind(email_address)
.bind(params.github_login) .bind(&params.github_login)
.bind(params.github_user_id) .bind(params.github_user_id)
.bind(admin) .bind(admin)
.fetch_one(&self.pool) .fetch_one(&mut tx)
.await?; .await?;
tx.commit().await?;
Ok(NewUserResult { Ok(NewUserResult {
user_id, user_id,
metrics_id, metrics_id,
@ -283,6 +300,7 @@ impl Db<sqlx::Postgres> {
inviting_user_id: None, inviting_user_id: None,
}) })
}) })
.await
} }
pub async fn create_user_from_invite( pub async fn create_user_from_invite(
@ -290,9 +308,7 @@ impl Db<sqlx::Postgres> {
invite: &Invite, invite: &Invite,
user: NewUserParams, user: NewUserParams,
) -> Result<Option<NewUserResult>> { ) -> Result<Option<NewUserResult>> {
test_support!(self, { self.transact(|mut tx| async {
let mut tx = self.pool.begin().await?;
let (signup_id, existing_user_id, inviting_user_id, signup_device_id): ( let (signup_id, existing_user_id, inviting_user_id, signup_device_id): (
i32, i32,
Option<UserId>, Option<UserId>,
@ -393,10 +409,11 @@ impl Db<sqlx::Postgres> {
signup_device_id, signup_device_id,
})) }))
}) })
.await
} }
pub async fn create_signup(&self, signup: Signup) -> Result<()> { pub async fn create_signup(&self, signup: Signup) -> Result<()> {
test_support!(self, { self.transact(|mut tx| async {
sqlx::query( sqlx::query(
" "
INSERT INTO signups INSERT INTO signups
@ -425,10 +442,12 @@ impl Db<sqlx::Postgres> {
.bind(&signup.editor_features) .bind(&signup.editor_features)
.bind(&signup.programming_languages) .bind(&signup.programming_languages)
.bind(&signup.device_id) .bind(&signup.device_id)
.execute(&self.pool) .execute(&mut tx)
.await?; .await?;
tx.commit().await?;
Ok(()) Ok(())
}) })
.await
} }
pub async fn create_invite_from_code( pub async fn create_invite_from_code(
@ -437,9 +456,7 @@ impl Db<sqlx::Postgres> {
email_address: &str, email_address: &str,
device_id: Option<&str>, device_id: Option<&str>,
) -> Result<Invite> { ) -> Result<Invite> {
test_support!(self, { self.transact(|mut tx| async {
let mut tx = self.pool.begin().await?;
let existing_user: Option<UserId> = sqlx::query_scalar( let existing_user: Option<UserId> = sqlx::query_scalar(
" "
SELECT id SELECT id
@ -516,10 +533,11 @@ impl Db<sqlx::Postgres> {
email_confirmation_code, email_confirmation_code,
}) })
}) })
.await
} }
pub async fn record_sent_invites(&self, invites: &[Invite]) -> Result<()> { pub async fn record_sent_invites(&self, invites: &[Invite]) -> Result<()> {
test_support!(self, { self.transact(|mut tx| async {
let emails = invites let emails = invites
.iter() .iter()
.map(|s| s.email_address.as_str()) .map(|s| s.email_address.as_str())
@ -532,15 +550,18 @@ impl Db<sqlx::Postgres> {
", ",
) )
.bind(&emails) .bind(&emails)
.execute(&self.pool) .execute(&mut tx)
.await?; .await?;
tx.commit().await?;
Ok(()) Ok(())
}) })
.await
} }
} }
impl<D> Db<D> impl<D> Db<D>
where where
Self: BeginTransaction<Database = D>,
D: sqlx::Database + sqlx::migrate::MigrateDatabase, D: sqlx::Database + sqlx::migrate::MigrateDatabase,
D::Connection: sqlx::migrate::Migrate, D::Connection: sqlx::migrate::Migrate,
for<'a> <D as sqlx::database::HasArguments<'a>>::Arguments: sqlx::IntoArguments<'a, D>, for<'a> <D as sqlx::database::HasArguments<'a>>::Arguments: sqlx::IntoArguments<'a, D>,
@ -627,18 +648,21 @@ where
// users // users
pub async fn get_all_users(&self, page: u32, limit: u32) -> Result<Vec<User>> { pub async fn get_all_users(&self, page: u32, limit: u32) -> Result<Vec<User>> {
test_support!(self, { self.transact(|tx| async {
let mut tx = tx;
let query = "SELECT * FROM users ORDER BY github_login ASC LIMIT $1 OFFSET $2"; let query = "SELECT * FROM users ORDER BY github_login ASC LIMIT $1 OFFSET $2";
Ok(sqlx::query_as(query) Ok(sqlx::query_as(query)
.bind(limit as i32) .bind(limit as i32)
.bind((page * limit) as i32) .bind((page * limit) as i32)
.fetch_all(&self.pool) .fetch_all(&mut tx)
.await?) .await?)
}) })
.await
} }
pub async fn get_user_by_id(&self, id: UserId) -> Result<Option<User>> { pub async fn get_user_by_id(&self, id: UserId) -> Result<Option<User>> {
test_support!(self, { self.transact(|tx| async {
let mut tx = tx;
let query = " let query = "
SELECT users.* SELECT users.*
FROM users FROM users
@ -647,16 +671,18 @@ where
"; ";
Ok(sqlx::query_as(query) Ok(sqlx::query_as(query)
.bind(&id) .bind(&id)
.fetch_optional(&self.pool) .fetch_optional(&mut tx)
.await?) .await?)
}) })
.await
} }
pub async fn get_users_with_no_invites( pub async fn get_users_with_no_invites(
&self, &self,
invited_by_another_user: bool, invited_by_another_user: bool,
) -> Result<Vec<User>> { ) -> Result<Vec<User>> {
test_support!(self, { self.transact(|tx| async {
let mut tx = tx;
let query = format!( let query = format!(
" "
SELECT users.* SELECT users.*
@ -667,8 +693,9 @@ where
if invited_by_another_user { " NOT" } else { "" } if invited_by_another_user { " NOT" } else { "" }
); );
Ok(sqlx::query_as(&query).fetch_all(&self.pool).await?) Ok(sqlx::query_as(&query).fetch_all(&mut tx).await?)
}) })
.await
} }
pub async fn get_user_by_github_account( pub async fn get_user_by_github_account(
@ -676,7 +703,8 @@ where
github_login: &str, github_login: &str,
github_user_id: Option<i32>, github_user_id: Option<i32>,
) -> Result<Option<User>> { ) -> Result<Option<User>> {
test_support!(self, { self.transact(|tx| async {
let mut tx = tx;
if let Some(github_user_id) = github_user_id { if let Some(github_user_id) = github_user_id {
let mut user = sqlx::query_as::<_, User>( let mut user = sqlx::query_as::<_, User>(
" "
@ -688,7 +716,7 @@ where
) )
.bind(github_login) .bind(github_login)
.bind(github_user_id) .bind(github_user_id)
.fetch_optional(&self.pool) .fetch_optional(&mut tx)
.await?; .await?;
if user.is_none() { if user.is_none() {
@ -702,7 +730,7 @@ where
) )
.bind(github_user_id) .bind(github_user_id)
.bind(github_login) .bind(github_login)
.fetch_optional(&self.pool) .fetch_optional(&mut tx)
.await?; .await?;
} }
@ -716,58 +744,62 @@ where
", ",
) )
.bind(github_login) .bind(github_login)
.fetch_optional(&self.pool) .fetch_optional(&mut tx)
.await?; .await?;
Ok(user) Ok(user)
} }
}) })
.await
} }
pub async fn set_user_is_admin(&self, id: UserId, is_admin: bool) -> Result<()> { pub async fn set_user_is_admin(&self, id: UserId, is_admin: bool) -> Result<()> {
test_support!(self, { self.transact(|mut tx| async {
let query = "UPDATE users SET admin = $1 WHERE id = $2"; let query = "UPDATE users SET admin = $1 WHERE id = $2";
Ok(sqlx::query(query) sqlx::query(query)
.bind(is_admin) .bind(is_admin)
.bind(id.0) .bind(id.0)
.execute(&self.pool) .execute(&mut tx)
.await .await?;
.map(drop)?) tx.commit().await?;
Ok(())
}) })
.await
} }
pub async fn set_user_connected_once(&self, id: UserId, connected_once: bool) -> Result<()> { pub async fn set_user_connected_once(&self, id: UserId, connected_once: bool) -> Result<()> {
test_support!(self, { self.transact(|mut tx| async move {
let query = "UPDATE users SET connected_once = $1 WHERE id = $2"; let query = "UPDATE users SET connected_once = $1 WHERE id = $2";
Ok(sqlx::query(query) sqlx::query(query)
.bind(connected_once) .bind(connected_once)
.bind(id.0) .bind(id.0)
.execute(&self.pool) .execute(&mut tx)
.await .await?;
.map(drop)?) tx.commit().await?;
Ok(())
}) })
.await
} }
pub async fn destroy_user(&self, id: UserId) -> Result<()> { pub async fn destroy_user(&self, id: UserId) -> Result<()> {
test_support!(self, { self.transact(|mut tx| async move {
let query = "DELETE FROM access_tokens WHERE user_id = $1;"; let query = "DELETE FROM access_tokens WHERE user_id = $1;";
sqlx::query(query) sqlx::query(query)
.bind(id.0) .bind(id.0)
.execute(&self.pool) .execute(&mut tx)
.await .await
.map(drop)?; .map(drop)?;
let query = "DELETE FROM users WHERE id = $1;"; let query = "DELETE FROM users WHERE id = $1;";
Ok(sqlx::query(query) sqlx::query(query).bind(id.0).execute(&mut tx).await?;
.bind(id.0) tx.commit().await?;
.execute(&self.pool) Ok(())
.await
.map(drop)?)
}) })
.await
} }
// signups // signups
pub async fn get_waitlist_summary(&self) -> Result<WaitlistSummary> { pub async fn get_waitlist_summary(&self) -> Result<WaitlistSummary> {
test_support!(self, { self.transact(|mut tx| async move {
Ok(sqlx::query_as( Ok(sqlx::query_as(
" "
SELECT SELECT
@ -784,13 +816,14 @@ where
) AS unsent ) AS unsent
", ",
) )
.fetch_one(&self.pool) .fetch_one(&mut tx)
.await?) .await?)
}) })
.await
} }
pub async fn get_unsent_invites(&self, count: usize) -> Result<Vec<Invite>> { pub async fn get_unsent_invites(&self, count: usize) -> Result<Vec<Invite>> {
test_support!(self, { self.transact(|mut tx| async move {
Ok(sqlx::query_as( Ok(sqlx::query_as(
" "
SELECT SELECT
@ -803,16 +836,16 @@ where
", ",
) )
.bind(count as i32) .bind(count as i32)
.fetch_all(&self.pool) .fetch_all(&mut tx)
.await?) .await?)
}) })
.await
} }
// invite codes // invite codes
pub async fn set_invite_count_for_user(&self, id: UserId, count: u32) -> Result<()> { pub async fn set_invite_count_for_user(&self, id: UserId, count: u32) -> Result<()> {
test_support!(self, { self.transact(|mut tx| async move {
let mut tx = self.pool.begin().await?;
if count > 0 { if count > 0 {
sqlx::query( sqlx::query(
" "
@ -841,10 +874,11 @@ where
tx.commit().await?; tx.commit().await?;
Ok(()) Ok(())
}) })
.await
} }
pub async fn get_invite_code_for_user(&self, id: UserId) -> Result<Option<(String, u32)>> { pub async fn get_invite_code_for_user(&self, id: UserId) -> Result<Option<(String, u32)>> {
test_support!(self, { self.transact(|mut tx| async move {
let result: Option<(String, i32)> = sqlx::query_as( let result: Option<(String, i32)> = sqlx::query_as(
" "
SELECT invite_code, invite_count SELECT invite_code, invite_count
@ -853,7 +887,7 @@ where
", ",
) )
.bind(id) .bind(id)
.fetch_optional(&self.pool) .fetch_optional(&mut tx)
.await?; .await?;
if let Some((code, count)) = result { if let Some((code, count)) = result {
Ok(Some((code, count.try_into().map_err(anyhow::Error::new)?))) Ok(Some((code, count.try_into().map_err(anyhow::Error::new)?)))
@ -861,10 +895,12 @@ where
Ok(None) Ok(None)
} }
}) })
.await
} }
pub async fn get_user_for_invite_code(&self, code: &str) -> Result<User> { pub async fn get_user_for_invite_code(&self, code: &str) -> Result<User> {
test_support!(self, { self.transact(|tx| async {
let mut tx = tx;
sqlx::query_as( sqlx::query_as(
" "
SELECT * SELECT *
@ -873,7 +909,7 @@ where
", ",
) )
.bind(code) .bind(code)
.fetch_optional(&self.pool) .fetch_optional(&mut tx)
.await? .await?
.ok_or_else(|| { .ok_or_else(|| {
Error::Http( Error::Http(
@ -882,6 +918,7 @@ where
) )
}) })
}) })
.await
} }
pub async fn create_room( pub async fn create_room(
@ -889,8 +926,7 @@ where
user_id: UserId, user_id: UserId,
connection_id: ConnectionId, connection_id: ConnectionId,
) -> Result<proto::Room> { ) -> Result<proto::Room> {
test_support!(self, { self.transact(|mut tx| async move {
let mut tx = self.pool.begin().await?;
let live_kit_room = nanoid::nanoid!(30); let live_kit_room = nanoid::nanoid!(30);
let room_id = sqlx::query_scalar( let room_id = sqlx::query_scalar(
" "
@ -920,7 +956,7 @@ where
.await?; .await?;
self.commit_room_transaction(room_id, tx).await self.commit_room_transaction(room_id, tx).await
}) }).await
} }
pub async fn call( pub async fn call(
@ -931,8 +967,7 @@ where
called_user_id: UserId, called_user_id: UserId,
initial_project_id: Option<ProjectId>, initial_project_id: Option<ProjectId>,
) -> Result<(proto::Room, proto::IncomingCall)> { ) -> Result<(proto::Room, proto::IncomingCall)> {
test_support!(self, { self.transact(|mut tx| async move {
let mut tx = self.pool.begin().await?;
sqlx::query( sqlx::query(
" "
INSERT INTO room_participants (room_id, user_id, calling_user_id, calling_connection_id, initial_project_id) INSERT INTO room_participants (room_id, user_id, calling_user_id, calling_connection_id, initial_project_id)
@ -951,15 +986,14 @@ where
let incoming_call = Self::build_incoming_call(&room, called_user_id) let incoming_call = Self::build_incoming_call(&room, called_user_id)
.ok_or_else(|| anyhow!("failed to build incoming call"))?; .ok_or_else(|| anyhow!("failed to build incoming call"))?;
Ok((room, incoming_call)) Ok((room, incoming_call))
}) }).await
} }
pub async fn incoming_call_for_user( pub async fn incoming_call_for_user(
&self, &self,
user_id: UserId, user_id: UserId,
) -> Result<Option<proto::IncomingCall>> { ) -> Result<Option<proto::IncomingCall>> {
test_support!(self, { self.transact(|mut tx| async move {
let mut tx = self.pool.begin().await?;
let room_id = sqlx::query_scalar::<_, RoomId>( let room_id = sqlx::query_scalar::<_, RoomId>(
" "
SELECT room_id SELECT room_id
@ -978,6 +1012,7 @@ where
Ok(None) Ok(None)
} }
}) })
.await
} }
fn build_incoming_call( fn build_incoming_call(
@ -1013,8 +1048,7 @@ where
room_id: RoomId, room_id: RoomId,
called_user_id: UserId, called_user_id: UserId,
) -> Result<proto::Room> { ) -> Result<proto::Room> {
test_support!(self, { self.transact(|mut tx| async move {
let mut tx = self.pool.begin().await?;
sqlx::query( sqlx::query(
" "
DELETE FROM room_participants DELETE FROM room_participants
@ -1028,6 +1062,7 @@ where
self.commit_room_transaction(room_id, tx).await self.commit_room_transaction(room_id, tx).await
}) })
.await
} }
pub async fn decline_call( pub async fn decline_call(
@ -1035,8 +1070,7 @@ where
expected_room_id: Option<RoomId>, expected_room_id: Option<RoomId>,
user_id: UserId, user_id: UserId,
) -> Result<proto::Room> { ) -> Result<proto::Room> {
test_support!(self, { self.transact(|mut tx| async move {
let mut tx = self.pool.begin().await?;
let room_id = sqlx::query_scalar( let room_id = sqlx::query_scalar(
" "
DELETE FROM room_participants DELETE FROM room_participants
@ -1053,6 +1087,7 @@ where
self.commit_room_transaction(room_id, tx).await self.commit_room_transaction(room_id, tx).await
}) })
.await
} }
pub async fn cancel_call( pub async fn cancel_call(
@ -1061,8 +1096,7 @@ where
calling_connection_id: ConnectionId, calling_connection_id: ConnectionId,
called_user_id: UserId, called_user_id: UserId,
) -> Result<proto::Room> { ) -> Result<proto::Room> {
test_support!(self, { self.transact(|mut tx| async move {
let mut tx = self.pool.begin().await?;
let room_id = sqlx::query_scalar( let room_id = sqlx::query_scalar(
" "
DELETE FROM room_participants DELETE FROM room_participants
@ -1079,7 +1113,7 @@ where
} }
self.commit_room_transaction(room_id, tx).await self.commit_room_transaction(room_id, tx).await
}) }).await
} }
pub async fn join_room( pub async fn join_room(
@ -1088,8 +1122,7 @@ where
user_id: UserId, user_id: UserId,
connection_id: ConnectionId, connection_id: ConnectionId,
) -> Result<proto::Room> { ) -> Result<proto::Room> {
test_support!(self, { self.transact(|mut tx| async move {
let mut tx = self.pool.begin().await?;
sqlx::query( sqlx::query(
" "
UPDATE room_participants UPDATE room_participants
@ -1105,15 +1138,14 @@ where
.await?; .await?;
self.commit_room_transaction(room_id, tx).await self.commit_room_transaction(room_id, tx).await
}) })
.await
} }
pub async fn leave_room_for_connection( pub async fn leave_room_for_connection(
&self, &self,
connection_id: ConnectionId, connection_id: ConnectionId,
) -> Result<Option<LeftRoom>> { ) -> Result<Option<LeftRoom>> {
test_support!(self, { self.transact(|mut tx| async move {
let mut tx = self.pool.begin().await?;
// Leave room. // Leave room.
let room_id = sqlx::query_scalar::<_, RoomId>( let room_id = sqlx::query_scalar::<_, RoomId>(
" "
@ -1198,6 +1230,7 @@ where
Ok(None) Ok(None)
} }
}) })
.await
} }
pub async fn update_room_participant_location( pub async fn update_room_participant_location(
@ -1206,13 +1239,13 @@ where
connection_id: ConnectionId, connection_id: ConnectionId,
location: proto::ParticipantLocation, location: proto::ParticipantLocation,
) -> Result<proto::Room> { ) -> Result<proto::Room> {
test_support!(self, { self.transact(|tx| async {
let mut tx = self.pool.begin().await?; let mut tx = tx;
let location_kind; let location_kind;
let location_project_id; let location_project_id;
match location match location
.variant .variant
.as_ref()
.ok_or_else(|| anyhow!("invalid location"))? .ok_or_else(|| anyhow!("invalid location"))?
{ {
proto::participant_location::Variant::SharedProject(project) => { proto::participant_location::Variant::SharedProject(project) => {
@ -1245,6 +1278,7 @@ where
self.commit_room_transaction(room_id, tx).await self.commit_room_transaction(room_id, tx).await
}) })
.await
} }
async fn commit_room_transaction( async fn commit_room_transaction(
@ -1375,8 +1409,7 @@ where
connection_id: ConnectionId, connection_id: ConnectionId,
worktrees: &[proto::WorktreeMetadata], worktrees: &[proto::WorktreeMetadata],
) -> Result<(ProjectId, proto::Room)> { ) -> Result<(ProjectId, proto::Room)> {
test_support!(self, { self.transact(|mut tx| async move {
let mut tx = self.pool.begin().await?;
let project_id = sqlx::query_scalar( let project_id = sqlx::query_scalar(
" "
INSERT INTO projects (room_id, host_user_id, host_connection_id) INSERT INTO projects (room_id, host_user_id, host_connection_id)
@ -1428,16 +1461,65 @@ where
let room = self.commit_room_transaction(room_id, tx).await?; let room = self.commit_room_transaction(room_id, tx).await?;
Ok((project_id, room)) Ok((project_id, room))
}) })
.await
} }
// pub async fn join_project( pub async fn update_project(
// &self, &self,
// user_id: UserId, project_id: ProjectId,
// connection_id: ConnectionId, connection_id: ConnectionId,
// project_id: ProjectId, worktrees: &[proto::WorktreeMetadata],
// ) -> Result<(Project, ReplicaId)> { ) -> Result<(proto::Room, Vec<ConnectionId>)> {
// todo!() self.transact(|mut tx| async move {
// } let room_id: RoomId = sqlx::query_scalar(
"
SELECT room_id
FROM projects
WHERE id = $1 AND host_connection_id = $2
",
)
.bind(project_id)
.bind(connection_id.0 as i32)
.fetch_one(&mut tx)
.await?;
for worktree in worktrees {
sqlx::query(
"
INSERT INTO worktrees (project_id, id, root_name)
VALUES ($1, $2, $3)
ON CONFLICT (project_id, id) DO UPDATE SET root_name = excluded.root_name
",
)
.bind(project_id)
.bind(worktree.id as i32)
.bind(&worktree.root_name)
.execute(&mut tx)
.await?;
}
let mut params = "?,".repeat(worktrees.len());
if !worktrees.is_empty() {
params.pop();
}
let query = format!(
"
DELETE FROM worktrees
WHERE id NOT IN ({params})
",
);
let mut query = sqlx::query(&query);
for worktree in worktrees {
query = query.bind(worktree.id as i32);
}
query.execute(&mut tx).await?;
let room = self.commit_room_transaction(room_id, tx).await?;
todo!()
})
.await
}
pub async fn unshare_project(&self, project_id: ProjectId) -> Result<()> { pub async fn unshare_project(&self, project_id: ProjectId) -> Result<()> {
todo!() todo!()
@ -1459,7 +1541,7 @@ where
// contacts // contacts
pub async fn get_contacts(&self, user_id: UserId) -> Result<Vec<Contact>> { pub async fn get_contacts(&self, user_id: UserId) -> Result<Vec<Contact>> {
test_support!(self, { self.transact(|mut tx| async move {
let query = " let query = "
SELECT user_id_a, user_id_b, a_to_b, accepted, should_notify SELECT user_id_a, user_id_b, a_to_b, accepted, should_notify
FROM contacts FROM contacts
@ -1468,7 +1550,7 @@ where
let mut rows = sqlx::query_as::<_, (UserId, UserId, bool, bool, bool)>(query) let mut rows = sqlx::query_as::<_, (UserId, UserId, bool, bool, bool)>(query)
.bind(user_id) .bind(user_id)
.fetch(&self.pool); .fetch(&mut tx);
let mut contacts = Vec::new(); let mut contacts = Vec::new();
while let Some(row) = rows.next().await { while let Some(row) = rows.next().await {
@ -1507,10 +1589,11 @@ where
Ok(contacts) Ok(contacts)
}) })
.await
} }
pub async fn has_contact(&self, user_id_1: UserId, user_id_2: UserId) -> Result<bool> { pub async fn has_contact(&self, user_id_1: UserId, user_id_2: UserId) -> Result<bool> {
test_support!(self, { self.transact(|mut tx| async move {
let (id_a, id_b) = if user_id_1 < user_id_2 { let (id_a, id_b) = if user_id_1 < user_id_2 {
(user_id_1, user_id_2) (user_id_1, user_id_2)
} else { } else {
@ -1525,14 +1608,15 @@ where
Ok(sqlx::query_scalar::<_, i32>(query) Ok(sqlx::query_scalar::<_, i32>(query)
.bind(id_a.0) .bind(id_a.0)
.bind(id_b.0) .bind(id_b.0)
.fetch_optional(&self.pool) .fetch_optional(&mut tx)
.await? .await?
.is_some()) .is_some())
}) })
.await
} }
pub async fn send_contact_request(&self, sender_id: UserId, receiver_id: UserId) -> Result<()> { pub async fn send_contact_request(&self, sender_id: UserId, receiver_id: UserId) -> Result<()> {
test_support!(self, { self.transact(|mut tx| async move {
let (id_a, id_b, a_to_b) = if sender_id < receiver_id { let (id_a, id_b, a_to_b) = if sender_id < receiver_id {
(sender_id, receiver_id, true) (sender_id, receiver_id, true)
} else { } else {
@ -1554,7 +1638,7 @@ where
.bind(id_a.0) .bind(id_a.0)
.bind(id_b.0) .bind(id_b.0)
.bind(a_to_b) .bind(a_to_b)
.execute(&self.pool) .execute(&mut tx)
.await?; .await?;
if result.rows_affected() == 1 { if result.rows_affected() == 1 {
@ -1562,11 +1646,11 @@ where
} else { } else {
Err(anyhow!("contact already requested"))? Err(anyhow!("contact already requested"))?
} }
}) }).await
} }
pub async fn remove_contact(&self, requester_id: UserId, responder_id: UserId) -> Result<()> { pub async fn remove_contact(&self, requester_id: UserId, responder_id: UserId) -> Result<()> {
test_support!(self, { self.transact(|mut tx| async move {
let (id_a, id_b) = if responder_id < requester_id { let (id_a, id_b) = if responder_id < requester_id {
(responder_id, requester_id) (responder_id, requester_id)
} else { } else {
@ -1579,7 +1663,7 @@ where
let result = sqlx::query(query) let result = sqlx::query(query)
.bind(id_a.0) .bind(id_a.0)
.bind(id_b.0) .bind(id_b.0)
.execute(&self.pool) .execute(&mut tx)
.await?; .await?;
if result.rows_affected() == 1 { if result.rows_affected() == 1 {
@ -1588,6 +1672,7 @@ where
Err(anyhow!("no such contact"))? Err(anyhow!("no such contact"))?
} }
}) })
.await
} }
pub async fn dismiss_contact_notification( pub async fn dismiss_contact_notification(
@ -1595,7 +1680,7 @@ where
user_id: UserId, user_id: UserId,
contact_user_id: UserId, contact_user_id: UserId,
) -> Result<()> { ) -> Result<()> {
test_support!(self, { self.transact(|mut tx| async move {
let (id_a, id_b, a_to_b) = if user_id < contact_user_id { let (id_a, id_b, a_to_b) = if user_id < contact_user_id {
(user_id, contact_user_id, true) (user_id, contact_user_id, true)
} else { } else {
@ -1617,7 +1702,7 @@ where
.bind(id_a.0) .bind(id_a.0)
.bind(id_b.0) .bind(id_b.0)
.bind(a_to_b) .bind(a_to_b)
.execute(&self.pool) .execute(&mut tx)
.await?; .await?;
if result.rows_affected() == 0 { if result.rows_affected() == 0 {
@ -1626,6 +1711,7 @@ where
Ok(()) Ok(())
}) })
.await
} }
pub async fn respond_to_contact_request( pub async fn respond_to_contact_request(
@ -1634,7 +1720,7 @@ where
requester_id: UserId, requester_id: UserId,
accept: bool, accept: bool,
) -> Result<()> { ) -> Result<()> {
test_support!(self, { self.transact(|mut tx| async move {
let (id_a, id_b, a_to_b) = if responder_id < requester_id { let (id_a, id_b, a_to_b) = if responder_id < requester_id {
(responder_id, requester_id, false) (responder_id, requester_id, false)
} else { } else {
@ -1650,7 +1736,7 @@ where
.bind(id_a.0) .bind(id_a.0)
.bind(id_b.0) .bind(id_b.0)
.bind(a_to_b) .bind(a_to_b)
.execute(&self.pool) .execute(&mut tx)
.await? .await?
} else { } else {
let query = " let query = "
@ -1661,7 +1747,7 @@ where
.bind(id_a.0) .bind(id_a.0)
.bind(id_b.0) .bind(id_b.0)
.bind(a_to_b) .bind(a_to_b)
.execute(&self.pool) .execute(&mut tx)
.await? .await?
}; };
if result.rows_affected() == 1 { if result.rows_affected() == 1 {
@ -1670,6 +1756,7 @@ where
Err(anyhow!("no such contact request"))? Err(anyhow!("no such contact request"))?
} }
}) })
.await
} }
// access tokens // access tokens
@ -1680,7 +1767,8 @@ where
access_token_hash: &str, access_token_hash: &str,
max_access_token_count: usize, max_access_token_count: usize,
) -> Result<()> { ) -> Result<()> {
test_support!(self, { self.transact(|tx| async {
let mut tx = tx;
let insert_query = " let insert_query = "
INSERT INTO access_tokens (user_id, hash) INSERT INTO access_tokens (user_id, hash)
VALUES ($1, $2); VALUES ($1, $2);
@ -1696,7 +1784,6 @@ where
) )
"; ";
let mut tx = self.pool.begin().await?;
sqlx::query(insert_query) sqlx::query(insert_query)
.bind(user_id.0) .bind(user_id.0)
.bind(access_token_hash) .bind(access_token_hash)
@ -1710,10 +1797,11 @@ where
.await?; .await?;
Ok(tx.commit().await?) Ok(tx.commit().await?)
}) })
.await
} }
pub async fn get_access_token_hashes(&self, user_id: UserId) -> Result<Vec<String>> { pub async fn get_access_token_hashes(&self, user_id: UserId) -> Result<Vec<String>> {
test_support!(self, { self.transact(|mut tx| async move {
let query = " let query = "
SELECT hash SELECT hash
FROM access_tokens FROM access_tokens
@ -1722,9 +1810,51 @@ where
"; ";
Ok(sqlx::query_scalar(query) Ok(sqlx::query_scalar(query)
.bind(user_id.0) .bind(user_id.0)
.fetch_all(&self.pool) .fetch_all(&mut tx)
.await?) .await?)
}) })
.await
}
async fn transact<F, Fut, T>(&self, f: F) -> Result<T>
where
F: Send + Fn(sqlx::Transaction<'static, D>) -> Fut,
Fut: Send + Future<Output = Result<T>>,
{
let body = async {
loop {
let tx = self.begin_transaction().await?;
match f(tx).await {
Ok(result) => return Ok(result),
Err(error) => match error {
Error::Database(error)
if error
.as_database_error()
.and_then(|error| error.code())
.as_deref()
== Some("hey") =>
{
// Retry (don't break the loop)
}
error @ _ => return Err(error),
},
}
}
};
#[cfg(test)]
{
if let Some(background) = self.background.as_ref() {
background.simulate_random_delay().await;
}
self.runtime.as_ref().unwrap().block_on(body)
}
#[cfg(not(test))]
{
body.await
}
} }
} }

8
crates/collab/src/lib.rs
View file

@ -4,6 +4,7 @@ pub type Result<T, E = Error> = std::result::Result<T, E>;
pub enum Error { pub enum Error {
Http(StatusCode, String), Http(StatusCode, String),
Database(sqlx::Error),
Internal(anyhow::Error), Internal(anyhow::Error),
} }
@ -15,7 +16,7 @@ impl From<anyhow::Error> for Error {
impl From<sqlx::Error> for Error { impl From<sqlx::Error> for Error {
fn from(error: sqlx::Error) -> Self { fn from(error: sqlx::Error) -> Self {
Self::Internal(error.into()) Self::Database(error)
} }
} }
@ -41,6 +42,9 @@ impl IntoResponse for Error {
fn into_response(self) -> axum::response::Response { fn into_response(self) -> axum::response::Response {
match self { match self {
Error::Http(code, message) => (code, message).into_response(), Error::Http(code, message) => (code, message).into_response(),
Error::Database(error) => {
(StatusCode::INTERNAL_SERVER_ERROR, format!("{}", &error)).into_response()
}
Error::Internal(error) => { Error::Internal(error) => {
(StatusCode::INTERNAL_SERVER_ERROR, format!("{}", &error)).into_response() (StatusCode::INTERNAL_SERVER_ERROR, format!("{}", &error)).into_response()
} }
@ -52,6 +56,7 @@ impl std::fmt::Debug for Error {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self { match self {
Error::Http(code, message) => (code, message).fmt(f), Error::Http(code, message) => (code, message).fmt(f),
Error::Database(error) => error.fmt(f),
Error::Internal(error) => error.fmt(f), Error::Internal(error) => error.fmt(f),
} }
} }
@ -61,6 +66,7 @@ impl std::fmt::Display for Error {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self { match self {
Error::Http(code, message) => write!(f, "{code}: {message}"), Error::Http(code, message) => write!(f, "{code}: {message}"),
Error::Database(error) => error.fmt(f),
Error::Internal(error) => error.fmt(f), Error::Internal(error) => error.fmt(f),
} }
} }

27
crates/collab/src/rpc.rs
View file

@ -1087,30 +1087,31 @@ impl Server {
response: Response<proto::UpdateProject>, response: Response<proto::UpdateProject>,
) -> Result<()> { ) -> Result<()> {
let project_id = ProjectId::from_proto(request.payload.project_id); let project_id = ProjectId::from_proto(request.payload.project_id);
{ let (room, guest_connection_ids) = self
let mut state = self.store().await; .app_state
let guest_connection_ids = state .db
.read_project(project_id, request.sender_connection_id)? .update_project(
.guest_connection_ids();
let room = state.update_project(
project_id, project_id,
&request.payload.worktrees,
request.sender_connection_id, request.sender_connection_id,
)?; &request.payload.worktrees,
)
.await?;
broadcast( broadcast(
request.sender_connection_id, request.sender_connection_id,
guest_connection_ids, guest_connection_ids,
|connection_id| { |connection_id| {
self.peer.forward_send( self.peer.send(
request.sender_connection_id,
connection_id, connection_id,
request.payload.clone(), proto::ProjectUpdated {
project_id: project_id.to_proto(),
worktrees: request.payload.worktrees.clone(),
room_version: room.version,
},
) )
}, },
); );
self.room_updated(room); self.room_updated(&room);
response.send(proto::Ack {})?; response.send(proto::Ack {})?;
};
Ok(()) Ok(())
} }