Fix adding/removing references on table content

In some places, we were adding/removing references but not taking into account
rollbacks/delays that should be performed when exceptions occur. In particular,
correct acquisition and release of write buffer blob files were not guarded by
`allocateTemp` and `freeTemp`.

In addition, to make our reference accounting slightly less error prone, we make
a few related changes:

* We add new  `addReferenceTableContent` and `removeReferenceTableContent`
functions, which replace most occurrences of manually increasing/decreasing
reference counts for the components of a  `TableContent`. For example, in
`duplicate` and `close` we now use the new functions.

* `newWith` now uses a `TempRegistry`, which fixes a double-free TODO.

* `newWith` now requires a full `TableContent` to be passed in.

Author: jorisdral

src/Database/LSMTree/Internal.hs

@@ -104,9 +104,7 @@ import           Database.LSMTree.Internal.Serialise (SerialisedBlob (..),
                      SerialisedKey, SerialisedValue)
 import           Database.LSMTree.Internal.UniqCounter
 import qualified Database.LSMTree.Internal.Vector as V
-import           Database.LSMTree.Internal.WriteBuffer (WriteBuffer)
 import qualified Database.LSMTree.Internal.WriteBuffer as WB
-import           Database.LSMTree.Internal.WriteBufferBlobs (WriteBufferBlobs)
 import qualified Database.LSMTree.Internal.WriteBufferBlobs as WBB
 import qualified System.FS.API as FS
 import           System.FS.API (FsError, FsErrorPath (..), FsPath, Handle,
@@ -620,53 +618,58 @@ withTable sesh conf = bracket (new sesh conf) close
   -> IO (TableHandle IO h) #-}
 -- | See 'Database.LSMTree.Normal.new'.
 new ::
-     (MonadSTM m, MonadThrow m, MonadMVar m, PrimMonad m)
+     (MonadSTM m, MonadMVar m, PrimMonad m, MonadMask m)
   => Session m h
   -> TableConfig
   -> m (TableHandle m h)
 new sesh conf = do
     traceWith (sessionTracer sesh) TraceNewTable
-    withOpenSession sesh $ \seshEnv -> do
-      am <- newArenaManager
-      blobpath <- Paths.tableBlobPath (sessionRoot seshEnv) <$>
-                    incrUniqCounter (sessionUniqCounter seshEnv)
-      wbblobs  <- WBB.new (sessionHasFS seshEnv) blobpath
-      newWith sesh seshEnv conf am WB.empty wbblobs V.empty
+    withOpenSession sesh $ \seshEnv ->
+      withTempRegistry $ \reg -> do
+        am <- newArenaManager
+        blobpath <- Paths.tableBlobPath (sessionRoot seshEnv) <$>
+                      incrUniqCounter (sessionUniqCounter seshEnv)
+        tableWriteBufferBlobs
+          <- allocateTemp reg
+               (WBB.new (sessionHasFS seshEnv) blobpath)
+               WBB.removeReference
+        let tableWriteBuffer = WB.empty
+            tableLevels = V.empty
+        tableCache <- mkLevelsCache tableLevels
+        let tc = TableContent {
+                tableWriteBuffer
+              , tableWriteBufferBlobs
+              , tableLevels
+              , tableCache
+              }
+        newWith reg sesh seshEnv conf am tc
 
 {-# SPECIALISE newWith ::
-     Session IO h
+     TempRegistry IO
+  -> Session IO h
   -> SessionEnv IO h
   -> TableConfig
   -> ArenaManager RealWorld
-  -> WriteBuffer
-  -> WriteBufferBlobs IO h
-  -> Levels IO h
+  -> TableContent IO h
   -> IO (TableHandle IO h) #-}
 newWith ::
      (MonadSTM m, MonadMVar m)
-  => Session m h
+  => TempRegistry m
+  -> Session m h
   -> SessionEnv m h
   -> TableConfig
   -> ArenaManager (PrimState m)
-  -> WriteBuffer
-  -> WriteBufferBlobs m h
-  -> Levels m h
+  -> TableContent m h
   -> m (TableHandle m h)
-newWith sesh seshEnv conf !am !wb !wbblobs !levels = do
+newWith reg sesh seshEnv conf !am !tc = do
     tableId <- incrUniqCounter (sessionUniqCounter seshEnv)
     let tr = TraceTable (uniqueToWord64 tableId) `contramap` sessionTracer sesh
     traceWith tr $ TraceCreateTableHandle conf
-    cache <- mkLevelsCache levels
     -- The session is kept open until we've updated the session's set of tracked
     -- tables. If 'closeSession' is called by another thread while this code
     -- block is being executed, that thread will block until it reads the
     -- /updated/ set of tracked tables.
-    contentVar <- RW.new $ TableContent
-        { tableWriteBuffer = wb
-        , tableWriteBufferBlobs = wbblobs
-        , tableLevels = levels
-        , tableCache = cache
-        }
+    contentVar <- RW.new $ tc
     tableVar <- RW.new $ TableHandleOpen $ TableHandleEnv {
           tableSession = sesh
         , tableSessionEnv = seshEnv
@@ -675,7 +678,8 @@ newWith sesh seshEnv conf !am !wb !wbblobs !levels = do
         }
     let !th = TableHandle conf tableVar am tr
     -- Track the current table
-    modifyMVar_ (sessionOpenTables seshEnv) $ pure . Map.insert (uniqueToWord64 tableId) th
+    freeTemp reg $ modifyMVar_ (sessionOpenTables seshEnv)
+                 $ pure . Map.insert (uniqueToWord64 tableId) th
     pure $! th
 
 {-# SPECIALISE close :: TableHandle IO h -> IO () #-}
@@ -686,17 +690,17 @@ close ::
   -> m ()
 close th = do
     traceWith (tableTracer th) TraceCloseTable
-    RW.withWriteAccess_ (tableHandleState th) $ \case
+    modifyWithTempRegistry_
+      (RW.unsafeAcquireWriteAccess (tableHandleState th))
+      (atomically . RW.unsafeReleaseWriteAccess (tableHandleState th)) $ \reg -> \case
       TableHandleClosed -> pure TableHandleClosed
       TableHandleOpen thEnv -> do
         -- Since we have a write lock on the table state, we know that we are the
         -- only thread currently closing the table. We can safely make the session
         -- forget about this table.
-        -- TODO: use TempRegistry
-        tableSessionUntrackTable thEnv
+        freeTemp reg (tableSessionUntrackTable thEnv)
         RW.withWriteAccess_ (tableContent thEnv) $ \tc -> do
-          forRunM_ (tableLevels tc) Run.removeReference
-          WBB.removeReference (tableWriteBufferBlobs tc)
+          removeReferenceTableContent reg tc
           pure tc
         pure TableHandleClosed
 
@@ -1289,12 +1293,22 @@ open sesh label override snap = do
         let runPaths = V.map (bimap (second $ RunFsPaths (Paths.activeDir $ sessionRoot seshEnv))
                                     (V.map (RunFsPaths (Paths.activeDir $ sessionRoot seshEnv))))
                             runNumbers
-        lvls <- openLevels reg hfs hbio (confDiskCachePolicy conf') runPaths
+
         am <- newArenaManager
         blobpath <- Paths.tableBlobPath (sessionRoot seshEnv) <$>
                       incrUniqCounter (sessionUniqCounter seshEnv)
-        wbblobs  <- WBB.new hfs blobpath
-        newWith sesh seshEnv conf' am WB.empty wbblobs lvls
+        tableWriteBufferBlobs
+          <- allocateTemp reg
+               (WBB.new hfs blobpath)
+               WBB.removeReference
+        tableLevels <- openLevels reg hfs hbio (confDiskCachePolicy conf') runPaths
+        tableCache <- mkLevelsCache tableLevels
+        newWith reg sesh seshEnv conf' am $! TableContent {
+            tableWriteBuffer = WB.empty
+          , tableWriteBufferBlobs
+          , tableLevels
+          , tableCache
+          }
 
 {-# SPECIALISE openLevels ::
      TempRegistry IO
@@ -1390,21 +1404,12 @@ duplicate th = do
           -- The table contents escape the read access, but we just added references
           -- to each run so it is safe.
           content <- RW.withReadAccess (tableContent thEnv) $ \content -> do
-            forRunM_ (tableLevels content) $ \r -> do
-              allocateTemp reg
-                (Run.addReference r)
-                (\_ -> Run.removeReference r)
+            addReferenceTableContent reg content
             pure content
-          WBB.addReference (tableWriteBufferBlobs content)
-          -- TODO: Fix possible double-free! See 'newCursor'.
-          -- In `newWith`, the table handle (in the open state) gets added to
-          -- `sessionOpenTables', even if later an async exception occurs and
-          -- the temp registry rolls back all allocations.
           newWith
+            reg
             (tableSession thEnv)
             (tableSessionEnv thEnv)
             (tableConfig th)
             (tableHandleArenaManager th)
-            (tableWriteBuffer content)
-            (tableWriteBufferBlobs content)
-            (tableLevels content)
+            content

src/Database/LSMTree/Internal/MergeSchedule.hs

@@ -7,6 +7,8 @@ module Database.LSMTree.Internal.MergeSchedule (
   , MergeTrace (..)
     -- * Table content
   , TableContent (..)
+  , addReferenceTableContent
+  , removeReferenceTableContent
     -- * Levels cache
   , LevelsCache (..)
   , mkLevelsCache
@@ -15,7 +17,6 @@ module Database.LSMTree.Internal.MergeSchedule (
   , Level (..)
   , MergingRun (..)
   , MergingRunState (..)
-  , forRunM_
     -- * Flushes and scheduled merges
   , updatesWithInterleavedFlushes
   , flushWriteBuffer
@@ -114,6 +115,28 @@ data TableContent m h = TableContent {
   , tableCache            :: !(LevelsCache m (Handle h))
   }
 
+{-# SPECIALISE addReferenceTableContent :: TempRegistry IO -> TableContent IO h -> IO () #-}
+addReferenceTableContent ::
+     (PrimMonad m, MonadMask m, MonadMVar m)
+  => TempRegistry m
+  -> TableContent m h
+  -> m ()
+addReferenceTableContent reg (TableContent _wb wbb levels _cache) = do
+    allocateTemp reg (WBB.addReference wbb) (\_ -> WBB.removeReference wbb)
+    addReferenceLevels reg levels
+    -- references on the cache are implicit
+
+{-# SPECIALISE removeReferenceTableContent :: TempRegistry IO -> TableContent IO h -> IO () #-}
+removeReferenceTableContent ::
+     (PrimMonad m, MonadMask m, MonadMVar m)
+  => TempRegistry m
+  -> TableContent m h
+  -> m ()
+removeReferenceTableContent reg (TableContent _wb wbb levels _cache) = do
+    freeTemp reg (WBB.removeReference wbb)
+    removeReferenceLevels reg levels
+    -- references on the cache are implicit
+
 {-------------------------------------------------------------------------------
   Levels cache
 -------------------------------------------------------------------------------}
@@ -170,23 +193,46 @@ data MergingRunState m h =
     CompletedMerge !(Run m (Handle h))
   | OngoingMerge !(V.Vector (Run m (Handle h))) !(Merge m h)
 
-{-# SPECIALISE forRunM_ ::
+{-# SPECIALISE addReferenceLevels :: TempRegistry IO -> Levels IO h -> IO () #-}
+addReferenceLevels ::
+     (PrimMonad m, MonadMVar m, MonadMask m)
+  => TempRegistry m
+  -> Levels m h
+  -> m ()
+addReferenceLevels reg levels =
+    forRunAndMergeM_ levels
+      (\r -> allocateTemp reg (Run.addReference r) (\_ -> Run.removeReference r))
+      (\m -> allocateTemp reg (Merge.addReference m) (\_ -> Merge.removeReference m))
+
+{-# SPECIALISE addReferenceLevels :: TempRegistry IO -> Levels IO h -> IO () #-}
+removeReferenceLevels ::
+     (PrimMonad m, MonadMVar m, MonadMask m)
+  => TempRegistry m
+  -> Levels m h
+  -> m ()
+removeReferenceLevels reg levels =
+    forRunAndMergeM_ levels
+      (\r -> freeTemp reg (Run.removeReference r))
+      (\m -> freeTemp reg (Merge.removeReference m))
+
+{-# SPECIALISE forRunAndMergeM_ ::
      Levels IO h
   -> (Run IO (Handle h) -> IO ())
+  -> (Merge IO h -> IO ())
   -> IO () #-}
-forRunM_ ::
+forRunAndMergeM_ ::
      MonadMVar m
   => Levels m h
   -> (Run m (Handle h) -> m ())
+  -> (Merge m h -> m ())
   -> m ()
-forRunM_ lvls k = V.forM_ lvls $ \(Level mr rs) -> do
+forRunAndMergeM_ lvls k1 k2 = V.forM_ lvls $ \(Level mr rs) -> do
     case mr of
-      SingleRun r    -> k r
+      SingleRun r    -> k1 r
       MergingRun var -> withMVar var $ \case
-        CompletedMerge r -> k r
-        OngoingMerge irs _m -> V.mapM_ k irs
-    V.mapM_ k rs
-
+        CompletedMerge r -> k1 r
+        OngoingMerge irs m -> V.mapM_ k1 irs >> k2 m
+    V.mapM_ k1 rs
 
 {-# SPECIALISE foldRunM ::
      (b -> Run IO (Handle h) -> IO b)
@@ -385,8 +431,8 @@ flushWriteBuffer tr conf@TableConfig{confDiskCachePolicy}
               (tableWriteBuffer tc)
               (tableWriteBufferBlobs tc))
             Run.removeReference
-    WBB.removeReference (tableWriteBufferBlobs tc)
-    wbblobs' <- WBB.new hfs (Paths.tableBlobPath root n)
+    freeTemp reg (WBB.removeReference (tableWriteBufferBlobs tc))
+    wbblobs' <- allocateTemp reg (WBB.new hfs (Paths.tableBlobPath root n)) WBB.removeReference
     levels' <- addRunToLevels tr conf resolve hfs hbio root uc r reg (tableLevels tc)
     cache' <- mkLevelsCache levels'
     pure $! TableContent {