Merge pull request #532 from IntersectMBO/jdral/fault-test-writebufferreader

Make `WriteBufferReader` functions exception safe

Author: wenkokke

lsm-tree.cabal

@@ -382,6 +382,7 @@ test-suite lsm-tree-test
     Test.Database.LSMTree.Internal.Vector
     Test.Database.LSMTree.Internal.Vector.Growing
     Test.Database.LSMTree.Internal.WriteBufferBlobs.FS
+    Test.Database.LSMTree.Internal.WriteBufferReader.FS
     Test.Database.LSMTree.Model.Table
     Test.Database.LSMTree.Monoidal
     Test.Database.LSMTree.StateMachine

src-extras/Database/LSMTree/Extras/RunData.hs

@@ -141,9 +141,7 @@ withSerialisedWriteBuffer hfs hbio wbPaths wb wbb =
     for_ [ Paths.writeBufferKOpsPath wbPaths
          , Paths.writeBufferBlobPath wbPaths
          , Paths.writeBufferChecksumsPath wbPaths
-         ] $ \fsPath -> do
-      fsPathExists <- FS.doesFileExist hfs fsPath
-      when fsPathExists $ FS.removeFile hfs fsPath
+         ] $ FS.removeFile hfs
 
 {-------------------------------------------------------------------------------
   RunData

src/Database/LSMTree/Internal/RunReader.hs

@@ -306,6 +306,10 @@ readDiskPage ::
 readDiskPage fs h = do
     mba <- newPinnedByteArray pageSize
     -- TODO: make sure no other exception type can be thrown
+    --
+    -- TODO: if FS.FsReachEOF is thrown as an injected disk fault, then we
+    -- incorrectly deduce that the file has no more contents. We should probably
+    -- use an explicit file pointer instead in the style of 'FilePointer'.
     handleJust (guard . FS.isFsErrorType FS.FsReachedEOF) (\_ -> pure Nothing) $ do
       bytesRead <- FS.hGetBufExactly fs h mba 0 (fromIntegral pageSize)
       assert (fromIntegral bytesRead == pageSize) $ pure ()

src/Database/LSMTree/Internal/WriteBufferReader.hs

@@ -7,7 +7,8 @@ module Database.LSMTree.Internal.WriteBufferReader (
 import           Control.Concurrent.Class.MonadMVar.Strict
 import           Control.Monad.Class.MonadST (MonadST (..))
 import           Control.Monad.Class.MonadSTM (MonadSTM (..))
-import           Control.Monad.Class.MonadThrow (MonadMask, MonadThrow (..))
+import           Control.Monad.Class.MonadThrow (MonadMask, MonadThrow (..),
+                     bracketOnError)
 import           Control.Monad.Primitive (PrimMonad (..))
 import           Control.RefCount (Ref, dupRef, releaseRef)
 import           Data.Primitive.MutVar (MutVar, newMutVar, readMutVar,
@@ -43,8 +44,8 @@ import           System.FS.BlockIO.API (HasBlockIO)
   #-}
 -- | Read a serialised `WriteBuffer` back into memory.
 --
---   NOTE: The `BlobFile` argument /must be/ the blob file associated with the
---         write buffer; @`readWriteBuffer`@ does not check this.
+-- The argument blob file ('BlobFile') must be the file associated with the
+-- argument key\/ops file ('ForKOps'). 'readWriteBuffer' does not check this.
 readWriteBuffer ::
      (MonadMVar m, MonadMask m, MonadSTM m, MonadST m)
   => ResolveSerialisedValue
@@ -54,7 +55,7 @@ readWriteBuffer ::
   -> Ref (BlobFile m h)
   -> m WriteBuffer
 readWriteBuffer resolve hfs hbio kOpsPath blobFile =
-  bracket (new hfs hbio kOpsPath blobFile) close $ readEntries
+    bracket (new hfs hbio kOpsPath blobFile) close $ readEntries
   where
     readEntries reader = readEntriesAcc WB.empty
       where
@@ -88,30 +89,42 @@ data WriteBufferReader m h = WriteBufferReader {
     -> IO (WriteBufferReader IO h)
   #-}
 -- | See 'Database.LSMTree.Internal.RunReader.new'.
+--
+-- REF: the resulting 'WriteBufferReader' must be closed once it is no longer
+-- used.
+--
+-- ASYNC: this should be called with asynchronous exceptions masked because it
+-- allocates/creates resources.
 new :: forall m h.
      (MonadMVar m, MonadST m, MonadMask m)
   => HasFS m h
   -> HasBlockIO m h
   -> ForKOps FS.FsPath
   -> Ref (BlobFile m h)
   -> m (WriteBufferReader m h)
-new readerHasFS readerHasBlockIO kOpsPath blobFile = do
-  readerKOpsHandle <- FS.hOpen readerHasFS (unForKOps kOpsPath) FS.ReadMode
-  -- Double the file readahead window (only applies to this file descriptor)
-  FS.hAdviseAll readerHasBlockIO readerKOpsHandle FS.AdviceSequential
-  readerBlobFile <- dupRef blobFile
-  -- Load first page from disk, if it exists.
-  readerCurrentEntryNo <- newPrimVar (0 :: Word16)
-  firstPage <- readDiskPage readerHasFS readerKOpsHandle
-  readerCurrentPage <- newMutVar firstPage
-  pure $ WriteBufferReader{..}
+new readerHasFS readerHasBlockIO kOpsPath blobFile =
+    bracketOnError openKOps (FS.hClose readerHasFS) $ \readerKOpsHandle -> do
+      -- Double the file readahead window (only applies to this file descriptor)
+      FS.hAdviseAll readerHasBlockIO readerKOpsHandle FS.AdviceSequential
+      bracketOnError (dupRef blobFile) releaseRef $ \readerBlobFile -> do
+        -- Load first page from disk, if it exists.
+        readerCurrentEntryNo <- newPrimVar (0 :: Word16)
+        firstPage <- readDiskPage readerHasFS readerKOpsHandle
+        readerCurrentPage <- newMutVar firstPage
+        pure $ WriteBufferReader{..}
+  where
+    openKOps = FS.hOpen readerHasFS (unForKOps kOpsPath) FS.ReadMode
 
 {-# SPECIALISE
   next ::
        WriteBufferReader IO h
     -> IO (Result IO h)
   #-}
 -- | See 'Database.LSMTree.Internal.RunReader.next'.
+--
+-- TODO: 'next' is currently only used in 'readWriteBuffer', where it is a safe
+-- use of an unsafe function. If this function is ever exported and used
+-- directly, the TODOs in the body of this function should be addressed first.
 next :: forall m h.
      (MonadSTM m, MonadST m, MonadMask m)
   => WriteBufferReader m h
@@ -124,13 +137,21 @@ next WriteBufferReader {..} = do
         entryNo <- readPrimVar readerCurrentEntryNo
         go entryNo page
   where
+    -- TODO: if 'readerCurrentEntryNo' is incremented but an exception is thrown
+    -- before the 'Result' is used by the caller of 'next', then we'll lose that
+    -- 'Result'. The following call to 'next' will not return the 'Result' we
+    -- missed.
     go :: Word16 -> RawPage -> m (Result m h)
     go !entryNo !page =
         -- take entry from current page (resolve blob if necessary)
         case rawPageIndex page entryNo of
           IndexNotPresent -> do
             -- if it is past the last one, load a new page from disk, try again
             newPage <- readDiskPage readerHasFS readerKOpsHandle
+            -- TODO: if the next disk page is read but an (async) exception is
+            -- thrown just before updating the MutVar below, then we lose the
+            -- disk page because 'readDiskPage' has already updated its file
+            -- pointer.
             stToIO $ writeMutVar readerCurrentPage newPage
             case newPage of
               Nothing -> do
@@ -154,10 +175,14 @@ next WriteBufferReader {..} = do
             return (ReadEntry key rawEntry)
 
 {-# SPECIALISE close :: WriteBufferReader IO h -> IO () #-}
+-- | Close the 'WriteBufferReader'.
+--
+-- ASYNC: this should be called with asynchronous exceptions masked because it
+-- releases/removes resources.
 close ::
      (MonadMask m, PrimMonad m)
   => WriteBufferReader m h
   -> m ()
 close WriteBufferReader{..} = do
   FS.hClose readerHasFS readerKOpsHandle
-  releaseRef readerBlobFile
+    `finally` releaseRef readerBlobFile

test/Main.hs

@@ -34,6 +34,7 @@ import qualified Test.Database.LSMTree.Internal.Snapshot.FS
 import qualified Test.Database.LSMTree.Internal.Vector
 import qualified Test.Database.LSMTree.Internal.Vector.Growing
 import qualified Test.Database.LSMTree.Internal.WriteBufferBlobs.FS
+import qualified Test.Database.LSMTree.Internal.WriteBufferReader.FS
 import qualified Test.Database.LSMTree.Model.Table
 import qualified Test.Database.LSMTree.Monoidal
 import qualified Test.Database.LSMTree.StateMachine
@@ -76,6 +77,7 @@ main = do
     , Test.Database.LSMTree.Internal.Vector.tests
     , Test.Database.LSMTree.Internal.Vector.Growing.tests
     , Test.Database.LSMTree.Internal.WriteBufferBlobs.FS.tests
+    , Test.Database.LSMTree.Internal.WriteBufferReader.FS.tests
     , Test.Database.LSMTree.Model.Table.tests
     , Test.Database.LSMTree.Monoidal.tests
     , Test.Database.LSMTree.UnitTests.tests

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

@@ -0,0 +1,80 @@
+module Test.Database.LSMTree.Internal.WriteBufferReader.FS (tests) where
+
+
+import           Control.Concurrent.Class.MonadSTM.Strict (MonadSTM (..))
+import           Control.Concurrent.Class.MonadSTM.Strict.TMVar
+import           Control.Monad.Class.MonadThrow
+import           Control.RefCount
+import           Database.LSMTree.Extras.Generators ()
+import           Database.LSMTree.Extras.RunData (RunData,
+                     withRunDataAsWriteBuffer, withSerialisedWriteBuffer)
+import           Database.LSMTree.Internal.Paths (ForKOps (ForKOps),
+                     WriteBufferFsPaths (WriteBufferFsPaths),
+                     writeBufferKOpsPath)
+import           Database.LSMTree.Internal.RunNumber (RunNumber (RunNumber))
+import           Database.LSMTree.Internal.Serialise (SerialisedBlob,
+                     SerialisedKey, SerialisedValue (..))
+import qualified Database.LSMTree.Internal.WriteBufferBlobs as WBB
+import           Database.LSMTree.Internal.WriteBufferReader
+import           System.FS.API
+import           System.FS.Sim.Error hiding (genErrors)
+import qualified System.FS.Sim.MockFS as MockFS
+import qualified System.FS.Sim.Stream as Stream
+import           Test.Tasty
+import           Test.Tasty.QuickCheck as QC
+import           Test.Util.FS
+
+tests :: TestTree
+tests = testGroup "Test.Database.LSMTree.Internal.WriteBufferReader.FS" [
+      testProperty "prop_fault_WriteBufferReader" prop_fault_WriteBufferReader
+    ]
+
+-- | Test that 'writeWriteBuffer' roundtrips with 'readWriteBuffer', and test
+-- that the presence of disk faults for the latter does not leak file handles
+-- and files.
+prop_fault_WriteBufferReader ::
+     NoCleanupErrors
+  -> RunData SerialisedKey SerialisedValue SerialisedBlob
+  -> Property
+prop_fault_WriteBufferReader (NoCleanupErrors readErrors) rdata =
+    ioProperty $
+    withSimErrorHasBlockIO propPost MockFS.empty emptyErrors $ \hfs hbio fsVar errsVar ->
+    withRunDataAsWriteBuffer hfs resolve inPath rdata $ \wb wbb ->
+    withSerialisedWriteBuffer hfs hbio outPath wb wbb $ do
+      fsBefore <- atomically $ readTMVar fsVar
+      eith <-
+        try @_ @FsError $
+          withErrors errsVar readErrors' $
+            withRef wbb $ \wbb' -> do
+              wb' <- readWriteBuffer resolve hfs hbio outKOpsPath (WBB.blobFile wbb')
+              pure (wb === wb')
+
+      fsAfter <- atomically $ readTMVar fsVar
+      pure $
+        case eith of
+          Left{}     -> do
+            label "FsError" $ property True
+          Right prop ->
+            label "Success" $ prop .&&. propEqNumDirEntries root fsBefore fsAfter
+  where
+    root = mkFsPath []
+    -- The run number for the original write buffer. Primarily used to name the
+    -- 'WriteBufferBlobs' corresponding to the write buffer.
+    inPath = WriteBufferFsPaths root (RunNumber 0)
+    -- The run number for the serialised write buffer. Used to name all files
+    -- that are the result of serialising the write buffer.
+    outPath = WriteBufferFsPaths root (RunNumber 1)
+    outKOpsPath = ForKOps (writeBufferKOpsPath outPath)
+    resolve (SerialisedValue x) (SerialisedValue y) = SerialisedValue (x <> y)
+    propPost fs = propNoOpenHandles fs .&&. propNoDirEntries root fs
+
+    -- TODO: fix, see the TODO on readDiskPage
+    readErrors' = readErrors {
+          hGetBufSomeE = Stream.filter (not . isFsReachedEOF) (hGetBufSomeE readErrors)
+        }
+
+    isFsReachedEOF Nothing = False
+    isFsReachedEOF (Just (Left e)) = case e of
+        FsReachedEOF -> True
+        _            -> False
+    isFsReachedEOF (Just (Right _)) = False

test/Test/Database/LSMTree/StateMachine.hs

@@ -134,8 +134,8 @@ import qualified Test.QuickCheck.StateModel.Lockstep.Defaults as Lockstep.Defaul
 import qualified Test.QuickCheck.StateModel.Lockstep.Run as Lockstep.Run
 import           Test.Tasty (TestTree, testGroup)
 import           Test.Tasty.QuickCheck (testProperty)
-import           Test.Util.FS (approximateEqStream, propNoOpenHandles,
-                     propNumOpenHandles)
+import           Test.Util.FS (approximateEqStream, noRemoveDirectoryRecursiveE,
+                     propNoOpenHandles, propNumOpenHandles)
 import           Test.Util.PrettyProxy
 import           Test.Util.QLS
 import           Test.Util.TypeFamilyWrappers (WrapBlob (..), WrapBlobRef (..),
@@ -1434,8 +1434,7 @@ arbitraryActionWithVars _ label ctx (ModelState st _stats) =
      ++ [ (1, fmap Some $ OpenSnapshot @k @v @b PrettyProxy <$>
                 genErrors <*> pure label <*> genUsedSnapshotName)
         | not (null usedSnapshotNames)
-          -- TODO: generate errors
-        , let genErrors = pure Nothing
+        , let genErrors = fmap noRemoveDirectoryRecursiveE <$> QC.arbitrary
         ]
 
      ++ [ (1, fmap Some $ DeleteSnapshot <$> genUsedSnapshotName)

test/Test/Util/FS.hs

@@ -16,8 +16,10 @@ module Test.Util.FS (
   , propTrivial
   , propNumOpenHandles
   , propNoOpenHandles
+  , numDirEntries
   , propNumDirEntries
   , propNoDirEntries
+  , propEqNumDirEntries
   , assertNoOpenHandles
   , assertNumOpenHandles
     -- * Equality
@@ -30,6 +32,10 @@ module Test.Util.FS (
     -- * Corruption
   , flipFileBit
   , hFlipBit
+    -- * Errors
+  , noHCloseE
+  , noRemoveFileE
+  , noRemoveDirectoryRecursiveE
     -- * Arbitrary
   , FsPathComponent (..)
   , fsPathComponentFsPath
@@ -190,6 +196,12 @@ propNumOpenHandles expected fs =
 propNoOpenHandles :: MockFS -> Property
 propNoOpenHandles fs = propNumOpenHandles 0 fs
 
+numDirEntries :: FsPath -> MockFS -> Int
+numDirEntries path fs = Set.size contents
+  where
+    (contents, _) =
+        runSimOrThrow $ runSimFS fs $ \hfs -> FS.listDirectory hfs path
+
 {-# INLINABLE propNumDirEntries #-}
 propNumDirEntries :: FsPath -> Int -> MockFS -> Property
 propNumDirEntries path expected fs =
@@ -199,19 +211,30 @@ propNumDirEntries path expected fs =
         (show path) actual) $
     printMockFSOnFailure fs $
     expected === actual
-  where
-    actual =
-      let (contents, _) = runSimOrThrow $
-            runSimFS fs $ \hfs ->
-              FS.listDirectory hfs path
-      in  Set.size contents
+  where actual = numDirEntries path fs
 
 {-# INLINABLE propNoDirEntries #-}
 propNoDirEntries :: FsPath -> MockFS -> Property
 propNoDirEntries path fs = propNumDirEntries path 0 fs
 
+{-# INLINABLE propEqNumDirEntries #-}
+propEqNumDirEntries :: FsPath -> MockFS -> MockFS -> Property
+propEqNumDirEntries path lhsFs rhsFs =
+    counterexample
+      (printf "The LHS has %d entries in the directory at %s, but the RHS has %d"
+        lhs (show path) rhs) $
+      printMockFSOnFailureWith "Mocked file system (LHS)" lhsFs $
+      printMockFSOnFailureWith "Mocked file system (RHS)" rhsFs $
+      lhs === rhs
+  where
+    lhs = numDirEntries path lhsFs
+    rhs = numDirEntries path rhsFs
+
 printMockFSOnFailure :: Testable prop => MockFS -> prop -> Property
-printMockFSOnFailure fs = counterexample ("Mocked file system: " <> pretty fs)
+printMockFSOnFailure = printMockFSOnFailureWith "Mocked file system"
+
+printMockFSOnFailureWith :: Testable prop => String -> MockFS -> prop -> Property
+printMockFSOnFailureWith s fs = counterexample (s <> ": " <> pretty fs)
 
 assertNoOpenHandles :: HasCallStack => MockFS -> a -> a
 assertNoOpenHandles fs = assertNumOpenHandles fs 0
@@ -354,6 +377,19 @@ hFlipBit hfs h bitOffset = do
     void $ hPutBufExactlyAt hfs h buf bufOff count off
 
 
+{-------------------------------------------------------------------------------
+  Errors
+-------------------------------------------------------------------------------}
+
+noHCloseE :: Errors -> Errors
+noHCloseE errs = errs { hCloseE = Stream.empty }
+
+noRemoveFileE :: Errors -> Errors
+noRemoveFileE errs = errs { removeFileE = Stream.empty }
+
+noRemoveDirectoryRecursiveE :: Errors -> Errors
+noRemoveDirectoryRecursiveE errs = errs { removeDirectoryRecursiveE = Stream.empty }
+
 {-------------------------------------------------------------------------------
   Arbitrary
 -------------------------------------------------------------------------------}
@@ -396,10 +432,11 @@ newtype NoCleanupErrors = NoCleanupErrors Errors
   deriving stock Show
 
 mkNoCleanupErrors :: Errors -> NoCleanupErrors
-mkNoCleanupErrors errs = NoCleanupErrors $ errs {
-      hCloseE = Stream.empty
-    , removeFileE = Stream.empty
-    }
+mkNoCleanupErrors errs = NoCleanupErrors $
+      noHCloseE
+    $ noRemoveFileE
+    $ noRemoveDirectoryRecursiveE
+    $ errs
 
 instance Arbitrary NoCleanupErrors where
   arbitrary = do