From 7739ad5fa7c8aa8d014a58498549d43795e1c294 Mon Sep 17 00:00:00 2001
From: Jon Shen <shenjonathan0@gmail.com>
Date: Wed, 22 Dec 2021 18:18:22 -0500
Subject: [PATCH 01/14] fix sql busy database is locked errors using retries

---
 ghcide/ghcide.cabal                           |  1 +
 .../session-loader/Development/IDE/Session.hs | 29 ++++++++-
 ghcide/src/Development/IDE/Core/Actions.hs    | 23 +++----
 ghcide/src/Development/IDE/Core/Rules.hs      |  8 +--
 ghcide/src/Development/IDE/Core/Service.hs    |  8 +--
 ghcide/src/Development/IDE/Core/Shake.hs      | 15 +++--
 .../src/Development/IDE/LSP/LanguageServer.hs | 13 ++--
 ghcide/src/Development/IDE/Main.hs            | 12 ++--
 ghcide/src/Development/IDE/Spans/AtPoint.hs   | 48 +++++++-------
 ghcide/test/exe/Main.hs                       | 65 ++++++++++++++++++-
 .../src/Ide/Plugin/CallHierarchy/Internal.hs  |  8 +--
 .../src/Ide/Plugin/Rename.hs                  | 18 ++---
 12 files changed, 172 insertions(+), 76 deletions(-)

diff --git a/ghcide/ghcide.cabal b/ghcide/ghcide.cabal
index e4095e239d..ed830a0ff8 100644
--- a/ghcide/ghcide.cabal
+++ b/ghcide/ghcide.cabal
@@ -397,6 +397,7 @@ test-suite ghcide-tests
         safe,
         safe-exceptions,
         shake,
+        sqlite-simple,
         stm,
         stm-containers,
         hls-graph,
diff --git a/ghcide/session-loader/Development/IDE/Session.hs b/ghcide/session-loader/Development/IDE/Session.hs
index 643bcf6303..a60e6dbcb3 100644
--- a/ghcide/session-loader/Development/IDE/Session.hs
+++ b/ghcide/session-loader/Development/IDE/Session.hs
@@ -1,3 +1,4 @@
+{-# LANGUAGE RankNTypes   #-}
 {-# LANGUAGE TypeFamilies #-}
 
 {-|
@@ -11,6 +12,7 @@ module Development.IDE.Session
   ,setInitialDynFlags
   ,getHieDbLoc
   ,runWithDb
+  ,retryOnSqliteBusy
   ) where
 
 -- Unfortunately, we cannot use loadSession with ghc-lib since hie-bios uses
@@ -41,7 +43,7 @@ import qualified Data.Text                            as T
 import           Data.Time.Clock
 import           Data.Version
 import           Development.IDE.Core.RuleTypes
-import           Development.IDE.Core.Shake
+import           Development.IDE.Core.Shake           hiding (withHieDb)
 import qualified Development.IDE.GHC.Compat           as Compat
 import           Development.IDE.GHC.Compat.Core      hiding (Target,
                                                        TargetFile, TargetModule,
@@ -165,11 +167,32 @@ setInitialDynFlags logger rootDir SessionLoadingOptions{..} = do
   mapM_ setUnsafeGlobalDynFlags dynFlags
   pure libdir
 
+
+-- | If the sqlite returns an SQLITE_BUSY then we sleep for a millisecond
+-- and try action again for a maximum of `maxRetryCount` times.
+-- `MonadIO`, `MonadCatch` are used as constraints because there are a few
+-- HieDb functions that don't return IO values.
+retryOnSqliteBusy :: (MonadIO m, MonadCatch m) => HieDb -> Int -> (HieDb -> m b) -> m b
+retryOnSqliteBusy hieDb !maxRetryCount f = do
+  let hieDbAction = f hieDb
+  let isErrorBusy e
+        | SQLError{ sqlError = ErrorBusy } <- e = Just e
+        | otherwise = Nothing
+  result <- tryJust isErrorBusy hieDbAction
+  case result of
+    Left e
+      | maxRetryCount > 0 ->
+        -- 1 millisecond
+        liftIO (threadDelay 1000) >> retryOnSqliteBusy hieDb (maxRetryCount - 1) f
+      | otherwise ->
+        liftIO $ throwIO e
+    Right b -> pure b
+
 -- | Wraps `withHieDb` to provide a database connection for reading, and a `HieWriterChan` for
 -- writing. Actions are picked off one by one from the `HieWriterChan` and executed in serial
 -- by a worker thread using a dedicated database connection.
 -- This is done in order to serialize writes to the database, or else SQLite becomes unhappy
-runWithDb :: Logger -> FilePath -> (HieDb -> IndexQueue -> IO ()) -> IO ()
+runWithDb :: Logger -> FilePath -> (WithHieDb -> IndexQueue -> IO ()) -> IO ()
 runWithDb logger fp k = do
   -- Delete the database if it has an incompatible schema version
   withHieDb fp (const $ pure ())
@@ -178,7 +201,7 @@ runWithDb logger fp k = do
     initConn writedb
     chan <- newTQueueIO
     withAsync (writerThread writedb chan) $ \_ -> do
-      withHieDb fp (flip k chan)
+      withHieDb fp (\readDb -> k (retryOnSqliteBusy readDb 10) chan)
   where
     writerThread db chan = do
       -- Clear the index of any files that might have been deleted since the last run
diff --git a/ghcide/src/Development/IDE/Core/Actions.hs b/ghcide/src/Development/IDE/Core/Actions.hs
index da7d310111..3d90640be5 100644
--- a/ghcide/src/Development/IDE/Core/Actions.hs
+++ b/ghcide/src/Development/IDE/Core/Actions.hs
@@ -1,4 +1,5 @@
-{-# LANGUAGE TypeFamilies    #-}
+{-# LANGUAGE RankNTypes   #-}
+{-# LANGUAGE TypeFamilies #-}
 module Development.IDE.Core.Actions
 ( getAtPoint
 , getDefinition
@@ -83,24 +84,20 @@ usesE k = MaybeT . fmap sequence . mapM (useWithStaleFast k)
 -- | Goto Definition.
 getDefinition :: NormalizedFilePath -> Position -> IdeAction (Maybe [Location])
 getDefinition file pos = runMaybeT $ do
-    ide <- ask
+    ide@ShakeExtras{ withHieDb, hiedbWriter } <- ask
     opts <- liftIO $ getIdeOptionsIO ide
     (HAR _ hf _ _ _, mapping) <- useE GetHieAst file
     (ImportMap imports, _) <- useE GetImportMap file
     !pos' <- MaybeT (pure $ fromCurrentPosition mapping pos)
-    hiedb <- lift $ asks hiedb
-    dbWriter <- lift $ asks hiedbWriter
-    toCurrentLocations mapping <$> AtPoint.gotoDefinition hiedb (lookupMod dbWriter) opts imports hf pos'
+    toCurrentLocations mapping <$> AtPoint.gotoDefinition withHieDb (lookupMod hiedbWriter) opts imports hf pos'
 
 getTypeDefinition :: NormalizedFilePath -> Position -> IdeAction (Maybe [Location])
 getTypeDefinition file pos = runMaybeT $ do
-    ide <- ask
+    ide@ShakeExtras{ withHieDb, hiedbWriter } <- ask
     opts <- liftIO $ getIdeOptionsIO ide
     (hf, mapping) <- useE GetHieAst file
     !pos' <- MaybeT (return $ fromCurrentPosition mapping pos)
-    hiedb <- lift $ asks hiedb
-    dbWriter <- lift $ asks hiedbWriter
-    toCurrentLocations mapping <$> AtPoint.gotoTypeDefinition hiedb (lookupMod dbWriter) opts hf pos'
+    toCurrentLocations mapping <$> AtPoint.gotoTypeDefinition withHieDb (lookupMod hiedbWriter) opts hf pos'
 
 highlightAtPoint :: NormalizedFilePath -> Position -> IdeAction (Maybe [DocumentHighlight])
 highlightAtPoint file pos = runMaybeT $ do
@@ -112,13 +109,13 @@ highlightAtPoint file pos = runMaybeT $ do
 -- Refs are not an IDE action, so it is OK to be slow and (more) accurate
 refsAtPoint :: NormalizedFilePath -> Position -> Action [Location]
 refsAtPoint file pos = do
-    ShakeExtras{hiedb} <- getShakeExtras
+    ShakeExtras{withHieDb} <- getShakeExtras
     fs <- HM.keys <$> getFilesOfInterestUntracked
     asts <- HM.fromList . mapMaybe sequence . zip fs <$> usesWithStale GetHieAst fs
-    AtPoint.referencesAtPoint hiedb file pos (AtPoint.FOIReferences asts)
+    liftIO $ AtPoint.referencesAtPoint withHieDb file pos (AtPoint.FOIReferences asts)
 
 workspaceSymbols :: T.Text -> IdeAction (Maybe [SymbolInformation])
 workspaceSymbols query = runMaybeT $ do
-  hiedb <- lift $ asks hiedb
-  res <- liftIO $ HieDb.searchDef hiedb $ T.unpack query
+  withHieDb <- lift $ asks withHieDb
+  res <- liftIO $ withHieDb (\hieDb -> HieDb.searchDef hieDb $ T.unpack query)
   pure $ mapMaybe AtPoint.defRowToSymbolInfo res
diff --git a/ghcide/src/Development/IDE/Core/Rules.hs b/ghcide/src/Development/IDE/Core/Rules.hs
index 14ff4a29fa..00c1180036 100644
--- a/ghcide/src/Development/IDE/Core/Rules.hs
+++ b/ghcide/src/Development/IDE/Core/Rules.hs
@@ -596,9 +596,9 @@ persistentDocMapRule = addPersistentRule GetDocMap $ \_ -> pure $ Just (DKMap me
 
 readHieFileForSrcFromDisk :: NormalizedFilePath -> MaybeT IdeAction Compat.HieFile
 readHieFileForSrcFromDisk file = do
-  db <- asks hiedb
+  withHieDb <- asks withHieDb
   log <- asks $ L.logDebug . logger
-  row <- MaybeT $ liftIO $ HieDb.lookupHieFileFromSource db $ fromNormalizedFilePath file
+  row <- MaybeT $ liftIO $ withHieDb (\hieDb -> HieDb.lookupHieFileFromSource hieDb $ fromNormalizedFilePath file)
   let hie_loc = HieDb.hieModuleHieFile row
   liftIO $ log $ "LOADING HIE FILE :" <> T.pack (show file)
   exceptToMaybeT $ readHieFileFromDisk hie_loc
@@ -770,13 +770,13 @@ getModIfaceFromDiskAndIndexRule =
   -- doesn't need early cutoff since all its dependencies already have it
   defineNoDiagnostics $ \GetModIfaceFromDiskAndIndex f -> do
   x <- use_ GetModIfaceFromDisk f
-  se@ShakeExtras{hiedb} <- getShakeExtras
+  se@ShakeExtras{withHieDb} <- getShakeExtras
 
   -- GetModIfaceFromDisk should have written a `.hie` file, must check if it matches version in db
   let ms = hirModSummary x
       hie_loc = Compat.ml_hie_file $ ms_location ms
   hash <- liftIO $ Util.getFileHash hie_loc
-  mrow <- liftIO $ HieDb.lookupHieFileFromSource hiedb (fromNormalizedFilePath f)
+  mrow <- liftIO $ withHieDb (\hieDb -> HieDb.lookupHieFileFromSource hieDb (fromNormalizedFilePath f))
   hie_loc' <- liftIO $ traverse (makeAbsolute . HieDb.hieModuleHieFile) mrow
   case mrow of
     Just row
diff --git a/ghcide/src/Development/IDE/Core/Service.hs b/ghcide/src/Development/IDE/Core/Service.hs
index 59f20deed1..b43a5cada2 100644
--- a/ghcide/src/Development/IDE/Core/Service.hs
+++ b/ghcide/src/Development/IDE/Core/Service.hs
@@ -30,7 +30,7 @@ import qualified Language.LSP.Types              as LSP
 
 import           Control.Monad
 import           Development.IDE.Core.Shake
-import           System.Environment             (lookupEnv)
+import           System.Environment              (lookupEnv)
 
 
 ------------------------------------------------------------
@@ -44,10 +44,10 @@ initialise :: Config
            -> Debouncer LSP.NormalizedUri
            -> IdeOptions
            -> VFSHandle
-           -> HieDb
+           -> WithHieDb
            -> IndexQueue
            -> IO IdeState
-initialise defaultConfig mainRule lspEnv logger debouncer options vfs hiedb hiedbChan = do
+initialise defaultConfig mainRule lspEnv logger debouncer options vfs withHieDb hiedbChan = do
     shakeProfiling <- do
         let fromConf = optShakeProfiling options
         fromEnv <- lookupEnv "GHCIDE_BUILD_PROFILING"
@@ -60,7 +60,7 @@ initialise defaultConfig mainRule lspEnv logger debouncer options vfs hiedb hied
         shakeProfiling
         (optReportProgress options)
         (optTesting options)
-        hiedb
+        withHieDb
         hiedbChan
         vfs
         (optShakeOptions options)
diff --git a/ghcide/src/Development/IDE/Core/Shake.hs b/ghcide/src/Development/IDE/Core/Shake.hs
index 0dda58478e..01766da74d 100644
--- a/ghcide/src/Development/IDE/Core/Shake.hs
+++ b/ghcide/src/Development/IDE/Core/Shake.hs
@@ -73,6 +73,7 @@ module Development.IDE.Core.Shake(
     IndexQueue,
     HieDb,
     HieDbWriter(..),
+    WithHieDb,
     VFSHandle(..),
     addPersistentRule,
     garbageCollectDirtyKeys,
@@ -169,6 +170,8 @@ import           Ide.Types                              (PluginId)
 import qualified "list-t" ListT
 import qualified StmContainers.Map                      as STM
 
+
+
 -- | We need to serialize writes to the database, so we send any function that
 -- needs to write to the database over the channel, where it will be picked up by
 -- a worker thread.
@@ -184,6 +187,10 @@ data HieDbWriter
 -- | Actions to queue up on the index worker thread
 type IndexQueue = TQueue (HieDb -> IO ())
 
+-- | Intended to represent HieDb calls wrapped with (currently) retry
+-- functionality
+type WithHieDb = forall a. (HieDb -> IO a) -> IO a
+
 -- information we stash inside the shakeExtra field
 data ShakeExtras = ShakeExtras
     { --eventer :: LSP.FromServerMessage -> IO ()
@@ -219,7 +226,7 @@ data ShakeExtras = ShakeExtras
     -- | A work queue for actions added via 'runInShakeSession'
     ,actionQueue :: ActionQueue
     ,clientCapabilities :: ClientCapabilities
-    , hiedb :: HieDb -- ^ Use only to read.
+    , withHieDb :: WithHieDb -- ^ Use only to read.
     , hiedbWriter :: HieDbWriter -- ^ use to write
     , persistentKeys :: TVar (HMap.HashMap Key GetStalePersistent)
       -- ^ Registery for functions that compute/get "stale" results for the rule
@@ -499,14 +506,14 @@ shakeOpen :: Maybe (LSP.LanguageContextEnv Config)
           -> Maybe FilePath
           -> IdeReportProgress
           -> IdeTesting
-          -> HieDb
+          -> WithHieDb
           -> IndexQueue
           -> VFSHandle
           -> ShakeOptions
           -> Rules ()
           -> IO IdeState
 shakeOpen lspEnv defaultConfig logger debouncer
-  shakeProfileDir (IdeReportProgress reportProgress) ideTesting@(IdeTesting testing) hiedb indexQueue vfs opts rules = mdo
+  shakeProfileDir (IdeReportProgress reportProgress) ideTesting@(IdeTesting testing) withHieDb indexQueue vfs opts rules = mdo
 
     us <- mkSplitUniqSupply 'r'
     ideNc <- newIORef (initNameCache us knownKeyNames)
@@ -528,7 +535,7 @@ shakeOpen lspEnv defaultConfig logger debouncer
         -- lazily initialize the exports map with the contents of the hiedb
         _ <- async $ do
             logDebug logger "Initializing exports map from hiedb"
-            em <- createExportsMapHieDb hiedb
+            em <- withHieDb createExportsMapHieDb
             atomically $ modifyTVar' exportsMap (<> em)
             logDebug logger $ "Done initializing exports map from hiedb (" <> pack(show (ExportsMap.size em)) <> ")"
 
diff --git a/ghcide/src/Development/IDE/LSP/LanguageServer.hs b/ghcide/src/Development/IDE/LSP/LanguageServer.hs
index ec2cf3f484..292de72a5f 100644
--- a/ghcide/src/Development/IDE/LSP/LanguageServer.hs
+++ b/ghcide/src/Development/IDE/LSP/LanguageServer.hs
@@ -44,6 +44,9 @@ import           System.IO.Unsafe                      (unsafeInterleaveIO)
 issueTrackerUrl :: T.Text
 issueTrackerUrl = "https://github.com/haskell/haskell-language-server/issues"
 
+-- used to smuggle RankNType WithHieDb through dbMVar
+newtype WithHieDbShield = WithHieDbShield WithHieDb
+
 runLanguageServer
     :: forall config. (Show config)
     => LSP.Options
@@ -53,7 +56,7 @@ runLanguageServer
     -> config
     -> (config -> Value -> Either T.Text config)
     -> LSP.Handlers (ServerM config)
-    -> (LSP.LanguageContextEnv config -> VFSHandle -> Maybe FilePath -> HieDb -> IndexQueue -> IO IdeState)
+    -> (LSP.LanguageContextEnv config -> VFSHandle -> Maybe FilePath -> WithHieDb -> IndexQueue -> IO IdeState)
     -> IO ()
 runLanguageServer options inH outH getHieDbLoc defaultConfig onConfigurationChange userHandlers getIdeState = do
 
@@ -136,9 +139,9 @@ runLanguageServer options inH outH getHieDbLoc defaultConfig onConfigurationChan
             -- The database needs to be open for the duration of the reactor thread, but we need to pass in a reference
             -- to 'getIdeState', so we use this dirty trick
             dbMVar <- newEmptyMVar
-            ~(hiedb,hieChan) <- unsafeInterleaveIO $ takeMVar dbMVar
+            ~(WithHieDbShield withHieDb,hieChan) <- unsafeInterleaveIO $ takeMVar dbMVar
 
-            ide <- getIdeState env (makeLSPVFSHandle env) root hiedb hieChan
+            ide <- getIdeState env (makeLSPVFSHandle env) root withHieDb hieChan
 
             let initConfig = parseConfiguration params
             logInfo (ideLogger ide) $ T.pack $ "Registering ide configuration: " <> show initConfig
@@ -183,8 +186,8 @@ runLanguageServer options inH outH getHieDbLoc defaultConfig onConfigurationChan
                             exceptionInHandler e
                             k $ ResponseError InternalError (T.pack $ show e) Nothing
             _ <- flip forkFinally handleServerException $ do
-                untilMVar lifetime $ runWithDb logger dbLoc $ \hiedb hieChan -> do
-                    putMVar dbMVar (hiedb,hieChan)
+                untilMVar lifetime $ runWithDb logger dbLoc $ \withHieDb hieChan -> do
+                    putMVar dbMVar (WithHieDbShield withHieDb,hieChan)
                     forever $ do
                         msg <- readChan clientMsgChan
                         -- We dispatch notifications synchronously and requests asynchronously
diff --git a/ghcide/src/Development/IDE/Main.hs b/ghcide/src/Development/IDE/Main.hs
index 541e7385b8..86a6007100 100644
--- a/ghcide/src/Development/IDE/Main.hs
+++ b/ghcide/src/Development/IDE/Main.hs
@@ -11,7 +11,8 @@ module Development.IDE.Main
 ,testing) where
 import           Control.Concurrent.Extra              (newLock, withLock,
                                                         withNumCapabilities)
-import           Control.Concurrent.STM.Stats          (atomically, dumpSTMStats)
+import           Control.Concurrent.STM.Stats          (atomically,
+                                                        dumpSTMStats)
 import           Control.Exception.Safe                (Exception (displayException),
                                                         catchAny)
 import           Control.Monad.Extra                   (concatMapM, unless,
@@ -56,6 +57,7 @@ import           Development.IDE.Core.Shake            (IdeState (shakeExtras),
 import           Development.IDE.Core.Tracing          (measureMemory)
 import           Development.IDE.Graph                 (action)
 import           Development.IDE.LSP.LanguageServer    (runLanguageServer)
+import           Development.IDE.Main.HeapStats        (withHeapStats)
 import           Development.IDE.Plugin                (Plugin (pluginHandlers, pluginModifyDynflags, pluginRules))
 import           Development.IDE.Plugin.HLS            (asGhcIdePlugin)
 import qualified Development.IDE.Plugin.HLS.GhcIde     as Ghcide
@@ -77,12 +79,10 @@ import           Development.IDE.Types.Options         (IdeGhcSession,
                                                         defaultIdeOptions,
                                                         optModifyDynFlags,
                                                         optTesting)
-import           Development.IDE.Types.Shake           (Key(Key),
-                                                        fromKeyType)
+import           Development.IDE.Types.Shake           (Key (Key), fromKeyType)
 import           GHC.Conc                              (getNumProcessors)
 import           GHC.IO.Encoding                       (setLocaleEncoding)
 import           GHC.IO.Handle                         (hDuplicate)
-import           Development.IDE.Main.HeapStats        (withHeapStats)
 import           HIE.Bios.Cradle                       (findCradle)
 import qualified HieDb.Run                             as HieDb
 import           Ide.Plugin.Config                     (CheckParents (NeverCheck),
@@ -272,7 +272,7 @@ defaultMain Arguments{..} = flip withHeapStats fun =<< argsLogger
             t <- offsetTime
             logInfo logger "Starting LSP server..."
             logInfo logger "If you are seeing this in a terminal, you probably should have run WITHOUT the --lsp option!"
-            runLanguageServer options inH outH argsGetHieDbLoc argsDefaultHlsConfig argsOnConfigChange (pluginHandlers plugins) $ \env vfs rootPath hiedb hieChan -> do
+            runLanguageServer options inH outH argsGetHieDbLoc argsDefaultHlsConfig argsOnConfigChange (pluginHandlers plugins) $ \env vfs rootPath withHieDb hieChan -> do
                 traverse_ IO.setCurrentDirectory rootPath
                 t <- t
                 logInfo logger $ T.pack $ "Started LSP server in " ++ showDuration t
@@ -313,7 +313,7 @@ defaultMain Arguments{..} = flip withHeapStats fun =<< argsLogger
                     debouncer
                     options
                     vfs
-                    hiedb
+                    withHieDb
                     hieChan
             dumpSTMStats
         Check argFiles -> do
diff --git a/ghcide/src/Development/IDE/Spans/AtPoint.hs b/ghcide/src/Development/IDE/Spans/AtPoint.hs
index 36bdd58303..060cac4019 100644
--- a/ghcide/src/Development/IDE/Spans/AtPoint.hs
+++ b/ghcide/src/Development/IDE/Spans/AtPoint.hs
@@ -1,8 +1,9 @@
 -- Copyright (c) 2019 The DAML Authors. All rights reserved.
 -- SPDX-License-Identifier: Apache-2.0
 
-{-# LANGUAGE CPP   #-}
-{-# LANGUAGE GADTs #-}
+{-# LANGUAGE CPP        #-}
+{-# LANGUAGE GADTs      #-}
+{-# LANGUAGE RankNTypes #-}
 
 -- | Gives information about symbols at a given point in DAML files.
 -- These are all pure functions that should execute quickly.
@@ -51,6 +52,7 @@ import           Data.List                            (isSuffixOf)
 import           Data.List.Extra                      (dropEnd1, nubOrd)
 
 import           Data.Version                         (showVersion)
+import           Development.IDE.Core.Shake           (WithHieDb)
 import           HieDb                                hiding (pointCommand)
 import           System.Directory                     (doesFileExist)
 
@@ -106,12 +108,12 @@ toCurrentLocation mapping (Location uri range) =
 
 referencesAtPoint
   :: MonadIO m
-  => HieDb
+  => WithHieDb
   -> NormalizedFilePath -- ^ The file the cursor is in
   -> Position -- ^ position in the file
   -> FOIReferences -- ^ references data for FOIs
   -> m [Location]
-referencesAtPoint hiedb nfp pos refs = do
+referencesAtPoint withHieDb nfp pos refs = do
   -- The database doesn't have up2date references data for the FOIs so we must collect those
   -- from the Shake graph.
   let (names, foiRefs, exclude) = foiReferencesAtPoint nfp pos refs
@@ -121,12 +123,12 @@ referencesAtPoint hiedb nfp pos refs = do
       Just mod -> do
          -- Look for references (strictly in project files, not dependencies),
          -- excluding the files in the FOIs (since those are in foiRefs)
-         rows <- liftIO $ findReferences hiedb True (nameOccName name) (Just $ moduleName mod) (Just $ moduleUnit mod) exclude
+         rows <- liftIO $ withHieDb (\hieDb -> findReferences hieDb True (nameOccName name) (Just $ moduleName mod) (Just $ moduleUnit mod) exclude)
          pure $ mapMaybe rowToLoc rows
   typeRefs <- forM names $ \name ->
     case nameModule_maybe name of
       Just mod | isTcClsNameSpace (occNameSpace $ nameOccName name) -> do
-        refs <- liftIO $ findTypeRefs hiedb True (nameOccName name) (Just $ moduleName mod) (Just $ moduleUnit mod) exclude
+        refs <- liftIO $ withHieDb (\hieDb -> findTypeRefs hieDb True (nameOccName name) (Just $ moduleName mod) (Just $ moduleUnit mod) exclude)
         pure $ mapMaybe typeRowToLoc refs
       _ -> pure []
   pure $ nubOrd $ foiRefs ++ concat nonFOIRefs ++ concat typeRefs
@@ -172,27 +174,27 @@ documentHighlight hf rf pos = pure highlights
 
 gotoTypeDefinition
   :: MonadIO m
-  => HieDb
+  => WithHieDb
   -> LookupModule m
   -> IdeOptions
   -> HieAstResult
   -> Position
   -> MaybeT m [Location]
-gotoTypeDefinition hiedb lookupModule ideOpts srcSpans pos
-  = lift $ typeLocationsAtPoint hiedb lookupModule ideOpts pos srcSpans
+gotoTypeDefinition withHieDb lookupModule ideOpts srcSpans pos
+  = lift $ typeLocationsAtPoint withHieDb lookupModule ideOpts pos srcSpans
 
 -- | Locate the definition of the name at a given position.
 gotoDefinition
   :: MonadIO m
-  => HieDb
+  => WithHieDb
   -> LookupModule m
   -> IdeOptions
   -> M.Map ModuleName NormalizedFilePath
   -> HieASTs a
   -> Position
   -> MaybeT m [Location]
-gotoDefinition hiedb getHieFile ideOpts imports srcSpans pos
-  = lift $ locationsAtPoint hiedb getHieFile ideOpts imports pos srcSpans
+gotoDefinition withHieDb getHieFile ideOpts imports srcSpans pos
+  = lift $ locationsAtPoint withHieDb getHieFile ideOpts imports pos srcSpans
 
 -- | Synopsis for the name at a given position.
 atPoint
@@ -252,13 +254,13 @@ atPoint IdeOptions{} (HAR _ hf _ _ kind) (DKMap dm km) env pos = listToMaybe $ p
 typeLocationsAtPoint
   :: forall m
    . MonadIO m
-  => HieDb
+  => WithHieDb
   -> LookupModule m
   -> IdeOptions
   -> Position
   -> HieAstResult
   -> m [Location]
-typeLocationsAtPoint hiedb lookupModule _ideOptions pos (HAR _ ast _ _ hieKind) =
+typeLocationsAtPoint withHieDb lookupModule _ideOptions pos (HAR _ ast _ _ hieKind) =
   case hieKind of
     HieFromDisk hf ->
       let arr = hie_types hf
@@ -283,12 +285,12 @@ typeLocationsAtPoint hiedb lookupModule _ideOptions pos (HAR _ ast _ _ hieKind)
             HQualTy a b -> getTypes [a,b]
             HCastTy a -> getTypes [a]
             _ -> []
-        in fmap nubOrd $ concatMapM (fmap (fromMaybe []) . nameToLocation hiedb lookupModule) (getTypes ts)
+        in fmap nubOrd $ concatMapM (fmap (fromMaybe []) . nameToLocation withHieDb lookupModule) (getTypes ts)
     HieFresh ->
       let ts = concat $ pointCommand ast pos getts
           getts x = nodeType ni  ++ (mapMaybe identType $ M.elems $ nodeIdentifiers ni)
             where ni = nodeInfo x
-        in fmap nubOrd $ concatMapM (fmap (fromMaybe []) . nameToLocation hiedb lookupModule) (getTypes ts)
+        in fmap nubOrd $ concatMapM (fmap (fromMaybe []) . nameToLocation withHieDb lookupModule) (getTypes ts)
 
 namesInType :: Type -> [Name]
 namesInType (TyVarTy n)      = [varName n]
@@ -306,23 +308,23 @@ getTypes ts = concatMap namesInType ts
 locationsAtPoint
   :: forall m a
    . MonadIO m
-  => HieDb
+  => WithHieDb
   -> LookupModule m
   -> IdeOptions
   -> M.Map ModuleName NormalizedFilePath
   -> Position
   -> HieASTs a
   -> m [Location]
-locationsAtPoint hiedb lookupModule _ideOptions imports pos ast =
+locationsAtPoint withHieDb lookupModule _ideOptions imports pos ast =
   let ns = concat $ pointCommand ast pos (M.keys . getNodeIds)
       zeroPos = Position 0 0
       zeroRange = Range zeroPos zeroPos
       modToLocation m = fmap (\fs -> pure $ Location (fromNormalizedUri $ filePathToUri' fs) zeroRange) $ M.lookup m imports
-    in fmap (nubOrd . concat) $ mapMaybeM (either (pure . modToLocation) $ nameToLocation hiedb lookupModule) ns
+    in fmap (nubOrd . concat) $ mapMaybeM (either (pure . modToLocation) $ nameToLocation withHieDb lookupModule) ns
 
 -- | Given a 'Name' attempt to find the location where it is defined.
-nameToLocation :: MonadIO m => HieDb -> LookupModule m -> Name -> m (Maybe [Location])
-nameToLocation hiedb lookupModule name = runMaybeT $
+nameToLocation :: MonadIO m => WithHieDb -> LookupModule m -> Name -> m (Maybe [Location])
+nameToLocation withHieDb lookupModule name = runMaybeT $
   case nameSrcSpan name of
     sp@(RealSrcSpan rsp _)
       -- Lookup in the db if we got a location in a boot file
@@ -344,14 +346,14 @@ nameToLocation hiedb lookupModule name = runMaybeT $
       -- In this case the interface files contain garbage source spans
       -- so we instead read the .hie files to get useful source spans.
       mod <- MaybeT $ return $ nameModule_maybe name
-      erow <- liftIO $ findDef hiedb (nameOccName name) (Just $ moduleName mod) (Just $ moduleUnit mod)
+      erow <- liftIO $ withHieDb (\hieDb -> findDef hieDb (nameOccName name) (Just $ moduleName mod) (Just $ moduleUnit mod))
       case erow of
         [] -> do
           -- If the lookup failed, try again without specifying a unit-id.
           -- This is a hack to make find definition work better with ghcide's nascent multi-component support,
           -- where names from a component that has been indexed in a previous session but not loaded in this
           -- session may end up with different unit ids
-          erow <- liftIO $ findDef hiedb (nameOccName name) (Just $ moduleName mod) Nothing
+          erow <- liftIO $ withHieDb (\hieDb -> findDef hieDb (nameOccName name) (Just $ moduleName mod) Nothing)
           case erow of
             [] -> MaybeT $ pure Nothing
             xs -> lift $ mapMaybeM (runMaybeT . defRowToLocation lookupModule) xs
diff --git a/ghcide/test/exe/Main.hs b/ghcide/test/exe/Main.hs
index 6a366ebbdd..a8f4bee00a 100644
--- a/ghcide/test/exe/Main.hs
+++ b/ghcide/test/exe/Main.hs
@@ -15,7 +15,8 @@
 module Main (main) where
 
 import           Control.Applicative.Combinators
-import           Control.Exception                        (bracket_, catch)
+import           Control.Concurrent.Extra                 as Concurrent
+import           Control.Exception                        (bracket_, catch, throwIO, tryJust, ErrorCall(ErrorCall), evaluate)
 import qualified Control.Lens                             as Lens
 import           Control.Monad
 import           Control.Monad.IO.Class                   (MonadIO, liftIO)
@@ -57,6 +58,7 @@ import           Development.IDE.Test                     (Cursor,
 import           Development.IDE.Test.Runfiles
 import qualified Development.IDE.Types.Diagnostics        as Diagnostics
 import           Development.IDE.Types.Location
+import           Development.IDE.Session                  (retryOnSqliteBusy)
 import qualified Language.LSP.Types.Lens                  as Lens (label)
 import           Development.Shake                        (getDirectoryFilesIO)
 import qualified Experiments                              as Bench
@@ -87,7 +89,6 @@ import           System.Process.Extra                     (CreateProcess (cwd),
                                                            readCreateProcessWithExitCode)
 import           Test.QuickCheck
 -- import Test.QuickCheck.Instances ()
-import           Control.Concurrent                       (threadDelay)
 import           Control.Concurrent.Async
 import           Control.Lens                             ((^.))
 import           Control.Monad.Extra                      (whenJust)
@@ -114,6 +115,7 @@ import           Test.Tasty.Ingredients.Rerun
 import           Test.Tasty.QuickCheck
 import           Text.Printf                              (printf)
 import           Text.Regex.TDFA                          ((=~))
+import           Database.SQLite.Simple                   as SQLite
 
 -- | Wait for the next progress begin step
 waitForProgressBegin :: Session ()
@@ -6141,6 +6143,65 @@ unitTests = do
            let msg = printf "Timestamps do not have millisecond resolution: %dus" resolution_us
            assertBool msg (resolution_us <= 1000)
      , Progress.tests
+     , testCase "retryOnSqliteBusy throws ErrorBusy after max retries" $ do
+        let isErrorBusy e
+              | SQLite.SQLError { sqlError = SQLite.ErrorBusy } <- e = Just e
+              | otherwise = Nothing
+        let expected = SQLite.SQLError{ sqlError = SQLite.ErrorBusy, sqlErrorDetails = "", sqlErrorContext = "" }
+        let hieDbAction _ = throwIO expected
+        let action = retryOnSqliteBusy undefined 1 hieDbAction
+           
+        result <- tryJust isErrorBusy action
+
+        case result of
+          Left actual -> actual @?= expected
+          Right _ -> assertFailure "Expected ErrorBusy exception"
+     , testCase "retryOnSqliteBusy doesn't throw if given function doesn't throw" $ do
+        let expected = 1 :: Int
+        let action = retryOnSqliteBusy undefined 0 (\_ -> pure expected)
+
+        actual <- action
+
+        actual @?= expected
+     , testCase "retryOnSqliteBusy retries the number of times it should" $ do
+        let isErrorBusy e
+              | SQLite.SQLError { sqlError = SQLite.ErrorBusy } <- e = Just e
+              | otherwise = Nothing
+        let errorBusy = SQLite.SQLError{ sqlError = SQLite.ErrorBusy, sqlErrorDetails = "", sqlErrorContext = "" }
+        let maxRetryCount = 3
+        let expected = maxRetryCount + 1
+
+        countVar <- newVar 0
+
+        let hieDbAction _ = Concurrent.modifyVar countVar (\count -> pure (dupe (count + 1))) >> throwIO errorBusy
+        let action = retryOnSqliteBusy undefined maxRetryCount hieDbAction
+
+        _ <- tryJust isErrorBusy action
+
+        actual <- Concurrent.readVar countVar
+
+        actual @?= expected
+     , testCase "retryOnSqliteBusy doesn't retry if exception is not ErrorBusy" $ do
+        let isErrorCall e
+              | ErrorCall _ <- e = Just e
+              | otherwise = Nothing
+
+        countVar <- newVar 0
+
+        let hieDbAction _ = do
+              count <- readVar countVar
+              if count == 0 then
+                evaluate (error "dummy exception")
+              else
+                Concurrent.modifyVar countVar (\count -> pure (dupe (count + 1)))
+
+        let action = retryOnSqliteBusy undefined 1 hieDbAction
+        let expected = 0 :: Int
+
+        void $ tryJust isErrorCall action
+        actual <- readVar countVar
+
+        actual @?= expected
      ]
 
 garbageCollectionTests :: TestTree
diff --git a/plugins/hls-call-hierarchy-plugin/src/Ide/Plugin/CallHierarchy/Internal.hs b/plugins/hls-call-hierarchy-plugin/src/Ide/Plugin/CallHierarchy/Internal.hs
index e54c7721ab..8641e79214 100644
--- a/plugins/hls-call-hierarchy-plugin/src/Ide/Plugin/CallHierarchy/Internal.hs
+++ b/plugins/hls-call-hierarchy-plugin/src/Ide/Plugin/CallHierarchy/Internal.hs
@@ -242,8 +242,8 @@ mkCallHierarchyCall mk v@Vertex{..} = do
     \case
       Just [item] -> pure $ Just $ mk item (List [range])
       _           -> do
-        ShakeExtras{hiedb} <- getShakeExtras
-        liftIO (Q.getSymbolPosition hiedb v) >>=
+        ShakeExtras{withHieDb} <- getShakeExtras
+        liftIO (withHieDb (\hieDb -> Q.getSymbolPosition hieDb v)) >>=
           \case
             (x:_) ->
               prepareCallHierarchyItem nfp (Position (psl x - 1) (psc x - 1)) >>=
@@ -263,12 +263,12 @@ queryCalls item queryFunc makeFunc merge
   | Just nfp <- uriToNormalizedFilePath $ toNormalizedUri uri = do
     refreshHieDb
 
-    ShakeExtras{hiedb} <- getShakeExtras
+    ShakeExtras{withHieDb} <- getShakeExtras
     maySymbol <- getSymbol nfp
     case maySymbol of
       Nothing -> error "CallHierarchy.Impossible"
       Just symbol -> do
-        vs <- liftIO $ queryFunc hiedb symbol
+        vs <- liftIO $ withHieDb (\hieDb -> queryFunc hieDb symbol)
         items <- Just . catMaybes <$> mapM makeFunc vs
         pure $ merge <$> items
   | otherwise = pure Nothing
diff --git a/plugins/hls-rename-plugin/src/Ide/Plugin/Rename.hs b/plugins/hls-rename-plugin/src/Ide/Plugin/Rename.hs
index fd85bccc5c..54470fc0d8 100644
--- a/plugins/hls-rename-plugin/src/Ide/Plugin/Rename.hs
+++ b/plugins/hls-rename-plugin/src/Ide/Plugin/Rename.hs
@@ -113,19 +113,21 @@ renameModRefs newNameText refs = everywhere $ mkT replace
 -- | Note: We only find exact name occurences (i.e. type reference "depth" is 0).
 refsAtName :: IdeState -> NormalizedFilePath -> Name -> ExceptT [Char] (LspT Config IO) [Location]
 refsAtName state nfp name = do
-    ShakeExtras{hiedb} <- liftIO $ runAction "Rename.HieDb" state getShakeExtras
+    ShakeExtras{withHieDb} <- liftIO $ runAction "Rename.HieDb" state getShakeExtras
     ast <- safeGetHieAst state nfp
     astRefs <- handleMaybe "Error: Could not get name AST references" $ getNameAstLocations name ast
     dbRefs <- case nameModule_maybe name of
         Nothing -> pure []
         Just mod -> liftIO $ mapMaybe rowToLoc <$>
-            findReferences
-                hiedb
-                True
-                (nameOccName name)
-                (Just $ moduleName mod)
-                (Just $ moduleUnit mod)
-                [fromNormalizedFilePath nfp]
+            withHieDb (\hieDb ->
+              findReferences
+                  hieDb
+                  True
+                  (nameOccName name)
+                  (Just $ moduleName mod)
+                  (Just $ moduleUnit mod)
+                  [fromNormalizedFilePath nfp]
+              )
     pure $ nubOrd $ astRefs ++ dbRefs
 
 getNameAstLocations :: Name -> (HieAstResult, PositionMapping) -> Maybe [Location]

From b97c74af42b844aff8b14f81d4634b930012b26d Mon Sep 17 00:00:00 2001
From: Jon Shen <shenjonathan0@gmail.com>
Date: Wed, 22 Dec 2021 20:48:20 -0500
Subject: [PATCH 02/14] fix ghc 9.0+

---
 ghcide/src/Development/IDE/Core/Actions.hs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ghcide/src/Development/IDE/Core/Actions.hs b/ghcide/src/Development/IDE/Core/Actions.hs
index 3d90640be5..d91fdc083d 100644
--- a/ghcide/src/Development/IDE/Core/Actions.hs
+++ b/ghcide/src/Development/IDE/Core/Actions.hs
@@ -116,6 +116,6 @@ refsAtPoint file pos = do
 
 workspaceSymbols :: T.Text -> IdeAction (Maybe [SymbolInformation])
 workspaceSymbols query = runMaybeT $ do
-  withHieDb <- lift $ asks withHieDb
+  ShakeExtras{withHieDb} <- ask
   res <- liftIO $ withHieDb (\hieDb -> HieDb.searchDef hieDb $ T.unpack query)
   pure $ mapMaybe AtPoint.defRowToSymbolInfo res

From 434a4fa64135f54a79d9266a88c6bfc4b03d341b Mon Sep 17 00:00:00 2001
From: Jon Shen <shenjonathan0@gmail.com>
Date: Wed, 22 Dec 2021 20:48:48 -0500
Subject: [PATCH 03/14] hlint fixes

---
 .../src/Ide/Plugin/CallHierarchy/Internal.hs                  | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/plugins/hls-call-hierarchy-plugin/src/Ide/Plugin/CallHierarchy/Internal.hs b/plugins/hls-call-hierarchy-plugin/src/Ide/Plugin/CallHierarchy/Internal.hs
index 8641e79214..299f23e4af 100644
--- a/plugins/hls-call-hierarchy-plugin/src/Ide/Plugin/CallHierarchy/Internal.hs
+++ b/plugins/hls-call-hierarchy-plugin/src/Ide/Plugin/CallHierarchy/Internal.hs
@@ -243,7 +243,7 @@ mkCallHierarchyCall mk v@Vertex{..} = do
       Just [item] -> pure $ Just $ mk item (List [range])
       _           -> do
         ShakeExtras{withHieDb} <- getShakeExtras
-        liftIO (withHieDb (\hieDb -> Q.getSymbolPosition hieDb v)) >>=
+        liftIO (withHieDb (`Q.getSymbolPosition` v)) >>=
           \case
             (x:_) ->
               prepareCallHierarchyItem nfp (Position (psl x - 1) (psc x - 1)) >>=
@@ -268,7 +268,7 @@ queryCalls item queryFunc makeFunc merge
     case maySymbol of
       Nothing -> error "CallHierarchy.Impossible"
       Just symbol -> do
-        vs <- liftIO $ withHieDb (\hieDb -> queryFunc hieDb symbol)
+        vs <- liftIO $ withHieDb (`queryFunc` symbol)
         items <- Just . catMaybes <$> mapM makeFunc vs
         pure $ merge <$> items
   | otherwise = pure Nothing

From 32ef8332a993c928c2d2ea01281c281ebc05136b Mon Sep 17 00:00:00 2001
From: Jon Shen <shenjonathan0@gmail.com>
Date: Wed, 22 Dec 2021 20:56:36 -0500
Subject: [PATCH 04/14] fix ghc 9.0+ again

---
 ghcide/src/Development/IDE/Core/Rules.hs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ghcide/src/Development/IDE/Core/Rules.hs b/ghcide/src/Development/IDE/Core/Rules.hs
index 00c1180036..5ce7017713 100644
--- a/ghcide/src/Development/IDE/Core/Rules.hs
+++ b/ghcide/src/Development/IDE/Core/Rules.hs
@@ -596,7 +596,7 @@ persistentDocMapRule = addPersistentRule GetDocMap $ \_ -> pure $ Just (DKMap me
 
 readHieFileForSrcFromDisk :: NormalizedFilePath -> MaybeT IdeAction Compat.HieFile
 readHieFileForSrcFromDisk file = do
-  withHieDb <- asks withHieDb
+  ShakeExtras{withHieDb} <- ask
   log <- asks $ L.logDebug . logger
   row <- MaybeT $ liftIO $ withHieDb (\hieDb -> HieDb.lookupHieFileFromSource hieDb $ fromNormalizedFilePath file)
   let hie_loc = HieDb.hieModuleHieFile row

From d3e275c326f705845b78aa9e5b096bcf83d177ec Mon Sep 17 00:00:00 2001
From: Jon Shen <shenjonathan0@gmail.com>
Date: Wed, 22 Dec 2021 23:52:32 -0500
Subject: [PATCH 05/14] remove accidentally added redundant liftIO, remove
 accidentally added empty lines

---
 ghcide/src/Development/IDE/Core/Actions.hs | 2 +-
 ghcide/src/Development/IDE/Core/Shake.hs   | 2 --
 2 files changed, 1 insertion(+), 3 deletions(-)

diff --git a/ghcide/src/Development/IDE/Core/Actions.hs b/ghcide/src/Development/IDE/Core/Actions.hs
index d91fdc083d..304dfd393e 100644
--- a/ghcide/src/Development/IDE/Core/Actions.hs
+++ b/ghcide/src/Development/IDE/Core/Actions.hs
@@ -112,7 +112,7 @@ refsAtPoint file pos = do
     ShakeExtras{withHieDb} <- getShakeExtras
     fs <- HM.keys <$> getFilesOfInterestUntracked
     asts <- HM.fromList . mapMaybe sequence . zip fs <$> usesWithStale GetHieAst fs
-    liftIO $ AtPoint.referencesAtPoint withHieDb file pos (AtPoint.FOIReferences asts)
+    AtPoint.referencesAtPoint withHieDb file pos (AtPoint.FOIReferences asts)
 
 workspaceSymbols :: T.Text -> IdeAction (Maybe [SymbolInformation])
 workspaceSymbols query = runMaybeT $ do
diff --git a/ghcide/src/Development/IDE/Core/Shake.hs b/ghcide/src/Development/IDE/Core/Shake.hs
index 01766da74d..0f9135bed4 100644
--- a/ghcide/src/Development/IDE/Core/Shake.hs
+++ b/ghcide/src/Development/IDE/Core/Shake.hs
@@ -170,8 +170,6 @@ import           Ide.Types                              (PluginId)
 import qualified "list-t" ListT
 import qualified StmContainers.Map                      as STM
 
-
-
 -- | We need to serialize writes to the database, so we send any function that
 -- needs to write to the database over the channel, where it will be picked up by
 -- a worker thread.

From 794d60962dea19e4708d8ae2986a21c38da2d536 Mon Sep 17 00:00:00 2001
From: Jon Shen <shenjonathan0@gmail.com>
Date: Thu, 23 Dec 2021 18:15:43 -0500
Subject: [PATCH 06/14] add logging to retries, add exponential backoff,
 refactor tests

---
 ghcide/ghcide.cabal                           |   3 +
 .../session-loader/Development/IDE/Session.hs |  70 +++++++++--
 ghcide/test/exe/HieDbRetry.hs                 | 117 ++++++++++++++++++
 ghcide/test/exe/Main.hs                       |  65 +---------
 4 files changed, 180 insertions(+), 75 deletions(-)
 create mode 100644 ghcide/test/exe/HieDbRetry.hs

diff --git a/ghcide/ghcide.cabal b/ghcide/ghcide.cabal
index ed830a0ff8..3f19b39133 100644
--- a/ghcide/ghcide.cabal
+++ b/ghcide/ghcide.cabal
@@ -74,6 +74,7 @@ library
         parallel,
         prettyprinter-ansi-terminal,
         prettyprinter,
+        random,
         regex-tdfa >= 1.3.1.0,
         retrie,
         rope-utf16-splay,
@@ -392,6 +393,7 @@ test-suite ghcide-tests
         process,
         QuickCheck,
         quickcheck-instances,
+        random,
         rope-utf16-splay,
         regex-tdfa ^>= 1.3.1,
         safe,
@@ -422,6 +424,7 @@ test-suite ghcide-tests
         Experiments
         Experiments.Types
         Progress
+        HieDbRetry
     default-extensions:
         BangPatterns
         DeriveFunctor
diff --git a/ghcide/session-loader/Development/IDE/Session.hs b/ghcide/session-loader/Development/IDE/Session.hs
index a60e6dbcb3..bfc311a547 100644
--- a/ghcide/session-loader/Development/IDE/Session.hs
+++ b/ghcide/session-loader/Development/IDE/Session.hs
@@ -87,6 +87,8 @@ import           Development.IDE.Core.Tracing         (withTrace)
 import           HieDb.Create
 import           HieDb.Types
 import           HieDb.Utils
+import           System.Random                        (RandomGen)
+import qualified System.Random                        as Random
 
 -- | Bump this version number when making changes to the format of the data stored in hiedb
 hiedbDataVersion :: String
@@ -168,12 +170,22 @@ setInitialDynFlags logger rootDir SessionLoadingOptions{..} = do
   pure libdir
 
 
--- | If the sqlite returns an SQLITE_BUSY then we sleep for a millisecond
--- and try action again for a maximum of `maxRetryCount` times.
+-- | If the sqlite db returns an SQLITE_BUSY then we sleep for a duration
+-- determined by the random exponential backoff formula,
+-- `uniformRandom(0, min (maxDelay, (baseDelay * 2) ^ retryAttempt))`, and try the
+-- action again for a maximum of `maxRetryCount` times.
 -- `MonadIO`, `MonadCatch` are used as constraints because there are a few
 -- HieDb functions that don't return IO values.
-retryOnSqliteBusy :: (MonadIO m, MonadCatch m) => HieDb -> Int -> (HieDb -> m b) -> m b
-retryOnSqliteBusy hieDb !maxRetryCount f = do
+retryOnSqliteBusy :: (MonadIO m, MonadCatch m, RandomGen g)
+                  => Logger
+                  -> HieDb -- ^ HieDb connection
+                  -> Int -- ^ maximum backoff delay in microseconds
+                  -> Int -- ^ base backoff delay in microseconds
+                  -> Int -- ^ maximum number of times to retry
+                  -> g -- ^ random number generator
+                  -> (HieDb -> m b) -- ^ function that uses HieDb connection
+                  -> m b
+retryOnSqliteBusy logger hieDb maxDelay !baseDelay !maxRetryCount rng f = do
   let hieDbAction = f hieDb
   let isErrorBusy e
         | SQLError{ sqlError = ErrorBusy } <- e = Just e
@@ -181,12 +193,38 @@ retryOnSqliteBusy hieDb !maxRetryCount f = do
   result <- tryJust isErrorBusy hieDbAction
   case result of
     Left e
-      | maxRetryCount > 0 ->
-        -- 1 millisecond
-        liftIO (threadDelay 1000) >> retryOnSqliteBusy hieDb (maxRetryCount - 1) f
-      | otherwise ->
-        liftIO $ throwIO e
+      | maxRetryCount > 0 -> do
+        -- multiply by 2 because baseDelay is midpoint of uniform range
+        let newBaseDelay = min maxDelay (baseDelay * 2)
+        let (delay, newRng) = Random.uniformR (0, newBaseDelay) rng
+        let newMaxRetryCount = maxRetryCount - 1
+        liftIO $ do
+          logInfo logger $ "Retrying - " <> makeLogMsgComponentsText (Right delay) newMaxRetryCount e
+          threadDelay delay
+        retryOnSqliteBusy logger hieDb maxDelay newBaseDelay newMaxRetryCount newRng f
+
+      | otherwise -> do
+        liftIO $ do
+          logDebug logger $ "Retries exhausted - " <> makeLogMsgComponentsText (Left baseDelay) maxRetryCount e
+          throwIO e
+
     Right b -> pure b
+  where
+    -- e.g. delay: 1010102, maximumDelay: 12010, maxRetryCount: 9, exception: SQLError { ... }
+    makeLogMsgComponentsText delay newMaxRetryCount e =
+      let
+        logMsgComponents =
+          [ either
+              (("base delay: " <>) . T.pack . show)
+              (("delay: " <>) . T.pack . show)
+              delay
+          , "maximumDelay: " <> T.pack (show maxDelay)
+          , "maxRetryCount: " <> T.pack (show newMaxRetryCount)
+          , "exception: " <> T.pack (show e)]
+      in
+        T.intercalate ", " logMsgComponents
+
+
 
 -- | Wraps `withHieDb` to provide a database connection for reading, and a `HieWriterChan` for
 -- writing. Actions are picked off one by one from the `HieWriterChan` and executed in serial
@@ -200,16 +238,22 @@ runWithDb logger fp k = do
   withHieDb fp $ \writedb -> do
     initConn writedb
     chan <- newTQueueIO
-    withAsync (writerThread writedb chan) $ \_ -> do
-      withHieDb fp (\readDb -> k (retryOnSqliteBusy readDb 10) chan)
+    -- use newStdGen because what if multiple HLS start at same time and send bursts of requests
+    rng <- Random.newStdGen
+    withAsync (writerThread writedb chan rng) $ \_ -> do
+      withHieDb fp (\readDb -> k (retryOnSqliteBusy logger readDb oneSecond oneMillisecond maxRetryCount rng) chan)
   where
-    writerThread db chan = do
+    oneSecond = 1000000
+    oneMillisecond = 1000
+    maxRetryCount = 10
+
+    writerThread db chan rng = do
       -- Clear the index of any files that might have been deleted since the last run
       deleteMissingRealFiles db
       _ <- garbageCollectTypeNames db
       forever $ do
         k <- atomically $ readTQueue chan
-        k db
+        retryOnSqliteBusy logger db oneSecond oneMillisecond maxRetryCount rng k
           `Safe.catch` \e@SQLError{} -> do
             logDebug logger $ T.pack $ "SQLite error in worker, ignoring: " ++ show e
           `Safe.catchAny` \e -> do
diff --git a/ghcide/test/exe/HieDbRetry.hs b/ghcide/test/exe/HieDbRetry.hs
new file mode 100644
index 0000000000..64aac586f7
--- /dev/null
+++ b/ghcide/test/exe/HieDbRetry.hs
@@ -0,0 +1,117 @@
+{-# LANGUAGE MultiWayIf #-}
+module HieDbRetry (tests) where
+
+import           Control.Concurrent.Extra     (Var, modifyVar, newVar, readVar,
+                                               withVar)
+import           Control.Exception            (ErrorCall (ErrorCall), evaluate,
+                                               throwIO, tryJust)
+import           Data.Text                    (Text)
+import           Data.Tuple.Extra             (dupe)
+import qualified Database.SQLite.Simple       as SQLite
+import           Development.IDE.Core.Shake   (HieDb)
+import           Development.IDE.Session      (retryOnSqliteBusy)
+import           Development.IDE.Types.Logger (Logger (Logger), Priority,
+                                               noLogging)
+import qualified System.Random                as Random
+import           Test.Tasty                   (TestTree, testGroup)
+import           Test.Tasty.HUnit             (assertFailure, testCase, (@?=))
+
+makeLogger :: Var [(Priority, Text)] -> Logger
+makeLogger msgsVar = Logger $ \priority msg -> modifyVar msgsVar (\msgs -> pure ((priority, msg) : msgs, ()))
+
+rng :: Random.StdGen
+rng = Random.mkStdGen 0
+
+defaultRetryOnSqliteBusy :: Logger -> Int -> (HieDb -> IO b) -> IO b
+defaultRetryOnSqliteBusy logger maxRetryCount = retryOnSqliteBusy logger undefined 1 1 maxRetryCount rng
+
+isErrorBusy :: SQLite.SQLError -> Maybe SQLite.SQLError
+isErrorBusy e
+  | SQLite.SQLError { sqlError = SQLite.ErrorBusy } <- e = Just e
+  | otherwise = Nothing
+
+errorBusy :: SQLite.SQLError
+errorBusy = SQLite.SQLError{ sqlError = SQLite.ErrorBusy, sqlErrorDetails = "", sqlErrorContext = "" }
+
+isErrorCall :: ErrorCall -> Maybe ErrorCall
+isErrorCall e
+  | ErrorCall _ <- e = Just e
+  | otherwise = Nothing
+
+tests :: TestTree
+tests = testGroup "RetryHieDb"
+  [ testCase "retryOnSqliteBusy throws ErrorBusy after max retries" $ do
+      logMsgsVar <- newVar []
+      let logger = makeLogger logMsgsVar
+      let maxRetryCount = 1
+      let action = defaultRetryOnSqliteBusy logger maxRetryCount (\_ -> throwIO errorBusy)
+
+      result <- tryJust isErrorBusy action
+
+      case result of
+        Left exception -> do
+          exception @?= errorBusy
+          withVar logMsgsVar $ \logMsgs ->
+            length logMsgs @?= 2
+            -- uncomment if want to compare log msgs
+            -- logMsgs @?= []
+        Right _ -> assertFailure "Expected ErrorBusy exception"
+
+   , testCase "retryOnSqliteBusy doesn't throw if given function doesn't throw" $ do
+      let expected = 1 :: Int
+      let maxRetryCount = 0
+      let action = defaultRetryOnSqliteBusy noLogging maxRetryCount (\_ -> pure expected)
+
+      actual <- action
+
+      actual @?= expected
+
+   , testCase "retryOnSqliteBusy retries the number of times it should" $ do
+      countVar <- newVar 0
+      let maxRetryCount = 3
+      let hieDbAction _ = modifyVar countVar (\count -> pure (dupe (count + 1))) >> throwIO errorBusy
+      let action = defaultRetryOnSqliteBusy noLogging maxRetryCount hieDbAction
+
+      _ <- tryJust isErrorBusy action
+
+      withVar countVar $ \count ->
+        count @?= maxRetryCount + 1
+
+   , testCase "retryOnSqliteBusy doesn't retry if exception is not ErrorBusy" $ do
+      countVar <- newVar (0 :: Int)
+      let maxRetryCount = 1
+
+      let hieDbAction _ = do
+            count <- readVar countVar
+            if count == 0 then
+              evaluate (error "dummy exception")
+            else
+              modifyVar countVar (\count -> pure (dupe (count + 1)))
+
+      let action = defaultRetryOnSqliteBusy noLogging maxRetryCount hieDbAction
+
+      _ <- tryJust isErrorCall action
+
+      withVar countVar $ \count ->
+        count @?= 0
+
+    , testCase "retryOnSqliteBusy exponentially backs off" $ do
+       logMsgsVar <- newVar ([] :: [(Priority, Text)])
+
+       let maxDelay = 100
+       let baseDelay = 1
+       let maxRetryCount = 6
+       let logger = makeLogger logMsgsVar
+       let action = retryOnSqliteBusy logger undefined maxDelay baseDelay maxRetryCount rng (\_ -> throwIO errorBusy)
+
+       result <- tryJust isErrorBusy action
+
+       case result of
+         Left _ -> do
+           withVar logMsgsVar $ \logMsgs ->
+             if | ((_, lastLogMsg) : _) <- logMsgs ->
+                  -- uses log messages to indirectly check backoff...
+                  lastLogMsg @?= "Retries exhausted - base delay: 64, maximumDelay: 100, maxRetryCount: 0, exception: SQLite3 returned ErrorBusy while attempting to perform : "
+                | otherwise -> assertFailure "Expected more than 0 log messages"
+         Right _ -> assertFailure "Expected ErrorBusy exception"
+  ]
diff --git a/ghcide/test/exe/Main.hs b/ghcide/test/exe/Main.hs
index a8f4bee00a..39b1925549 100644
--- a/ghcide/test/exe/Main.hs
+++ b/ghcide/test/exe/Main.hs
@@ -16,7 +16,7 @@ module Main (main) where
 
 import           Control.Applicative.Combinators
 import           Control.Concurrent.Extra                 as Concurrent
-import           Control.Exception                        (bracket_, catch, throwIO, tryJust, ErrorCall(ErrorCall), evaluate)
+import           Control.Exception                        (bracket_, catch)
 import qualified Control.Lens                             as Lens
 import           Control.Monad
 import           Control.Monad.IO.Class                   (MonadIO, liftIO)
@@ -58,7 +58,6 @@ import           Development.IDE.Test                     (Cursor,
 import           Development.IDE.Test.Runfiles
 import qualified Development.IDE.Types.Diagnostics        as Diagnostics
 import           Development.IDE.Types.Location
-import           Development.IDE.Session                  (retryOnSqliteBusy)
 import qualified Language.LSP.Types.Lens                  as Lens (label)
 import           Development.Shake                        (getDirectoryFilesIO)
 import qualified Experiments                              as Bench
@@ -115,7 +114,7 @@ import           Test.Tasty.Ingredients.Rerun
 import           Test.Tasty.QuickCheck
 import           Text.Printf                              (printf)
 import           Text.Regex.TDFA                          ((=~))
-import           Database.SQLite.Simple                   as SQLite
+import qualified HieDbRetry
 
 -- | Wait for the next progress begin step
 waitForProgressBegin :: Session ()
@@ -184,6 +183,7 @@ main = do
     , codeActionHelperFunctionTests
     , referenceTests
     , garbageCollectionTests
+    , HieDbRetry.tests
     ]
 
 initializeResponseTests :: TestTree
@@ -6143,65 +6143,6 @@ unitTests = do
            let msg = printf "Timestamps do not have millisecond resolution: %dus" resolution_us
            assertBool msg (resolution_us <= 1000)
      , Progress.tests
-     , testCase "retryOnSqliteBusy throws ErrorBusy after max retries" $ do
-        let isErrorBusy e
-              | SQLite.SQLError { sqlError = SQLite.ErrorBusy } <- e = Just e
-              | otherwise = Nothing
-        let expected = SQLite.SQLError{ sqlError = SQLite.ErrorBusy, sqlErrorDetails = "", sqlErrorContext = "" }
-        let hieDbAction _ = throwIO expected
-        let action = retryOnSqliteBusy undefined 1 hieDbAction
-           
-        result <- tryJust isErrorBusy action
-
-        case result of
-          Left actual -> actual @?= expected
-          Right _ -> assertFailure "Expected ErrorBusy exception"
-     , testCase "retryOnSqliteBusy doesn't throw if given function doesn't throw" $ do
-        let expected = 1 :: Int
-        let action = retryOnSqliteBusy undefined 0 (\_ -> pure expected)
-
-        actual <- action
-
-        actual @?= expected
-     , testCase "retryOnSqliteBusy retries the number of times it should" $ do
-        let isErrorBusy e
-              | SQLite.SQLError { sqlError = SQLite.ErrorBusy } <- e = Just e
-              | otherwise = Nothing
-        let errorBusy = SQLite.SQLError{ sqlError = SQLite.ErrorBusy, sqlErrorDetails = "", sqlErrorContext = "" }
-        let maxRetryCount = 3
-        let expected = maxRetryCount + 1
-
-        countVar <- newVar 0
-
-        let hieDbAction _ = Concurrent.modifyVar countVar (\count -> pure (dupe (count + 1))) >> throwIO errorBusy
-        let action = retryOnSqliteBusy undefined maxRetryCount hieDbAction
-
-        _ <- tryJust isErrorBusy action
-
-        actual <- Concurrent.readVar countVar
-
-        actual @?= expected
-     , testCase "retryOnSqliteBusy doesn't retry if exception is not ErrorBusy" $ do
-        let isErrorCall e
-              | ErrorCall _ <- e = Just e
-              | otherwise = Nothing
-
-        countVar <- newVar 0
-
-        let hieDbAction _ = do
-              count <- readVar countVar
-              if count == 0 then
-                evaluate (error "dummy exception")
-              else
-                Concurrent.modifyVar countVar (\count -> pure (dupe (count + 1)))
-
-        let action = retryOnSqliteBusy undefined 1 hieDbAction
-        let expected = 0 :: Int
-
-        void $ tryJust isErrorCall action
-        actual <- readVar countVar
-
-        actual @?= expected
      ]
 
 garbageCollectionTests :: TestTree

From 4106410165fd04d4cb96fda5ff28c6b89f5fbcca Mon Sep 17 00:00:00 2001
From: Jon Shen <shenjonathan0@gmail.com>
Date: Thu, 23 Dec 2021 18:55:41 -0500
Subject: [PATCH 07/14] add random-1.2.1 to older stack.yamls

---
 stack-8.6.5.yaml | 1 +
 stack-8.8.4.yaml | 1 +
 2 files changed, 2 insertions(+)

diff --git a/stack-8.6.5.yaml b/stack-8.6.5.yaml
index bf6c913636..952742b16d 100644
--- a/stack-8.6.5.yaml
+++ b/stack-8.6.5.yaml
@@ -77,6 +77,7 @@ extra-deps:
   - ormolu-0.1.4.1
   - parser-combinators-1.2.1
   - primitive-0.7.1.0
+  - random-1.2.1
   - refinery-0.4.0.0
   - regex-base-0.94.0.0
   - regex-pcre-builtin-0.95.1.1.8.43
diff --git a/stack-8.8.4.yaml b/stack-8.8.4.yaml
index 0954312fab..1c384a4604 100644
--- a/stack-8.8.4.yaml
+++ b/stack-8.8.4.yaml
@@ -62,6 +62,7 @@ extra-deps:
   - monad-dijkstra-0.1.1.2
   - opentelemetry-0.6.1
   - opentelemetry-extra-0.6.1
+  - random-1.2.1
   - refinery-0.4.0.0
   - retrie-1.1.0.0
   - semigroups-0.18.5

From 91fdff8aad0135d9b6353be39824965327c3e7ae Mon Sep 17 00:00:00 2001
From: Jon Shen <shenjonathan0@gmail.com>
Date: Thu, 23 Dec 2021 19:11:25 -0500
Subject: [PATCH 08/14] use Random typeclass instead of uniformR, revert
 stack.yamls

---
 ghcide/session-loader/Development/IDE/Session.hs | 2 +-
 stack-8.6.5.yaml                                 | 1 -
 stack-8.8.4.yaml                                 | 1 -
 3 files changed, 1 insertion(+), 3 deletions(-)

diff --git a/ghcide/session-loader/Development/IDE/Session.hs b/ghcide/session-loader/Development/IDE/Session.hs
index bfc311a547..f6c5ae0d67 100644
--- a/ghcide/session-loader/Development/IDE/Session.hs
+++ b/ghcide/session-loader/Development/IDE/Session.hs
@@ -196,7 +196,7 @@ retryOnSqliteBusy logger hieDb maxDelay !baseDelay !maxRetryCount rng f = do
       | maxRetryCount > 0 -> do
         -- multiply by 2 because baseDelay is midpoint of uniform range
         let newBaseDelay = min maxDelay (baseDelay * 2)
-        let (delay, newRng) = Random.uniformR (0, newBaseDelay) rng
+        let (delay, newRng) = Random.randomR (0, newBaseDelay) rng
         let newMaxRetryCount = maxRetryCount - 1
         liftIO $ do
           logInfo logger $ "Retrying - " <> makeLogMsgComponentsText (Right delay) newMaxRetryCount e
diff --git a/stack-8.6.5.yaml b/stack-8.6.5.yaml
index 952742b16d..bf6c913636 100644
--- a/stack-8.6.5.yaml
+++ b/stack-8.6.5.yaml
@@ -77,7 +77,6 @@ extra-deps:
   - ormolu-0.1.4.1
   - parser-combinators-1.2.1
   - primitive-0.7.1.0
-  - random-1.2.1
   - refinery-0.4.0.0
   - regex-base-0.94.0.0
   - regex-pcre-builtin-0.95.1.1.8.43
diff --git a/stack-8.8.4.yaml b/stack-8.8.4.yaml
index 1c384a4604..0954312fab 100644
--- a/stack-8.8.4.yaml
+++ b/stack-8.8.4.yaml
@@ -62,7 +62,6 @@ extra-deps:
   - monad-dijkstra-0.1.1.2
   - opentelemetry-0.6.1
   - opentelemetry-extra-0.6.1
-  - random-1.2.1
   - refinery-0.4.0.0
   - retrie-1.1.0.0
   - semigroups-0.18.5

From 7574560ffe92bd4b8a2e51ee489235477ed7d5b6 Mon Sep 17 00:00:00 2001
From: Jon Shen <shenjonathan0@gmail.com>
Date: Fri, 24 Dec 2021 13:42:06 -0500
Subject: [PATCH 09/14] logInfo instead of logDebug

---
 ghcide/session-loader/Development/IDE/Session.hs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ghcide/session-loader/Development/IDE/Session.hs b/ghcide/session-loader/Development/IDE/Session.hs
index f6c5ae0d67..650901e68b 100644
--- a/ghcide/session-loader/Development/IDE/Session.hs
+++ b/ghcide/session-loader/Development/IDE/Session.hs
@@ -205,7 +205,7 @@ retryOnSqliteBusy logger hieDb maxDelay !baseDelay !maxRetryCount rng f = do
 
       | otherwise -> do
         liftIO $ do
-          logDebug logger $ "Retries exhausted - " <> makeLogMsgComponentsText (Left baseDelay) maxRetryCount e
+          logInfo logger $ "Retries exhausted - " <> makeLogMsgComponentsText (Left baseDelay) maxRetryCount e
           throwIO e
 
     Right b -> pure b

From c3fd6491b17e67ceb7788e1ade03d8255610eda0 Mon Sep 17 00:00:00 2001
From: Jon Shen <shenjonathan0@gmail.com>
Date: Fri, 24 Dec 2021 14:49:09 -0500
Subject: [PATCH 10/14] dont wrap action with hiedb retry on the writer side,
 give hiedb retry wrapper to action instead

---
 ghcide/session-loader/Development/IDE/Session.hs | 2 +-
 ghcide/src/Development/IDE/Core/Compile.hs       | 4 ++--
 ghcide/src/Development/IDE/Core/FileStore.hs     | 2 +-
 ghcide/src/Development/IDE/Core/Shake.hs         | 4 +++-
 4 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/ghcide/session-loader/Development/IDE/Session.hs b/ghcide/session-loader/Development/IDE/Session.hs
index 650901e68b..fc9c24310b 100644
--- a/ghcide/session-loader/Development/IDE/Session.hs
+++ b/ghcide/session-loader/Development/IDE/Session.hs
@@ -253,7 +253,7 @@ runWithDb logger fp k = do
       _ <- garbageCollectTypeNames db
       forever $ do
         k <- atomically $ readTQueue chan
-        retryOnSqliteBusy logger db oneSecond oneMillisecond maxRetryCount rng k
+        k (retryOnSqliteBusy logger db oneSecond oneMillisecond maxRetryCount rng)
           `Safe.catch` \e@SQLError{} -> do
             logDebug logger $ T.pack $ "SQLite error in worker, ignoring: " ++ show e
           `Safe.catchAny` \e -> do
diff --git a/ghcide/src/Development/IDE/Core/Compile.hs b/ghcide/src/Development/IDE/Core/Compile.hs
index 7b002f08fa..a1d5244038 100644
--- a/ghcide/src/Development/IDE/Core/Compile.hs
+++ b/ghcide/src/Development/IDE/Core/Compile.hs
@@ -521,7 +521,7 @@ indexHieFile se mod_summary srcPath !hash hf = do
       -- hiedb doesn't use the Haskell src, so we clear it to avoid unnecessarily keeping it around
       let !hf' = hf{hie_hs_src = mempty}
       modifyTVar' indexPending $ HashMap.insert srcPath hash
-      writeTQueue indexQueue $ \db -> do
+      writeTQueue indexQueue $ \withHieDb -> do
         -- We are now in the worker thread
         -- Check if a newer index of this file has been scheduled, and if so skip this one
         newerScheduled <- atomically $ do
@@ -532,7 +532,7 @@ indexHieFile se mod_summary srcPath !hash hf = do
             Just pendingHash -> pendingHash /= hash
         unless newerScheduled $ do
           pre optProgressStyle
-          addRefsFromLoaded db targetPath (RealFile $ fromNormalizedFilePath srcPath) hash hf'
+          withHieDb (\db -> HieDb.addRefsFromLoaded db targetPath (HieDb.RealFile $ fromNormalizedFilePath srcPath) hash hf')
           post
   where
     mod_location    = ms_location mod_summary
diff --git a/ghcide/src/Development/IDE/Core/FileStore.hs b/ghcide/src/Development/IDE/Core/FileStore.hs
index 9a37663f37..f6f93d3c02 100644
--- a/ghcide/src/Development/IDE/Core/FileStore.hs
+++ b/ghcide/src/Development/IDE/Core/FileStore.hs
@@ -294,7 +294,7 @@ setSomethingModified state keys reason = do
         fail "setSomethingModified can't be called on this type of VFSHandle"
     -- Update database to remove any files that might have been renamed/deleted
     atomically $ do
-        writeTQueue (indexQueue $ hiedbWriter $ shakeExtras state) deleteMissingRealFiles
+        writeTQueue (indexQueue $ hiedbWriter $ shakeExtras state) (\withHieDb -> withHieDb deleteMissingRealFiles)
         modifyTVar' (dirtyKeys $ shakeExtras state) $ \x ->
             foldl' (flip HSet.insert) x keys
     void $ restartShakeSession (shakeExtras state) reason []
diff --git a/ghcide/src/Development/IDE/Core/Shake.hs b/ghcide/src/Development/IDE/Core/Shake.hs
index 0f9135bed4..8938945a61 100644
--- a/ghcide/src/Development/IDE/Core/Shake.hs
+++ b/ghcide/src/Development/IDE/Core/Shake.hs
@@ -183,7 +183,9 @@ data HieDbWriter
   }
 
 -- | Actions to queue up on the index worker thread
-type IndexQueue = TQueue (HieDb -> IO ())
+-- The inner `(HieDb -> IO ()) -> IO ()` wraps `HieDb -> IO ()`
+-- with (currently) retry functionality
+type IndexQueue = TQueue (((HieDb -> IO ()) -> IO ()) -> IO ())
 
 -- | Intended to represent HieDb calls wrapped with (currently) retry
 -- functionality

From 1bc3adc80e6676adbf5ab3521a973de54b00b351 Mon Sep 17 00:00:00 2001
From: Jon Shen <shenjonathan0@gmail.com>
Date: Sun, 26 Dec 2021 21:41:35 -0500
Subject: [PATCH 11/14] bump log priorities to warning, wrap all hiedb calls in
 runWithDb and writerThread with retries, promote time duration and
 maxRetryCount constants to top level

---
 .../session-loader/Development/IDE/Session.hs | 47 +++++++++++++------
 1 file changed, 32 insertions(+), 15 deletions(-)

diff --git a/ghcide/session-loader/Development/IDE/Session.hs b/ghcide/session-loader/Development/IDE/Session.hs
index fc9c24310b..52e6927691 100644
--- a/ghcide/session-loader/Development/IDE/Session.hs
+++ b/ghcide/session-loader/Development/IDE/Session.hs
@@ -199,13 +199,13 @@ retryOnSqliteBusy logger hieDb maxDelay !baseDelay !maxRetryCount rng f = do
         let (delay, newRng) = Random.randomR (0, newBaseDelay) rng
         let newMaxRetryCount = maxRetryCount - 1
         liftIO $ do
-          logInfo logger $ "Retrying - " <> makeLogMsgComponentsText (Right delay) newMaxRetryCount e
+          logWarning logger $ "Retrying - " <> makeLogMsgComponentsText (Right delay) newMaxRetryCount e
           threadDelay delay
         retryOnSqliteBusy logger hieDb maxDelay newBaseDelay newMaxRetryCount newRng f
 
       | otherwise -> do
         liftIO $ do
-          logInfo logger $ "Retries exhausted - " <> makeLogMsgComponentsText (Left baseDelay) maxRetryCount e
+          logWarning logger $ "Retries exhausted - " <> makeLogMsgComponentsText (Left baseDelay) maxRetryCount e
           throwIO e
 
     Right b -> pure b
@@ -224,7 +224,17 @@ retryOnSqliteBusy logger hieDb maxDelay !baseDelay !maxRetryCount rng f = do
       in
         T.intercalate ", " logMsgComponents
 
+-- | in microseconds
+oneSecond :: Int
+oneSecond = 1000000
 
+-- | in microseconds
+oneMillisecond :: Int
+oneMillisecond = 1000
+
+-- | default maximum number of times to retry hiedb call
+maxRetryCount :: Int
+maxRetryCount = 10
 
 -- | Wraps `withHieDb` to provide a database connection for reading, and a `HieWriterChan` for
 -- writing. Actions are picked off one by one from the `HieWriterChan` and executed in serial
@@ -232,28 +242,35 @@ retryOnSqliteBusy logger hieDb maxDelay !baseDelay !maxRetryCount rng f = do
 -- This is done in order to serialize writes to the database, or else SQLite becomes unhappy
 runWithDb :: Logger -> FilePath -> (WithHieDb -> IndexQueue -> IO ()) -> IO ()
 runWithDb logger fp k = do
+  -- use non-deterministic seed because maybe multiple HLS start at same time
+  -- and send bursts of requests
+  rng <- Random.newStdGen
   -- Delete the database if it has an incompatible schema version
-  withHieDb fp (const $ pure ())
+  withHieDb fp (const (pure ()) . makeWithRetryableHieDb rng)
     `Safe.catch` \IncompatibleSchemaVersion{} -> removeFile fp
   withHieDb fp $ \writedb -> do
-    initConn writedb
+    -- the type signature is necessary to avoid concretizing the RankNType
+    -- e.g. using it with initConn will set tyvar a to ()
+    let withRetryableWriteDb :: WithHieDb
+        withRetryableWriteDb = makeWithRetryableHieDb rng writedb
+    withRetryableWriteDb initConn
+
     chan <- newTQueueIO
-    -- use newStdGen because what if multiple HLS start at same time and send bursts of requests
-    rng <- Random.newStdGen
-    withAsync (writerThread writedb chan rng) $ \_ -> do
-      withHieDb fp (\readDb -> k (retryOnSqliteBusy logger readDb oneSecond oneMillisecond maxRetryCount rng) chan)
+
+    withAsync (writerThread withRetryableWriteDb chan) $ \_ -> do
+      withHieDb fp (\readDb -> k (makeWithRetryableHieDb rng readDb) chan)
   where
-    oneSecond = 1000000
-    oneMillisecond = 1000
-    maxRetryCount = 10
+    makeWithRetryableHieDb :: RandomGen g => g -> HieDb -> WithHieDb
+    makeWithRetryableHieDb rng hieDb = retryOnSqliteBusy logger hieDb oneSecond oneMillisecond maxRetryCount rng
 
-    writerThread db chan rng = do
+    writerThread :: WithHieDb -> IndexQueue -> IO ()
+    writerThread withRetryableWriteDb chan = do
       -- Clear the index of any files that might have been deleted since the last run
-      deleteMissingRealFiles db
-      _ <- garbageCollectTypeNames db
+      _ <- withRetryableWriteDb deleteMissingRealFiles
+      _ <- withRetryableWriteDb garbageCollectTypeNames
       forever $ do
         k <- atomically $ readTQueue chan
-        k (retryOnSqliteBusy logger db oneSecond oneMillisecond maxRetryCount rng)
+        k withRetryableWriteDb
           `Safe.catch` \e@SQLError{} -> do
             logDebug logger $ T.pack $ "SQLite error in worker, ignoring: " ++ show e
           `Safe.catchAny` \e -> do

From 3296e8ba66207f5ceb3ee143c5ed6b37a5a4049f Mon Sep 17 00:00:00 2001
From: Jon Shen <shenjonathan0@gmail.com>
Date: Sun, 26 Dec 2021 22:43:13 -0500
Subject: [PATCH 12/14] fix ghc 9.0.1

---
 ghcide/session-loader/Development/IDE/Session.hs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ghcide/session-loader/Development/IDE/Session.hs b/ghcide/session-loader/Development/IDE/Session.hs
index 52e6927691..a7c6547b41 100644
--- a/ghcide/session-loader/Development/IDE/Session.hs
+++ b/ghcide/session-loader/Development/IDE/Session.hs
@@ -246,7 +246,7 @@ runWithDb logger fp k = do
   -- and send bursts of requests
   rng <- Random.newStdGen
   -- Delete the database if it has an incompatible schema version
-  withHieDb fp (const (pure ()) . makeWithRetryableHieDb rng)
+  withHieDb fp (\hieDb -> makeWithRetryableHieDb rng hieDb $ const (pure ()))
     `Safe.catch` \IncompatibleSchemaVersion{} -> removeFile fp
   withHieDb fp $ \writedb -> do
     -- the type signature is necessary to avoid concretizing the RankNType

From 6acc965907f6e633c394081dc5a85b7c0b02c059 Mon Sep 17 00:00:00 2001
From: Jon Shen <shenjonathan0@gmail.com>
Date: Mon, 27 Dec 2021 05:25:57 -0500
Subject: [PATCH 13/14] refactor retryOnSqliteBusy into retryOnException et al,
 wrap Hiedb.runCommand with retry, fix tests

---
 .../session-loader/Development/IDE/Session.hs | 85 +++++++++++--------
 ghcide/src/Development/IDE/Main.hs            |  5 +-
 ghcide/test/exe/HieDbRetry.hs                 | 52 +++++++-----
 3 files changed, 83 insertions(+), 59 deletions(-)

diff --git a/ghcide/session-loader/Development/IDE/Session.hs b/ghcide/session-loader/Development/IDE/Session.hs
index a7c6547b41..33e51f11d1 100644
--- a/ghcide/session-loader/Development/IDE/Session.hs
+++ b/ghcide/session-loader/Development/IDE/Session.hs
@@ -13,6 +13,7 @@ module Development.IDE.Session
   ,getHieDbLoc
   ,runWithDb
   ,retryOnSqliteBusy
+  ,retryOnException
   ) where
 
 -- Unfortunately, we cannot use loadSession with ghc-lib since hie-bios uses
@@ -169,28 +170,24 @@ setInitialDynFlags logger rootDir SessionLoadingOptions{..} = do
   mapM_ setUnsafeGlobalDynFlags dynFlags
   pure libdir
 
-
--- | If the sqlite db returns an SQLITE_BUSY then we sleep for a duration
--- determined by the random exponential backoff formula,
--- `uniformRandom(0, min (maxDelay, (baseDelay * 2) ^ retryAttempt))`, and try the
--- action again for a maximum of `maxRetryCount` times.
+-- | If the action throws exception that satisfies predicate then we sleep for
+-- a duration determined by the random exponential backoff formula,
+-- `uniformRandom(0, min (maxDelay, (baseDelay * 2) ^ retryAttempt))`, and try
+-- the action again for a maximum of `maxRetryCount` times.
 -- `MonadIO`, `MonadCatch` are used as constraints because there are a few
 -- HieDb functions that don't return IO values.
-retryOnSqliteBusy :: (MonadIO m, MonadCatch m, RandomGen g)
-                  => Logger
-                  -> HieDb -- ^ HieDb connection
-                  -> Int -- ^ maximum backoff delay in microseconds
-                  -> Int -- ^ base backoff delay in microseconds
-                  -> Int -- ^ maximum number of times to retry
-                  -> g -- ^ random number generator
-                  -> (HieDb -> m b) -- ^ function that uses HieDb connection
-                  -> m b
-retryOnSqliteBusy logger hieDb maxDelay !baseDelay !maxRetryCount rng f = do
-  let hieDbAction = f hieDb
-  let isErrorBusy e
-        | SQLError{ sqlError = ErrorBusy } <- e = Just e
-        | otherwise = Nothing
-  result <- tryJust isErrorBusy hieDbAction
+retryOnException
+  :: (MonadIO m, MonadCatch m, RandomGen g, Exception e)
+  => (e -> Maybe e) -- ^ only retry on exception if this predicate returns Just
+  -> Logger
+  -> Int -- ^ maximum backoff delay in microseconds
+  -> Int -- ^ base backoff delay in microseconds
+  -> Int -- ^ maximum number of times to retry
+  -> g -- ^ random number generator
+  -> m a -- ^ action that may throw exception
+  -> m a
+retryOnException exceptionPred logger maxDelay !baseDelay !maxRetryCount rng action = do
+  result <- tryJust exceptionPred action
   case result of
     Left e
       | maxRetryCount > 0 -> do
@@ -201,7 +198,7 @@ retryOnSqliteBusy logger hieDb maxDelay !baseDelay !maxRetryCount rng f = do
         liftIO $ do
           logWarning logger $ "Retrying - " <> makeLogMsgComponentsText (Right delay) newMaxRetryCount e
           threadDelay delay
-        retryOnSqliteBusy logger hieDb maxDelay newBaseDelay newMaxRetryCount newRng f
+        retryOnException exceptionPred logger maxDelay newBaseDelay newMaxRetryCount newRng action
 
       | otherwise -> do
         liftIO $ do
@@ -236,6 +233,19 @@ oneMillisecond = 1000
 maxRetryCount :: Int
 maxRetryCount = 10
 
+retryOnSqliteBusy :: (MonadIO m, MonadCatch m, RandomGen g)
+                  => Logger -> g -> m a -> m a
+retryOnSqliteBusy logger rng action =
+  let isErrorBusy e
+        | SQLError{ sqlError = ErrorBusy } <- e = Just e
+        | otherwise = Nothing
+  in
+    retryOnException isErrorBusy logger oneSecond oneMillisecond maxRetryCount rng action
+
+makeWithHieDbRetryable :: RandomGen g => Logger -> g -> HieDb -> WithHieDb
+makeWithHieDbRetryable logger rng hieDb f =
+  retryOnSqliteBusy logger rng (f hieDb)
+
 -- | Wraps `withHieDb` to provide a database connection for reading, and a `HieWriterChan` for
 -- writing. Actions are picked off one by one from the `HieWriterChan` and executed in serial
 -- by a worker thread using a dedicated database connection.
@@ -246,31 +256,32 @@ runWithDb logger fp k = do
   -- and send bursts of requests
   rng <- Random.newStdGen
   -- Delete the database if it has an incompatible schema version
-  withHieDb fp (\hieDb -> makeWithRetryableHieDb rng hieDb $ const (pure ()))
-    `Safe.catch` \IncompatibleSchemaVersion{} -> removeFile fp
+  retryOnSqliteBusy
+    logger
+    rng
+    (withHieDb fp (const $ pure ()) `Safe.catch` \IncompatibleSchemaVersion{} -> removeFile fp)
+
   withHieDb fp $ \writedb -> do
-    -- the type signature is necessary to avoid concretizing the RankNType
-    -- e.g. using it with initConn will set tyvar a to ()
-    let withRetryableWriteDb :: WithHieDb
-        withRetryableWriteDb = makeWithRetryableHieDb rng writedb
-    withRetryableWriteDb initConn
+    -- the type signature is necessary to avoid concretizing the tyvar
+    -- e.g. `withWriteDbRetrable initConn` without type signature will
+    -- instantiate tyvar `a` to `()`
+    let withWriteDbRetryable :: WithHieDb
+        withWriteDbRetryable = makeWithHieDbRetryable logger rng writedb
+    withWriteDbRetryable initConn
 
     chan <- newTQueueIO
 
-    withAsync (writerThread withRetryableWriteDb chan) $ \_ -> do
-      withHieDb fp (\readDb -> k (makeWithRetryableHieDb rng readDb) chan)
+    withAsync (writerThread withWriteDbRetryable chan) $ \_ -> do
+      withHieDb fp (\readDb -> k (makeWithHieDbRetryable logger rng readDb) chan)
   where
-    makeWithRetryableHieDb :: RandomGen g => g -> HieDb -> WithHieDb
-    makeWithRetryableHieDb rng hieDb = retryOnSqliteBusy logger hieDb oneSecond oneMillisecond maxRetryCount rng
-
     writerThread :: WithHieDb -> IndexQueue -> IO ()
-    writerThread withRetryableWriteDb chan = do
+    writerThread withHieDbRetryable chan = do
       -- Clear the index of any files that might have been deleted since the last run
-      _ <- withRetryableWriteDb deleteMissingRealFiles
-      _ <- withRetryableWriteDb garbageCollectTypeNames
+      _ <- withHieDbRetryable deleteMissingRealFiles
+      _ <- withHieDbRetryable garbageCollectTypeNames
       forever $ do
         k <- atomically $ readTQueue chan
-        k withRetryableWriteDb
+        k withHieDbRetryable
           `Safe.catch` \e@SQLError{} -> do
             logDebug logger $ T.pack $ "SQLite error in worker, ignoring: " ++ show e
           `Safe.catchAny` \e -> do
diff --git a/ghcide/src/Development/IDE/Main.hs b/ghcide/src/Development/IDE/Main.hs
index 86a6007100..e708670810 100644
--- a/ghcide/src/Development/IDE/Main.hs
+++ b/ghcide/src/Development/IDE/Main.hs
@@ -65,6 +65,7 @@ import qualified Development.IDE.Plugin.Test           as Test
 import           Development.IDE.Session               (SessionLoadingOptions,
                                                         getHieDbLoc,
                                                         loadSessionWithOptions,
+                                                        retryOnSqliteBusy,
                                                         runWithDb,
                                                         setInitialDynFlags)
 import           Development.IDE.Types.Location        (NormalizedUri,
@@ -116,6 +117,7 @@ import           System.IO                             (BufferMode (LineBufferin
                                                         hSetBuffering,
                                                         hSetEncoding, stderr,
                                                         stdin, stdout, utf8)
+import           System.Random                         (newStdGen)
 import           System.Time.Extra                     (offsetTime,
                                                         showDuration)
 import           Text.Printf                           (printf)
@@ -384,9 +386,10 @@ defaultMain Arguments{..} = flip withHeapStats fun =<< argsLogger
             dbLoc <- getHieDbLoc dir
             hPutStrLn stderr $ "Using hiedb at: " ++ dbLoc
             mlibdir <- setInitialDynFlags logger dir def
+            rng <- newStdGen
             case mlibdir of
                 Nothing     -> exitWith $ ExitFailure 1
-                Just libdir -> HieDb.runCommand libdir opts{HieDb.database = dbLoc} cmd
+                Just libdir -> retryOnSqliteBusy logger rng (HieDb.runCommand libdir opts{HieDb.database = dbLoc} cmd)
 
         Custom projectRoot (IdeCommand c) -> do
           dbLoc <- getHieDbLoc projectRoot
diff --git a/ghcide/test/exe/HieDbRetry.hs b/ghcide/test/exe/HieDbRetry.hs
index 64aac586f7..f3a29cea39 100644
--- a/ghcide/test/exe/HieDbRetry.hs
+++ b/ghcide/test/exe/HieDbRetry.hs
@@ -8,8 +8,8 @@ import           Control.Exception            (ErrorCall (ErrorCall), evaluate,
 import           Data.Text                    (Text)
 import           Data.Tuple.Extra             (dupe)
 import qualified Database.SQLite.Simple       as SQLite
-import           Development.IDE.Core.Shake   (HieDb)
-import           Development.IDE.Session      (retryOnSqliteBusy)
+import           Development.IDE.Session      (retryOnException,
+                                               retryOnSqliteBusy)
 import           Development.IDE.Types.Logger (Logger (Logger), Priority,
                                                noLogging)
 import qualified System.Random                as Random
@@ -22,8 +22,8 @@ makeLogger msgsVar = Logger $ \priority msg -> modifyVar msgsVar (\msgs -> pure
 rng :: Random.StdGen
 rng = Random.mkStdGen 0
 
-defaultRetryOnSqliteBusy :: Logger -> Int -> (HieDb -> IO b) -> IO b
-defaultRetryOnSqliteBusy logger maxRetryCount = retryOnSqliteBusy logger undefined 1 1 maxRetryCount rng
+retryOnSqliteBusyForTest :: Logger -> Int -> IO a -> IO a
+retryOnSqliteBusyForTest logger maxRetryCount = retryOnException isErrorBusy logger 1 1 maxRetryCount rng
 
 isErrorBusy :: SQLite.SQLError -> Maybe SQLite.SQLError
 isErrorBusy e
@@ -40,13 +40,12 @@ isErrorCall e
 
 tests :: TestTree
 tests = testGroup "RetryHieDb"
-  [ testCase "retryOnSqliteBusy throws ErrorBusy after max retries" $ do
+  [ testCase "retryOnException throws exception after max retries" $ do
       logMsgsVar <- newVar []
       let logger = makeLogger logMsgsVar
       let maxRetryCount = 1
-      let action = defaultRetryOnSqliteBusy logger maxRetryCount (\_ -> throwIO errorBusy)
 
-      result <- tryJust isErrorBusy action
+      result <- tryJust isErrorBusy (retryOnSqliteBusyForTest logger maxRetryCount (throwIO errorBusy))
 
       case result of
         Left exception -> do
@@ -57,54 +56,65 @@ tests = testGroup "RetryHieDb"
             -- logMsgs @?= []
         Right _ -> assertFailure "Expected ErrorBusy exception"
 
-   , testCase "retryOnSqliteBusy doesn't throw if given function doesn't throw" $ do
+   , testCase "retryOnException doesn't throw if given function doesn't throw" $ do
       let expected = 1 :: Int
       let maxRetryCount = 0
-      let action = defaultRetryOnSqliteBusy noLogging maxRetryCount (\_ -> pure expected)
 
-      actual <- action
+      actual <- retryOnSqliteBusyForTest noLogging maxRetryCount (pure expected)
 
       actual @?= expected
 
-   , testCase "retryOnSqliteBusy retries the number of times it should" $ do
+   , testCase "retryOnException retries the number of times it should" $ do
       countVar <- newVar 0
       let maxRetryCount = 3
-      let hieDbAction _ = modifyVar countVar (\count -> pure (dupe (count + 1))) >> throwIO errorBusy
-      let action = defaultRetryOnSqliteBusy noLogging maxRetryCount hieDbAction
+      let incrementThenThrow = modifyVar countVar (\count -> pure (dupe (count + 1))) >> throwIO errorBusy
 
-      _ <- tryJust isErrorBusy action
+      _ <- tryJust isErrorBusy (retryOnSqliteBusyForTest noLogging maxRetryCount incrementThenThrow)
 
       withVar countVar $ \count ->
         count @?= maxRetryCount + 1
 
-   , testCase "retryOnSqliteBusy doesn't retry if exception is not ErrorBusy" $ do
+   , testCase "retryOnException doesn't retry if exception is not ErrorBusy" $ do
       countVar <- newVar (0 :: Int)
       let maxRetryCount = 1
 
-      let hieDbAction _ = do
+      let throwThenIncrement = do
             count <- readVar countVar
             if count == 0 then
               evaluate (error "dummy exception")
             else
               modifyVar countVar (\count -> pure (dupe (count + 1)))
 
-      let action = defaultRetryOnSqliteBusy noLogging maxRetryCount hieDbAction
 
-      _ <- tryJust isErrorCall action
+      _ <- tryJust isErrorCall (retryOnSqliteBusyForTest noLogging maxRetryCount throwThenIncrement)
 
       withVar countVar $ \count ->
         count @?= 0
 
-    , testCase "retryOnSqliteBusy exponentially backs off" $ do
+   , testCase "retryOnSqliteBusy retries on ErrorBusy" $ do
+      countVar <- newVar (0 :: Int)
+
+      let incrementThenThrowThenIncrement = do
+            count <- readVar countVar
+            if count == 0 then
+              modifyVar countVar (\count -> pure (dupe (count + 1))) >> throwIO errorBusy
+            else
+              modifyVar countVar (\count -> pure (dupe (count + 1)))
+
+      _ <- retryOnSqliteBusy noLogging rng incrementThenThrowThenIncrement
+
+      withVar countVar $ \count ->
+        count @?= 2
+
+    , testCase "retryOnException exponentially backs off" $ do
        logMsgsVar <- newVar ([] :: [(Priority, Text)])
 
        let maxDelay = 100
        let baseDelay = 1
        let maxRetryCount = 6
        let logger = makeLogger logMsgsVar
-       let action = retryOnSqliteBusy logger undefined maxDelay baseDelay maxRetryCount rng (\_ -> throwIO errorBusy)
 
-       result <- tryJust isErrorBusy action
+       result <- tryJust isErrorBusy (retryOnException isErrorBusy logger maxDelay baseDelay maxRetryCount rng (throwIO errorBusy))
 
        case result of
          Left _ -> do

From db72aebfa7079aa8afb18c7675d83e1a0b79f373 Mon Sep 17 00:00:00 2001
From: Jon Shen <shenjonathan0@gmail.com>
Date: Mon, 27 Dec 2021 05:48:30 -0500
Subject: [PATCH 14/14] push WithHieDb into createExportsMapHieDb to
 potentially retry less stuff, move WithHieDb to Types.Shake to avoid circular
 dependency

---
 ghcide/session-loader/Development/IDE/Session.hs |  1 +
 ghcide/src/Development/IDE/Core/Service.hs       |  1 +
 ghcide/src/Development/IDE/Core/Shake.hs         |  7 +------
 ghcide/src/Development/IDE/LSP/LanguageServer.hs |  1 +
 ghcide/src/Development/IDE/Spans/AtPoint.hs      |  2 +-
 ghcide/src/Development/IDE/Types/Exports.hs      | 10 ++++++----
 ghcide/src/Development/IDE/Types/Shake.hs        |  8 +++++++-
 7 files changed, 18 insertions(+), 12 deletions(-)

diff --git a/ghcide/session-loader/Development/IDE/Session.hs b/ghcide/session-loader/Development/IDE/Session.hs
index 33e51f11d1..bf93c060cc 100644
--- a/ghcide/session-loader/Development/IDE/Session.hs
+++ b/ghcide/session-loader/Development/IDE/Session.hs
@@ -85,6 +85,7 @@ import           Data.Foldable                        (for_)
 import qualified Data.HashSet                         as Set
 import           Database.SQLite.Simple
 import           Development.IDE.Core.Tracing         (withTrace)
+import           Development.IDE.Types.Shake          (WithHieDb)
 import           HieDb.Create
 import           HieDb.Types
 import           HieDb.Utils
diff --git a/ghcide/src/Development/IDE/Core/Service.hs b/ghcide/src/Development/IDE/Core/Service.hs
index b43a5cada2..a98f80bfb4 100644
--- a/ghcide/src/Development/IDE/Core/Service.hs
+++ b/ghcide/src/Development/IDE/Core/Service.hs
@@ -30,6 +30,7 @@ import qualified Language.LSP.Types              as LSP
 
 import           Control.Monad
 import           Development.IDE.Core.Shake
+import           Development.IDE.Types.Shake     (WithHieDb)
 import           System.Environment              (lookupEnv)
 
 
diff --git a/ghcide/src/Development/IDE/Core/Shake.hs b/ghcide/src/Development/IDE/Core/Shake.hs
index 8938945a61..86df150faa 100644
--- a/ghcide/src/Development/IDE/Core/Shake.hs
+++ b/ghcide/src/Development/IDE/Core/Shake.hs
@@ -73,7 +73,6 @@ module Development.IDE.Core.Shake(
     IndexQueue,
     HieDb,
     HieDbWriter(..),
-    WithHieDb,
     VFSHandle(..),
     addPersistentRule,
     garbageCollectDirtyKeys,
@@ -187,10 +186,6 @@ data HieDbWriter
 -- with (currently) retry functionality
 type IndexQueue = TQueue (((HieDb -> IO ()) -> IO ()) -> IO ())
 
--- | Intended to represent HieDb calls wrapped with (currently) retry
--- functionality
-type WithHieDb = forall a. (HieDb -> IO a) -> IO a
-
 -- information we stash inside the shakeExtra field
 data ShakeExtras = ShakeExtras
     { --eventer :: LSP.FromServerMessage -> IO ()
@@ -535,7 +530,7 @@ shakeOpen lspEnv defaultConfig logger debouncer
         -- lazily initialize the exports map with the contents of the hiedb
         _ <- async $ do
             logDebug logger "Initializing exports map from hiedb"
-            em <- withHieDb createExportsMapHieDb
+            em <- createExportsMapHieDb withHieDb
             atomically $ modifyTVar' exportsMap (<> em)
             logDebug logger $ "Done initializing exports map from hiedb (" <> pack(show (ExportsMap.size em)) <> ")"
 
diff --git a/ghcide/src/Development/IDE/LSP/LanguageServer.hs b/ghcide/src/Development/IDE/LSP/LanguageServer.hs
index 292de72a5f..f3e4f4d9e8 100644
--- a/ghcide/src/Development/IDE/LSP/LanguageServer.hs
+++ b/ghcide/src/Development/IDE/LSP/LanguageServer.hs
@@ -39,6 +39,7 @@ import           Development.IDE.LSP.HoverDefinition
 import           Development.IDE.Types.Logger
 
 import           Control.Monad.IO.Unlift               (MonadUnliftIO)
+import           Development.IDE.Types.Shake           (WithHieDb)
 import           System.IO.Unsafe                      (unsafeInterleaveIO)
 
 issueTrackerUrl :: T.Text
diff --git a/ghcide/src/Development/IDE/Spans/AtPoint.hs b/ghcide/src/Development/IDE/Spans/AtPoint.hs
index 060cac4019..2f1a359d81 100644
--- a/ghcide/src/Development/IDE/Spans/AtPoint.hs
+++ b/ghcide/src/Development/IDE/Spans/AtPoint.hs
@@ -52,7 +52,7 @@ import           Data.List                            (isSuffixOf)
 import           Data.List.Extra                      (dropEnd1, nubOrd)
 
 import           Data.Version                         (showVersion)
-import           Development.IDE.Core.Shake           (WithHieDb)
+import           Development.IDE.Types.Shake          (WithHieDb)
 import           HieDb                                hiding (pointCommand)
 import           System.Directory                     (doesFileExist)
 
diff --git a/ghcide/src/Development/IDE/Types/Exports.hs b/ghcide/src/Development/IDE/Types/Exports.hs
index bb08bb416e..becd09a6b2 100644
--- a/ghcide/src/Development/IDE/Types/Exports.hs
+++ b/ghcide/src/Development/IDE/Types/Exports.hs
@@ -1,5 +1,6 @@
 {-# LANGUAGE DeriveAnyClass     #-}
 {-# LANGUAGE DerivingStrategies #-}
+{-# LANGUAGE RankNTypes         #-}
 module Development.IDE.Types.Exports
 (
     IdentInfo(..),
@@ -26,6 +27,7 @@ import           Data.Text                   (Text, pack)
 import           Development.IDE.GHC.Compat
 import           Development.IDE.GHC.Orphans ()
 import           Development.IDE.GHC.Util
+import           Development.IDE.Types.Shake (WithHieDb)
 import           GHC.Generics                (Generic)
 import           HieDb
 
@@ -154,13 +156,13 @@ createExportsMapTc modIface = do
 nonInternalModules :: ModuleName -> Bool
 nonInternalModules = not . (".Internal" `isSuffixOf`) . moduleNameString
 
-createExportsMapHieDb :: HieDb -> IO ExportsMap
-createExportsMapHieDb hiedb = do
-    mods <- getAllIndexedMods hiedb
+createExportsMapHieDb :: WithHieDb -> IO ExportsMap
+createExportsMapHieDb withHieDb = do
+    mods <- withHieDb getAllIndexedMods
     idents <- forM (filter (nonInternalModules . modInfoName . hieModInfo) mods) $ \m -> do
         let mn = modInfoName $ hieModInfo m
             mText = pack $ moduleNameString mn
-        fmap (wrap . unwrap mText) <$> getExportsForModule hiedb mn
+        fmap (wrap . unwrap mText) <$> withHieDb (\hieDb -> getExportsForModule hieDb mn)
     let exportsMap = Map.fromListWith (<>) (concat idents)
     return $ ExportsMap exportsMap $ buildModuleExportMap (concat idents)
   where
diff --git a/ghcide/src/Development/IDE/Types/Shake.hs b/ghcide/src/Development/IDE/Types/Shake.hs
index 5df0de9a76..32a9959991 100644
--- a/ghcide/src/Development/IDE/Types/Shake.hs
+++ b/ghcide/src/Development/IDE/Types/Shake.hs
@@ -1,6 +1,7 @@
 {-# LANGUAGE DerivingStrategies        #-}
 {-# LANGUAGE ExistentialQuantification #-}
 {-# LANGUAGE PatternSynonyms           #-}
+{-# LANGUAGE RankNTypes                #-}
 {-# LANGUAGE TypeFamilies              #-}
 module Development.IDE.Types.Shake
   ( Q (..),
@@ -13,7 +14,7 @@ module Development.IDE.Types.Shake
     ShakeValue(..),
     currentValue,
     isBadDependency,
-  toShakeValue,encodeShakeValue,decodeShakeValue,toKey,toNoFileKey,fromKey,fromKeyType)
+  toShakeValue,encodeShakeValue,decodeShakeValue,toKey,toNoFileKey,fromKey,fromKeyType,WithHieDb)
 where
 
 import           Control.DeepSeq
@@ -29,6 +30,7 @@ import qualified Development.IDE.Graph                as Shake
 import           Development.IDE.Types.Diagnostics
 import           Development.IDE.Types.Location
 import           GHC.Generics
+import           HieDb.Types                          (HieDb)
 import           Language.LSP.Types
 import qualified StmContainers.Map                    as STM
 import           Type.Reflection                      (SomeTypeRep (SomeTypeRep),
@@ -37,6 +39,10 @@ import           Type.Reflection                      (SomeTypeRep (SomeTypeRep)
                                                        typeRepTyCon)
 import           Unsafe.Coerce                        (unsafeCoerce)
 
+-- | Intended to represent HieDb calls wrapped with (currently) retry
+-- functionality
+type WithHieDb = forall a. (HieDb -> IO a) -> IO a
+
 data Value v
     = Succeeded TextDocumentVersion v
     | Stale (Maybe PositionDelta) TextDocumentVersion v