[BUG] Cursor reads and blob retrieval return keys/values/blobs with retained unused memory

[jorisdral]

lsm-tree can be used in an "untyped" way by setting types of keys, values, and blobs as raw bytes. For example, a user could decide on either of the two following options:

type Typed = Table IO Word64 Word64 Word64
type Untyped = Table IO RawBytes RawBytes RawBytes

The latter option bypasses lsm-tree's serialisation framework, which is perfectly valid to do, and it is useful in some cases where lsm-tree's serialisation classes are too restrictive. The caveat is that the onus for (de-)serialisation now falls on the user of lsm-tree. For example ouroboros-consensus does this because they want to (de-)serialise keys and values with additional context.

This unearths a problem, however, since some public API functions applied to untyped tables will return keys, values, and blobs that retain unused memory. Inspecting RawBytes, we see that it is a newtype wrapper around a primitive vector:

-- | Unboxed vectors of primitive types.
data Vector a = Vector {-# UNPACK #-} !Int       -- ^ offset
                       {-# UNPACK #-} !Int       -- ^ length
                       {-# UNPACK #-} !ByteArray -- ^ underlying byte array

A RawBytes that retains unused memory is slice of a primitive vector. For example:

Vector 0 8 (replicate 8 0) -- not a slice
Vector 1 8 (replicate 8 0) -- a slice
Vector 0 7 (replicate 8 0) -- a slice
Vector 3 4 (replicate 8 0) -- a slice

Ideally, public API functions should never return slices, but some do.

We wrote a reproducer program that inserts a bunch of keys/values/blobs into an untyped table, then performs lookups and range lookups (i.e., cursor reads) while retrieving blobs, and prints the results. The reproducer program and its output is included below:

Inserting keys, values, and blobs
Printing inserted keys, values, and blobs as (bytes, bytes, Maybe bytes)
([0,0,0,0,0,0,0,0],[0,0,0,0,0,0,0,0],Just [0,0,0,0,0,0,0,0])
([0,0,0,0,0,0,0,1],[1,0,0,0,0,0,0,0],Just [1,0,0,0,0,0,0,0])
([0,0,0,0,0,0,0,2],[2,0,0,0,0,0,0,0],Just [2,0,0,0,0,0,0,0])
([0,0,0,0,0,0,0,3],[3,0,0,0,0,0,0,0],Just [3,0,0,0,0,0,0,0])
([0,0,0,0,0,0,0,4],[4,0,0,0,0,0,0,0],Just [4,0,0,0,0,0,0,0])
([0,0,0,0,0,0,0,5],[5,0,0,0,0,0,0,0],Just [5,0,0,0,0,0,0,0])
([0,0,0,0,0,0,0,6],[6,0,0,0,0,0,0,0],Just [6,0,0,0,0,0,0,0])
([0,0,0,0,0,0,0,7],[7,0,0,0,0,0,0,0],Just [7,0,0,0,0,0,0,0])
([0,0,0,0,0,0,0,8],[8,0,0,0,0,0,0,0],Just [8,0,0,0,0,0,0,0])
([0,0,0,0,0,0,0,9],[9,0,0,0,0,0,0,0],Just [9,0,0,0,0,0,0,0])
Performing lookups
Printing looked up values as (bytes, offset, size)
([0,0,0,0,0,0,0,0],0,8)
([1,0,0,0,0,0,0,0],0,8)
([2,0,0,0,0,0,0,0],0,8)
([3,0,0,0,0,0,0,0],0,8)
([4,0,0,0,0,0,0,0],0,8)
([5,0,0,0,0,0,0,0],0,8)
([6,0,0,0,0,0,0,0],0,8)
([7,0,0,0,0,0,0,0],0,8)
([8,0,0,0,0,0,0,0],0,8)
([9,0,0,0,0,0,0,0],0,8)
Printing looked up blobs as (bytes, offset, size)
([0,0,0,0,0,0,0,0],0,8)
([1,0,0,0,0,0,0,0],8,8)
([2,0,0,0,0,0,0,0],16,8)
([3,0,0,0,0,0,0,0],24,8)
([4,0,0,0,0,0,0,0],32,8)
([5,0,0,0,0,0,0,0],40,8)
([6,0,0,0,0,0,0,0],48,8)
([7,0,0,0,0,0,0,0],56,8)
([8,0,0,0,0,0,0,0],64,8)
([9,0,0,0,0,0,0,0],72,8)
Performing a range lookup
Printing looked up keys as (bytes, offset, size)
([0,0,0,0,0,0,0,0],106,8)
([0,0,0,0,0,0,0,1],114,8)
([0,0,0,0,0,0,0,2],122,8)
([0,0,0,0,0,0,0,3],130,8)
([0,0,0,0,0,0,0,4],138,8)
([0,0,0,0,0,0,0,5],106,8)
([0,0,0,0,0,0,0,6],114,8)
([0,0,0,0,0,0,0,7],122,8)
([0,0,0,0,0,0,0,8],130,8)
([0,0,0,0,0,0,0,9],138,8)
Printing looked up values as (bytes, offset, size)
([0,0,0,0,0,0,0,0],146,8)
([1,0,0,0,0,0,0,0],154,8)
([2,0,0,0,0,0,0,0],162,8)
([3,0,0,0,0,0,0,0],170,8)
([4,0,0,0,0,0,0,0],178,8)
([5,0,0,0,0,0,0,0],146,8)
([6,0,0,0,0,0,0,0],154,8)
([7,0,0,0,0,0,0,0],162,8)
([8,0,0,0,0,0,0,0],170,8)
([9,0,0,0,0,0,0,0],178,8)
Printing looked up blobs as (bytes, offset, size)
([0,0,0,0,0,0,0,0],0,8)
([1,0,0,0,0,0,0,0],8,8)
([2,0,0,0,0,0,0,0],16,8)
([3,0,0,0,0,0,0,0],24,8)
([4,0,0,0,0,0,0,0],32,8)
([5,0,0,0,0,0,0,0],40,8)
([6,0,0,0,0,0,0,0],48,8)
([7,0,0,0,0,0,0,0],56,8)
([8,0,0,0,0,0,0,0],64,8)
([9,0,0,0,0,0,0,0],72,8)

What we can see here is that lookups do not return slices, but blob retrieval and range lookups (i.e., cursor reads) do. We should fix this, presumably by forceing primitive vectors before returning them as raw bytes from public API functions. We should also write tests for this.

Reproducer

{-# OPTIONS_GHC -Wno-orphans #-}

module Main (main) where

import           Control.Tracer
import           Data.Maybe
import qualified Data.Vector as V
import qualified Data.Vector.Primitive as VP
import           Data.Word
import           Database.LSMTree

main :: IO ()
main = do
  foo

foo :: IO ()
foo =
    withOpenSessionIO nullTracer "_temp" $ \s ->
    withTableWith conf s $ \(t :: Table IO RawBytes RawBytes RawBytes) -> do
      do
        putStrLn "Inserting keys, values, and blobs"
        let entries = V.fromList [
                ( serialiseKey k
                , serialiseValue v
                , fmap serialiseValue mb
                )
              | (k, v, mb) <- contents
              ]
        inserts t entries
        putStrLn "Printing inserted keys, values, and blobs as (bytes, bytes, Maybe bytes)"
        mapM_ print entries
      do
        putStrLn "Performing lookups"
        res <- lookups t $ V.fromList [
            serialiseKey k
          | (k, _, _) <- contents
          ]
        let values = V.map (fromJust . getValue) res
        let blobRefs = V.mapMaybe getBlob res
        blobs <- retrieveBlobs s blobRefs
        putStrLn "Printing looked up values as (bytes, offset, size)"
        mapM_ print $ fmap deconstruct values
        putStrLn "Printing looked up blobs as (bytes, offset, size)"
        mapM_ print $ fmap deconstruct blobs
      do
        putStrLn "Performing a range lookup"
        res <- rangeLookup t $
          FromToIncluding (serialiseKey @Word64 minBound)
                          (serialiseKey @Word64 maxBound)
        let keys = V.map getEntryKey res
        let values = V.map getEntryValue res
        let blobRefs = V.mapMaybe getEntryBlob res
        blobs <- retrieveBlobs s blobRefs
        putStrLn "Printing looked up keys as (bytes, offset, size)"
        mapM_ print $ fmap deconstruct keys
        putStrLn "Printing looked up values as (bytes, offset, size)"
        mapM_ print $ fmap deconstruct values
        putStrLn "Printing looked up blobs as (bytes, offset, size)"
        mapM_ print $ fmap deconstruct blobs
  where
    conf = defaultTableConfig {
        confWriteBufferAlloc = AllocNumEntries 5
      }

    contents :: [(Word64, Word64, Maybe Word64)]
    contents = [(i, i, Just i) | i <- [0..9]]

instance SerialiseKey RawBytes where
  serialiseKey = id
  deserialiseKey = id

instance SerialiseValue RawBytes where
  serialiseValue = id
  deserialiseValue = id

instance ResolveValue RawBytes where
  resolve = const

deconstruct :: RawBytes -> (RawBytes, Int, Int)
deconstruct rb@(RawBytes (VP.Vector off sz _)) = (rb, off, sz)

[jorisdral]

Presumably, we didn't see this problem before because deserialisation functions normally produce a new Haskell value from the serialised form, instead of using id as the deserialisation function. For example deserialise :: RawBytes -> Word64 produces a new memory location for Word64, while deserialise = id does not

[jorisdral]

Alternatively, we could reconsider whether the current definition of RawBytes is the right user-facing type for the serialised form of keys/values/blobs. Currently, we have:

newtype RawBytes = RawBytes (VP.Vector Word8)

Should the user-facing type rather be just:

newtype RawBytes = RawBytes ByteArray

So that the problem of slices retaining unused memory goes away entirely?

NOTE: we do require the internal serialised form to be sliceable, but the user-facing serialised form could be different as long as we can easily/cheaply map between the two

[dcoutts]

Perhaps the rule should be that the deserialisation must not retain the input. That way we don't need extra copies for deserialisation functions that do copy.

So essentially the deserialisation instance for RawBytes would be force rather than id. And perhaps also provide an instance for ByteArray that does the same.

[dcoutts]

Related notes:

1. the lookup (in intraPageLookupsOn) already uses copySerialisedValue to force the raw bytes used in deserialising values, though it keeps it as a RawBytes type.
2. in intraPageLookupsOn, for the LookupEntryOverflow case, we do not force appending the buffers together. The problem is if the value deserialisation itself is lazy then we will end up retaining the original buffers.

We ought to be consistent and treat the keys and values the same, since both are small portions of larger pages.

Indeed in lookups since we use arena allocation for pages we must (and do) copy to avoid retaining, since the I/O pages are not retained. So slicing there is not possible.

For cursors, we use ordinary GC'd buffers so retaining is sound but undesirable.

On the other hand, no slicing or copying is needed for blobs because they don't use page allocation and always get fresh GC'd allocations. So retaining and non-copying there is fine.

So arguably the right thing to do here is to just pay the cost of copying the keys and values in cursors too, and make sure we're consistent with copying in the large/overflow value case too. Then the user deserialisation instances do not need to worry about copying.

Perhaps the deserialiseKey and deserialiseValue should be given a simple ByteArray to drive home the point that we guarantee no slicing and that retaining the array without copying is fine and indeed encouraged.

[jorisdral]

2. in intraPageLookupsOn, for the LookupEntryOverflow case, we do not force appending the buffers together. The problem is if the value deserialisation itself is lazy then we will end up retaining the original buffers.

We use (<>) to append buffers together in this case, that should take care of copying the input buffers to a single output buffer without retaining additional memory. Maybe in an edge case when all buffers except one are empty vectors, it wouldn't copy, but that sounds like an "impossible" case in lsm-tree

EDIT: I misread. Indeed, yes when deserialisation is lazy then we could run into retaining problems

[dcoutts]

Right, the <> is fine, but we need to force the result.

Arguably, we should make the page reading data structure live in ST so that we can enforce these things more reliably.