Merge pull request #5785 from unisonweb/cp/check-and-set-branch-updates

Add check-and-set when running branch updates

Author: aryairani

unison-cli/src/Unison/Cli/Monad.hs

@@ -189,7 +189,7 @@ data Env = Env
 --
 -- There's an additional pseudo @"currentPath"@ field lens, for convenience.
 data LoopState = LoopState
-  { -- the current position in the codebase, with the head being the most recent lcoation.
+  { -- the current position in the codebase, with the head being the most recent location.
     projectPathStack :: List.NonEmpty PP.ProjectPathIds,
     -- TBD
     -- , _activeEdits :: Set Branch.EditGuid

unison-cli/src/Unison/Cli/MonadUtils.hs

@@ -423,16 +423,26 @@ updateProjectBranchRoot :: ProjectBranch -> Text -> (Branch IO -> Cli (Branch IO
 updateProjectBranchRoot projectBranch reason f = do
   env <- ask
   Cli.time "updateProjectBranchRoot" do
-    old <- getProjectBranchRoot projectBranch
-    (new, result) <- f old
-    when (old /= new) do
+    beforeUpdates <- getProjectBranchRoot projectBranch
+    (new, result) <- f beforeUpdates
+    when (beforeUpdates /= new) do
       liftIO $ Codebase.putBranch env.codebase new
-      Cli.runTransaction do
-        -- TODO: If we transactionally check that the project branch hasn't changed while we were computing the new
-        -- branch, and if it has, abort the transaction and return an error, then we can
-        -- remove the single UCM per codebase restriction.
-        causalHashId <- Q.expectCausalHashIdByCausalHash (Branch.headHash new)
-        Q.setProjectBranchHead reason projectBranch.projectId projectBranch.branchId causalHashId
+      Cli.runTransactionWithRollback \rollback -> do
+        causalHashId <- Q.expectProjectBranchHead projectBranch.projectId projectBranch.branchId
+        currentHeadHash <- Q.expectCausalHash causalHashId
+        -- Inside the transaction we ensure that the branch from before the updates matches the current head of the
+        -- project branch, like a check-and-set operation.
+        -- If it doesn't, then some other process has updated the branch between when we read it and computed the
+        -- updates. We should abort and ask the user to try again.
+        if
+          | (currentHeadHash == Branch.headHash new) -> do
+              -- Someone else updated the branch, but they set it to what we wanted to anyways.
+              pure ()
+          | (currentHeadHash /= Branch.headHash beforeUpdates) -> do
+              rollback Output.BranchUpdate'BranchChanged
+          | otherwise -> do
+              causalHashId <- Q.expectCausalHashIdByCausalHash (Branch.headHash new)
+              Q.setProjectBranchHead reason projectBranch.projectId projectBranch.branchId causalHashId
       -- The input to this function isn't necessarily the *current* project branch, which is what LSP cares about. But
       -- it might be! There's no harm in unconditionally notifying the LSP that the current project branch may have
       -- changed, but it is slightly more efficient for us to just do the == comparison here (since otherwise the LSP

unison-cli/src/Unison/Codebase/Editor/Output.hs

@@ -452,6 +452,7 @@ data Output
   | OpenCodebaseError CodebasePath OpenCodebaseError
   | UCMServerNotRunning
   | BranchSquashSuccess ({- source -} ProjectAndBranch Project ProjectBranch) ({- dest branch -} ProjectAndBranch Project ProjectBranch)
+  | BranchUpdate'BranchChanged
 
 data MoreEntriesThanShown = MoreEntriesThanShown | AllEntriesShown
   deriving (Eq, Show)
@@ -694,6 +695,7 @@ isFailure o = case o of
   OpenCodebaseError {} -> True
   UCMServerNotRunning -> True
   BranchSquashSuccess {} -> False
+  BranchUpdate'BranchChanged {} -> True
 
 isNumberedFailure :: NumberedOutput -> Bool
 isNumberedFailure = \case

unison-cli/src/Unison/CommandLine/OutputMessages.hs

@@ -2346,6 +2346,9 @@ notifyUser dir = \case
       P.lines
         [ P.wrap $ "I squashed " <> sourceName <> " into " <> destName
         ]
+  BranchUpdate'BranchChanged -> do
+    pure $
+      P.wrap "Another process updated the codebase while your command was running, so I didn't apply the update. Please run the command again."
 
 prettyShareError :: ShareError -> Pretty
 prettyShareError =