LMDB.hs

{-# OPTIONS_HADDOCK prune #-}
{-# LANGUAGE BangPatterns   #-}
{-# LANGUAGE EmptyDataDecls #-}

-- |
-- Module: Database.LMDB
-- Copyright: (c) 2026 Jared Tobin
-- License: MIT
-- Maintainer: Jared Tobin <jared@ppad.tech>
--
-- Minimal bindings to
-- [LMDB](https://www.symas.com/lmdb), an embedded ACID key-value
-- store. The upstream LMDB C source is vendored under @cbits\/@ at
-- release @LMDB_0.9.33@; no external @liblmdb@ is needed.
--
-- LMDB has a single-writer, many-reader transaction model and uses
-- memory mapping for zero-copy reads. This module presents that as a
-- bracketed @IO@-only interface.
--
-- Read-only and read-write transactions are distinguished at the type
-- level via a phantom parameter on 'Txn', so using a write primitive
-- (e.g. 'put') on a read transaction is a compile-time error.
--
-- Most operations throw 'LMDBException' on error. Lookups ('get' and
-- the cursor functions) and 'del' are the exceptions: they report a
-- missing key as 'Nothing' or 'False', since absence is normal control
-- flow.
--
-- Transactions may be used from any thread, under either RTS; see
-- 'withReadTxn' and 'withWriteTxn' for the details.
--
-- Lookups, writes and cursor moves are unsafe foreign calls, for
-- speed. On a database much larger than memory, a page fault in one
-- stalls garbage collection under the threaded RTS until the read
-- completes.
--
-- Handles are not checked for validity. A 'Txn' or 'Cursor' is valid
-- only within the bracket that supplied it, and a 'Dbi' only until its
-- 'Env' is closed; an 'Env' must be closed at most once, after its
-- transactions have ended. Using a handle outside those bounds is
-- undefined behaviour. As LMDB requires, a process must not have the
-- same environment open twice at once.

module Database.LMDB (
  -- * Environments
    Env
  , EnvFlags(..)
  , defaultEnvFlags
  , open
  , close
  , withEnv
  , sync

  -- * Transactions
  , Txn
  , RO
  , RW
  , withReadTxn
  , withWriteTxn

  -- * Databases
  , Dbi
  , openDbi

  -- * Key-value operations
  , get
  , put
  , del

  -- * Cursors
  , Cursor
  , withCursor
  , cursorFirst
  , cursorLast
  , cursorNext
  , cursorPrev
  , cursorSeek

  -- * Errors
  , LMDBException(..)
  ) where

import Control.Concurrent (rtsSupportsBoundThreads, runInBoundThread)
import Control.Concurrent.MVar (MVar, newMVar, withMVar)
import Control.Exception
  (Exception, bracket, mask, onException, throwIO)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Unsafe as BU
import Foreign.C.Error (Errno(..), eINVAL)
import Foreign.C.String (CString, peekCString, withCString)
import Foreign.C.Types (CChar, CInt, CSize, CUInt)
import Foreign.Marshal.Alloc (alloca, allocaBytes)
import Foreign.Ptr (Ptr, nullPtr, plusPtr)
import Foreign.Storable (peek, peekByteOff, pokeByteOff, sizeOf)
import Database.LMDB.Internal

-- handles --------------------------------------------------------------------

-- | An LMDB environment. Roughly, an open database file (or, with
--   'envNoSubdir' off, a directory containing one).
data Env = Env
  {-# UNPACK #-} !(Ptr MDB_env)
  {-# UNPACK #-} !(MVar ()) -- held for the duration of a write txn

-- | Phantom tag for read-only transactions.
data RO

-- | Phantom tag for read-write transactions.
data RW

-- | A transaction. The phantom @s@ distinguishes 'RO' from 'RW' so
--   that write operations refuse to type-check against read-only
--   transactions. Valid only within the 'withReadTxn' or
--   'withWriteTxn' call that supplied it.
newtype Txn s = Txn (Ptr MDB_txn)

-- | A database handle, valid until its 'Env' is closed.
newtype Dbi = Dbi MDB_dbi

-- | A cursor for range-scanning a 'Dbi'. Valid only within the
--   'withCursor' call that supplied it.
newtype Cursor s = Cursor (Ptr MDB_cursor)

-- environment flags ----------------------------------------------------------

-- | Configuration for 'open'.
data EnvFlags = EnvFlags
  { envMapSize  :: !Int
    -- ^ map size in bytes; LMDB will refuse writes beyond this. The
    --   default of 10 MiB is intentionally small — set it to something
    --   appropriate for your workload before opening.
  , envMaxDbs   :: !Int
    -- ^ maximum number of named sub-databases. Default 1.
  , envNoSubdir :: !Bool
    -- ^ if 'True', treat the path as a regular file rather than a
    --   directory containing @data.mdb@ and @lock.mdb@.
  , envReadOnly :: !Bool
    -- ^ open the environment read-only.
  , envNoSync   :: !Bool
    -- ^ skip @fsync@ at commit time. Faster, but an OS crash or power
    --   loss may undo the most recent commits, and corrupt the
    --   database unless the filesystem preserves write order. A
    --   process crash loses nothing. Use 'sync' to flush explicitly.
  } deriving Show

-- | Default 'EnvFlags': 10 MiB map, one database, directory layout,
--   read-write, syncing at commit.
--
--   >>> envMapSize defaultEnvFlags
--   10485760
defaultEnvFlags :: EnvFlags
defaultEnvFlags = EnvFlags
  { envMapSize  = 10 * 1024 * 1024
  , envMaxDbs   = 1
  , envNoSubdir = False
  , envReadOnly = False
  , envNoSync   = False
  }

-- MDB_NOTLS is always set: GHC multiplexes Haskell threads over OS
-- threads, so reader slots must belong to txns rather than threads.
envFlagBits :: EnvFlags -> CUInt
envFlagBits EnvFlags { envNoSubdir = nsd, envReadOnly = ro, envNoSync = ns } =
  _MDB_NOTLS
    + (if nsd then _MDB_NOSUBDIR else 0)
    + (if ro then _MDB_RDONLY   else 0)
    + (if ns then _MDB_NOSYNC   else 0)

-- error handling -------------------------------------------------------------

-- | A typed wrapper around LMDB error codes. 'LMDBOther' carries the
--   raw error code plus the message from @mdb_strerror@.
data LMDBException
  = LMDBKeyExist
  | LMDBNotFound          -- ^ Thrown by 'openDbi' for a missing named
                          --   database. Lookups, cursor moves and
                          --   'del' report absence as 'Nothing' or
                          --   'False' instead.
  | LMDBMapFull
  | LMDBCorrupted
  | LMDBPanic
  | LMDBVersionMismatch
  | LMDBOther !Int !String
  deriving Show

instance Exception LMDBException

throwLmdb :: CInt -> IO a
throwLmdb code
  | code == _MDB_KEYEXIST         = throwIO LMDBKeyExist
  | code == _MDB_NOTFOUND         = throwIO LMDBNotFound
  | code == _MDB_MAP_FULL         = throwIO LMDBMapFull
  | code == _MDB_CORRUPTED        = throwIO LMDBCorrupted
  | code == _MDB_PANIC            = throwIO LMDBPanic
  | code == _MDB_VERSION_MISMATCH = throwIO LMDBVersionMismatch
  | otherwise = do
      msg <- mdb_strerror code >>= peekCString
      throwIO (LMDBOther (fromIntegral code) msg)

check :: CInt -> IO ()
check 0 = pure ()
check c = throwLmdb c
{-# INLINE check #-}

-- Reject an argument before it reaches LMDB, with the EINVAL LMDB
-- itself would use.
invalid :: String -> IO a
invalid msg = throwIO (LMDBOther (fromIntegral e) msg)
  where
    Errno e = eINVAL

-- environment lifecycle ------------------------------------------------------

-- | Open an LMDB environment at the given path. The path must already
--   exist (as a directory unless 'envNoSubdir' is set, in which case
--   the file's parent directory must exist).
--
--   Throws 'LMDBException' on failure, including for a negative
--   'envMapSize' or an out-of-range 'envMaxDbs'.
--
--   >>> env <- open "/tmp/mydb" defaultEnvFlags { envNoSubdir = True }
open :: FilePath -> EnvFlags -> IO Env
open path flags = do
  checkFlags flags
  lock <- newMVar ()
  mask $ \restore -> do
    envp <- alloca $ \pp -> do
      check =<< mdb_env_create pp
      peek pp
    let env = Env envp lock
    flip onException (close env) . restore $ do
      check =<< mdb_env_set_mapsize envp
                    (fromIntegral (envMapSize flags))
      check =<< mdb_env_set_maxdbs envp
                    (fromIntegral (envMaxDbs flags))
      withCString path $ \cpath ->
        check =<< mdb_env_open envp cpath (envFlagBits flags) 0o644
    pure env

-- Reject values that would wrap when converted to LMDB's types.
checkFlags :: EnvFlags -> IO ()
checkFlags EnvFlags { envMapSize = sz, envMaxDbs = dbs }
  | sz < 0 = invalid "open: negative envMapSize"
  | dbs < 0 || toInteger dbs > toInteger (maxBound :: MDB_dbi) =
      invalid "open: envMaxDbs out of range"
  | otherwise = pure ()

-- | Close an environment. Must be called at most once, after all of
--   its transactions have ended.
close :: Env -> IO ()
close (Env envp _) = mdb_env_close envp

-- | Bracketed environment lifecycle: 'open' the environment, run the
--   action, 'close' on exit (including async exceptions).
--
--   >>> withEnv "/tmp/mydb" defaultEnvFlags $ \env -> ...
withEnv :: FilePath -> EnvFlags -> (Env -> IO a) -> IO a
withEnv path flags = bracket (open path flags) close

-- | Flush the environment's buffers to disk. Commits already do so
--   unless 'envNoSync' is set.
--
--   >>> sync env
sync :: Env -> IO ()
sync (Env envp _) = check =<< mdb_env_sync envp 1

-- transactions ---------------------------------------------------------------

beginTxn :: Env -> CUInt -> IO (Ptr MDB_txn)
beginTxn (Env envp _) flags = alloca $ \pp -> do
  check =<< mdb_txn_begin envp nullPtr flags pp
  peek pp

-- | Run an action in a read-only transaction. The transaction is
--   aborted on exit; read transactions have no work to commit.
--
--   Read transactions don't wait on writers. They may be nested, or
--   run from any number of threads at once, up to LMDB's limit of
--   126 concurrent readers.
--
--   >>> withReadTxn env $ \txn -> ...
withReadTxn :: Env -> (Txn RO -> IO a) -> IO a
withReadTxn env act =
  bracket (beginTxn env _MDB_RDONLY) mdb_txn_abort (act . Txn)

-- | Run an action in a read-write transaction. Commits if the action
--   returns normally; aborts on exception.
--
--   Write transactions are serialized: one waits for any other in
--   progress to finish. They must not be nested.
--
--   LMDB requires a write transaction to begin and end on the same
--   OS thread. Under the threaded RTS, if the calling thread is
--   unbound, the action therefore runs in a temporary bound thread
--   (see 'Control.Concurrent.runInBoundThread'), and asynchronous
--   exceptions thrown to the caller are deferred until the
--   transaction has ended.
--
--   >>> withWriteTxn env $ \txn -> ...
withWriteTxn :: Env -> (Txn RW -> IO a) -> IO a
withWriteTxn env act = locked env 0 (act . Txn)

-- Run a txn under the env lock, committing if the action returns
-- normally and aborting if it throws.
locked :: Env -> CUInt -> (Ptr MDB_txn -> IO a) -> IO a
locked env@(Env _ lock) flags act =
  -- writers queue on the lock, not on LMDB's own mutex, which would
  -- block an OS thread: under the non-threaded RTS, the only one
  withMVar lock $ \_ -> pinned $ mask $ \restore -> do
    txnp <- beginTxn env flags
    r <- restore (act txnp) `onException` mdb_txn_abort txnp
    -- the commit runs masked, and must not be followed by an abort:
    -- mdb_txn_commit ends the txn whether or not it succeeds
    check =<< mdb_txn_commit txnp
    pure r

-- Run an action on a single OS thread. Under the non-threaded RTS,
-- all Haskell code already does.
pinned :: IO a -> IO a
pinned act
  | rtsSupportsBoundThreads = runInBoundThread act
  | otherwise               = act

-- databases ------------------------------------------------------------------

-- | Open a database, in a transaction of its own. Pass 'Nothing' for
--   the unnamed default database. If @create@ is 'True', the database
--   is created if it does not already exist, which requires a
--   writable environment; otherwise a missing named database throws
--   'LMDBNotFound'.
--
--   The handle is valid until the 'Env' is closed. Open each database
--   once, up front, and reuse the handle: a transaction that began
--   before 'openDbi' returned can't use it. Like a write transaction,
--   'openDbi' must not be called within 'withWriteTxn'.
--
--   >>> dbi <- openDbi env Nothing True
openDbi
  :: Env
  -> Maybe BS.ByteString -- ^ database name; 'Nothing' for the default
  -> Bool                -- ^ create if not present
  -> IO Dbi
openDbi env mname create = withName mname $ \cname ->
  -- LMDB forbids concurrent mdb_dbi_open calls, and keeps a handle
  -- only once the txn that opened it commits; 'locked' ensures both
  locked env txnFlags $ \txnp -> alloca $ \pp -> do
    check =<< mdb_dbi_open txnp cname dbiFlags pp
    Dbi <$> peek pp
  where
    (txnFlags, dbiFlags)
      | create    = (0, _MDB_CREATE)
      | otherwise = (_MDB_RDONLY, 0)

withName :: Maybe BS.ByteString -> (CString -> IO a) -> IO a
withName Nothing k = k nullPtr
withName (Just nm) k
  | BS.elem 0 nm = invalid "openDbi: database name contains NUL"
  | otherwise    = BS.useAsCString nm k

-- bytestring <-> MDB_val helpers --------------------------------------------

-- An MDB_val is a size_t length followed by a data pointer, each one
-- word wide on every GHC target. We poke fields directly rather than
-- going through the Storable instance so hot paths don't depend on it
-- being inlined.

-- Offset of an MDB_val's data pointer, and the size of an MDB_val.
valData, valSize :: Int
valData = sizeOf (nullPtr :: Ptr ())
valSize = 2 * valData

pokeVal :: Ptr MDB_val -> Ptr CChar -> Int -> IO ()
pokeVal vp p n = do
  pokeByteOff vp 0 (fromIntegral n :: CSize)
  pokeByteOff vp valData p
{-# INLINE pokeVal #-}

withBSAsVal :: BS.ByteString -> (Ptr MDB_val -> IO a) -> IO a
withBSAsVal bs k =
  BU.unsafeUseAsCStringLen bs $ \(p, n) ->
  allocaBytes valSize $ \vp -> do
    pokeVal vp p n
    k vp
{-# INLINE withBSAsVal #-}

peekVal :: Ptr MDB_val -> IO BS.ByteString
peekVal vp = do
  sz <- peekByteOff vp 0 :: IO CSize
  pd <- peekByteOff vp valData
  BS.packCStringLen (pd, fromIntegral sz)
{-# INLINE peekVal #-}

-- key-value ops --------------------------------------------------------------

-- | Look up a key. Returns 'Nothing' if not found. Other LMDB errors
--   throw 'LMDBException'.
--
--   The returned 'BS.ByteString' is a copy of the data; LMDB's
--   memory-mapped buffer is not aliased into Haskell.
--
--   >>> get txn dbi "hello"
--   Just "world"
get
  :: Txn s
  -> Dbi
  -> BS.ByteString               -- ^ key
  -> IO (Maybe BS.ByteString)    -- ^ value, if present
get (Txn txnp) (Dbi dbi) k =
  BU.unsafeUseAsCStringLen k $ \(kdata, klen) ->
  allocaBytes (2 * valSize) $ \kp -> do
    let !vp = kp `plusPtr` valSize
    pokeVal kp kdata klen
    code <- mdb_get txnp dbi kp vp
    if   code == _MDB_NOTFOUND
    then pure Nothing
    else do
      check code
      Just <$> peekVal vp
{-# INLINE get #-}

-- | Insert or replace a key. Throws 'LMDBException' on failure.
--
--   >>> put txn dbi "hello" "world"
put
  :: Txn RW
  -> Dbi
  -> BS.ByteString -- ^ key
  -> BS.ByteString -- ^ value
  -> IO ()
put (Txn txnp) (Dbi dbi) k v =
  BU.unsafeUseAsCStringLen k $ \(kdata, klen) ->
  BU.unsafeUseAsCStringLen v $ \(vdata, vlen) ->
  allocaBytes (2 * valSize) $ \kp -> do
    let !vp = kp `plusPtr` valSize
    pokeVal kp kdata klen
    pokeVal vp vdata vlen
    check =<< mdb_put txnp dbi kp vp 0
{-# INLINE put #-}

-- | Delete a key. Returns 'True' if the key existed and was removed,
--   'False' if it was already absent.
--
--   >>> del txn dbi "hello"
--   True
del
  :: Txn RW
  -> Dbi
  -> BS.ByteString -- ^ key
  -> IO Bool       -- ^ 'True' if the key was present and removed
del (Txn txnp) (Dbi dbi) k =
  withBSAsVal k $ \kp -> do
    code <- mdb_del txnp dbi kp nullPtr
    if   code == _MDB_NOTFOUND then pure False
    else do
      check code
      pure True
{-# INLINE del #-}

-- cursors --------------------------------------------------------------------

-- | Bracketed cursor lifecycle: open a cursor over the given 'Dbi',
--   run the action, close the cursor on exit (including async
--   exceptions).
--
--   >>> withCursor txn dbi $ \cur -> ...
withCursor :: Txn s -> Dbi -> (Cursor s -> IO a) -> IO a
withCursor (Txn txnp) (Dbi dbi) act = bracket open' mdb_cursor_close use'
  where
    open' = alloca $ \pp -> do
      check =<< mdb_cursor_open txnp dbi pp
      peek pp
    use' p = act (Cursor p)

-- Scan-position helper (no input key): one buffer for both output
-- MDB_vals.
cursorMove
  :: Cursor s
  -> CInt              -- MDB_cursor_op
  -> IO (Maybe (BS.ByteString, BS.ByteString))
cursorMove (Cursor cp) op =
  allocaBytes (2 * valSize) $ \kp -> do
    let !vp = kp `plusPtr` valSize
    code <- mdb_cursor_get cp kp vp op
    if   code == _MDB_NOTFOUND
    then pure Nothing
    else do
      check code
      k <- peekVal kp
      v <- peekVal vp
      pure (Just (k, v))
{-# INLINE cursorMove #-}

-- | Position at the first key/value pair, or 'Nothing' on an empty
--   database.
--
--   >>> cursorFirst cur
--   Just ("a","alpha")
cursorFirst :: Cursor s -> IO (Maybe (BS.ByteString, BS.ByteString))
cursorFirst c = cursorMove c _MDB_FIRST
{-# INLINE cursorFirst #-}

-- | Position at the last key/value pair, or 'Nothing' on an empty
--   database.
cursorLast :: Cursor s -> IO (Maybe (BS.ByteString, BS.ByteString))
cursorLast c = cursorMove c _MDB_LAST
{-# INLINE cursorLast #-}

-- | Advance to the next key/value pair, or 'Nothing' if already at
--   the end.
cursorNext :: Cursor s -> IO (Maybe (BS.ByteString, BS.ByteString))
cursorNext c = cursorMove c _MDB_NEXT
{-# INLINE cursorNext #-}

-- | Step back to the previous key/value pair, or 'Nothing' if already
--   at the start.
cursorPrev :: Cursor s -> IO (Maybe (BS.ByteString, BS.ByteString))
cursorPrev c = cursorMove c _MDB_PREV
{-# INLINE cursorPrev #-}

-- | Position at the first key /greater than or equal to/ the given
--   key, or 'Nothing' if no such key exists.
--
--   >>> cursorSeek cur "m"
--   Just ("n","november")
cursorSeek
  :: Cursor s
  -> BS.ByteString
  -> IO (Maybe (BS.ByteString, BS.ByteString))
cursorSeek (Cursor cp) k =
  BU.unsafeUseAsCStringLen k $ \(kdata, klen) ->
  allocaBytes (2 * valSize) $ \kp -> do
    let !vp = kp `plusPtr` valSize
    pokeVal kp kdata klen
    code <- mdb_cursor_get cp kp vp _MDB_SET_RANGE
    if   code == _MDB_NOTFOUND
    then pure Nothing
    else do
      check code
      ok <- peekVal kp
      ov <- peekVal vp
      pure (Just (ok, ov))