From 9b5e6eec2b2c8e778e40fc4dda723799f99cd3e5 Mon Sep 17 00:00:00 2001
From: Nikita Galaiko <nikita@galaiko.rocks>
Date: Mon, 26 Jun 2023 12:19:46 +0200
Subject: [PATCH] protect vbranches with a semaphore

---
 src-tauri/src/app.rs                  | 37 ++++++++++++++++++++++++---
 src-tauri/src/main.rs                 | 10 ++++++--
 src-tauri/src/virtual_branches/mod.rs |  1 +
 3 files changed, 42 insertions(+), 6 deletions(-)

diff --git a/src-tauri/src/app.rs b/src-tauri/src/app.rs
index b5667784c..4dc54a225 100644
--- a/src-tauri/src/app.rs
+++ b/src-tauri/src/app.rs
@@ -1,7 +1,7 @@
 use std::{collections::HashMap, ops, sync};
 
 use anyhow::{Context, Result};
-use tokio::sync::mpsc;
+use tokio::sync::{mpsc, Semaphore};
 use tokio_util::sync::CancellationToken;
 
 use crate::{
@@ -29,6 +29,8 @@ pub struct App {
     files_database: files::Database,
     deltas_database: deltas::Database,
     bookmarks_database: bookmarks::Database,
+
+    vbranch_semaphores: sync::Arc<tokio::sync::Mutex<HashMap<String, Semaphore>>>,
 }
 
 #[derive(Debug, thiserror::Error)]
@@ -63,6 +65,7 @@ impl App {
             deltas_database: deltas::Database::new(database.clone()),
             files_database: files::Database::new(database.clone()),
             bookmarks_database: bookmarks::Database::new(database),
+            vbranch_semaphores: sync::Arc::new(tokio::sync::Mutex::new(HashMap::new())),
         })
     }
 
@@ -333,7 +336,7 @@ impl App {
         Ok(())
     }
 
-    pub fn list_virtual_branches(
+    pub async fn list_virtual_branches(
         &self,
         project_id: &str,
     ) -> Result<Vec<virtual_branches::VirtualBranch>> {
@@ -341,11 +344,30 @@ impl App {
         let project = self.gb_project(project_id)?;
         let project_repository = project_repository::Repository::open(&project)
             .context("failed to open project repository")?;
+
+        let mut semaphores = self.vbranch_semaphores.lock().await;
+        let semaphore = semaphores
+            .entry(project_id.to_string())
+            .or_insert_with(|| Semaphore::new(1));
+        let _permit = semaphore.acquire().await?;
+
         virtual_branches::list_virtual_branches(&gb_repository, &project_repository)
     }
 
-    pub fn create_virtual_branch(&self, project_id: &str, name: &str, path: &str) -> Result<()> {
+    pub async fn create_virtual_branch(
+        &self,
+        project_id: &str,
+        name: &str,
+        path: &str,
+    ) -> Result<()> {
         let gb_repository = self.gb_repository(project_id)?;
+
+        let mut semaphores = self.vbranch_semaphores.lock().await;
+        let semaphore = semaphores
+            .entry(project_id.to_string())
+            .or_insert_with(|| Semaphore::new(1));
+        let _permit = semaphore.acquire().await?;
+
         let branch_id = virtual_branches::create_virtual_branch(&gb_repository, name)?;
         virtual_branches::move_files(&gb_repository, &branch_id, &vec![path.try_into()?])?;
         Ok(())
@@ -361,7 +383,7 @@ impl App {
         Ok(())
     }
 
-    pub fn move_virtual_branch_files(
+    pub async fn move_virtual_branch_files(
         &self,
         project_id: &str,
         branch: &str,
@@ -372,6 +394,13 @@ impl App {
             .into_iter()
             .map(|p| p.try_into())
             .collect::<Result<Vec<branch::Ownership>, _>>()?;
+
+        let mut semaphores = self.vbranch_semaphores.lock().await;
+        let semaphore = semaphores
+            .entry(project_id.to_string())
+            .or_insert_with(|| Semaphore::new(1));
+        let _permit = semaphore.acquire().await?;
+
         virtual_branches::move_files(&gb_repository, branch, &paths)?;
         Ok(())
     }
diff --git a/src-tauri/src/main.rs b/src-tauri/src/main.rs
index feff236b2..86f9db988 100644
--- a/src-tauri/src/main.rs
+++ b/src-tauri/src/main.rs
@@ -577,6 +577,7 @@ async fn list_virtual_branches(
     let app = handle.state::<app::App>();
     let branches = app
         .list_virtual_branches(project_id)
+        .await
         .context("failed to list virtual branches")?;
     Ok(branches)
 }
@@ -590,7 +591,9 @@ async fn create_virtual_branch(
     path: &str,
 ) -> Result<(), Error> {
     let app = handle.state::<app::App>();
-    app.create_virtual_branch(project_id, name, path)?;
+    app.create_virtual_branch(project_id, name, path)
+        .await
+        .context("failed to create virtual branch")?;
     Ok(())
 }
 
@@ -615,7 +618,10 @@ async fn move_virtual_branch_files(
     paths: Vec<&str>,
 ) -> Result<(), Error> {
     let app = handle.state::<app::App>();
-    let target = app.move_virtual_branch_files(project_id, branch, paths)?;
+    let target = app
+        .move_virtual_branch_files(project_id, branch, paths)
+        .await
+        .context("failed to move virtual branch files")?;
     Ok(target)
 }
 
diff --git a/src-tauri/src/virtual_branches/mod.rs b/src-tauri/src/virtual_branches/mod.rs
index 541f53727..771411064 100644
--- a/src-tauri/src/virtual_branches/mod.rs
+++ b/src-tauri/src/virtual_branches/mod.rs
@@ -13,6 +13,7 @@ use serde::Serialize;
 
 pub use branch::Branch;
 pub use iterator::BranchIterator as Iterator;
+use tokio::sync::Semaphore;
 use uuid::Uuid;
 
 use crate::{gb_repository, project_repository, reader, sessions};