|
| 1 | +{- |
| 2 | + Copyright (c) Meta Platforms, Inc. and affiliates. |
| 3 | + All rights reserved. |
| 4 | +
|
| 5 | + This source code is licensed under the BSD-style license found in the |
| 6 | + LICENSE file in the root directory of this source tree. |
| 7 | +-} |
| 8 | + |
| 9 | +module Glean.Database.Storage.LMDB |
| 10 | + ( LMDB(..) |
| 11 | + , newStorage |
| 12 | + ) where |
| 13 | + |
| 14 | +import Control.Exception |
| 15 | +import Data.Int |
| 16 | +import Foreign.C.String |
| 17 | +import Foreign.C.Types |
| 18 | +import Foreign.ForeignPtr |
| 19 | +import Foreign.Ptr |
| 20 | +import System.Directory |
| 21 | +import System.IO.Temp (withTempDirectory) |
| 22 | +import System.FilePath |
| 23 | +import System.Process |
| 24 | + |
| 25 | +import Util.FFI |
| 26 | +import Util.IO (safeRemovePathForcibly) |
| 27 | + |
| 28 | +import Glean.Database.Backup.Backend (Data(Data)) |
| 29 | +import Glean.Database.Repo (databasePath) |
| 30 | +import Glean.Database.Storage |
| 31 | +import Glean.Database.Storage.DB |
| 32 | +import Glean.FFI |
| 33 | +import Glean.RTS.Foreign.Lookup (CanLookup(..)) |
| 34 | +import Glean.RTS.Foreign.Ownership as Ownership |
| 35 | + hiding (computeDerivedOwnership) |
| 36 | +import Glean.RTS.Types (Fid(..), invalidFid) |
| 37 | +import qualified Glean.ServerConfig.Types as ServerConfig |
| 38 | +import Glean.Types (Repo) |
| 39 | +import Glean.Util.Disk |
| 40 | +import Glean.Impl.MemoryReader |
| 41 | + |
| 42 | +data LMDB = LMDB |
| 43 | + { lmdbRoot :: FilePath |
| 44 | + , lmdbMaxDiskSize :: Maybe Int |
| 45 | + -- ^ virtual limit to report capped disk capacities. The limit is |
| 46 | + -- not enforced. It's up to each io usage to check diskspace before writing. |
| 47 | + -- We're using this to avoid serving too many dbs on query servers, |
| 48 | + -- and smarter sharding. |
| 49 | + } |
| 50 | + |
| 51 | +newStorage :: FilePath -> ServerConfig.Config -> IO LMDB |
| 52 | +newStorage root ServerConfig.Config{..} = do |
| 53 | + mem_capacity <- totalMemCapacityKB |
| 54 | + return LMDB |
| 55 | + { lmdbRoot = root |
| 56 | + , lmdbMaxDiskSize = case mem_capacity of |
| 57 | + Just mem -> (* (mem * 1024)) . fromIntegral <$> |
| 58 | + config_db_rocksdb_disk_mem_capacity_ratio_limit |
| 59 | + Nothing -> Nothing |
| 60 | + } |
| 61 | + |
| 62 | +newtype instance Database LMDB = Database DB |
| 63 | + deriving (CanLookup) |
| 64 | + |
| 65 | +instance Storage LMDB where |
| 66 | + describe db = "lmdb:" <> lmdbRoot db |
| 67 | + |
| 68 | + open lmdb repo mode (DBVersion version) = do |
| 69 | + (cmode, start, ownership) <- case mode of |
| 70 | + ReadOnly -> return (0, invalidFid, Nothing) |
| 71 | + ReadWrite -> return (1, invalidFid, Nothing) |
| 72 | + Create start ownership _ -> do |
| 73 | + createDirectoryIfMissing True path |
| 74 | + return (2, start, ownership) |
| 75 | + withCString path $ \cpath -> |
| 76 | + using (invoke $ glean_lmdb_container_open cpath cmode) |
| 77 | + $ \container -> do |
| 78 | + fp <- mask_ $ do |
| 79 | + first_unit_id <- maybe (return firstUsetId) nextUsetId ownership |
| 80 | + p <- invoke $ |
| 81 | + glean_lmdb_container_open_database container start |
| 82 | + first_unit_id version |
| 83 | + newForeignPtr glean_rocksdb_database_free p |
| 84 | + return (Database (DB (castForeignPtr fp) repo)) |
| 85 | + where |
| 86 | + path = containerPath lmdb repo |
| 87 | + |
| 88 | + delete lmdb = safeRemovePathForcibly . containerPath lmdb |
| 89 | + |
| 90 | + safeRemoveForcibly lmdb = |
| 91 | + safeRemovePathForcibly . databasePath (lmdbRoot lmdb) |
| 92 | + |
| 93 | + getTotalCapacity lmdb = do |
| 94 | + exists <- doesDirectoryExist (lmdbRoot lmdb) |
| 95 | + if exists |
| 96 | + then do |
| 97 | + fullDiskCapacity <- getDiskSize (lmdbRoot lmdb) |
| 98 | + return $ Just $ case lmdbMaxDiskSize lmdb of |
| 99 | + Just maxDiskSize -> min maxDiskSize fullDiskCapacity |
| 100 | + Nothing -> fullDiskCapacity |
| 101 | + else return Nothing |
| 102 | + |
| 103 | + getUsedCapacity lmdb = do |
| 104 | + exists <- doesDirectoryExist (lmdbRoot lmdb) |
| 105 | + if exists |
| 106 | + then Just <$> getUsedDiskSpace (lmdbRoot lmdb) |
| 107 | + else return Nothing |
| 108 | + |
| 109 | + getFreeCapacity lmdb = do |
| 110 | + used <- getUsedCapacity lmdb |
| 111 | + total <- getTotalCapacity lmdb |
| 112 | + case (used,total) of |
| 113 | + (Just used, Just total) -> return $ total - used |
| 114 | + _ -> getFreeDiskSpace (lmdbRoot lmdb) -- not aware of disk limit |
| 115 | + |
| 116 | + withScratchRoot rocks f = f $ lmdbRoot rocks </> ".scratch" |
| 117 | + |
| 118 | + restore lmdb repo scratch scratch_file = do |
| 119 | + withTempDirectory scratch "restore" $ \scratch_restore -> do |
| 120 | + let db = scratch_restore </> "db" |
| 121 | + createDirectoryIfMissing True db |
| 122 | + callProcess "unsquashfs" ["-d", db, scratch_file ] |
| 123 | + -- to avoid retaining an extra copy of the DB during restore, |
| 124 | + -- delete the input file now. |
| 125 | + let target = containerPath lmdb repo |
| 126 | + createDirectoryIfMissing True $ takeDirectory target |
| 127 | + renameDirectory db target |
| 128 | + |
| 129 | + |
| 130 | +containerPath :: LMDB -> Repo -> FilePath |
| 131 | +containerPath LMDB{..} repo = databasePath lmdbRoot repo </> "db" |
| 132 | + |
| 133 | +instance DatabaseOps (Database LMDB) where |
| 134 | + close (Database db) = close db |
| 135 | + predicateStats (Database db) = predicateStats db |
| 136 | + store (Database db) = store db |
| 137 | + retrieve (Database db) = retrieve db |
| 138 | + commit (Database db) = commit db |
| 139 | + addOwnership (Database db) = addOwnership db |
| 140 | + optimize (Database db) = optimize db |
| 141 | + computeOwnership (Database db) = computeOwnership db |
| 142 | + storeOwnership (Database db) = storeOwnership db |
| 143 | + getOwnership (Database db) = getOwnership db |
| 144 | + getUnitId (Database db) = getUnitId db |
| 145 | + getUnit (Database db) = getUnit db |
| 146 | + addDefineOwnership (Database db) = addDefineOwnership db |
| 147 | + computeDerivedOwnership (Database db) = computeDerivedOwnership db |
| 148 | + cacheOwnership (Database db) = cacheOwnership db |
| 149 | + prepareFactOwnerCache (Database db) = prepareFactOwnerCache db |
| 150 | + |
| 151 | + backup (Database db) scratch process = do |
| 152 | + createDirectoryIfMissing True (scratch </> "backup") |
| 153 | + backup db scratch $ \_ _ -> do |
| 154 | + withTempDirectory scratch "out" $ \tmpdir -> do |
| 155 | + let out = tmpdir </> "db.squashfs" |
| 156 | + callProcess "mksquashfs" [ |
| 157 | + scratch </> "backup", out, |
| 158 | + "-comp", "zstd", "-Xcompression-level", "8" ] |
| 159 | + size <- getFileSize out |
| 160 | + process out (Data $ fromIntegral size) |
| 161 | + |
| 162 | +foreign import ccall safe glean_lmdb_container_open |
| 163 | + :: CString |
| 164 | + -> CInt |
| 165 | + -> Ptr Container |
| 166 | + -> IO CString |
| 167 | + |
| 168 | +foreign import ccall safe glean_lmdb_container_open_database |
| 169 | + :: Container |
| 170 | + -> Fid |
| 171 | + -> UsetId |
| 172 | + -> Int64 |
| 173 | + -> Ptr (Ptr (Database LMDB)) |
| 174 | + -> IO CString |
| 175 | +foreign import ccall safe "&glean_rocksdb_database_free" |
| 176 | + glean_rocksdb_database_free :: Destroy (Database LMDB) |
0 commit comments