module Unison.Codebase.Watch
  ( watchDirectory,
  )
where

import Control.Concurrent (threadDelay)
import Control.Concurrent.STM qualified as STM
import Control.Exception (MaskingState (..))
import Data.IORef (newIORef, readIORef, writeIORef)
import Data.Map qualified as Map
import Data.Time.Clock (UTCTime, diffUTCTime)
import GHC.Conc (registerDelay)
import GHC.IO (unsafeUnmask)
import Ki qualified
import System.FSNotify (Event (Added, Modified))
import System.FSNotify qualified as FSNotify
import Unison.Prelude
import UnliftIO.Exception (finally, tryAny)
import UnliftIO.STM (atomically)

watchDirectory :: Ki.Scope -> FSNotify.WatchManager -> FilePath -> (FilePath -> Bool) -> IO (IO (FilePath, Text))
watchDirectory :: Scope
-> WatchManager
-> FilePath
-> (FilePath -> Bool)
-> IO (IO (FilePath, Text))
watchDirectory Scope
scope WatchManager
mgr FilePath
dir FilePath -> Bool
allow = do
  TQueue (FilePath, UTCTime)
eventQueue <- Scope
-> WatchManager
-> FilePath
-> (FilePath -> Bool)
-> IO (TQueue (FilePath, UTCTime))
forkDirWatcherThread Scope
scope WatchManager
mgr FilePath
dir FilePath -> Bool
allow

  -- Await an event from the event queue with the following simple debounce logic, which is intended to work around the
  -- tendency for modern editors to create a flurry of rapid filesystem events when a file is saved:
  --
  -- 1. Block until an event arrives.
  -- 2. Keep consuming events until 50ms elapse without an event.
  -- 3. Return only the last event.
  --
  -- Note we don't have any smarts here for a flurry of events that are related to more than one file; we just throw
  -- everything away except the last event. In practice, this has seemed to work fine.
  let awaitEvent0 :: IO (FilePath, UTCTime)
      awaitEvent0 :: IO (FilePath, UTCTime)
awaitEvent0 = do
        let go :: (FilePath, UTCTime) -> IO (FilePath, UTCTime)
            go :: (FilePath, UTCTime) -> IO (FilePath, UTCTime)
go (FilePath, UTCTime)
event0 = do
              TVar Bool
var <- Int -> IO (TVar Bool)
registerDelay Int
50_000
              (IO (IO (FilePath, UTCTime)) -> IO (FilePath, UTCTime)
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (IO (IO (FilePath, UTCTime)) -> IO (FilePath, UTCTime))
-> ([STM (IO (FilePath, UTCTime))] -> IO (IO (FilePath, UTCTime)))
-> [STM (IO (FilePath, UTCTime))]
-> IO (FilePath, UTCTime)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. STM (IO (FilePath, UTCTime)) -> IO (IO (FilePath, UTCTime))
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM (IO (FilePath, UTCTime)) -> IO (IO (FilePath, UTCTime)))
-> ([STM (IO (FilePath, UTCTime))] -> STM (IO (FilePath, UTCTime)))
-> [STM (IO (FilePath, UTCTime))]
-> IO (IO (FilePath, UTCTime))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [STM (IO (FilePath, UTCTime))] -> STM (IO (FilePath, UTCTime))
forall (t :: * -> *) (f :: * -> *) a.
(Foldable t, Alternative f) =>
t (f a) -> f a
asum)
                [ do
                    (FilePath, UTCTime)
event1 <- TQueue (FilePath, UTCTime) -> STM (FilePath, UTCTime)
forall a. TQueue a -> STM a
STM.readTQueue TQueue (FilePath, UTCTime)
eventQueue
                    IO (FilePath, UTCTime) -> STM (IO (FilePath, UTCTime))
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((FilePath, UTCTime) -> IO (FilePath, UTCTime)
go (FilePath, UTCTime)
event1),
                  do
                    TVar Bool -> STM Bool
forall a. TVar a -> STM a
STM.readTVar TVar Bool
var STM Bool -> (Bool -> STM ()) -> STM ()
forall a b. STM a -> (a -> STM b) -> STM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Bool -> STM ()
STM.check
                    IO (FilePath, UTCTime) -> STM (IO (FilePath, UTCTime))
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((FilePath, UTCTime) -> IO (FilePath, UTCTime)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (FilePath, UTCTime)
event0)
                ]
        (FilePath, UTCTime)
event <- STM (FilePath, UTCTime) -> IO (FilePath, UTCTime)
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (TQueue (FilePath, UTCTime) -> STM (FilePath, UTCTime)
forall a. TQueue a -> STM a
STM.readTQueue TQueue (FilePath, UTCTime)
eventQueue)
        (FilePath, UTCTime) -> IO (FilePath, UTCTime)
go (FilePath, UTCTime)
event

  -- Enhance the previous "await event" action with a small file cache that serves as a second debounce implementation.
  -- We keep in memory the file contents of previously-saved files, so that we can avoid emitting events for files that
  -- last changed less than 500ms ago, and whose contents haven't changed.
  IORef (Map FilePath (Text, UTCTime))
previousFilesRef <- Map FilePath (Text, UTCTime)
-> IO (IORef (Map FilePath (Text, UTCTime)))
forall a. a -> IO (IORef a)
newIORef Map FilePath (Text, UTCTime)
forall k a. Map k a
Map.empty
  let awaitEvent1 :: IO (FilePath, Text)
      awaitEvent1 :: IO (FilePath, Text)
awaitEvent1 = do
        (FilePath
file, UTCTime
t) <- IO (FilePath, UTCTime)
awaitEvent0
        IO Text -> IO (Either SomeException Text)
forall (m :: * -> *) a.
MonadUnliftIO m =>
m a -> m (Either SomeException a)
tryAny (FilePath -> IO Text
readUtf8 FilePath
file) IO (Either SomeException Text)
-> (Either SomeException Text -> IO (FilePath, Text))
-> IO (FilePath, Text)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
          -- Somewhat-expected read error from a file that was just written. Just ignore the event and try again.
          Left SomeException
_ -> IO (FilePath, Text)
awaitEvent1
          Right Text
contents -> do
            Map FilePath (Text, UTCTime)
previousFiles <- IORef (Map FilePath (Text, UTCTime))
-> IO (Map FilePath (Text, UTCTime))
forall a. IORef a -> IO a
readIORef IORef (Map FilePath (Text, UTCTime))
previousFilesRef
            case FilePath -> Map FilePath (Text, UTCTime) -> Maybe (Text, UTCTime)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup FilePath
file Map FilePath (Text, UTCTime)
previousFiles of
              Just (Text
contents0, UTCTime
t0) | Text
contents Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
contents0 Bool -> Bool -> Bool
&& (UTCTime
t UTCTime -> UTCTime -> NominalDiffTime
`diffUTCTime` UTCTime
t0) NominalDiffTime -> NominalDiffTime -> Bool
forall a. Ord a => a -> a -> Bool
< NominalDiffTime
0.5 -> IO (FilePath, Text)
awaitEvent1
              Maybe (Text, UTCTime)
_ -> do
                IORef (Map FilePath (Text, UTCTime))
-> Map FilePath (Text, UTCTime) -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Map FilePath (Text, UTCTime))
previousFilesRef (Map FilePath (Text, UTCTime) -> IO ())
-> Map FilePath (Text, UTCTime) -> IO ()
forall a b. (a -> b) -> a -> b
$! FilePath
-> (Text, UTCTime)
-> Map FilePath (Text, UTCTime)
-> Map FilePath (Text, UTCTime)
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert FilePath
file (Text
contents, UTCTime
t) Map FilePath (Text, UTCTime)
previousFiles
                (FilePath, Text) -> IO (FilePath, Text)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (FilePath
file, Text
contents)

  -- Enhance the previous "await" event action by first clearing the whole event queue (tossing old filesystem events
  -- we may have accumulated while e.g. running a long-running IO action), and *then* waiting.
  let awaitEvent2 :: IO (FilePath, Text)
      awaitEvent2 :: IO (FilePath, Text)
awaitEvent2 = do
        [(FilePath, UTCTime)]
_ <- STM [(FilePath, UTCTime)] -> IO [(FilePath, UTCTime)]
forall a. STM a -> IO a
STM.atomically (TQueue (FilePath, UTCTime) -> STM [(FilePath, UTCTime)]
forall a. TQueue a -> STM [a]
STM.flushTQueue TQueue (FilePath, UTCTime)
eventQueue)
        IO (FilePath, Text)
awaitEvent1

  IO (FilePath, Text) -> IO (IO (FilePath, Text))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure IO (FilePath, Text)
awaitEvent2

-- | `forkDirWatcherThread scope mgr dir allow` forks a background thread into `scope` that, using "file watcher
-- manager" `mgr` (just a boilerplate argument the caller is responsible for creating), watches directory `dir` for
-- all "added" and "modified" filesystem events that occur on files that pass the `allow` predicate. It returns a queue
-- of such event that is (obviously) meant to be read or flushed, never written.
forkDirWatcherThread :: Ki.Scope -> FSNotify.WatchManager -> FilePath -> (FilePath -> Bool) -> IO (STM.TQueue (FilePath, UTCTime))
forkDirWatcherThread :: Scope
-> WatchManager
-> FilePath
-> (FilePath -> Bool)
-> IO (TQueue (FilePath, UTCTime))
forkDirWatcherThread Scope
scope WatchManager
mgr FilePath
dir FilePath -> Bool
allow = do
  TQueue (FilePath, UTCTime)
queue <- IO (TQueue (FilePath, UTCTime))
forall a. IO (TQueue a)
STM.newTQueueIO

  let handler :: Event -> IO ()
      handler :: Event -> IO ()
handler = \case
        Added FilePath
fp UTCTime
t EventIsDirectory
FSNotify.IsFile | FilePath -> Bool
allow FilePath
fp -> STM () -> IO ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (TQueue (FilePath, UTCTime) -> (FilePath, UTCTime) -> STM ()
forall a. TQueue a -> a -> STM ()
STM.writeTQueue TQueue (FilePath, UTCTime)
queue (FilePath
fp, UTCTime
t))
        Modified FilePath
fp UTCTime
t EventIsDirectory
FSNotify.IsFile | FilePath -> Bool
allow FilePath
fp -> STM () -> IO ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (TQueue (FilePath, UTCTime) -> (FilePath, UTCTime) -> STM ()
forall a. TQueue a -> a -> STM ()
STM.writeTQueue TQueue (FilePath, UTCTime)
queue (FilePath
fp, UTCTime
t))
        Event
_ -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

  -- A bit of a "one too many threads" situation but there's not much we can easily do about it. The `fsnotify` API
  -- doesn't expose any synchronous API; the only option is to fork a background thread with a callback. So, we spawn
  -- a thread that spawns *that* thread, then waits forever. The purpose here is to simply leverage `ki` exception
  -- propagation machinery to ensure that the `fsnotify` thread is properly cleaned up.
  Scope -> ThreadOptions -> IO Void -> IO ()
Ki.forkWith_ Scope
scope ThreadOptions
Ki.defaultThreadOptions {Ki.maskingState = MaskedUninterruptible} do
    -- The goal here is to prevent spawning this background watching thread before installing an exception handler that
    -- guarantees it's killed. Unfortunately the fsnotify API doesn't seem to make that possible (hence the first
    -- `unsafeUnmask` here), since we do need the thread *it* spawns to be killable, and (at least as of version
    -- 0.4.2.0) they don't take care to guarantee that; it just inherits the masking state.
    IO ()
stopListening <- IO (IO ()) -> IO (IO ())
forall a. IO a -> IO a
unsafeUnmask (WatchManager
-> FilePath -> ActionPredicate -> (Event -> IO ()) -> IO (IO ())
FSNotify.watchDir WatchManager
mgr FilePath
dir (Bool -> ActionPredicate
forall a b. a -> b -> a
const Bool
True) Event -> IO ()
handler) IO (IO ()) -> IO (IO ()) -> IO (IO ())
forall a. IO a -> IO a -> IO a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> IO () -> IO (IO ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (() -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())
    IO Void -> IO Void
forall a. IO a -> IO a
unsafeUnmask (IO () -> IO Void
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (Int -> IO ()
threadDelay Int
forall a. Bounded a => a
maxBound)) IO Void -> IO () -> IO Void
forall (m :: * -> *) a b. MonadUnliftIO m => m a -> m b -> m a
`finally` IO ()
stopListening

  TQueue (FilePath, UTCTime) -> IO (TQueue (FilePath, UTCTime))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TQueue (FilePath, UTCTime)
queue