RocksDB.hs

{-
  Copyright (c) Meta Platforms, Inc. and affiliates.
  All rights reserved.

  This source code is licensed under the BSD-style license found in the
  LICENSE file in the root directory of this source tree.
-}

module Glean.Database.Storage.RocksDB
  ( RocksDB(..)
  , newStorage
  ) where

import Control.Exception
import Control.Monad
import qualified Data.HashMap.Strict as HashMap
import Data.Int
import Data.List (unzip4)
import qualified Data.Vector.Storable as VS
import Data.Word
import Foreign.C.String
import Foreign.C.Types
import Foreign.ForeignPtr
import Foreign.Marshal.Array
import Foreign.Ptr
import Foreign.Storable
import System.Directory
import System.IO.Temp (withTempDirectory)
import System.FilePath
import System.Process (readProcessWithExitCode)
import System.Exit (ExitCode(ExitSuccess))

import Util.FFI
import Util.IO (safeRemovePathForcibly)

import Glean.Database.Backup.Backend (Data(Data))
import Glean.Database.Repo (databasePath)
import Glean.Database.Storage
import Glean.FFI
import Glean.Repo.Text
import Glean.RTS.Foreign.FactSet (FactSet)
import Glean.RTS.Foreign.Lookup
  (CanLookup(..), Lookup(..))
import Glean.RTS.Foreign.Ownership as Ownership
import Glean.RTS.Foreign.Stats (marshalPredicateStats)
import Glean.RTS.Types (Fid(..), invalidFid, Pid(..))
import qualified Glean.ServerConfig.Types as ServerConfig
import Glean.Types (Repo)
import Glean.Util.Disk
import Glean.Impl.MemoryReader
import System.IO.Extra (withTempFile)

newtype Cache = Cache (ForeignPtr Cache)

instance Object Cache where
  wrap = Cache
  unwrap (Cache p) = p
  destroy = glean_rocksdb_free_cache

newCache :: Int -> IO Cache
newCache size =
  construct $ invoke $ glean_rocksdb_new_cache (fromIntegral size)

withCache :: Maybe Cache -> (Ptr Cache -> IO a) -> IO a
withCache (Just cache) f = with cache f
withCache Nothing f = f nullPtr

data RocksDB = RocksDB
  { rocksRoot :: FilePath
  , rocksCache :: Maybe Cache
  , rocksCacheIndexAndFilterBlocks :: Bool
  , rocksMaxDiskSize :: Maybe Int
    -- ^ virtual limit to report capped disk capacities. The limit is
    -- not enforced. It's up to each io usage to check diskspace before writing.
    -- We're using this to avoid serving too many dbs on query servers,
    -- and smarter sharding.
  }

newStorage :: FilePath -> ServerConfig.Config -> IO RocksDB
newStorage root ServerConfig.Config{..} = do
  cache <- if config_db_rocksdb_cache_mb > 0
    then
      Just <$> newCache (fromIntegral config_db_rocksdb_cache_mb * 1024 * 1024)
    else return Nothing
  mem_capacity <- totalMemCapacity
  return RocksDB
    { rocksRoot = root
    , rocksCache = cache
    , rocksCacheIndexAndFilterBlocks =
        config_db_rocksdb_cache_index_and_filter_blocks
    , rocksMaxDiskSize = case mem_capacity of
        Just mem -> (* mem) . fromIntegral <$>
          config_db_rocksdb_disk_mem_capacity_ratio_limit
        Nothing -> Nothing
    }

newtype Container = Container (Ptr Container)
  deriving(Storable)

instance Static Container where
  destroyStatic = glean_rocksdb_container_free

instance Storage RocksDB where
  data Database RocksDB = Database
    { dbPtr :: ForeignPtr (Database RocksDB)
    , dbRepo :: Repo
    }

  describe rocks = "rocksdb:" <> rocksRoot rocks

  open rocks repo mode (DBVersion version) = do
    (cmode, start, ownership) <- case mode of
      ReadOnly -> return (0, invalidFid, Nothing)
      ReadWrite -> return (1, invalidFid, Nothing)
      Create start ownership _ -> do
        createDirectoryIfMissing True path
        return (2, start, ownership)
    withCString path $ \cpath ->
      withCache (rocksCache rocks) $ \cache_ptr ->
      using (invoke $
          glean_rocksdb_container_open cpath
            cmode
            (fromIntegral (fromEnum (rocksCacheIndexAndFilterBlocks rocks)))
            cache_ptr)
        $ \container -> do
      fp <- mask_ $ do
        first_unit_id <- maybe (return firstUsetId) nextUsetId ownership
        p <- invoke $
          glean_rocksdb_container_open_database container start
            first_unit_id version
        newForeignPtr glean_rocksdb_database_free p
      return (Database fp repo)
    where
      path = containerPath rocks repo

  close db = withContainer db glean_rocksdb_container_close

  delete rocks = safeRemovePathForcibly . containerPath rocks

  safeRemoveForcibly rocks =
      safeRemovePathForcibly . databasePath (rocksRoot rocks)

  predicateStats db = unsafeWithForeignPtr (dbPtr db)
    $ marshalPredicateStats . glean_rocksdb_database_predicateStats

  store db key value =
    withContainer db $ \s_ptr ->
    unsafeWithBytes key $ \key_ptr key_size ->
    unsafeWithBytes value $ \value_ptr value_size ->
    invoke $ glean_rocksdb_container_write_data
      s_ptr
      key_ptr
      key_size
      value_ptr
      value_size

  retrieve db key =
    withContainer db $ \s_ptr ->
    unsafeWithBytes key $ \key_ptr key_size -> do
      (value_ptr, value_size, found)
        <- invoke $ glean_rocksdb_container_read_data s_ptr key_ptr key_size
      if found /= 0
        then Just <$> unsafeMallocedByteString value_ptr value_size
        else return Nothing

  commit db facts = unsafeWithForeignPtr (dbPtr db) $ \db_ptr -> do
    with facts $ \facts_ptr -> invoke $ glean_rocksdb_commit db_ptr facts_ptr

  addOwnership db _ owned =
    unsafeWithForeignPtr (dbPtr db) $ \db_ptr ->
    when (not $ HashMap.null owned) $
      withMany entry (HashMap.toList owned) $ \xs ->
      let (unit_ptrs, unit_sizes, facts_ptrs, facts_sizes) = unzip4 xs
      in
      withArray unit_ptrs $ \p_unit_ptrs ->
      withArray unit_sizes $ \p_unit_sizes ->
      withArray facts_ptrs $ \p_facts_ptrs ->
      withArray facts_sizes $ \p_facts_sizes ->
      invoke $ glean_rocksdb_add_ownership
        db_ptr
        (fromIntegral $ HashMap.size owned)
        p_unit_ptrs
        p_unit_sizes
        p_facts_ptrs
        p_facts_sizes
    where
      entry (unit, facts) f =
        unsafeWithBytes unit $ \unit_ptr unit_size ->
        VS.unsafeWith facts $ \facts_ptr ->
        f (unit_ptr, unit_size, facts_ptr, fromIntegral $ VS.length facts)

  optimize db compact = withContainer db $ \s_ptr ->
    invoke $ glean_rocksdb_container_optimize s_ptr
      (fromIntegral (fromEnum compact))

  computeOwnership db base inv =
    unsafeWithForeignPtr (dbPtr db) $ \db_ptr ->
    using (invoke $ glean_rocksdb_get_ownership_unit_iterator db_ptr) $
    Ownership.compute inv db base

  storeOwnership db _ own =
    unsafeWithForeignPtr (dbPtr db) $ \db_ptr ->
    with own $ \own_ptr ->
    invoke $ glean_rocksdb_store_ownership db_ptr own_ptr

  getOwnership db = fmap Just $
    unsafeWithForeignPtr (dbPtr db) $ \db_ptr ->
    construct $ invoke $ glean_rocksdb_get_ownership db_ptr

  getUnitId db unit =
    unsafeWithForeignPtr (dbPtr db) $ \db_ptr ->
    unsafeWithBytes unit $ \unit_ptr unit_size -> do
      w64 <- invoke $ glean_rocksdb_get_unit_id db_ptr unit_ptr unit_size
      if w64 == 0xffffffffffffffff
        then return Nothing
        else return (Just (UnitId (fromIntegral w64)))

  getUnit db unit =
    unsafeWithForeignPtr (dbPtr db) $ \db_ptr -> do
      (unit_ptr, unit_size) <- invoke $ glean_rocksdb_get_unit db_ptr unit
      if unit_size /= 0
        then Just <$> unsafeMallocedByteString unit_ptr unit_size
        else return Nothing

  addDefineOwnership db _ define =
    unsafeWithForeignPtr (dbPtr db) $ \db_ptr ->
    with define $ \define_ptr ->
      invoke $ glean_rocksdb_add_define_ownership db_ptr define_ptr

  computeDerivedOwnership db _ ownership base (Pid pid) =
    unsafeWithForeignPtr (dbPtr db) $ \db_ptr ->
    using
      (invoke $
        glean_rocksdb_get_derived_fact_ownership_iterator
          db_ptr
          (fromIntegral pid)) $
      Ownership.computeDerivedOwnership ownership base

  cacheOwnership db =
    unsafeWithForeignPtr (dbPtr db) $ \db_ptr ->
      invoke $ glean_rocksdb_cache_ownership db_ptr

  prepareFactOwnerCache db =
    unsafeWithForeignPtr (dbPtr db) $ \db_ptr ->
      invoke $ glean_rocksdb_prepare_fact_owner_cache db_ptr

  getTotalCapacity rocksdb = do
    exists <- doesDirectoryExist (rocksRoot rocksdb)
    if exists
      then do
        fullDiskCapacity <- getDiskSize (rocksRoot rocksdb)
        return $ Just $ case rocksMaxDiskSize rocksdb of
          Just maxDiskSize -> min maxDiskSize fullDiskCapacity
          Nothing -> fullDiskCapacity
      else return Nothing

  getUsedCapacity rocksdb = do
    exists <- doesDirectoryExist (rocksRoot rocksdb)
    if exists
      then Just <$> getUsedDiskSpace (rocksRoot rocksdb)
      else return Nothing

  getFreeCapacity rocksdb = do
    used <- getUsedCapacity rocksdb
    total <- getTotalCapacity rocksdb
    case (used,total) of
      (Just used, Just total) -> return $ total - used
      _ -> getFreeDiskSpace (rocksRoot rocksdb) -- not aware of disk limit

  withScratchRoot rocks f = f $ rocksRoot rocks </> ".scratch"

  backup db scratch process = do
    createDirectoryIfMissing True path
    withContainer db $ \s_ptr ->
      withCString path $ invoke . glean_rocksdb_container_backup s_ptr
    withTempFile $ \tarFile -> do
      tar ["-cf", tarFile, "-C", scratch, "backup"]
      size <- getFileSize tarFile
      process tarFile (Data $ fromIntegral size)
    where
      path = scratch </> "backup"

  restore rocks repo scratch scratch_file =
    withTempDirectory scratch "restore" $ \scratch_restore -> do
      unTar scratch_file scratch_restore
      -- to avoid retaining an extra copy of the DB during restore,
      -- delete the input file now.
      removeFile scratch_file

      -- If the tarfile contains "backup/.." then it is a RocksDB backup
      -- If the tarfile contains "db/.." then it is a plain tarball of the DB
      let scratch_restore_backup = scratch_restore </> "backup"
      is_rocksdb_backup <- doesDirectoryExist scratch_restore_backup
      db <-
        if is_rocksdb_backup
          then do
            let scratch_db = scratch </> "db"
            createDirectoryIfMissing True scratch_db
            withCString scratch_db $ \p_target ->
              withCString (scratch_restore </> "backup") $ \p_source ->
                invoke $ glean_rocksdb_restore p_target p_source
            return scratch_db
          else do
            let scratch_restore_db = scratch_restore </> "db"
            is_copy <- doesDirectoryExist scratch_restore_db
            if is_copy
              then return scratch_restore_db
              else throwIO $ userError "unrecognised backup"

      let target = containerPath rocks repo
      createDirectoryIfMissing True $ takeDirectory target
      renameDirectory db target


unTar :: FilePath -> FilePath -> IO ()
unTar scratch_file scratch_restore =
  tar ["-xf", scratch_file, "-C", scratch_restore]

tar :: [String] -> IO ()
tar args = do
  tarPath <- findExecutable "tar"
  case tarPath of
    Nothing  -> throwIO $ userError "Cannot find tar executable"
    Just path -> do
      (ec, _, err) <- readProcessWithExitCode path args ""
      unless (ec == ExitSuccess) $ throwIO $ userError err

containerPath :: RocksDB -> Repo -> FilePath
containerPath RocksDB{..} repo = databasePath rocksRoot repo </> "db"

instance CanLookup (Database RocksDB) where
  lookupName db = "rocksdb:" <> repoToText (dbRepo db)
  withLookup db f = unsafeWithForeignPtr (dbPtr db) $
    f . glean_rocksdb_database_lookup

withContainer :: Database RocksDB -> (Container -> IO a) -> IO a
withContainer db f = unsafeWithForeignPtr (dbPtr db) $
  f . glean_rocksdb_database_container

foreign import ccall unsafe glean_rocksdb_new_cache
  :: CSize -> Ptr (Ptr Cache) -> IO CString
foreign import ccall unsafe "&glean_rocksdb_free_cache"
  glean_rocksdb_free_cache :: Destroy Cache


foreign import ccall safe glean_rocksdb_container_open
  :: CString
  -> CInt
  -> CBool
  -> Ptr Cache
  -> Ptr Container
  -> IO CString
foreign import ccall safe glean_rocksdb_container_free
  :: Container -> IO ()

foreign import ccall safe glean_rocksdb_container_close
  :: Container -> IO ()

foreign import ccall unsafe glean_rocksdb_container_write_data
  :: Container
  -> Ptr ()
  -> CSize
  -> Ptr ()
  -> CSize
  -> IO CString

foreign import ccall unsafe glean_rocksdb_container_read_data
  :: Container
  -> Ptr ()
  -> CSize
  -> Ptr (Ptr ())
  -> Ptr CSize
  -> Ptr CChar
  -> IO CString

foreign import ccall safe glean_rocksdb_container_optimize
  :: Container -> CBool -> IO CString

foreign import ccall safe glean_rocksdb_container_backup
  :: Container -> CString -> IO CString

foreign import ccall safe glean_rocksdb_container_open_database
  :: Container
  -> Fid
  -> UsetId
  -> Int64
  -> Ptr (Ptr (Database RocksDB))
  -> IO CString
foreign import ccall safe "&glean_rocksdb_database_free"
  glean_rocksdb_database_free :: Destroy (Database RocksDB)

foreign import ccall unsafe glean_rocksdb_database_container
  :: Ptr (Database RocksDB) -> Container

foreign import ccall unsafe glean_rocksdb_database_lookup
  :: Ptr (Database RocksDB) -> Ptr Lookup

foreign import ccall safe glean_rocksdb_commit
  :: Ptr (Database RocksDB)
  -> Ptr FactSet
  -> IO CString

foreign import ccall safe glean_rocksdb_add_ownership
  :: Ptr (Database RocksDB)
  -> CSize
  -> Ptr (Ptr ())
  -> Ptr CSize
  -> Ptr (Ptr Fid)
  -> Ptr CSize
  -> IO CString

foreign import ccall safe glean_rocksdb_get_ownership_unit_iterator
  :: Ptr (Database RocksDB)
  -> Ptr Ownership.UnitIterator
  -> IO CString

foreign import ccall unsafe glean_rocksdb_database_predicateStats
  :: Ptr (Database RocksDB)
  -> Ptr CSize
  -> Ptr (Ptr Int64)
  -> Ptr (Ptr Word64)
  -> Ptr (Ptr Word64)
  -> IO CString

foreign import ccall safe glean_rocksdb_restore
  :: CString
  -> CString
  -> IO CString

foreign import ccall unsafe glean_rocksdb_get_unit_id
  :: Ptr (Database RocksDB)
  -> Ptr ()
  -> CSize
  -> Ptr Word64
  -> IO CString

foreign import ccall unsafe glean_rocksdb_get_unit
  :: Ptr (Database RocksDB)
  -> UnitId
  -> Ptr (Ptr ())
  -> Ptr CSize
  -> IO CString

foreign import ccall safe glean_rocksdb_store_ownership
  :: Ptr (Database RocksDB)
  -> Ptr ComputedOwnership
  -> IO CString

foreign import ccall unsafe glean_rocksdb_get_ownership
  :: Ptr (Database RocksDB)
  -> Ptr (Ptr Ownership)
  -> IO CString

foreign import ccall safe glean_rocksdb_add_define_ownership
  :: Ptr (Database RocksDB)
  -> Ptr DefineOwnership
  -> IO CString

foreign import ccall unsafe glean_rocksdb_get_derived_fact_ownership_iterator
  :: Ptr (Database RocksDB)
  -> Word64
  -> Ptr Ownership.DerivedFactOwnershipIterator
  -> IO CString

foreign import ccall unsafe glean_rocksdb_cache_ownership
  :: Ptr (Database RocksDB)
  -> IO CString

foreign import ccall safe glean_rocksdb_prepare_fact_owner_cache
  :: Ptr (Database RocksDB)
  -> IO CString