diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,5 @@ +# Revision history for persistent-stm + +## 0.1.0.0 -- YYYY-mm-dd + +* First version. Released on an unsuspecting world. diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,26 @@ +Copyright (c) 2022, Chris Smith + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +3. Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/persistent-stm.cabal b/persistent-stm.cabal new file mode 100644 --- /dev/null +++ b/persistent-stm.cabal @@ -0,0 +1,53 @@ +cabal-version: 2.4 +name: persistent-stm +version: 0.1.0.0 +synopsis: STM transactions involving persistent storage +description: While Haskell's STM monad allows you to execute code + transactionally, it does not allow you to persist the state + on disk. This package implements a persistent storage + bridge that runs inside the already existing STM monad. +category: Database +homepage: https://github.com/cdsmith/persistent-stm +bug-reports: https://github.com/cdsmith/persistent-stm/issues +license: BSD-3-Clause +license-file: LICENSE + +author: Chris Smith +maintainer: Chris Smith + +extra-source-files: CHANGELOG.md + +tested-with: GHC == 8.6.5 || == 8.8.4 || == 8.10.7 || == 9.0.1 || == 9.2.2 + +source-repository head + type: git + location: git://github.com/cdsmith/persistent-stm.git + +library + exposed-modules: PersistentSTM + build-depends: base >= 4.12 && < 4.17, + binary >= 0.8 && < 0.9, + bytestring >= 0.10 && < 0.12, + containers >= 0.6 && < 0.7, + directory >= 1.3 && < 1.4, + filepath >= 1.4 && < 1.5, + extra >= 1.7 && < 1.8, + filelock >= 0.1 && < 0.2, + focus >= 1.0 && < 1.1, + stm >= 2.5 && < 2.6, + stm-containers >= 1.2 && < 1.3, + hs-source-dirs: src + default-language: Haskell2010 + ghc-options: -Wall + +test-suite tests + type: exitcode-stdio-1.0 + main-is: Main.hs + build-depends: persistent-stm, + base >= 4.12 && < 4.17, + hspec >= 2.10 && < 2.11, + stm >= 2.5 && < 2.6, + temporary >= 1.3 && < 1.4, + hs-source-dirs: test + default-language: Haskell2010 + ghc-options: -Wall diff --git a/src/PersistentSTM.hs b/src/PersistentSTM.hs new file mode 100644 --- /dev/null +++ b/src/PersistentSTM.hs @@ -0,0 +1,407 @@ +{-# 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 _ = pure . Binary.decode + + encode :: a -> ByteString + default encode :: (Binary a) => a -> ByteString + encode = 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 = Binary.encode . fmap encode + decode db = traverse (decode db) . 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. + 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. + persistentWrite :: Map String (Maybe ByteString) -> IO (), + -- | Perform any cleanup that is needed after the 'DB' is closed. This can + -- include releasing locks, for example. + 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. + dbRefs :: SMap.Map String (TypeRep, Weak SomeTVar), + -- | The last written generation number, used to find out when writes are + -- committed. + 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. + dbDirty :: TVar (Map String (SomeTVar, Maybe ByteString)), + -- | The persistence that is used for this database. + 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. + dbClosing :: TVar Bool, + -- | If True, the writer thread as finished writing all dirty values, and + -- it's okay for the process to exit. + 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 _ k1 _ == DBRef _ k2 _ = k1 == k2 + +-- | Only 'DBRef's in the same 'DB' should be compared. +instance Ord (DBRef a) where + compare (DBRef _ k1 _) (DBRef _ k2 _) = compare k1 k2 + +instance Show (DBRef a) where + show (DBRef _ s _) = s + +instance DBStorable a => DBStorable (DBRef a) where + decode db bs = getDBRef db (Binary.decode bs) + encode (DBRef _ dbkey _) = Binary.encode 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 dir = do + createDirectoryIfMissing True dir + tryLockFile (dir ".lock") Exclusive >>= \case + Nothing -> error "Directory is already in use" + Just lock -> + return $ + Persistence + { persistentRead = \key -> do + ex <- doesFileExist (dir key) + if ex + then Just <$> LBS.fromStrict <$> BS.readFile (dir key) + else return Nothing, + persistentWrite = \dirtyMap -> forM_ (Map.toList dirtyMap) $ + \(key, mbs) -> case mbs of + Just bs -> BS.writeFile (dir key) (LBS.toStrict bs) + Nothing -> removeFile (dir key), + persistentFinish = unlockFile 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 = do + refs <- SMap.newIO + generation <- newTVarIO 0 + dirty <- newTVarIO Map.empty + closing <- newTVarIO False + closed <- newTVarIO False + _ <- forkIO $ do + whileM $ do + (d, c) <- atomically $ do + c <- readTVar closing + d <- readTVar dirty + when (not c && Map.null d) retry + when (not (Map.null d)) $ writeTVar dirty Map.empty + return (d, c) + when (not (Map.null d)) $ persistentWrite persistence (snd <$> d) + atomically $ modifyTVar generation (+ 1) + return (not c) + atomically $ writeTVar closed True + let db = + DB + { dbRefs = refs, + dbGeneration = generation, + dbDirty = dirty, + dbPersistence = persistence, + dbClosing = closing, + dbClosed = closed + } + return 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 = do + atomically $ writeTVar (dbClosing db) True + atomically $ readTVar (dbClosed db) >>= bool retry (return ()) + persistentFinish (dbPersistence 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 f = bracket (openDB persistence) closeDB 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 maxLen = do + dirty <- readTVar (dbDirty db) + when (Map.size dirty > maxLen) 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 = do + c <- readTVar (dbClosing db) + when c $ error "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 txn = do + (result, gen) <- atomically $ do + result <- txn + gen <- readTVar (dbGeneration db) + dirty <- readTVar (dbDirty db) + if Map.null dirty + then return (result, Nothing) + else return (result, Just (gen + 1)) + case gen of + Just n -> + atomically $ + readTVar (dbGeneration db) >>= bool retry (return ()) . (>= n) + _ -> return () + return 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 key = do + failIfClosing db + SMap.lookup key (dbRefs db) >>= \case + Just (tr, weakRef) + | tr == typeRep (Proxy @a) -> + unsafeIOToSTM (deRefWeak weakRef) >>= \case + Just (SomeTVar ref) -> return (DBRef db key (unsafeCoerce ref)) + Nothing -> insert + | otherwise -> error "Type mismatch in DBRef" + Nothing -> insert + where + insert = do + ref <- newTVar Loading + ptr <- unsafeIOToSTM $ mkWeak ref (SomeTVar ref) (Just cleanupKey) + SMap.insert (typeRep (Proxy @a), ptr) key (dbRefs db) + v <- unsafeIOToSTM $ do + mvar <- newEmptyMVar + _ <- forkIO $ putMVar mvar =<< readKey + takeMVar mvar + writeTVar ref v + return (DBRef db key ref) + + readKey = do + readResult <- persistentRead (dbPersistence db) key + case readResult of + Just bs -> Present <$> atomically (decode db bs) + Nothing -> return Missing + + cleanupKey = + atomically $ + SMap.focus + ( Focus.updateM + ( \(tr, p) -> + unsafeIOToSTM (deRefWeak p) >>= \case + Nothing -> return Nothing + Just _ -> return (Just (tr, p)) + ) + ) + key + (dbRefs 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 db _ ref) = do + failIfClosing db + readTVar ref >>= \case + Loading -> retry + Missing -> return Nothing + Present a -> return (Just 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 db dbkey ref) a = do + failIfClosing db + writeTVar ref (Present a) + d <- readTVar (dbDirty db) + writeTVar (dbDirty db) (Map.insert dbkey (SomeTVar ref, Just (encode a)) 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 db dbkey ref) = do + failIfClosing db + writeTVar ref Missing + d <- readTVar (dbDirty db) + writeTVar (dbDirty db) (Map.insert dbkey (SomeTVar ref, Nothing) d) diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,68 @@ +{-# LANGUAGE LambdaCase #-} + +module Main where + +import Control.Concurrent (threadDelay) +import Control.Concurrent.STM (atomically) +import Control.Monad (replicateM) +import Data.IORef (newIORef, readIORef, writeIORef) +import PersistentSTM + ( Persistence (..), + filePersistence, + getDBRef, + readDBRef, + synchronously, + withDB, + writeDBRef, + ) +import System.IO.Temp (withSystemTempDirectory) +import Test.Hspec (Spec, describe, hspec, it, shouldBe) + +endToEndSpec :: Spec +endToEndSpec = + describe "end-to-end" $ do + it "maintains a persistent counter" $ do + withSystemTempDirectory "persistent-stm" $ \dir -> do + nums <- replicateM 100 $ do + p <- filePersistence dir + withDB p $ \db -> do + atomically $ do + ref <- getDBRef db "my-key" + readDBRef ref >>= \case + Nothing -> do + writeDBRef ref (1 :: Int) + return 1 + Just n -> do + writeDBRef ref (n + 1) + return (n + 1) + nums `shouldBe` [1 .. 100] + + it "waits for writes before synchronously returns" $ do + written <- newIORef False + let dummyPersistence = + Persistence + { persistentRead = const (return Nothing), + persistentWrite = const $ do + threadDelay 250000 + writeIORef written True, + persistentFinish = return () + } + withDB dummyPersistence $ \db -> do + synchronously db $ do + ref <- getDBRef db "my-key" + writeDBRef ref () + isWritten <- readIORef written + isWritten `shouldBe` True + + writeIORef written False + + withDB dummyPersistence $ \db -> do + atomically $ do + ref <- getDBRef db "my-key" + writeDBRef ref () + isWritten <- readIORef written + isWritten `shouldBe` False + +main :: IO () +main = hspec $ do + endToEndSpec