-- | Handles to save/load files across the hydra-node. We use a simple JSON
-- encoding and two modes of operation to store things: Full and Incremental.
module Hydra.Persistence where

import Hydra.Prelude

import Control.Concurrent.Class.MonadSTM (newTVarIO, throwSTM, writeTVar)
import Control.Lens.Combinators (iforM)
import Control.Monad.Class.MonadFork (myThreadId)
import Data.Aeson qualified as Aeson
import Data.ByteString qualified as BS
import Data.ByteString.Char8 qualified as C8
import System.Directory (createDirectoryIfMissing, doesFileExist)
import System.FilePath (takeDirectory)
import UnliftIO.IO.File (withBinaryFile, writeBinaryFileDurableAtomic)

data PersistenceException
  = PersistenceException String
  | IncorrectAccessException String
  deriving stock (PersistenceException -> PersistenceException -> Bool
(PersistenceException -> PersistenceException -> Bool)
-> (PersistenceException -> PersistenceException -> Bool)
-> Eq PersistenceException
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: PersistenceException -> PersistenceException -> Bool
== :: PersistenceException -> PersistenceException -> Bool
$c/= :: PersistenceException -> PersistenceException -> Bool
/= :: PersistenceException -> PersistenceException -> Bool
Eq, Int -> PersistenceException -> ShowS
[PersistenceException] -> ShowS
PersistenceException -> String
(Int -> PersistenceException -> ShowS)
-> (PersistenceException -> String)
-> ([PersistenceException] -> ShowS)
-> Show PersistenceException
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> PersistenceException -> ShowS
showsPrec :: Int -> PersistenceException -> ShowS
$cshow :: PersistenceException -> String
show :: PersistenceException -> String
$cshowList :: [PersistenceException] -> ShowS
showList :: [PersistenceException] -> ShowS
Show)

instance Exception PersistenceException

-- | Handle to save and load files to/from disk using JSON encoding.
data Persistence a m = Persistence
  { forall a (m :: * -> *). Persistence a m -> ToJSON a => a -> m ()
save :: ToJSON a => a -> m ()
  , forall a (m :: * -> *).
Persistence a m -> FromJSON a => m (Maybe a)
load :: FromJSON a => m (Maybe a)
  }

-- | Initialize persistence handle for given type 'a' at given file path.
createPersistence ::
  (MonadIO m, MonadThrow m) =>
  FilePath ->
  m (Persistence a m)
createPersistence :: forall (m :: * -> *) a.
(MonadIO m, MonadThrow m) =>
String -> m (Persistence a m)
createPersistence String
fp = do
  IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> (String -> IO ()) -> String -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> String -> IO ()
createDirectoryIfMissing Bool
True (String -> m ()) -> String -> m ()
forall a b. (a -> b) -> a -> b
$ ShowS
takeDirectory String
fp
  Persistence a m -> m (Persistence a m)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Persistence a m -> m (Persistence a m))
-> Persistence a m -> m (Persistence a m)
forall a b. (a -> b) -> a -> b
$
    Persistence
      { $sel:save:Persistence :: ToJSON a => a -> m ()
save = \a
a -> do
          String -> ByteString -> m ()
forall (m :: * -> *). MonadIO m => String -> ByteString -> m ()
writeBinaryFileDurableAtomic String
fp (ByteString -> m ())
-> (ByteString -> ByteString) -> ByteString -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
forall l s. LazyStrict l s => l -> s
toStrict (ByteString -> m ()) -> ByteString -> m ()
forall a b. (a -> b) -> a -> b
$ a -> ByteString
forall a. ToJSON a => a -> ByteString
Aeson.encode a
a
      , $sel:load:Persistence :: FromJSON a => m (Maybe a)
load =
          IO Bool -> m Bool
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (String -> IO Bool
doesFileExist String
fp) m Bool -> (Bool -> m (Maybe a)) -> m (Maybe a)
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
            Bool
False -> Maybe a -> m (Maybe a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
Nothing
            Bool
True -> do
              ByteString
bs <- String -> m ByteString
forall (m :: * -> *). MonadIO m => String -> m ByteString
readFileBS String
fp
              if ByteString -> Bool
BS.null ByteString
bs
                then Maybe a -> m (Maybe a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
Nothing
                else case ByteString -> Either String a
forall a. FromJSON a => ByteString -> Either String a
Aeson.eitherDecodeStrict' ByteString
bs of
                  Left String
e -> PersistenceException -> m (Maybe a)
forall e a. Exception e => e -> m a
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwIO (PersistenceException -> m (Maybe a))
-> PersistenceException -> m (Maybe a)
forall a b. (a -> b) -> a -> b
$ String -> PersistenceException
PersistenceException String
e
                  Right a
a -> Maybe a -> m (Maybe a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> Maybe a
forall a. a -> Maybe a
Just a
a)
      }

-- | Handle to save incrementally and load files to/from disk using JSON encoding.
data PersistenceIncremental a m = PersistenceIncremental
  { forall a (m :: * -> *).
PersistenceIncremental a m -> ToJSON a => a -> m ()
append :: ToJSON a => a -> m ()
  , forall a (m :: * -> *).
PersistenceIncremental a m -> FromJSON a => m [a]
loadAll :: FromJSON a => m [a]
  }

-- | Initialize persistence handle for given type 'a' at given file path.
--
-- This instance of `PersistenceIncremental` is "thread-safe" in the sense that
-- it prevents loading from a different thread once one starts `append`ing
-- through the handle. If another thread attempts to `loadAll` after this point,
-- an `IncorrectAccessException` will be raised.
createPersistenceIncremental ::
  forall a m.
  (MonadIO m, MonadThrow m, MonadSTM m, MonadThread m, MonadThrow (STM m)) =>
  FilePath ->
  m (PersistenceIncremental a m)
createPersistenceIncremental :: forall a (m :: * -> *).
(MonadIO m, MonadThrow m, MonadSTM m, MonadThread m,
 MonadThrow (STM m)) =>
String -> m (PersistenceIncremental a m)
createPersistenceIncremental String
fp = do
  IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> (String -> IO ()) -> String -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> String -> IO ()
createDirectoryIfMissing Bool
True (String -> m ()) -> String -> m ()
forall a b. (a -> b) -> a -> b
$ ShowS
takeDirectory String
fp
  TVar m (Maybe (ThreadId m))
authorizedThread <- Maybe (ThreadId m) -> m (TVar m (Maybe (ThreadId m)))
forall a. a -> m (TVar m a)
forall (m :: * -> *) a. MonadSTM m => a -> m (TVar m a)
newTVarIO Maybe (ThreadId m)
forall a. Maybe a
Nothing
  PersistenceIncremental a m -> m (PersistenceIncremental a m)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PersistenceIncremental a m -> m (PersistenceIncremental a m))
-> PersistenceIncremental a m -> m (PersistenceIncremental a m)
forall a b. (a -> b) -> a -> b
$
    PersistenceIncremental
      { $sel:append:PersistenceIncremental :: ToJSON a => a -> m ()
append = \a
a -> do
          ThreadId m
tid <- m (ThreadId m)
forall (m :: * -> *). MonadThread m => m (ThreadId m)
myThreadId
          STM m () -> m ()
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (STM m () -> m ()) -> STM m () -> m ()
forall a b. (a -> b) -> a -> b
$ TVar m (Maybe (ThreadId m)) -> Maybe (ThreadId m) -> STM m ()
forall a. TVar m a -> a -> STM m ()
forall (m :: * -> *) a. MonadSTM m => TVar m a -> a -> STM m ()
writeTVar TVar m (Maybe (ThreadId m))
authorizedThread (Maybe (ThreadId m) -> STM m ()) -> Maybe (ThreadId m) -> STM m ()
forall a b. (a -> b) -> a -> b
$ ThreadId m -> Maybe (ThreadId m)
forall a. a -> Maybe a
Just ThreadId m
tid
          let bytes :: ByteString
bytes = ByteString -> ByteString
forall l s. LazyStrict l s => l -> s
toStrict (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ a -> ByteString
forall a. ToJSON a => a -> ByteString
Aeson.encode a
a ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
"\n"
          IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ String -> IOMode -> (Handle -> IO ()) -> IO ()
forall (m :: * -> *) a.
MonadUnliftIO m =>
String -> IOMode -> (Handle -> m a) -> m a
withBinaryFile String
fp IOMode
AppendMode (Handle -> ByteString -> IO ()
`BS.hPut` ByteString
bytes)
      , $sel:loadAll:PersistenceIncremental :: FromJSON a => m [a]
loadAll = do
          ThreadId m
tid <- m (ThreadId m)
forall (m :: * -> *). MonadThread m => m (ThreadId m)
myThreadId
          STM m () -> m ()
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (STM m () -> m ()) -> STM m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
            Maybe (ThreadId m)
authTid <- TVar m (Maybe (ThreadId m)) -> STM m (Maybe (ThreadId m))
forall a. TVar m a -> STM m a
forall (m :: * -> *) a. MonadSTM m => TVar m a -> STM m a
readTVar TVar m (Maybe (ThreadId m))
authorizedThread
            Bool -> STM m () -> STM m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Maybe (ThreadId m) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (ThreadId m)
authTid Bool -> Bool -> Bool
&& Maybe (ThreadId m)
authTid Maybe (ThreadId m) -> Maybe (ThreadId m) -> Bool
forall a. Eq a => a -> a -> Bool
/= ThreadId m -> Maybe (ThreadId m)
forall a. a -> Maybe a
Just ThreadId m
tid) (STM m () -> STM m ()) -> STM m () -> STM m ()
forall a b. (a -> b) -> a -> b
$
              PersistenceException -> STM m ()
forall (m :: * -> *) e a.
(MonadSTM m, MonadThrow (STM m), Exception e) =>
e -> STM m a
throwSTM (String -> PersistenceException
IncorrectAccessException (String -> PersistenceException) -> String -> PersistenceException
forall a b. (a -> b) -> a -> b
$ String
"Trying to load persisted data in " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
fp String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
" from different thread")

          IO Bool -> m Bool
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (String -> IO Bool
doesFileExist String
fp) m Bool -> (Bool -> m [a]) -> m [a]
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
            Bool
False -> [a] -> m [a]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
            Bool
True -> do
              ByteString
bs <- String -> m ByteString
forall (m :: * -> *). MonadIO m => String -> m ByteString
readFileBS String
fp
              -- NOTE: We require the whole file to be loadable. It might
              -- happen that the data written by 'append' is only there
              -- partially and then this will fail (which we accept now).
              [ByteString] -> (Int -> ByteString -> m a) -> m [a]
forall i (t :: * -> *) (m :: * -> *) a b.
(TraversableWithIndex i t, Monad m) =>
t a -> (i -> a -> m b) -> m (t b)
iforM (ByteString -> [ByteString]
C8.lines ByteString
bs) ((Int -> ByteString -> m a) -> m [a])
-> (Int -> ByteString -> m a) -> m [a]
forall a b. (a -> b) -> a -> b
$ \Int
i ByteString
o ->
                case ByteString -> Either String a
forall a. FromJSON a => ByteString -> Either String a
Aeson.eitherDecodeStrict' ByteString
o of
                  Left String
e -> PersistenceException -> m a
forall e a. Exception e => e -> m a
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwIO (PersistenceException -> m a) -> PersistenceException -> m a
forall a b. (a -> b) -> a -> b
$ String -> PersistenceException
PersistenceException (String
"Error at line: " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall b a. (Show a, IsString b) => a -> b
show (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
" in file " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
fp String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
" - " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
e)
                  Right a
decoded -> a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
decoded
      }