{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MonoLocalBinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE UndecidableInstances #-}

-- | A scheme for adding persistence to Haskell's STM transactions.  A @'DBRef'
-- a@ is like a @'TVar' ('Maybe' a)@, except that it exists (or not) in
-- persistent storage as well as in memory.
--
-- The choice of persistent storage is up to the user, and is specified with a
-- 'Persistence'.  There is a default implementation called 'filePersistence'
-- that uses files on disk.  Note that 'filePersistence' doesn't guarantee
-- transactional atomicity in the presence of sudden termination of the process,
-- such as in a power outage or system crash.  Therefore, for serious use,
-- it's recommended that you use a different 'Persistence' implementation based
-- on a storage layer with stronger transactional guarantees.
--
-- For this scheme to work at all, this process must be the only entity to
-- access the persistent storage.  You may not even use a single-writer,
-- multiple-reader architecture, because consistency guarantees for reads, as
-- well, depend on all writes happening in the current process.
module PersistentSTM
  ( DB,
    openDB,
    closeDB,
    withDB,
    waitForMaxBacklog,
    synchronously,
    DBRef,
    DBStorable (..),
    getDBRef,
    readDBRef,
    writeDBRef,
    deleteDBRef,
    Persistence (..),
    filePersistence,
  )
where

import Control.Concurrent
  ( forkIO,
    newEmptyMVar,
    putMVar,
    takeMVar,
  )
import Control.Concurrent.STM
  ( STM,
    TVar,
    atomically,
    modifyTVar,
    newTVar,
    newTVarIO,
    readTVar,
    retry,
    writeTVar,
  )
import Control.Exception (bracket)
import Control.Monad (forM_, when)
import Control.Monad.Extra (whileM)
import Data.Binary (Binary)
import qualified Data.Binary as Binary
import Data.Bool (bool)
import qualified Data.ByteString as BS
import Data.ByteString.Lazy (ByteString)
import qualified Data.ByteString.Lazy as LBS
import Data.ByteString.Short (ShortByteString)
import Data.Int (Int16, Int32, Int64, Int8)
import Data.Map (Map)
import qualified Data.Map.Strict as Map
import Data.Proxy (Proxy (..))
import Data.Typeable (TypeRep, Typeable, typeRep)
import Data.Word (Word16, Word32, Word64, Word8)
import qualified Focus
import GHC.Conc (unsafeIOToSTM)
import Numeric.Natural (Natural)
import qualified StmContainers.Map as SMap
import System.Directory (createDirectoryIfMissing, doesFileExist, removeFile)
import System.FileLock (SharedExclusive (..), tryLockFile, unlockFile)
import System.FilePath ((</>))
import System.Mem.Weak (Weak, deRefWeak, mkWeak)
import Unsafe.Coerce (unsafeCoerce)

-- | A type class for things that can be stored in a DBRef.  This is similar to
-- a serialization class like 'Binary', but reads have access to the 'DB' and
-- the STM monad, which is important because it allows for one 'DBRef' to be
-- stored inside the value of another.  (In this case, 'decode' will call
-- 'getDBRef'.)
class Typeable a => DBStorable a where
  decode :: DB -> ByteString -> STM a
  default decode :: (Binary a) => DB -> ByteString -> STM a
  decode DB
_ = a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> STM a) -> (ByteString -> a) -> ByteString -> STM a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> a
forall a. Binary a => ByteString -> a
Binary.decode

  encode :: a -> ByteString
  default encode :: (Binary a) => a -> ByteString
  encode = a -> ByteString
forall a. Binary a => a -> ByteString
Binary.encode

instance DBStorable ()

instance DBStorable Bool

instance DBStorable Char

instance DBStorable Double

instance DBStorable Float

instance DBStorable Int

instance DBStorable Int8

instance DBStorable Int16

instance DBStorable Int32

instance DBStorable Int64

instance DBStorable Integer

instance DBStorable Natural

instance DBStorable Ordering

instance DBStorable Word

instance DBStorable Word8

instance DBStorable Word16

instance DBStorable Word32

instance DBStorable Word64

instance DBStorable BS.ByteString

instance DBStorable ByteString

instance DBStorable ShortByteString

instance DBStorable a => DBStorable [a] where
  encode :: [a] -> ByteString
encode = [ByteString] -> ByteString
forall a. Binary a => a -> ByteString
Binary.encode ([ByteString] -> ByteString)
-> ([a] -> [ByteString]) -> [a] -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> ByteString) -> [a] -> [ByteString]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> ByteString
forall a. DBStorable a => a -> ByteString
encode
  decode :: DB -> ByteString -> STM [a]
decode DB
db = (ByteString -> STM a) -> [ByteString] -> STM [a]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (DB -> ByteString -> STM a
forall a. DBStorable a => DB -> ByteString -> STM a
decode DB
db) ([ByteString] -> STM [a])
-> (ByteString -> [ByteString]) -> ByteString -> STM [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> [ByteString]
forall a. Binary a => ByteString -> a
Binary.decode

-- | Internal state of a 'DBRef'.  Loading means that the value is already
-- being loaded from persistent storage in a different thread, so the current
-- transaction can just retry to wait for it to load.
data Possible a = Loading | Missing | Present a

-- | Existential wrapper around 'TVar' that lets 'TVar's of various types be
-- cached.
data SomeTVar = forall a. SomeTVar (TVar (Possible a))

-- | A strategy for persisting values from 'DBRef' to some persistent storage.
-- The 'filePersistence' implementation is provided as a quick way to get
-- started, but note the weaknesses in its documentation.
--
-- A 'Persistence' can read one value at a time, but should be able to atomically
-- write/delete an entire set of keys at once, preferably atomically.
data Persistence = Persistence
  { -- | Read a single value from persistent storage.  Return the serialized
    -- representation if it exists, and Nothing otherwise.
    Persistence -> String -> IO (Maybe ByteString)
persistentRead :: String -> IO (Maybe ByteString),
    -- | Write (for 'Just' values) or delete (for 'Nothing' values) an entire
    -- set of values to persistent storage.  The values should ideally be
    -- written atomically, and if they are not then the implementation will be
    -- vulnerable to inconsistent data and corruption if the process is suddenly
    -- terminated.
    Persistence -> Map String (Maybe ByteString) -> IO ()
persistentWrite :: Map String (Maybe ByteString) -> IO (),
    -- | Perform any cleanup that is needed after the 'DB' is closed.  This can
    -- include releasing locks, for example.
    Persistence -> IO ()
persistentFinish :: IO ()
  }

-- | A currently open database in which 'DBRef's can be read and written.  See
-- 'openDB', 'closeDB', and 'withDB' to manage 'DB' values.
data DB = DB
  { -- | Cached 'TVar's corresponding to 'DBRef's that are already loading or
    -- loaded.
    DB -> Map String (TypeRep, Weak SomeTVar)
dbRefs :: SMap.Map String (TypeRep, Weak SomeTVar),
    -- | The last written generation number, used to find out when writes are
    -- committed.
    DB -> TVar Natural
dbGeneration :: TVar Natural,
    -- | Collection of dirty values that need to be written.  Only the
    -- 'ByteString' from the value is needed, but keeping the 'TVar' as well
    -- ensures that the 'TVar' won't be garbage collected and removed from
    -- dbRefs, which guarantees the value won't be read again until after the
    -- write is complete.  This is needed for consistency.
    DB -> TVar (Map String (SomeTVar, Maybe ByteString))
dbDirty :: TVar (Map String (SomeTVar, Maybe ByteString)),
    -- | The persistence that is used for this database.
    DB -> Persistence
dbPersistence :: Persistence,
    -- | If True, then 'closeDB' has been called, and the no new accesses to the
    -- 'DBRef's should be allowed.  This also triggers the writer thread to exit
    -- as soon as it has finished writing all dirty values.
    DB -> TVar Bool
dbClosing :: TVar Bool,
    -- | If True, the writer thread as finished writing all dirty values, and
    -- it's okay for the process to exit.
    DB -> TVar Bool
dbClosed :: TVar Bool
  }

-- | A reference to persistent data from some 'DB' that can be accessed in 'STM'
-- transaction.  @'DBRef' a@ is similar to @'TVar ('Maybe' a)@, except that
-- values exist in persistent storage as well as in memory.
data DBRef a = DBRef DB String (TVar (Possible a))

-- | Only 'DBRef's in the same 'DB' should be compared.
instance Eq (DBRef a) where
  DBRef DB
_ String
k1 TVar (Possible a)
_ == :: DBRef a -> DBRef a -> Bool
== DBRef DB
_ String
k2 TVar (Possible a)
_ = String
k1 String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
k2

-- | Only 'DBRef's in the same 'DB' should be compared.
instance Ord (DBRef a) where
  compare :: DBRef a -> DBRef a -> Ordering
compare (DBRef DB
_ String
k1 TVar (Possible a)
_) (DBRef DB
_ String
k2 TVar (Possible a)
_) = String -> String -> Ordering
forall a. Ord a => a -> a -> Ordering
compare String
k1 String
k2

instance Show (DBRef a) where
  show :: DBRef a -> String
show (DBRef DB
_ String
s TVar (Possible a)
_) = String
s

instance DBStorable a => DBStorable (DBRef a) where
  decode :: DB -> ByteString -> STM (DBRef a)
decode DB
db ByteString
bs = DB -> String -> STM (DBRef a)
forall a. DBStorable a => DB -> String -> STM (DBRef a)
getDBRef DB
db (ByteString -> String
forall a. Binary a => ByteString -> a
Binary.decode ByteString
bs)
  encode :: DBRef a -> ByteString
encode (DBRef DB
_ String
dbkey TVar (Possible a)
_) = String -> ByteString
forall a. Binary a => a -> ByteString
Binary.encode String
dbkey

-- | A simple 'Persistence' that stores data in a directory in the local
-- filesystem.  This is an easy way to get started.  However, note that because
-- writes are not atomic, your data can be corrupted during a crash or power
-- outage.  For this reason, it's recommended that you use a different
-- 'Persistence' for most applications.
filePersistence :: FilePath -> IO Persistence
filePersistence :: String -> IO Persistence
filePersistence String
dir = do
  Bool -> String -> IO ()
createDirectoryIfMissing Bool
True String
dir
  String -> SharedExclusive -> IO (Maybe FileLock)
tryLockFile (String
dir String -> ShowS
</> String
".lock") SharedExclusive
Exclusive IO (Maybe FileLock)
-> (Maybe FileLock -> IO Persistence) -> IO Persistence
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Maybe FileLock
Nothing -> String -> IO Persistence
forall a. HasCallStack => String -> a
error String
"Directory is already in use"
    Just FileLock
lock ->
      Persistence -> IO Persistence
forall (m :: * -> *) a. Monad m => a -> m a
return (Persistence -> IO Persistence) -> Persistence -> IO Persistence
forall a b. (a -> b) -> a -> b
$
        Persistence :: (String -> IO (Maybe ByteString))
-> (Map String (Maybe ByteString) -> IO ()) -> IO () -> Persistence
Persistence
          { persistentRead :: String -> IO (Maybe ByteString)
persistentRead = \String
key -> do
              Bool
ex <- String -> IO Bool
doesFileExist (String
dir String -> ShowS
</> String
key)
              if Bool
ex
                then ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just (ByteString -> Maybe ByteString)
-> (ByteString -> ByteString) -> ByteString -> Maybe ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ByteString -> ByteString
LBS.fromStrict (ByteString -> Maybe ByteString)
-> IO ByteString -> IO (Maybe ByteString)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> IO ByteString
BS.readFile (String
dir String -> ShowS
</> String
key)
                else Maybe ByteString -> IO (Maybe ByteString)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe ByteString
forall a. Maybe a
Nothing,
            persistentWrite :: Map String (Maybe ByteString) -> IO ()
persistentWrite = \Map String (Maybe ByteString)
dirtyMap -> [(String, Maybe ByteString)]
-> ((String, Maybe ByteString) -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (Map String (Maybe ByteString) -> [(String, Maybe ByteString)]
forall k a. Map k a -> [(k, a)]
Map.toList Map String (Maybe ByteString)
dirtyMap) (((String, Maybe ByteString) -> IO ()) -> IO ())
-> ((String, Maybe ByteString) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$
              \(String
key, Maybe ByteString
mbs) -> case Maybe ByteString
mbs of
                Just ByteString
bs -> String -> ByteString -> IO ()
BS.writeFile (String
dir String -> ShowS
</> String
key) (ByteString -> ByteString
LBS.toStrict ByteString
bs)
                Maybe ByteString
Nothing -> String -> IO ()
removeFile (String
dir String -> ShowS
</> String
key),
            persistentFinish :: IO ()
persistentFinish = FileLock -> IO ()
unlockFile FileLock
lock
          }

-- | Opens a 'DB' using the given 'Persistence'.  The caller should guarantee
-- that 'closeDB' is called when the 'DB' is no longer needed.
openDB :: Persistence -> IO DB
openDB :: Persistence -> IO DB
openDB Persistence
persistence = do
  Map String (TypeRep, Weak SomeTVar)
refs <- IO (Map String (TypeRep, Weak SomeTVar))
forall key value. IO (Map key value)
SMap.newIO
  TVar Natural
generation <- Natural -> IO (TVar Natural)
forall a. a -> IO (TVar a)
newTVarIO Natural
0
  TVar (Map String (SomeTVar, Maybe ByteString))
dirty <- Map String (SomeTVar, Maybe ByteString)
-> IO (TVar (Map String (SomeTVar, Maybe ByteString)))
forall a. a -> IO (TVar a)
newTVarIO Map String (SomeTVar, Maybe ByteString)
forall k a. Map k a
Map.empty
  TVar Bool
closing <- Bool -> IO (TVar Bool)
forall a. a -> IO (TVar a)
newTVarIO Bool
False
  TVar Bool
closed <- Bool -> IO (TVar Bool)
forall a. a -> IO (TVar a)
newTVarIO Bool
False
  ThreadId
_ <- IO () -> IO ThreadId
forkIO (IO () -> IO ThreadId) -> IO () -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ do
    IO Bool -> IO ()
forall (m :: * -> *). Monad m => m Bool -> m ()
whileM (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ do
      (Map String (SomeTVar, Maybe ByteString)
d, Bool
c) <- STM (Map String (SomeTVar, Maybe ByteString), Bool)
-> IO (Map String (SomeTVar, Maybe ByteString), Bool)
forall a. STM a -> IO a
atomically (STM (Map String (SomeTVar, Maybe ByteString), Bool)
 -> IO (Map String (SomeTVar, Maybe ByteString), Bool))
-> STM (Map String (SomeTVar, Maybe ByteString), Bool)
-> IO (Map String (SomeTVar, Maybe ByteString), Bool)
forall a b. (a -> b) -> a -> b
$ do
        Bool
c <- TVar Bool -> STM Bool
forall a. TVar a -> STM a
readTVar TVar Bool
closing
        Map String (SomeTVar, Maybe ByteString)
d <- TVar (Map String (SomeTVar, Maybe ByteString))
-> STM (Map String (SomeTVar, Maybe ByteString))
forall a. TVar a -> STM a
readTVar TVar (Map String (SomeTVar, Maybe ByteString))
dirty
        Bool -> STM () -> STM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Bool -> Bool
not Bool
c Bool -> Bool -> Bool
&& Map String (SomeTVar, Maybe ByteString) -> Bool
forall k a. Map k a -> Bool
Map.null Map String (SomeTVar, Maybe ByteString)
d) STM ()
forall a. STM a
retry
        Bool -> STM () -> STM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Bool -> Bool
not (Map String (SomeTVar, Maybe ByteString) -> Bool
forall k a. Map k a -> Bool
Map.null Map String (SomeTVar, Maybe ByteString)
d)) (STM () -> STM ()) -> STM () -> STM ()
forall a b. (a -> b) -> a -> b
$ TVar (Map String (SomeTVar, Maybe ByteString))
-> Map String (SomeTVar, Maybe ByteString) -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (Map String (SomeTVar, Maybe ByteString))
dirty Map String (SomeTVar, Maybe ByteString)
forall k a. Map k a
Map.empty
        (Map String (SomeTVar, Maybe ByteString), Bool)
-> STM (Map String (SomeTVar, Maybe ByteString), Bool)
forall (m :: * -> *) a. Monad m => a -> m a
return (Map String (SomeTVar, Maybe ByteString)
d, Bool
c)
      Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Bool -> Bool
not (Map String (SomeTVar, Maybe ByteString) -> Bool
forall k a. Map k a -> Bool
Map.null Map String (SomeTVar, Maybe ByteString)
d)) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Persistence -> Map String (Maybe ByteString) -> IO ()
persistentWrite Persistence
persistence ((SomeTVar, Maybe ByteString) -> Maybe ByteString
forall a b. (a, b) -> b
snd ((SomeTVar, Maybe ByteString) -> Maybe ByteString)
-> Map String (SomeTVar, Maybe ByteString)
-> Map String (Maybe ByteString)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map String (SomeTVar, Maybe ByteString)
d)
      STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar Natural -> (Natural -> Natural) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar TVar Natural
generation (Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
+ Natural
1)
      Bool -> IO Bool
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> Bool
not Bool
c)
    STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar Bool -> Bool -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar Bool
closed Bool
True
  let db :: DB
db =
        DB :: Map String (TypeRep, Weak SomeTVar)
-> TVar Natural
-> TVar (Map String (SomeTVar, Maybe ByteString))
-> Persistence
-> TVar Bool
-> TVar Bool
-> DB
DB
          { dbRefs :: Map String (TypeRep, Weak SomeTVar)
dbRefs = Map String (TypeRep, Weak SomeTVar)
refs,
            dbGeneration :: TVar Natural
dbGeneration = TVar Natural
generation,
            dbDirty :: TVar (Map String (SomeTVar, Maybe ByteString))
dbDirty = TVar (Map String (SomeTVar, Maybe ByteString))
dirty,
            dbPersistence :: Persistence
dbPersistence = Persistence
persistence,
            dbClosing :: TVar Bool
dbClosing = TVar Bool
closing,
            dbClosed :: TVar Bool
dbClosed = TVar Bool
closed
          }
  DB -> IO DB
forall (m :: * -> *) a. Monad m => a -> m a
return DB
db

-- | Closes a 'DB'.  When this call returns, all data will be written to
-- persistent storage, and the program can exit without possibly losing data.
closeDB :: DB -> IO ()
closeDB :: DB -> IO ()
closeDB DB
db = do
  STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar Bool -> Bool -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (DB -> TVar Bool
dbClosing DB
db) Bool
True
  STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar Bool -> STM Bool
forall a. TVar a -> STM a
readTVar (DB -> TVar Bool
dbClosed DB
db) STM Bool -> (Bool -> STM ()) -> STM ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= STM () -> STM () -> Bool -> STM ()
forall a. a -> a -> Bool -> a
bool STM ()
forall a. STM a
retry (() -> STM ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
  Persistence -> IO ()
persistentFinish (DB -> Persistence
dbPersistence DB
db)

-- | Runs an action with a 'DB' open.  The 'DB' will be closed when the action
-- is finished.  The 'DB' value should not be used after the action has
-- returned.
withDB :: Persistence -> (DB -> IO a) -> IO a
withDB :: Persistence -> (DB -> IO a) -> IO a
withDB Persistence
persistence DB -> IO a
f = IO DB -> (DB -> IO ()) -> (DB -> IO a) -> IO a
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket (Persistence -> IO DB
openDB Persistence
persistence) DB -> IO ()
closeDB DB -> IO a
f

-- | Check that there are at most the given number of queued writes to the
-- database, and retries the transaction if so.  Adding this to the beginning of
-- your transactions can help prevent writes from falling too far behind the
-- live data.  Prioritizing writes this way can also reduce memory usage,
-- because unreachable 'DBRef's no longer need to be retained once they are
-- written to disk.
waitForMaxBacklog :: DB -> Int -> STM ()
waitForMaxBacklog :: DB -> Int -> STM ()
waitForMaxBacklog DB
db Int
maxLen = do
  Map String (SomeTVar, Maybe ByteString)
dirty <- TVar (Map String (SomeTVar, Maybe ByteString))
-> STM (Map String (SomeTVar, Maybe ByteString))
forall a. TVar a -> STM a
readTVar (DB -> TVar (Map String (SomeTVar, Maybe ByteString))
dbDirty DB
db)
  Bool -> STM () -> STM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Map String (SomeTVar, Maybe ByteString) -> Int
forall k a. Map k a -> Int
Map.size Map String (SomeTVar, Maybe ByteString)
dirty Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
maxLen) STM ()
forall a. STM a
retry

-- | Throws an error if the given 'DB' is closing.  This prevents more work from
-- being added to the queue when we're supposed to be waiting for the last
-- writes to flush out.
failIfClosing :: DB -> STM ()
failIfClosing :: DB -> STM ()
failIfClosing DB
db = do
  Bool
c <- TVar Bool -> STM Bool
forall a. TVar a -> STM a
readTVar (DB -> TVar Bool
dbClosing DB
db)
  Bool -> STM () -> STM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
c (STM () -> STM ()) -> STM () -> STM ()
forall a b. (a -> b) -> a -> b
$ String -> STM ()
forall a. HasCallStack => String -> a
error String
"DB is closing"

-- | Atomically performs an STM transaction just like 'atomically', but also
-- waits for any changes it might have observed in the 'DB' to be written to
-- persistent storage before returning.  This guarantees that a transaction
-- whose results were observed will not be rolled back if the program crashes.
synchronously :: DB -> STM a -> IO a
synchronously :: DB -> STM a -> IO a
synchronously DB
db STM a
txn = do
  (a
result, Maybe Natural
gen) <- STM (a, Maybe Natural) -> IO (a, Maybe Natural)
forall a. STM a -> IO a
atomically (STM (a, Maybe Natural) -> IO (a, Maybe Natural))
-> STM (a, Maybe Natural) -> IO (a, Maybe Natural)
forall a b. (a -> b) -> a -> b
$ do
    a
result <- STM a
txn
    Natural
gen <- TVar Natural -> STM Natural
forall a. TVar a -> STM a
readTVar (DB -> TVar Natural
dbGeneration DB
db)
    Map String (SomeTVar, Maybe ByteString)
dirty <- TVar (Map String (SomeTVar, Maybe ByteString))
-> STM (Map String (SomeTVar, Maybe ByteString))
forall a. TVar a -> STM a
readTVar (DB -> TVar (Map String (SomeTVar, Maybe ByteString))
dbDirty DB
db)
    if Map String (SomeTVar, Maybe ByteString) -> Bool
forall k a. Map k a -> Bool
Map.null Map String (SomeTVar, Maybe ByteString)
dirty
      then (a, Maybe Natural) -> STM (a, Maybe Natural)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
result, Maybe Natural
forall a. Maybe a
Nothing)
      else (a, Maybe Natural) -> STM (a, Maybe Natural)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
result, Natural -> Maybe Natural
forall a. a -> Maybe a
Just (Natural
gen Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
+ Natural
1))
  case Maybe Natural
gen of
    Just Natural
n ->
      STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$
        TVar Natural -> STM Natural
forall a. TVar a -> STM a
readTVar (DB -> TVar Natural
dbGeneration DB
db) STM Natural -> (Natural -> STM ()) -> STM ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= STM () -> STM () -> Bool -> STM ()
forall a. a -> a -> Bool -> a
bool STM ()
forall a. STM a
retry (() -> STM ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()) (Bool -> STM ()) -> (Natural -> Bool) -> Natural -> STM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Natural -> Natural -> Bool
forall a. Ord a => a -> a -> Bool
>= Natural
n)
    Maybe Natural
_ -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
  a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
result

-- | Retrieves a 'DBRef' from a 'DB' for the given key.  Throws an exception if
-- the 'DBRef' requested has a different type from a previous time the key was
-- used in this process, or if a serialized value in persistent storage cannot
-- be parsed.
getDBRef :: forall a. DBStorable a => DB -> String -> STM (DBRef a)
getDBRef :: DB -> String -> STM (DBRef a)
getDBRef DB
db String
key = do
  DB -> STM ()
failIfClosing DB
db
  String
-> Map String (TypeRep, Weak SomeTVar)
-> STM (Maybe (TypeRep, Weak SomeTVar))
forall key value.
(Eq key, Hashable key) =>
key -> Map key value -> STM (Maybe value)
SMap.lookup String
key (DB -> Map String (TypeRep, Weak SomeTVar)
dbRefs DB
db) STM (Maybe (TypeRep, Weak SomeTVar))
-> (Maybe (TypeRep, Weak SomeTVar) -> STM (DBRef a))
-> STM (DBRef a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Just (TypeRep
tr, Weak SomeTVar
weakRef)
      | TypeRep
tr TypeRep -> TypeRep -> Bool
forall a. Eq a => a -> a -> Bool
== Proxy a -> TypeRep
forall k (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep (Proxy a
forall k (t :: k). Proxy t
Proxy @a) ->
          IO (Maybe SomeTVar) -> STM (Maybe SomeTVar)
forall a. IO a -> STM a
unsafeIOToSTM (Weak SomeTVar -> IO (Maybe SomeTVar)
forall v. Weak v -> IO (Maybe v)
deRefWeak Weak SomeTVar
weakRef) STM (Maybe SomeTVar)
-> (Maybe SomeTVar -> STM (DBRef a)) -> STM (DBRef a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
            Just (SomeTVar TVar (Possible a)
ref) -> DBRef a -> STM (DBRef a)
forall (m :: * -> *) a. Monad m => a -> m a
return (DB -> String -> TVar (Possible a) -> DBRef a
forall a. DB -> String -> TVar (Possible a) -> DBRef a
DBRef DB
db String
key (TVar (Possible a) -> TVar (Possible a)
forall a b. a -> b
unsafeCoerce TVar (Possible a)
ref))
            Maybe SomeTVar
Nothing -> STM (DBRef a)
insert
      | Bool
otherwise -> String -> STM (DBRef a)
forall a. HasCallStack => String -> a
error String
"Type mismatch in DBRef"
    Maybe (TypeRep, Weak SomeTVar)
Nothing -> STM (DBRef a)
insert
  where
    insert :: STM (DBRef a)
insert = do
      TVar (Possible a)
ref <- Possible a -> STM (TVar (Possible a))
forall a. a -> STM (TVar a)
newTVar Possible a
forall a. Possible a
Loading
      Weak SomeTVar
ptr <- IO (Weak SomeTVar) -> STM (Weak SomeTVar)
forall a. IO a -> STM a
unsafeIOToSTM (IO (Weak SomeTVar) -> STM (Weak SomeTVar))
-> IO (Weak SomeTVar) -> STM (Weak SomeTVar)
forall a b. (a -> b) -> a -> b
$ TVar (Possible a)
-> SomeTVar -> Maybe (IO ()) -> IO (Weak SomeTVar)
forall k v. k -> v -> Maybe (IO ()) -> IO (Weak v)
mkWeak TVar (Possible a)
ref (TVar (Possible a) -> SomeTVar
forall a. TVar (Possible a) -> SomeTVar
SomeTVar TVar (Possible a)
ref) (IO () -> Maybe (IO ())
forall a. a -> Maybe a
Just IO ()
cleanupKey)
      (TypeRep, Weak SomeTVar)
-> String -> Map String (TypeRep, Weak SomeTVar) -> STM ()
forall key value.
(Eq key, Hashable key) =>
value -> key -> Map key value -> STM ()
SMap.insert (Proxy a -> TypeRep
forall k (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep (Proxy a
forall k (t :: k). Proxy t
Proxy @a), Weak SomeTVar
ptr) String
key (DB -> Map String (TypeRep, Weak SomeTVar)
dbRefs DB
db)
      Possible a
v <- IO (Possible a) -> STM (Possible a)
forall a. IO a -> STM a
unsafeIOToSTM (IO (Possible a) -> STM (Possible a))
-> IO (Possible a) -> STM (Possible a)
forall a b. (a -> b) -> a -> b
$ do
        MVar (Possible a)
mvar <- IO (MVar (Possible a))
forall a. IO (MVar a)
newEmptyMVar
        ThreadId
_ <- IO () -> IO ThreadId
forkIO (IO () -> IO ThreadId) -> IO () -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ MVar (Possible a) -> Possible a -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar (Possible a)
mvar (Possible a -> IO ()) -> IO (Possible a) -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< IO (Possible a)
readKey
        MVar (Possible a) -> IO (Possible a)
forall a. MVar a -> IO a
takeMVar MVar (Possible a)
mvar
      TVar (Possible a) -> Possible a -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (Possible a)
ref Possible a
v
      DBRef a -> STM (DBRef a)
forall (m :: * -> *) a. Monad m => a -> m a
return (DB -> String -> TVar (Possible a) -> DBRef a
forall a. DB -> String -> TVar (Possible a) -> DBRef a
DBRef DB
db String
key TVar (Possible a)
ref)

    readKey :: IO (Possible a)
readKey = do
      Maybe ByteString
readResult <- Persistence -> String -> IO (Maybe ByteString)
persistentRead (DB -> Persistence
dbPersistence DB
db) String
key
      case Maybe ByteString
readResult of
        Just ByteString
bs -> a -> Possible a
forall a. a -> Possible a
Present (a -> Possible a) -> IO a -> IO (Possible a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> STM a -> IO a
forall a. STM a -> IO a
atomically (DB -> ByteString -> STM a
forall a. DBStorable a => DB -> ByteString -> STM a
decode DB
db ByteString
bs)
        Maybe ByteString
Nothing -> Possible a -> IO (Possible a)
forall (m :: * -> *) a. Monad m => a -> m a
return Possible a
forall a. Possible a
Missing

    cleanupKey :: IO ()
cleanupKey =
      STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$
        Focus (TypeRep, Weak SomeTVar) STM ()
-> String -> Map String (TypeRep, Weak SomeTVar) -> STM ()
forall key value result.
(Eq key, Hashable key) =>
Focus value STM result -> key -> Map key value -> STM result
SMap.focus
          ( ((TypeRep, Weak SomeTVar) -> STM (Maybe (TypeRep, Weak SomeTVar)))
-> Focus (TypeRep, Weak SomeTVar) STM ()
forall (m :: * -> *) a.
Monad m =>
(a -> m (Maybe a)) -> Focus a m ()
Focus.updateM
              ( \(TypeRep
tr, Weak SomeTVar
p) ->
                  IO (Maybe SomeTVar) -> STM (Maybe SomeTVar)
forall a. IO a -> STM a
unsafeIOToSTM (Weak SomeTVar -> IO (Maybe SomeTVar)
forall v. Weak v -> IO (Maybe v)
deRefWeak Weak SomeTVar
p) STM (Maybe SomeTVar)
-> (Maybe SomeTVar -> STM (Maybe (TypeRep, Weak SomeTVar)))
-> STM (Maybe (TypeRep, Weak SomeTVar))
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
                    Maybe SomeTVar
Nothing -> Maybe (TypeRep, Weak SomeTVar)
-> STM (Maybe (TypeRep, Weak SomeTVar))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (TypeRep, Weak SomeTVar)
forall a. Maybe a
Nothing
                    Just SomeTVar
_ -> Maybe (TypeRep, Weak SomeTVar)
-> STM (Maybe (TypeRep, Weak SomeTVar))
forall (m :: * -> *) a. Monad m => a -> m a
return ((TypeRep, Weak SomeTVar) -> Maybe (TypeRep, Weak SomeTVar)
forall a. a -> Maybe a
Just (TypeRep
tr, Weak SomeTVar
p))
              )
          )
          String
key
          (DB -> Map String (TypeRep, Weak SomeTVar)
dbRefs DB
db)

-- | Gets the value stored in a 'DBRef'.  The value is @'Just' x@ if @x@ was
-- last value stored in the database using this key, or 'Nothing' if there is no
-- value stored in the database.
readDBRef :: DBRef a -> STM (Maybe a)
readDBRef :: DBRef a -> STM (Maybe a)
readDBRef (DBRef DB
db String
_ TVar (Possible a)
ref) = do
  DB -> STM ()
failIfClosing DB
db
  TVar (Possible a) -> STM (Possible a)
forall a. TVar a -> STM a
readTVar TVar (Possible a)
ref STM (Possible a) -> (Possible a -> STM (Maybe a)) -> STM (Maybe a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Possible a
Loading -> STM (Maybe a)
forall a. STM a
retry
    Possible a
Missing -> Maybe a -> STM (Maybe a)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
forall a. Maybe a
Nothing
    Present a
a -> Maybe a -> STM (Maybe a)
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> Maybe a
forall a. a -> Maybe a
Just a
a)

-- | Updates the value stored in a 'DBRef'.  The update will be persisted to
-- storage soon, but not synchronously.
writeDBRef :: DBStorable a => DBRef a -> a -> STM ()
writeDBRef :: DBRef a -> a -> STM ()
writeDBRef (DBRef DB
db String
dbkey TVar (Possible a)
ref) a
a = do
  DB -> STM ()
failIfClosing DB
db
  TVar (Possible a) -> Possible a -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (Possible a)
ref (a -> Possible a
forall a. a -> Possible a
Present a
a)
  Map String (SomeTVar, Maybe ByteString)
d <- TVar (Map String (SomeTVar, Maybe ByteString))
-> STM (Map String (SomeTVar, Maybe ByteString))
forall a. TVar a -> STM a
readTVar (DB -> TVar (Map String (SomeTVar, Maybe ByteString))
dbDirty DB
db)
  TVar (Map String (SomeTVar, Maybe ByteString))
-> Map String (SomeTVar, Maybe ByteString) -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (DB -> TVar (Map String (SomeTVar, Maybe ByteString))
dbDirty DB
db) (String
-> (SomeTVar, Maybe ByteString)
-> Map String (SomeTVar, Maybe ByteString)
-> Map String (SomeTVar, Maybe ByteString)
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert String
dbkey (TVar (Possible a) -> SomeTVar
forall a. TVar (Possible a) -> SomeTVar
SomeTVar TVar (Possible a)
ref, ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just (a -> ByteString
forall a. DBStorable a => a -> ByteString
encode a
a)) Map String (SomeTVar, Maybe ByteString)
d)

-- | Deletes the value stored in a 'DBRef'.  The delete will be persisted to
-- storage soon, but not synchronously.
deleteDBRef :: DBStorable a => DBRef a -> STM ()
deleteDBRef :: DBRef a -> STM ()
deleteDBRef (DBRef DB
db String
dbkey TVar (Possible a)
ref) = do
  DB -> STM ()
failIfClosing DB
db
  TVar (Possible a) -> Possible a -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (Possible a)
ref Possible a
forall a. Possible a
Missing
  Map String (SomeTVar, Maybe ByteString)
d <- TVar (Map String (SomeTVar, Maybe ByteString))
-> STM (Map String (SomeTVar, Maybe ByteString))
forall a. TVar a -> STM a
readTVar (DB -> TVar (Map String (SomeTVar, Maybe ByteString))
dbDirty DB
db)
  TVar (Map String (SomeTVar, Maybe ByteString))
-> Map String (SomeTVar, Maybe ByteString) -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (DB -> TVar (Map String (SomeTVar, Maybe ByteString))
dbDirty DB
db) (String
-> (SomeTVar, Maybe ByteString)
-> Map String (SomeTVar, Maybe ByteString)
-> Map String (SomeTVar, Maybe ByteString)
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert String
dbkey (TVar (Possible a) -> SomeTVar
forall a. TVar (Possible a) -> SomeTVar
SomeTVar TVar (Possible a)
ref, Maybe ByteString
forall a. Maybe a
Nothing) Map String (SomeTVar, Maybe ByteString)
d)