Skip to content

Commit 9cf968f

Browse files
authored
Merge pull request #472 from IntersectMBO/wenkokke/checksum-handle
Move `ChecksumHandle` and its methods into their own module
2 parents d0ae185 + 588171b commit 9cf968f

File tree

7 files changed

+346
-230
lines changed

7 files changed

+346
-230
lines changed

lsm-tree.cabal

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,7 @@ library
124124
Database.LSMTree.Internal.BloomFilter
125125
Database.LSMTree.Internal.BloomFilterQuery1
126126
Database.LSMTree.Internal.ByteString
127+
Database.LSMTree.Internal.ChecksumHandle
127128
Database.LSMTree.Internal.Chunk
128129
Database.LSMTree.Internal.Config
129130
Database.LSMTree.Internal.CRC32C

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ import Data.Word
3636
import Database.LSMTree.Internal as Internal
3737
import Database.LSMTree.Internal.BlobFile
3838
import Database.LSMTree.Internal.BlobRef
39+
import Database.LSMTree.Internal.ChecksumHandle
3940
import Database.LSMTree.Internal.Config
4041
import Database.LSMTree.Internal.CRC32C
4142
import Database.LSMTree.Internal.Entry
@@ -206,6 +207,18 @@ deriving anyclass instance NoThunks SessionRoot
206207
deriving stock instance Generic RunFsPaths
207208
deriving anyclass instance NoThunks RunFsPaths
208209

210+
deriving stock instance Generic (ForKOps a)
211+
deriving anyclass instance NoThunks a => NoThunks (ForKOps a)
212+
213+
deriving stock instance Generic (ForBlob a)
214+
deriving anyclass instance NoThunks a => NoThunks (ForBlob a)
215+
216+
deriving stock instance Generic (ForFilter a)
217+
deriving anyclass instance NoThunks a => NoThunks (ForFilter a)
218+
219+
deriving stock instance Generic (ForIndex a)
220+
deriving anyclass instance NoThunks a => NoThunks (ForIndex a)
221+
209222
deriving stock instance Generic (ForRunFiles a)
210223
deriving anyclass instance NoThunks a => NoThunks (ForRunFiles a)
211224

src/Database/LSMTree/Internal/CRC32C.hs

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -348,4 +348,3 @@ formatChecksumsFile checksums =
348348
<> BS.word32HexFixed crc
349349
<> BS.char8 '\n'
350350
| (ChecksumsFileName name, CRC32C crc) <- Map.toList checksums ]
351-
Lines changed: 247 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,247 @@
1+
module Database.LSMTree.Internal.ChecksumHandle
2+
(
3+
-- * Checksum handles
4+
-- $checksum-handles
5+
ChecksumHandle (..),
6+
makeHandle,
7+
readChecksum,
8+
dropCache,
9+
closeHandle,
10+
writeToHandle,
11+
-- * Specialised writers
12+
writeRawPage,
13+
writeRawOverflowPages,
14+
writeBlob,
15+
copyBlob,
16+
writeFilter,
17+
writeIndexHeader,
18+
writeIndexChunk,
19+
writeIndexFinal,
20+
) where
21+
22+
import Control.Monad.Class.MonadSTM (MonadSTM (..))
23+
import Control.Monad.Class.MonadThrow (MonadThrow)
24+
import Control.Monad.Primitive
25+
import Data.BloomFilter (Bloom)
26+
import qualified Data.ByteString.Lazy as BSL
27+
import Data.Primitive.PrimVar
28+
import Data.Word (Word64)
29+
import Database.LSMTree.Internal.BlobRef (BlobSpan (..), RawBlobRef)
30+
import qualified Database.LSMTree.Internal.BlobRef as BlobRef
31+
import Database.LSMTree.Internal.BloomFilter (bloomFilterToLBS)
32+
import Database.LSMTree.Internal.Chunk (Chunk)
33+
import qualified Database.LSMTree.Internal.Chunk as Chunk (toByteString)
34+
import Database.LSMTree.Internal.CRC32C (CRC32C)
35+
import qualified Database.LSMTree.Internal.CRC32C as CRC
36+
import Database.LSMTree.Internal.Entry
37+
import Database.LSMTree.Internal.IndexCompact (IndexCompact)
38+
import qualified Database.LSMTree.Internal.IndexCompact as Index
39+
import Database.LSMTree.Internal.Paths (ForBlob (..), ForFilter (..),
40+
ForIndex (..), ForKOps (..))
41+
import qualified Database.LSMTree.Internal.RawBytes as RB
42+
import Database.LSMTree.Internal.RawOverflowPage (RawOverflowPage)
43+
import qualified Database.LSMTree.Internal.RawOverflowPage as RawOverflowPage
44+
import Database.LSMTree.Internal.RawPage (RawPage)
45+
import qualified Database.LSMTree.Internal.RawPage as RawPage
46+
import Database.LSMTree.Internal.Serialise
47+
import qualified System.FS.API as FS
48+
import System.FS.API
49+
import qualified System.FS.BlockIO.API as FS
50+
import System.FS.BlockIO.API (HasBlockIO)
51+
52+
{-------------------------------------------------------------------------------
53+
ChecksumHandle
54+
-------------------------------------------------------------------------------}
55+
56+
{- $checksum-handles
57+
A handle ('ChecksumHandle') that maintains a running CRC32 checksum.
58+
-}
59+
60+
-- | Tracks the checksum of a (write mode) file handle.
61+
data ChecksumHandle s h = ChecksumHandle !(FS.Handle h) !(PrimVar s CRC32C)
62+
63+
{-# SPECIALISE makeHandle ::
64+
HasFS IO h
65+
-> FS.FsPath
66+
-> IO (ChecksumHandle RealWorld h) #-}
67+
makeHandle ::
68+
(MonadSTM m, PrimMonad m)
69+
=> HasFS m h
70+
-> FS.FsPath
71+
-> m (ChecksumHandle (PrimState m) h)
72+
makeHandle fs path =
73+
ChecksumHandle
74+
<$> FS.hOpen fs path (FS.WriteMode FS.MustBeNew)
75+
<*> newPrimVar CRC.initialCRC32C
76+
77+
{-# SPECIALISE readChecksum ::
78+
ChecksumHandle RealWorld h
79+
-> IO CRC32C #-}
80+
readChecksum ::
81+
PrimMonad m
82+
=> ChecksumHandle (PrimState m) h
83+
-> m CRC32C
84+
readChecksum (ChecksumHandle _h checksum) = readPrimVar checksum
85+
86+
dropCache :: HasBlockIO m h -> ChecksumHandle (PrimState m) h -> m ()
87+
dropCache hbio (ChecksumHandle h _) = FS.hDropCacheAll hbio h
88+
89+
closeHandle :: HasFS m h -> ChecksumHandle (PrimState m) h -> m ()
90+
closeHandle fs (ChecksumHandle h _checksum) = FS.hClose fs h
91+
92+
{-# SPECIALISE writeToHandle ::
93+
HasFS IO h
94+
-> ChecksumHandle RealWorld h
95+
-> BSL.ByteString
96+
-> IO () #-}
97+
writeToHandle ::
98+
(MonadSTM m, PrimMonad m)
99+
=> HasFS m h
100+
-> ChecksumHandle (PrimState m) h
101+
-> BSL.ByteString
102+
-> m ()
103+
writeToHandle fs (ChecksumHandle h checksum) lbs = do
104+
crc <- readPrimVar checksum
105+
(_, crc') <- CRC.hPutAllChunksCRC32C fs h lbs crc
106+
writePrimVar checksum crc'
107+
108+
{-------------------------------------------------------------------------------
109+
Specialised Writers for ChecksumHandle
110+
-------------------------------------------------------------------------------}
111+
112+
{-# SPECIALISE writeRawPage ::
113+
HasFS IO h
114+
-> ForKOps (ChecksumHandle RealWorld h)
115+
-> RawPage
116+
-> IO () #-}
117+
writeRawPage ::
118+
(MonadSTM m, PrimMonad m)
119+
=> HasFS m h
120+
-> ForKOps (ChecksumHandle (PrimState m) h)
121+
-> RawPage
122+
-> m ()
123+
writeRawPage hfs kOpsHandle =
124+
writeToHandle hfs (unForKOps kOpsHandle)
125+
. BSL.fromStrict
126+
. RB.unsafePinnedToByteString -- 'RawPage' is guaranteed to be pinned
127+
. RawPage.rawPageRawBytes
128+
129+
{-# SPECIALISE writeRawOverflowPages ::
130+
HasFS IO h
131+
-> ForKOps (ChecksumHandle RealWorld h)
132+
-> [RawOverflowPage]
133+
-> IO () #-}
134+
writeRawOverflowPages ::
135+
(MonadSTM m, PrimMonad m)
136+
=> HasFS m h
137+
-> ForKOps (ChecksumHandle (PrimState m) h)
138+
-> [RawOverflowPage]
139+
-> m ()
140+
writeRawOverflowPages hfs kOpsHandle =
141+
writeToHandle hfs (unForKOps kOpsHandle)
142+
. BSL.fromChunks
143+
. map (RawOverflowPage.rawOverflowPageToByteString)
144+
145+
{-# SPECIALISE writeBlob ::
146+
HasFS IO h
147+
-> PrimVar RealWorld Word64
148+
-> ForBlob (ChecksumHandle RealWorld h)
149+
-> SerialisedBlob
150+
-> IO BlobSpan #-}
151+
writeBlob ::
152+
(MonadSTM m, PrimMonad m)
153+
=> HasFS m h
154+
-> PrimVar (PrimState m) Word64
155+
-> ForBlob (ChecksumHandle (PrimState m) h)
156+
-> SerialisedBlob
157+
-> m BlobSpan
158+
writeBlob hfs blobOffset blobHandle blob = do
159+
-- NOTE: This is different from BlobFile.writeBlob. This is because BlobFile
160+
-- internalises a regular Handle, rather than a ChecksumHandle. These two
161+
-- functions cannot be easily unified, because BlobFile.writeBlob permits
162+
-- writing blobs to arbitrary positions in the blob file, whereas, by the
163+
-- very nature of CRC32 checksums, ChecksumHandle.writeBlob only supports
164+
-- sequential writes.
165+
let size = sizeofBlob64 blob
166+
offset <- readPrimVar blobOffset
167+
modifyPrimVar blobOffset (+size)
168+
let SerialisedBlob rb = blob
169+
let lbs = BSL.fromStrict $ RB.toByteString rb
170+
writeToHandle hfs (unForBlob blobHandle) lbs
171+
return (BlobSpan offset (fromIntegral size))
172+
173+
{-# SPECIALISE copyBlob ::
174+
HasFS IO h
175+
-> PrimVar RealWorld Word64
176+
-> ForBlob (ChecksumHandle RealWorld h)
177+
-> RawBlobRef IO h
178+
-> IO BlobSpan #-}
179+
copyBlob ::
180+
(MonadSTM m, MonadThrow m, PrimMonad m)
181+
=> HasFS m h
182+
-> PrimVar (PrimState m) Word64
183+
-> ForBlob (ChecksumHandle (PrimState m) h)
184+
-> RawBlobRef m h
185+
-> m BlobSpan
186+
copyBlob hfs blobOffset blobHandle blobref = do
187+
blob <- BlobRef.readRawBlobRef hfs blobref
188+
writeBlob hfs blobOffset blobHandle blob
189+
190+
{-# SPECIALISE writeFilter ::
191+
HasFS IO h
192+
-> ForFilter (ChecksumHandle RealWorld h)
193+
-> Bloom SerialisedKey
194+
-> IO () #-}
195+
writeFilter ::
196+
(MonadSTM m, PrimMonad m)
197+
=> HasFS m h
198+
-> ForFilter (ChecksumHandle (PrimState m) h)
199+
-> Bloom SerialisedKey
200+
-> m ()
201+
writeFilter hfs filterHandle bf =
202+
writeToHandle hfs (unForFilter filterHandle) (bloomFilterToLBS bf)
203+
204+
{-# SPECIALISE writeIndexHeader ::
205+
HasFS IO h
206+
-> ForIndex (ChecksumHandle RealWorld h)
207+
-> IO () #-}
208+
writeIndexHeader ::
209+
(MonadSTM m, PrimMonad m)
210+
=> HasFS m h
211+
-> ForIndex (ChecksumHandle (PrimState m) h)
212+
-> m ()
213+
writeIndexHeader hfs indexHandle =
214+
writeToHandle hfs (unForIndex indexHandle) $
215+
Index.headerLBS
216+
217+
{-# SPECIALISE writeIndexChunk ::
218+
HasFS IO h
219+
-> ForIndex (ChecksumHandle RealWorld h)
220+
-> Chunk
221+
-> IO () #-}
222+
writeIndexChunk ::
223+
(MonadSTM m, PrimMonad m)
224+
=> HasFS m h
225+
-> ForIndex (ChecksumHandle (PrimState m) h)
226+
-> Chunk
227+
-> m ()
228+
writeIndexChunk hfs indexHandle chunk =
229+
writeToHandle hfs (unForIndex indexHandle) $
230+
BSL.fromStrict $ Chunk.toByteString chunk
231+
232+
{-# SPECIALISE writeIndexFinal ::
233+
HasFS IO h
234+
-> ForIndex (ChecksumHandle RealWorld h)
235+
-> NumEntries
236+
-> IndexCompact
237+
-> IO () #-}
238+
writeIndexFinal ::
239+
(MonadSTM m, PrimMonad m)
240+
=> HasFS m h
241+
-> ForIndex (ChecksumHandle (PrimState m) h)
242+
-> NumEntries
243+
-> IndexCompact
244+
-> m ()
245+
writeIndexFinal hfs indexHandle numEntries index =
246+
writeToHandle hfs (unForIndex indexHandle) $
247+
Index.finalLBS numEntries index

0 commit comments

Comments
 (0)