Make BlobFile functions exception safe

lsm-tree.cabal

@@ -356,6 +356,7 @@ test-suite lsm-tree-test
     Test.Database.LSMTree.Class
     Test.Database.LSMTree.Generators
     Test.Database.LSMTree.Internal
+    Test.Database.LSMTree.Internal.BlobFile.FS
     Test.Database.LSMTree.Internal.BloomFilter
     Test.Database.LSMTree.Internal.Chunk
     Test.Database.LSMTree.Internal.CRC32C

src/Database/LSMTree/Internal/BlobFile.hs

@@ -9,7 +9,8 @@ module Database.LSMTree.Internal.BlobFile (
   ) where
 
 import           Control.DeepSeq (NFData (..))
-import           Control.Monad.Class.MonadThrow (MonadThrow)
+import           Control.Monad.Class.MonadThrow (MonadCatch (bracketOnError),
+                     MonadThrow (..))
 import           Control.Monad.Primitive (PrimMonad)
 import           Control.RefCount
 import qualified Data.Primitive.ByteArray as P
@@ -51,24 +52,35 @@ instance NFData BlobSpan where
 
 -- | Open the given file to make a 'BlobFile'. The finaliser will close and
 -- delete the file.
+--
+-- REF: the resulting reference must be released once it is no longer used.
+--
+-- ASYNC: this should be called with asynchronous exceptions masked.
 {-# SPECIALISE openBlobFile :: HasCallStack => HasFS IO h -> FS.FsPath -> FS.OpenMode -> IO (Ref (BlobFile IO h)) #-}
 openBlobFile ::
-     PrimMonad m
+     (PrimMonad m, MonadCatch m)
   => HasCallStack
   => HasFS m h
   -> FS.FsPath
   -> FS.OpenMode
   -> m (Ref (BlobFile m h))
-openBlobFile fs path mode = do
-    blobFileHandle <- FS.hOpen fs path mode
-    let finaliser = do
-          FS.hClose fs blobFileHandle
-          FS.removeFile fs (FS.handlePath blobFileHandle)
-    newRef finaliser $ \blobFileRefCounter ->
-      BlobFile {
-        blobFileHandle,
-        blobFileRefCounter
-      }
+openBlobFile fs path mode =
+    bracketOnError (FS.hOpen fs path mode) (FS.hClose fs) $ \blobFileHandle -> do
+      let finaliser = do
+            FS.hClose fs blobFileHandle
+            -- If we fail to close the file handle, then we won't try to remove
+            -- the file.
+            --
+            -- TODO: this function takes ownership of the file path. The file is
+            -- removed when the blob file is finalised, which may lead to
+            -- surprise errors when the file is also deleted elsewhere. Maybe
+            -- file paths should be guarded by 'Ref's as well?
+            FS.removeFile fs (FS.handlePath blobFileHandle)
+      newRef finaliser $ \blobFileRefCounter ->
+        BlobFile {
+          blobFileHandle,
+          blobFileRefCounter
+        }
 
 {-# INLINE readBlob #-}
 readBlob ::

src/Database/LSMTree/Internal/Run.hs

@@ -183,6 +183,7 @@ setRunDataCaching hbio runKOpsFile NoCacheRunData = do
      RunDataCaching
   -> RunBuilder IO h
   -> IO (Ref (Run IO h)) #-}
+-- TODO: make exception safe
 fromMutable ::
      (MonadST m, MonadSTM m, MonadMask m)
   => RunDataCaching
@@ -192,6 +193,7 @@ fromMutable runRunDataCaching builder = do
     (runHasFS, runHasBlockIO, runRunFsPaths, runFilter, runIndex, runNumEntries) <-
       Builder.unsafeFinalise (runRunDataCaching == NoCacheRunData) builder
     runKOpsFile <- FS.hOpen runHasFS (runKOpsPath runRunFsPaths) FS.ReadMode
+    -- TODO: openBlobFile should be called with exceptions masked
     runBlobFile <- openBlobFile runHasFS (runBlobPath runRunFsPaths) FS.ReadMode
     setRunDataCaching runHasBlockIO runKOpsFile runRunDataCaching
     newRef (finaliser runHasFS runKOpsFile runBlobFile runRunFsPaths)
@@ -263,6 +265,7 @@ openFromDisk ::
   -> RunDataCaching
   -> RunFsPaths
   -> m (Ref (Run m h))
+-- TODO: make exception safe
 openFromDisk fs hbio runRunDataCaching runRunFsPaths = do
     expectedChecksums <-
        expectValidFile (runChecksumsPath runRunFsPaths) . fromChecksumsFile
@@ -282,6 +285,7 @@ openFromDisk fs hbio runRunDataCaching runRunFsPaths = do
         =<< readCRC (forRunIndexRaw expectedChecksums) (forRunIndexRaw paths)
 
     runKOpsFile <- FS.hOpen fs (runKOpsPath runRunFsPaths) FS.ReadMode
+    -- TODO: openBlobFile should be called with exceptions masked
     runBlobFile <- openBlobFile fs (runBlobPath runRunFsPaths) FS.ReadMode
     setRunDataCaching hbio runKOpsFile runRunDataCaching
     newRef (finaliser fs runKOpsFile runBlobFile runRunFsPaths) $ \runRefCounter ->

src/Database/LSMTree/Internal/WriteBufferBlobs.hs

@@ -134,9 +134,12 @@ open ::
   -> FS.FsPath
   -> FS.AllowExisting
   -> m (Ref (WriteBufferBlobs m h))
+-- TODO: make exception safe
 open fs blobFileName blobFileAllowExisting = do
     -- Must use read/write mode because we write blobs when adding, but
     -- we can also be asked to retrieve blobs at any time.
+    --
+    -- TODO: openBlobFile should be called with exceptions masked
     fromBlobFile fs =<< openBlobFile fs blobFileName (FS.ReadWriteMode blobFileAllowExisting)
 
 {-# SPECIALISE fromBlobFile :: HasFS IO h -> Ref (BlobFile IO h) -> IO (Ref (WriteBufferBlobs IO h)) #-}

test/Main.hs

@@ -8,6 +8,7 @@ import qualified Test.Data.Arena
 import qualified Test.Database.LSMTree.Class
 import qualified Test.Database.LSMTree.Generators
 import qualified Test.Database.LSMTree.Internal
+import qualified Test.Database.LSMTree.Internal.BlobFile.FS
 import qualified Test.Database.LSMTree.Internal.BloomFilter
 import qualified Test.Database.LSMTree.Internal.Chunk
 import qualified Test.Database.LSMTree.Internal.CRC32C
@@ -48,6 +49,7 @@ main = do
     , Test.Database.LSMTree.Class.tests
     , Test.Database.LSMTree.Generators.tests
     , Test.Database.LSMTree.Internal.tests
+    , Test.Database.LSMTree.Internal.BlobFile.FS.tests
     , Test.Database.LSMTree.Internal.BloomFilter.tests
     , Test.Database.LSMTree.Internal.Chunk.tests
     , Test.Database.LSMTree.Internal.CRC32C.tests

test/Test/Database/LSMTree/Internal/BlobFile/FS.hs

@@ -0,0 +1,64 @@
+module Test.Database.LSMTree.Internal.BlobFile.FS (tests) where
+
+import           Control.Concurrent.Class.MonadSTM.Strict
+import           Control.Monad
+import           Control.Monad.Class.MonadThrow
+import           Control.RefCount
+import           Database.LSMTree.Internal.BlobFile
+import           System.FS.API
+import           System.FS.Sim.Error hiding (genErrors)
+import qualified System.FS.Sim.MockFS as MockFS
+import           Test.Tasty
+import           Test.Tasty.QuickCheck as QC
+import           Test.Util.FS
+
+tests :: TestTree
+tests = testGroup "Test.Database.LSMTree.Internal.BlobFile.FS" [
+      testProperty "prop_fault_openRelease" prop_fault_openRelease
+    ]
+
+-- Test that opening and releasing a blob file properly cleans handles and files
+-- in the presence of disk faults.
+prop_fault_openRelease ::
+     Bool -- ^ create the file or not
+  -> OpenMode
+  -> NoCleanupErrors
+  -> NoCleanupErrors
+  -> Property
+prop_fault_openRelease doCreateFile om
+  (NoCleanupErrors openErrors)
+  (NoCleanupErrors releaseErrors) =
+    ioProperty $
+    withSimErrorHasFS propPost MockFS.empty emptyErrors $ \hfs fsVar errsVar -> do
+      when doCreateFile $
+        withFile hfs path (WriteMode MustBeNew) $ \_ -> pure ()
+      eith <- try @_ @FsError $
+        bracket (acquire hfs errsVar) (release errsVar) $ \_blobFile -> do
+          fs' <- atomically $ readTMVar fsVar
+          pure $ propNumOpenHandles 1 fs' .&&. propNumDirEntries (mkFsPath []) 1 fs'
+      pure $ case eith of
+        Left{}     ->
+          label "FsError" $ property True
+        Right prop ->
+          label "Success" $ prop
+  where
+    root = mkFsPath []
+    path = mkFsPath ["blobfile"]
+
+    acquire hfs errsVar =
+        withErrors errsVar openErrors $ openBlobFile hfs path om
+
+    release errsVar blobFile =
+        withErrors errsVar releaseErrors $ releaseRef blobFile
+
+    propPost fs = propNoOpenHandles fs .&&.
+        if doCreateFile then
+          case allowExisting om of
+            AllowExisting ->
+              -- TODO: fix, see the TODO on openBlobFile
+              propNoDirEntries root fs .||. propNumDirEntries root 1 fs
+            MustBeNew ->
+              propNumDirEntries root 1 fs
+        else
+          propNoDirEntries root fs
+

test/Test/Util/FS.hs

@@ -1,3 +1,4 @@
+{-# OPTIONS_GHC -Wno-orphans #-}
 
 {- HLINT ignore "Redundant if" -}
 
@@ -21,6 +22,8 @@ module Test.Util.FS (
   , assertNumOpenHandles
     -- * Equality
   , approximateEqStream
+    -- * Arbitrary
+  , NoCleanupErrors (..)
   ) where
 
 import           Control.Concurrent.Class.MonadMVar
@@ -39,6 +42,7 @@ import           System.FS.IO
 import           System.FS.Sim.Error
 import           System.FS.Sim.MockFS
 import           System.FS.Sim.STM
+import qualified System.FS.Sim.Stream as Stream
 import           System.FS.Sim.Stream (InternalInfo (..), Stream (..))
 import           System.IO.Temp
 import           Test.QuickCheck
@@ -216,3 +220,80 @@ approximateEqStream (UnsafeStream infoXs xs) (UnsafeStream infoYs ys) =
       (Infinite, Infinite) -> True
       (Finite, Finite)     ->  xs == ys
       (_, _)               -> False
+
+{-------------------------------------------------------------------------------
+  Arbitrary
+-------------------------------------------------------------------------------}
+
+-- | No errors on closing file handles and removing files
+newtype NoCleanupErrors = NoCleanupErrors Errors
+  deriving stock Show
+
+mkNoCleanupErrors :: Errors -> NoCleanupErrors
+mkNoCleanupErrors errs = NoCleanupErrors $ errs {
+      hCloseE = Stream.empty
+    , removeFileE = Stream.empty
+    }
+
+instance Arbitrary NoCleanupErrors where
+  arbitrary = do
+      errs <- arbitrary
+      pure $ mkNoCleanupErrors errs
+
+  -- The shrinker for 'Errors' does not re-introduce 'hCloseE' and 'removeFile'.
+  shrink (NoCleanupErrors errs) = NoCleanupErrors <$> shrink errs
+
+newtype TestOpenMode = TestOpenMode OpenMode
+  deriving stock Show
+
+instance Arbitrary OpenMode where
+  arbitrary = genOpenMode
+  shrink = shrinkOpenMode
+
+genOpenMode :: Gen OpenMode
+genOpenMode = oneof [
+      pure ReadMode
+    , WriteMode <$> genAllowExisting
+    , ReadWriteMode <$> genAllowExisting
+    , AppendMode <$> genAllowExisting
+    ]
+  where
+    _coveredAllCases x = case x of
+        ReadMode{}      -> ()
+        WriteMode{}     -> ()
+        ReadWriteMode{} -> ()
+        AppendMode{}    -> ()
+
+shrinkOpenMode :: OpenMode -> [OpenMode]
+shrinkOpenMode = \case
+    ReadMode -> []
+    WriteMode ae ->
+        ReadMode
+      : (WriteMode <$> shrinkAllowExisting ae)
+    ReadWriteMode ae ->
+        ReadMode
+      : WriteMode ae
+      : (ReadWriteMode <$> shrinkAllowExisting ae)
+    AppendMode ae ->
+        ReadMode
+      : WriteMode ae
+      : ReadWriteMode ae
+      : (AppendMode <$> shrinkAllowExisting ae)
+
+instance Arbitrary AllowExisting where
+  arbitrary = genAllowExisting
+  shrink = shrinkAllowExisting
+
+genAllowExisting :: Gen AllowExisting
+genAllowExisting = elements [
+      AllowExisting
+    , MustBeNew
+    ]
+  where
+    _coveredAllCases x = case x of
+        AllowExisting -> ()
+        MustBeNew     -> ()
+
+shrinkAllowExisting :: AllowExisting -> [AllowExisting]
+shrinkAllowExisting AllowExisting = []
+shrinkAllowExisting MustBeNew     = [AllowExisting]