Reconstruct bloom filters upon loading snapshots

lsm-tree/lsm-tree.cabal

@@ -570,6 +570,7 @@ library core
     Database.LSMTree.Internal.BlobFile
     Database.LSMTree.Internal.BlobRef
     Database.LSMTree.Internal.BloomFilter
+    Database.LSMTree.Internal.BloomFilterAcc
     Database.LSMTree.Internal.ByteString
     Database.LSMTree.Internal.ChecksumHandle
     Database.LSMTree.Internal.Chunk
@@ -607,10 +608,12 @@ library core
     Database.LSMTree.Internal.RunBuilder
     Database.LSMTree.Internal.RunNumber
     Database.LSMTree.Internal.RunReader
+    Database.LSMTree.Internal.RunReader.Keys
     Database.LSMTree.Internal.Serialise
     Database.LSMTree.Internal.Serialise.Class
     Database.LSMTree.Internal.Snapshot
     Database.LSMTree.Internal.Snapshot.Codec
+    Database.LSMTree.Internal.Snapshot.Codec.Monad
     Database.LSMTree.Internal.Types
     Database.LSMTree.Internal.UniqCounter
     Database.LSMTree.Internal.Unsafe
@@ -639,6 +642,7 @@ library core
     , io-classes:strict-mvar
     , lsm-tree:control
     , lsm-tree:kmerge
+    , mtl                     ^>=2.2      || ^>=2.3
     , primitive               ^>=0.9
     , serialise               ^>=0.2
     , text                    ^>=2.1.1

lsm-tree/src-core/Database/LSMTree/Internal/BloomFilter.hs

@@ -14,35 +14,20 @@ module Database.LSMTree.Internal.BloomFilter (
   bloomQueries,
   RunIxKeyIx(RunIxKeyIx),
   RunIx, KeyIx,
-
-  -- * Serialisation
-  bloomFilterVersion,
-  bloomFilterToLBS,
-  bloomFilterFromFile,
 ) where
 
 import           Data.Bits
-import qualified Data.ByteString as BS
-import qualified Data.ByteString.Builder.Extra as B
-import qualified Data.ByteString.Lazy as LBS
 import qualified Data.Primitive as P
 import qualified Data.Vector as V
 import qualified Data.Vector.Primitive as VP
-import           Data.Word (Word32, Word64, byteSwap32)
-import           Text.Printf (printf)
+import           Data.Word (Word32)
 
 import           Control.Exception (assert)
-import           Control.Monad (void, when)
-import           Control.Monad.Class.MonadThrow
-import           Control.Monad.Primitive (PrimMonad)
+
 import           Control.Monad.ST (ST, runST)
-import           System.FS.API
 
 import           Data.BloomFilter.Blocked (Bloom)
 import qualified Data.BloomFilter.Blocked as Bloom
-import           Database.LSMTree.Internal.ByteString (byteArrayToByteString)
-import           Database.LSMTree.Internal.CRC32C (FileCorruptedError (..),
-                     FileFormat (..))
 import           Database.LSMTree.Internal.Serialise (SerialisedKey)
 import qualified Database.LSMTree.Internal.Vector as P
 
@@ -206,121 +191,3 @@ toFiltersArray     = id
 indexFiltersArray  = V.unsafeIndex
 lengthFiltersArray = V.length
 #endif
-
-
--- serialising
------------------------------------------------------------
-
--- | By writing out the version in host endianness, we also indicate endianness.
--- During deserialisation, we would discover an endianness mismatch.
---
--- We base our version number on the 'Bloom.formatVersion' from the @bloomfilter@
--- library, plus our own version here. This accounts both for changes in the
--- format code here, and changes in the library.
---
-bloomFilterVersion :: Word32
-bloomFilterVersion = 1 + fromIntegral Bloom.formatVersion
-
-bloomFilterToLBS :: Bloom a -> LBS.ByteString
-bloomFilterToLBS bf =
-    let (size, salt, ba, off, len) = Bloom.serialise bf
-     in header size salt <> byteArrayToLBS ba off len
-  where
-    header Bloom.BloomSize { sizeBits, sizeHashes } salt =
-        -- creates a single 24 byte chunk
-        B.toLazyByteStringWith (B.safeStrategy 24 B.smallChunkSize) mempty $
-             B.word32Host bloomFilterVersion
-          <> B.word32Host (fromIntegral sizeHashes)
-          <> B.word64Host (fromIntegral sizeBits)
-          <> B.word64Host salt
-
-    byteArrayToLBS :: P.ByteArray -> Int -> Int -> LBS.ByteString
-    byteArrayToLBS ba off len =
-      LBS.fromStrict (byteArrayToByteString off len ba)
-
--- deserialising
------------------------------------------------------------
-
-{-# SPECIALISE bloomFilterFromFile ::
-     HasFS IO h
-  -> Bloom.Salt
-  -> Handle h
-  -> IO (Bloom a) #-}
--- | Read a 'Bloom' from a file.
---
-bloomFilterFromFile ::
-     (PrimMonad m, MonadCatch m)
-  => HasFS m h
-  -> Bloom.Salt -- ^ Expected salt
-  -> Handle h  -- ^ The open file, in read mode
-  -> m (Bloom a)
-bloomFilterFromFile hfs expectedSalt h = do
-    header <- rethrowEOFError "Doesn't contain a header" $
-              hGetByteArrayExactly hfs h 24
-
-    let !version = P.indexByteArray header 0 :: Word32
-        !nhashes = P.indexByteArray header 1 :: Word32
-        !nbits   = P.indexByteArray header 1 :: Word64
-        !salt    = P.indexByteArray header 2 :: Bloom.Salt
-
-    when (version /= bloomFilterVersion) $ throwFormatError $
-      if byteSwap32 version == bloomFilterVersion
-      then "Different byte order"
-      else "Unsupported version"
-
-    when (nbits <= 0) $ throwFormatError "Length is zero"
-
-    -- limit to 2^48 bits
-    when (nbits >= fromIntegral Bloom.maxSizeBits) $ throwFormatError "Too large bloomfilter"
-
-    when (expectedSalt /= salt) $ throwFormatError $
-      printf "Expected salt does not match actual salt: %d /= %d"
-        expectedSalt
-        salt
-
-    -- read the filter data from the file directly into the bloom filter
-    bloom <-
-      Bloom.deserialise
-        Bloom.BloomSize {
-          Bloom.sizeBits   = fromIntegral nbits,
-          Bloom.sizeHashes = fromIntegral nhashes
-        }
-        salt
-        (\buf off len ->
-            rethrowEOFError "bloom filter file too short" $
-              void $ hGetBufExactly hfs
-                       h buf (BufferOffset off) (fromIntegral len))
-
-    -- check we're now at EOF
-    trailing <- hGetSome hfs h 1
-    when (not (BS.null trailing)) $
-      throwFormatError "Byte array is too large for components"
-    pure bloom
-  where
-    throwFormatError = throwIO
-                     . ErrFileFormatInvalid
-                         (mkFsErrorPath hfs (handlePath h))
-                         FormatBloomFilterFile
-    rethrowEOFError msg =
-      handleJust
-        (\e -> if isFsErrorType FsReachedEOF e then Just e else Nothing)
-        (\e -> throwIO $
-                 ErrFileFormatInvalid
-                   (fsErrorPath e) FormatBloomFilterFile msg)
-
-{-# SPECIALISE hGetByteArrayExactly ::
-     HasFS IO h
-  -> Handle h
-  -> Int
-  -> IO P.ByteArray #-}
-hGetByteArrayExactly ::
-     (PrimMonad m, MonadThrow m)
-  => HasFS m h
-  -> Handle h
-  -> Int
-  -> m P.ByteArray
-hGetByteArrayExactly hfs h len = do
-    buf <- P.newByteArray len
-    _   <- hGetBufExactly hfs h buf 0 (fromIntegral len)
-    P.unsafeFreezeByteArray buf
-

lsm-tree/src-core/Database/LSMTree/Internal/BloomFilterAcc.hs

@@ -0,0 +1,58 @@
+{-# LANGUAGE RecordWildCards #-}
+{-# OPTIONS_HADDOCK not-home #-}
+
+-- | Incremental, in-memory bloom filter construction
+--
+module Database.LSMTree.Internal.BloomFilterAcc (
+    -- * Bloom filter allocation
+    RunBloomFilterAlloc (..)
+    -- * Incremental, in-memory bloom filter construction
+  , newMBloom
+  , bloomInserts
+  ) where
+
+import           Control.DeepSeq (NFData (..))
+import           Control.Monad.ST.Strict
+import qualified Data.BloomFilter.Blocked as Bloom
+import           Database.LSMTree.Internal.BloomFilter (MBloom)
+import           Database.LSMTree.Internal.Entry (NumEntries (..))
+import           Database.LSMTree.Internal.PageAcc (PageAcc)
+import qualified Database.LSMTree.Internal.PageAcc as PageAcc
+import           Database.LSMTree.Internal.Serialise (SerialisedKey)
+
+{-------------------------------------------------------------------------------
+  Bloom filter allocation
+-------------------------------------------------------------------------------}
+
+-- | See 'Database.LSMTree.Internal.Config.BloomFilterAlloc'
+data RunBloomFilterAlloc =
+    -- | Bits per element in a filter
+    RunAllocFixed      !Double
+  | RunAllocRequestFPR !Double
+  deriving stock (Show, Eq)
+
+instance NFData RunBloomFilterAlloc where
+    rnf (RunAllocFixed a)      = rnf a
+    rnf (RunAllocRequestFPR a) = rnf a
+
+{-------------------------------------------------------------------------------
+  Incremental, in-memory bloom filter construction
+-------------------------------------------------------------------------------}
+
+newMBloom :: NumEntries -> RunBloomFilterAlloc -> Bloom.Salt -> ST s (MBloom s a)
+newMBloom (NumEntries nentries) alloc salt =
+    Bloom.new (Bloom.sizeForPolicy (policy alloc) nentries) salt
+  where
+    --TODO: it'd be possible to turn the RunBloomFilterAlloc into a BloomPolicy
+    -- without the NumEntries, and cache the policy, avoiding recalculating the
+    -- policy every time.
+    policy (RunAllocFixed bitsPerEntry) = Bloom.policyForBits bitsPerEntry
+    policy (RunAllocRequestFPR fpr)     = Bloom.policyForFPR fpr
+
+-- An instance of insertMany specialised to SerialisedKey and indexKeyPageAcc.
+-- This is a performance-sensitive function. It is marked NOINLINE so we can
+-- easily inspect the core and check all the specialisations worked as expected.
+{-# NOINLINE bloomInserts #-}
+bloomInserts :: MBloom s SerialisedKey -> PageAcc s -> Int -> ST s ()
+bloomInserts !mbloom !mpageacc !nkeys =
+    Bloom.insertMany mbloom (PageAcc.indexKeyPageAcc mpageacc) nkeys

lsm-tree/src-core/Database/LSMTree/Internal/ChecksumHandle.hs

@@ -15,7 +15,6 @@ module Database.LSMTree.Internal.ChecksumHandle
     writeRawOverflowPages,
     writeBlob,
     copyBlob,
-    writeFilter,
     writeIndexHeader,
     writeIndexChunk,
     writeIndexFinal,
@@ -29,16 +28,15 @@ import           Data.Primitive.PrimVar
 import           Data.Word (Word64)
 import           Database.LSMTree.Internal.BlobRef (BlobSpan (..), RawBlobRef)
 import qualified Database.LSMTree.Internal.BlobRef as BlobRef
-import           Database.LSMTree.Internal.BloomFilter (Bloom, bloomFilterToLBS)
 import           Database.LSMTree.Internal.Chunk (Chunk)
 import qualified Database.LSMTree.Internal.Chunk as Chunk (toByteString)
 import           Database.LSMTree.Internal.CRC32C (CRC32C)
 import qualified Database.LSMTree.Internal.CRC32C as CRC
 import           Database.LSMTree.Internal.Entry
 import           Database.LSMTree.Internal.Index (Index, IndexType)
 import qualified Database.LSMTree.Internal.Index as Index (finalLBS, headerLBS)
-import           Database.LSMTree.Internal.Paths (ForBlob (..), ForFilter (..),
-                     ForIndex (..), ForKOps (..))
+import           Database.LSMTree.Internal.Paths (ForBlob (..), ForIndex (..),
+                     ForKOps (..))
 import qualified Database.LSMTree.Internal.RawBytes as RB
 import           Database.LSMTree.Internal.RawOverflowPage (RawOverflowPage)
 import qualified Database.LSMTree.Internal.RawOverflowPage as RawOverflowPage
@@ -188,20 +186,6 @@ copyBlob hfs blobOffset blobHandle blobref = do
     blob <- BlobRef.readRawBlobRef hfs blobref
     writeBlob hfs blobOffset blobHandle blob
 
-{-# SPECIALISE writeFilter ::
-     HasFS IO h
-  -> ForFilter (ChecksumHandle RealWorld h)
-  -> Bloom SerialisedKey
-  -> IO () #-}
-writeFilter ::
-     (MonadSTM m, PrimMonad m)
-  => HasFS m h
-  -> ForFilter (ChecksumHandle (PrimState m) h)
-  -> Bloom SerialisedKey
-  -> m ()
-writeFilter hfs filterHandle bf =
-    writeToHandle hfs (unForFilter filterHandle) (bloomFilterToLBS bf)
-
 {-# SPECIALISE writeIndexHeader ::
      HasFS IO h
   -> ForIndex (ChecksumHandle RealWorld h)

lsm-tree/src-core/Database/LSMTree/Internal/Config.hs

@@ -16,7 +16,7 @@ module Database.LSMTree.Internal.Config (
   , WriteBufferAlloc (..)
     -- * Bloom filter allocation
   , BloomFilterAlloc (..)
-  , bloomFilterAllocForLevel
+  , bloomFilterAllocForRun
     -- * Fence pointer index
   , FencePointerIndexType (..)
   , indexTypeForRun
@@ -153,7 +153,7 @@ runParamsForLevel :: TableConfig -> RunLevelNo -> RunParams
 runParamsForLevel conf@TableConfig {..} levelNo =
     RunParams
       { runParamCaching = diskCachePolicyForLevel confDiskCachePolicy levelNo
-      , runParamAlloc   = bloomFilterAllocForLevel conf levelNo
+      , runParamAlloc   = bloomFilterAllocForRun conf
       , runParamIndex   = indexTypeForRun confFencePointerIndex
       }
 
@@ -291,8 +291,8 @@ instance NFData BloomFilterAlloc where
   rnf (AllocFixed n)        = rnf n
   rnf (AllocRequestFPR fpr) = rnf fpr
 
-bloomFilterAllocForLevel :: TableConfig -> RunLevelNo -> RunBloomFilterAlloc
-bloomFilterAllocForLevel conf _levelNo =
+bloomFilterAllocForRun :: TableConfig -> RunBloomFilterAlloc
+bloomFilterAllocForRun conf =
     case confBloomFilterAlloc conf of
       AllocFixed n        -> RunAllocFixed n
       AllocRequestFPR fpr -> RunAllocRequestFPR fpr

lsm-tree/src-core/Database/LSMTree/Internal/Config/Override.hs

@@ -175,8 +175,8 @@ instance Override RunDataCaching RunParams where
 
 instance Override RunDataCaching SnapshotRun where
   override rdc
-    (SnapshotRun (rn :: RunNumber) (_rdc :: RunDataCaching) (it ::IndexType))
-    = SnapshotRun rn rdc it
+    (SnapshotRun (rn :: RunNumber) (_rdc :: RunDataCaching) (rbfa :: RunBloomFilterAlloc) (it ::IndexType))
+    = SnapshotRun rn rdc rbfa it
 
 instance Override RunDataCaching (SnapMergingTree SnapshotRun) where
   override rdc (SnapMergingTree (smts ::  SnapMergingTreeState SnapshotRun))