Attempt to fix memory leaks

unison-cli/src/Unison/Codebase/TranscriptParser.hs

@@ -235,9 +235,14 @@ run dir stanzas codebase runtime sbRuntime config ucmVersion baseURL = UnliftIO.
       ]
   initialRootCausalHash <- Codebase.getRootCausalHash codebase
   rootVar <- newEmptyTMVarIO
+  currentBranchVar <- newEmptyTMVarIO
+  currentPPEDVar <- newEmptyTMVarIO
   void $ Ki.fork scope do
     root <- Codebase.getRootBranch codebase
-    atomically $ putTMVar rootVar root
+    atomically $ do
+      putTMVar rootVar root
+      putTMVar currentBranchVar root
+      putTMVar currentPPEDVar root
   mayShareAccessToken <- fmap Text.pack <$> lookupEnv accessTokenEnvVarKey
   credMan <- AuthN.newCredentialManager
   let tokenProvider :: AuthN.TokenProvider
@@ -485,7 +490,8 @@ run dir stanzas codebase runtime sbRuntime config ucmVersion baseURL = UnliftIO.
             texts <- readIORef out
             pure $ Text.concat (Text.pack <$> toList (texts :: Seq String))
 
-  loop (Cli.loopState0 initialRootCausalHash rootVar initialPath)
+  let initState = Cli.loopState0 initialRootCausalHash rootVar initialPath
+  loop initState
 
 transcriptFailure :: IORef (Seq String) -> Text -> IO b
 transcriptFailure out msg = do

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

@@ -23,6 +23,7 @@ import Unison.Auth.HTTPClient (AuthenticatedHttpClient)
 import qualified Unison.Auth.HTTPClient as AuthN
 import qualified Unison.Auth.Tokens as AuthN
 import qualified Unison.Cli.Monad as Cli
+import qualified Unison.Cli.MonadUtils as Cli
 import Unison.Codebase (Codebase)
 import qualified Unison.Codebase as Codebase
 import Unison.Codebase.Branch (Branch)
@@ -102,9 +103,8 @@ main ::
   Maybe Server.BaseUrl ->
   UCMVersion ->
   (Branch IO -> STM ()) ->
-  (Path.Absolute -> STM ()) ->
   IO ()
-main dir welcome initialPath (config, cancelConfig) initialInputs runtime sbRuntime codebase serverBaseUrl ucmVersion notifyBranchChange notifyPathChange = Ki.scoped \scope -> do
+main dir welcome initialPath (config, cancelConfig) initialInputs runtime sbRuntime codebase serverBaseUrl ucmVersion notifyBranchChange = Ki.scoped \scope -> do
   rootVar <- newEmptyTMVarIO
   initialRootCausalHash <- Codebase.getRootCausalHash codebase
   _ <- Ki.fork scope $ do
@@ -217,7 +217,9 @@ main dir welcome initialPath (config, cancelConfig) initialInputs runtime sbRunt
         loop0 s0 = do
           let step = do
                 input <- awaitInput s0
-                (result, resultState) <- Cli.runCli env s0 (HandleInput.loop input)
+                (result, resultState) <- Cli.runCli env s0 $ do
+                  HandleInput.loop input
+                  Cli.getCurrentBranch
                 let sNext = case input of
                       Left _ -> resultState
                       Right inp -> resultState & #lastInput ?~ inp
@@ -232,10 +234,13 @@ main dir welcome initialPath (config, cancelConfig) initialInputs runtime sbRunt
               Text.Lazy.hPutStrLn stderr ("Encountered exception:\n" <> pShow e)
               loop0 s0
             Right (Right (result, s1)) -> do
-              when ((s0 ^. #currentPath) /= (s1 ^. #currentPath :: Path.Absolute)) (atomically . notifyPathChange $ s1 ^. #currentPath)
               case result of
-                Cli.Success () -> loop0 s1
-                Cli.Continue -> loop0 s1
+                Cli.Success curBranch -> do
+                  atomically $ notifyBranchChange curBranch
+                  loop0 s1
+                Cli.Continue ->
+                  -- TODO: we should notify on branch changes here.
+                  loop0 s1
                 Cli.HaltRepl -> pure ()
 
     withInterruptHandler onInterrupt (loop0 initialState `finally` cleanup)

unison-cli/src/Unison/LSP.hs

@@ -20,8 +20,8 @@ import Language.LSP.VFS
 import qualified Network.Simple.TCP as TCP
 import Network.Socket (socketToHandle)
 import System.Environment (lookupEnv)
+import U.Codebase.HashTags (BranchHash)
 import Unison.Codebase
-import Unison.Codebase.Branch (Branch)
 import qualified Unison.Codebase.Path as Path
 import Unison.Codebase.Runtime (Runtime)
 import qualified Unison.Debug as Debug
@@ -36,6 +36,7 @@ import Unison.LSP.Orphans ()
 import Unison.LSP.Types
 import Unison.LSP.UCMWorker (ucmWorker)
 import qualified Unison.LSP.VFS as VFS
+import Unison.Names (Names)
 import Unison.Parser.Ann
 import Unison.Prelude
 import qualified Unison.PrettyPrintEnvDecl as PPED
@@ -47,8 +48,8 @@ getLspPort :: IO String
 getLspPort = fromMaybe "5757" <$> lookupEnv "UNISON_LSP_PORT"
 
 -- | Spawn an LSP server on the configured port.
-spawnLsp :: Codebase IO Symbol Ann -> Runtime Symbol -> STM (Branch IO) -> STM (Path.Absolute) -> IO ()
-spawnLsp codebase runtime latestBranch latestPath = TCP.withSocketsDo do
+spawnLsp :: Codebase IO Symbol Ann -> Runtime Symbol -> STM (BranchHash, Names) -> IO ()
+spawnLsp codebase runtime latestNames = TCP.withSocketsDo do
   lspPort <- getLspPort
   UnliftIO.handleIO (handleFailure lspPort) $ do
     TCP.serve (TCP.Host "127.0.0.1") lspPort $ \(sock, _sockaddr) -> do
@@ -58,14 +59,14 @@ spawnLsp codebase runtime latestBranch latestPath = TCP.withSocketsDo do
         -- different un-saved state for the same file.
         initVFS $ \vfs -> do
           vfsVar <- newMVar vfs
-          void $ runServerWithHandles lspServerLogger lspClientLogger sockHandle sockHandle (serverDefinition vfsVar codebase runtime scope latestBranch latestPath)
+          void $ runServerWithHandles lspServerLogger lspClientLogger sockHandle sockHandle (serverDefinition vfsVar codebase runtime scope latestNames)
   where
     handleFailure :: String -> IOException -> IO ()
     handleFailure lspPort ioerr =
       case Errno <$> ioe_errno ioerr of
         Just errNo
           | errNo == eADDRINUSE -> do
-            putStrLn $ "Note: Port " <> lspPort <> " is already bound by another process or another UCM. The LSP server will not be started."
+              putStrLn $ "Note: Port " <> lspPort <> " is already bound by another process or another UCM. The LSP server will not be started."
         _ -> do
           Debug.debugM Debug.LSP "LSP Exception" ioerr
           Debug.debugM Debug.LSP "LSP Errno" (ioe_errno ioerr)
@@ -80,14 +81,13 @@ serverDefinition ::
   Codebase IO Symbol Ann ->
   Runtime Symbol ->
   Ki.Scope ->
-  STM (Branch IO) ->
-  STM (Path.Absolute) ->
+  STM (BranchHash, Names) ->
   ServerDefinition Config
-serverDefinition vfsVar codebase runtime scope latestBranch latestPath =
+serverDefinition vfsVar codebase runtime scope latestNames =
   ServerDefinition
     { defaultConfig = lspDefaultConfig,
       onConfigurationChange = lspOnConfigurationChange,
-      doInitialize = lspDoInitialize vfsVar codebase runtime scope latestBranch latestPath,
+      doInitialize = lspDoInitialize vfsVar codebase runtime scope latestNames,
       staticHandlers = lspStaticHandlers,
       interpretHandler = lspInterpretHandler,
       options = lspOptions
@@ -106,12 +106,11 @@ lspDoInitialize ::
   Codebase IO Symbol Ann ->
   Runtime Symbol ->
   Ki.Scope ->
-  STM (Branch IO) ->
-  STM (Path.Absolute) ->
+  STM (BranchHash, Names) ->
   LanguageContextEnv Config ->
   Message 'Initialize ->
   IO (Either ResponseError Env)
-lspDoInitialize vfsVar codebase runtime scope latestBranch latestPath lspContext _initMsg = do
+lspDoInitialize vfsVar codebase runtime scope latestNames lspContext _initMsg = do
   -- TODO: some of these should probably be MVars so that we correctly wait for names and
   -- things to be generated before serving requests.
   checkedFilesVar <- newTVarIO mempty
@@ -123,7 +122,7 @@ lspDoInitialize vfsVar codebase runtime scope latestBranch latestPath lspContext
   let env = Env {ppeCache = readTVarIO ppeCacheVar, parseNamesCache = readTVarIO parseNamesCacheVar, currentPathCache = readTVarIO currentPathCacheVar, ..}
   let lspToIO = flip runReaderT lspContext . unLspT . flip runReaderT env . runLspM
   Ki.fork scope (lspToIO Analysis.fileAnalysisWorker)
-  Ki.fork scope (lspToIO $ ucmWorker ppeCacheVar parseNamesCacheVar latestBranch latestPath)
+  Ki.fork scope (lspToIO $ ucmWorker ppeCacheVar parseNamesCacheVar latestNames)
   pure $ Right $ env
 
 -- | LSP request handlers that don't register/unregister dynamically

unison-cli/src/Unison/LSP/Types.hs

@@ -74,8 +74,8 @@ type FileVersion = Int32
 type LexedSource = (Text, [Lexer.Token Lexer.Lexeme])
 
 data FileAnalysis = FileAnalysis
-  { fileUri :: Uri,
-    fileVersion :: FileVersion,
+  { fileUri :: !Uri,
+    fileVersion :: !FileVersion,
     lexedSource :: LexedSource,
     parsedFile :: Maybe (UF.UnisonFile Symbol Ann),
     typecheckedFile :: Maybe (UF.TypecheckedUnisonFile Symbol Ann),

unison-cli/src/Unison/LSP/UCMWorker.hs

@@ -1,46 +1,39 @@
 module Unison.LSP.UCMWorker where
 
 import Control.Monad.Reader
+import U.Codebase.HashTags
 import qualified Unison.Codebase as Codebase
-import Unison.Codebase.Branch (Branch)
-import qualified Unison.Codebase.Path as Path
-import qualified Unison.Debug as Debug
 import Unison.LSP.Types
 import qualified Unison.LSP.VFS as VFS
+import Unison.Names (Names)
 import Unison.NamesWithHistory (NamesWithHistory)
+import qualified Unison.NamesWithHistory as NamesWithHistory
 import Unison.PrettyPrintEnvDecl
 import qualified Unison.PrettyPrintEnvDecl.Names as PPE
-import qualified Unison.Server.Backend as Backend
 import UnliftIO.STM
 
 -- | Watches for state changes in UCM and updates cached LSP state accordingly
 ucmWorker ::
   TVar PrettyPrintEnvDecl ->
   TVar NamesWithHistory ->
-  STM (Branch IO) ->
-  STM Path.Absolute ->
+  STM (BranchHash, Names) ->
   Lsp ()
-ucmWorker ppeVar parseNamesVar getLatestRoot getLatestPath = do
+ucmWorker ppeVar parseNamesVar getCurrentBranchNames = do
   Env {codebase} <- ask
-  let loop :: (Branch IO, Path.Absolute) -> Lsp a
-      loop (currentRoot, currentPath) = do
-        Debug.debugM Debug.LSP "LSP path: " currentPath
-        let parseNames = Backend.getCurrentParseNames (Backend.Within (Path.unabsolute currentPath)) currentRoot
-        hl <- liftIO $ Codebase.hashLength codebase
-        let ppe = PPE.fromNamesDecl hl parseNames
-        atomically $ do
-          writeTVar parseNamesVar parseNames
-          writeTVar ppeVar ppe
+  let loop :: Maybe (BranchHash, NamesWithHistory) -> Lsp a
+      loop previous = do
         -- Re-check everything with the new names and ppe
         VFS.markAllFilesDirty
-        latest <- atomically $ do
-          latestRoot <- getLatestRoot
-          latestPath <- getLatestPath
-          guard $ (currentRoot /= latestRoot || currentPath /= latestPath)
-          pure (latestRoot, latestPath)
-        loop latest
-
-  -- Bootstrap manually from codebase just in case we're in headless mode and don't get any
-  -- updates from UCM
-  rootBranch <- liftIO $ Codebase.getRootBranch codebase
-  loop (rootBranch, Path.absoluteEmpty)
+        latest@(_latestHash, latestNames) <- atomically $ do
+          (hash, names) <- getCurrentBranchNames
+          case previous of
+            Nothing -> pure ()
+            Just (currentHash, _) -> guard $ (currentHash /= hash)
+          pure (hash, NamesWithHistory.fromCurrentNames names)
+        hl <- liftIO $ Codebase.hashLength codebase
+        let ppe = PPE.fromNamesDecl hl latestNames
+        atomically $ do
+          writeTVar parseNamesVar $! latestNames
+          writeTVar ppeVar $! ppe
+        loop (Just latest)
+  loop Nothing

unison-cli/src/Unison/LSP/VFS.hs

@@ -15,7 +15,6 @@ import qualified Data.Set as Set
 import Data.Set.Lens (setOf)
 import qualified Data.Text as Text
 import qualified Data.Text.Utf16.Rope as Rope
-import Data.Tuple (swap)
 import qualified Language.LSP.Logging as LSP
 import Language.LSP.Types
 import Language.LSP.Types.Lens (HasCharacter (character), HasParams (params), HasPosition (position), HasTextDocument (textDocument), HasUri (uri))
@@ -31,7 +30,9 @@ import UnliftIO
 usingVFS :: forall a. StateT VFS Lsp a -> Lsp a
 usingVFS m = do
   vfsVar' <- asks vfsVar
-  modifyMVar vfsVar' $ \vfs -> swap <$> runStateT m vfs
+  modifyMVar vfsVar' $ \vfs -> do
+    (!a, !b) <- runStateT m vfs
+    pure $! (b, a)
 
 getVirtualFile :: (HasUri doc Uri) => doc -> Lsp (Maybe VirtualFile)
 getVirtualFile p = do

unison-cli/unison/Main.hs

@@ -53,6 +53,9 @@ import Text.Pretty.Simple (pHPrint)
 import Unison.Codebase (Codebase, CodebasePath)
 import qualified Unison.Codebase as Codebase
 import Unison.Codebase.Branch (Branch)
+import qualified Unison.Codebase.Branch.Names as Branch
+import qualified Unison.Codebase.Branch.Type as Branch
+import qualified Unison.Codebase.Causal.Type as Causal
 import qualified Unison.Codebase.Editor.Input as Input
 import Unison.Codebase.Editor.RemoteRepo (ReadShareRemoteNamespace)
 import qualified Unison.Codebase.Editor.VersionParser as VP
@@ -63,6 +66,7 @@ import Unison.Codebase.Init.OpenCodebaseError (OpenCodebaseError (..))
 import qualified Unison.Codebase.Path as Path
 import qualified Unison.Codebase.Runtime as Rt
 import qualified Unison.Codebase.SqliteCodebase as SC
+import qualified Unison.Codebase.SqliteCodebase.Conversions as Cv
 import qualified Unison.Codebase.TranscriptParser as TR
 import Unison.CommandLine (plural', watchConfig)
 import qualified Unison.CommandLine.Main as CommandLine
@@ -132,11 +136,10 @@ main = withCP65001 . Ki.scoped $ \scope -> do
                   rt <- RTI.startRuntime False RTI.OneOff Version.gitDescribeWithDate
                   sbrt <- RTI.startRuntime True RTI.OneOff Version.gitDescribeWithDate
                   let fileEvent = Input.UnisonFileChanged (Text.pack file) contents
-                  let noOpRootNotifier _ = pure ()
-                  let noOpPathNotifier _ = pure ()
+                  let noOpBranchNotifier _ = pure ()
                   let serverUrl = Nothing
                   let startPath = Nothing
-                  launch currentDir config rt sbrt theCodebase [Left fileEvent, Right $ Input.ExecuteI mainName args, Right Input.QuitI] serverUrl startPath ShouldNotDownloadBase initRes noOpRootNotifier noOpPathNotifier
+                  launch currentDir config rt sbrt theCodebase [Left fileEvent, Right $ Input.ExecuteI mainName args, Right Input.QuitI] serverUrl startPath ShouldNotDownloadBase initRes noOpBranchNotifier
       Run (RunFromPipe mainName) args -> do
         e <- safeReadUtf8StdIn
         case e of
@@ -146,8 +149,7 @@ main = withCP65001 . Ki.scoped $ \scope -> do
               rt <- RTI.startRuntime False RTI.OneOff Version.gitDescribeWithDate
               sbrt <- RTI.startRuntime True RTI.OneOff Version.gitDescribeWithDate
               let fileEvent = Input.UnisonFileChanged (Text.pack "<standard input>") contents
-              let noOpRootNotifier _ = pure ()
-              let noOpPathNotifier _ = pure ()
+              let noOpBranchNotifier _ = pure ()
               let serverUrl = Nothing
               let startPath = Nothing
               launch
@@ -161,8 +163,7 @@ main = withCP65001 . Ki.scoped $ \scope -> do
                 startPath
                 ShouldNotDownloadBase
                 initRes
-                noOpRootNotifier
-                noOpPathNotifier
+                noOpBranchNotifier
       Run (RunCompiled file) args ->
         BL.readFile file >>= \bs ->
           try (evaluate $ RTI.decodeStandalone bs) >>= \case
@@ -232,23 +233,22 @@ main = withCP65001 . Ki.scoped $ \scope -> do
       Launch isHeadless codebaseServerOpts downloadBase mayStartingPath -> do
         getCodebaseOrExit mCodePathOption SC.MigrateAfterPrompt \(initRes, _, theCodebase) -> do
           runtime <- RTI.startRuntime False RTI.Persistent Version.gitDescribeWithDate
-          rootVar <- newEmptyTMVarIO
-          pathVar <- newTVarIO initialPath
-          let notifyOnRootChanges :: Branch IO -> STM ()
-              notifyOnRootChanges b = do
-                isEmpty <- isEmptyTMVar rootVar
+          namesVar <- newEmptyTMVarIO
+          let notifyOnBranchChanges :: Branch IO -> STM ()
+              notifyOnBranchChanges b@(Branch.Branch c) = do
+                let bh = Cv.branchHash1to2 $ Causal.valueHash c
+                let names = Branch.toNames (Branch.head b)
+                isEmpty <- isEmptyTMVar namesVar
                 if isEmpty
-                  then putTMVar rootVar b
-                  else void $ swapTMVar rootVar b
-          let notifyOnPathChanges :: Path.Absolute -> STM ()
-              notifyOnPathChanges = writeTVar pathVar
+                  then putTMVar namesVar $! (bh, names)
+                  else void . swapTMVar namesVar $! (bh, names)
           sbRuntime <- RTI.startRuntime True RTI.Persistent Version.gitDescribeWithDate
           -- Unfortunately, the windows IO manager on GHC 8.* is prone to just hanging forever
           -- when waiting for input on handles, so if we listen for LSP connections it will
           -- prevent UCM from shutting down properly. Hopefully we can re-enable LSP on
           -- Windows when we move to GHC 9.*
           -- https://gitlab.haskell.org/ghc/ghc/-/merge_requests/1224
-          when (not onWindows) . void . Ki.fork scope $ LSP.spawnLsp theCodebase runtime (readTMVar rootVar) (readTVar pathVar)
+          when (not onWindows) . void . Ki.fork scope $ LSP.spawnLsp theCodebase runtime (readTMVar namesVar)
           Server.startServer (Backend.BackendEnv {Backend.useNamesIndex = False}) codebaseServerOpts sbRuntime theCodebase $ \baseUrl -> do
             case exitOption of
               DoNotExit -> do
@@ -271,7 +271,7 @@ main = withCP65001 . Ki.scoped $ \scope -> do
                     takeMVar mvar
                   WithCLI -> do
                     PT.putPrettyLn $ P.string "Now starting the Unison Codebase Manager (UCM)..."
-                    launch currentDir config runtime sbRuntime theCodebase [] (Just baseUrl) mayStartingPath downloadBase initRes notifyOnRootChanges notifyOnPathChanges
+                    launch currentDir config runtime sbRuntime theCodebase [] (Just baseUrl) mayStartingPath downloadBase initRes notifyOnBranchChanges
               Exit -> do Exit.exitSuccess
 
 -- | Set user agent and configure TLS on global http client.
@@ -421,9 +421,8 @@ launch ::
   ShouldDownloadBase ->
   InitResult ->
   (Branch IO -> STM ()) ->
-  (Path.Absolute -> STM ()) ->
   IO ()
-launch dir config runtime sbRuntime codebase inputs serverBaseUrl mayStartingPath shouldDownloadBase initResult notifyRootChange notifyPathChange =
+launch dir config runtime sbRuntime codebase inputs serverBaseUrl mayStartingPath shouldDownloadBase initResult notifyBranchChange =
   let downloadBase = case defaultBaseLib of
         Just remoteNS | shouldDownloadBase == ShouldDownloadBase -> Welcome.DownloadBase remoteNS
         _ -> Welcome.DontDownloadBase
@@ -444,8 +443,7 @@ launch dir config runtime sbRuntime codebase inputs serverBaseUrl mayStartingPat
         codebase
         serverBaseUrl
         ucmVersion
-        notifyRootChange
-        notifyPathChange
+        notifyBranchChange
 
 newtype MarkdownFile = MarkdownFile FilePath