[lmdb facebookincubator#7] Optionally mount restored DBs using squashfs

Author: simonmar

glean/config/server/server_config.thrift

@@ -324,6 +324,12 @@ struct Config {
   // Args to pass to mksquashfs when creating a backup of an LMDB
   39: list<string> db_lmdb_mksquashfs_args =
     [ "-comp", "zstd", "-Xcompression-level", "8" ];
+
+  // LMDB: unpack a restored DB. If false, the DB will be mounted using
+  // squashfs instead, which is more space-efficient but incurs a small
+  // runtime cost to decompress the DB on demand. Mounting requires
+  // squashfs-tools and squashfuse to be installed on Linux.
+  40: bool db_lmdb_restore_unpack = true;
 }
 
 // The following were automatically generated and may benefit from renaming.

glean/db/Glean/Database/Backup.hs

@@ -310,13 +310,13 @@ doRestore env@Env{..} repo meta
   , Just (_, Some site, r_repo) <- fromRepoLocator envBackupBackends loc
   , r_repo == repo =
     loggingAction (runLogRepo "restore" env repo) (const mempty) $ do
-      ServerConfig.Config{..} <- Observed.get envServerConfig
+      cfg@ServerConfig.Config{..} <- Observed.get envServerConfig
       let maybeTimeout =
             case config_restore_timeout of
               Just seconds -> void . timeout (fromIntegral seconds * 1000000)
               Nothing -> id
       withStorageFor env repo meta $ \storage ->
-        maybeTimeout $ restore site storage size `catch` handler storage
+        maybeTimeout $ restore site storage cfg size `catch` handler storage
 
       -- NOTE: No point in adding the repo to the sinbin if there was
       -- an exception, the handler removed it from the list of known DBs
@@ -333,8 +333,15 @@ doRestore env@Env{..} repo meta
   where
   say log s = log $ inRepo repo $ "restore: " ++ s
 
-  restore :: (Storage st, Site s) => s -> st -> Maybe Int64 -> IO ()
-  restore site storage bytes = traceMsg envTracer (GleanTraceDownload repo) $ do
+  restore
+    :: (Storage st, Site s)
+    => s
+    -> st
+    -> ServerConfig.Config
+    -> Maybe Int64
+    -> IO ()
+  restore site storage cfg bytes =
+      traceMsg envTracer (GleanTraceDownload repo) $ do
     atomically $ notify envListener $ RestoreStarted repo
     mbFreeBytes <- (Just <$> Storage.getFreeCapacity storage)
                   `catch` \(_ :: IOException) -> return Nothing
@@ -359,7 +366,7 @@ doRestore env@Env{..} repo meta
     say logInfo "restoring"
     createDirectoryIfMissing True scratch_restore
     traceMsg envTracer GleanTraceStorageRestore $
-      Storage.restore storage repo scratch_restore scratch_file
+      Storage.restore storage cfg repo scratch_restore scratch_file
     say logInfo "adding"
     traceMsg envTracer GleanTraceFinishRestore $
       Catalog.finishRestoring envCatalog repo

glean/db/Glean/Database/Storage.hs

@@ -121,6 +121,7 @@ class DatabaseOps (Database s) => Storage s where
   -- to reduce the number of copies of the DB on disk during a restore.
   restore
     :: s   -- ^ storage
+    -> ServerConfig.Config  -- ^ server config
     -> Repo  -- ^ repo
     -> FilePath  -- ^ scratch directory
     -> FilePath  -- ^ file containing the serialiased database (produced by 'backup')

glean/db/Glean/Database/Storage/LMDB.hs

@@ -12,6 +12,7 @@ module Glean.Database.Storage.LMDB
   ) where
 
 import Control.Exception
+import Control.Monad
 import Data.Int
 import qualified Data.Text as Text
 import Foreign.C.String
@@ -60,15 +61,25 @@ newStorage root ServerConfig.Config{..} = do
         Nothing -> Nothing
     }
 
-newtype instance Database LMDB = Database DB
-  deriving (CanLookup)
+data instance Database LMDB = Database DB LMDB
+
+instance CanLookup (Database LMDB) where
+  withLookup (Database db _) = withLookup db
+  lookupName (Database db _) = lookupName db
 
 instance Storage LMDB where
   describe db = "lmdb:" <> lmdbRoot db
 
   open lmdb repo mode (DBVersion version) = do
     (cmode, start, ownership) <- case mode of
-      ReadOnly -> return (0, invalidFid, Nothing)
+      ReadOnly -> do
+        exists <- doesDirectoryExist (path </> "data.mdb")
+        when (not exists) $ do
+          haveSquash <- doesFileExist squash
+          when haveSquash $ do
+            createDirectoryIfMissing True path
+            callProcess "squashfuse_ll" [squash, path]
+        return (0, invalidFid, Nothing)
       ReadWrite -> return (1, invalidFid, Nothing)
       Create start ownership _ -> do
         createDirectoryIfMissing True path
@@ -82,9 +93,10 @@ instance Storage LMDB where
           glean_lmdb_container_open_database container start
             first_unit_id version
         newForeignPtr glean_rocksdb_database_free p
-      return (Database (DB (castForeignPtr fp) repo))
+      return (Database (DB (castForeignPtr fp) repo) lmdb)
     where
       path = containerPath lmdb repo
+      squash = path <.> "squashfs"
 
   delete lmdb = safeRemovePathForcibly . containerPath lmdb
 
@@ -116,39 +128,49 @@ instance Storage LMDB where
 
   withScratchRoot rocks f = f $ lmdbRoot rocks </> ".scratch"
 
-  restore lmdb repo scratch scratch_file = do
+  restore lmdb cfg repo scratch scratch_file = do
     withTempDirectory scratch "restore" $ \scratch_restore -> do
-      let db = scratch_restore </> "db"
-      createDirectoryIfMissing True db
-      callProcess "unsquashfs" ["-d", db, scratch_file ]
-        -- to avoid retaining an extra copy of the DB during restore,
-        -- delete the input file now.
       let target = containerPath lmdb repo
       createDirectoryIfMissing True $ takeDirectory target
-      renameDirectory db target
+      if ServerConfig.config_db_lmdb_restore_unpack cfg
+        then do
+          let db = scratch_restore </> "db"
+          createDirectoryIfMissing True db
+          callProcess "unsquashfs" ["-d", db, scratch_file ]
+            -- to avoid retaining an extra copy of the DB during restore,
+            -- delete the input file now.
+          renameDirectory db target
+        else
+          renameFile scratch_file (target <.> "squashfs")
 
 containerPath :: LMDB -> Repo -> FilePath
 containerPath LMDB{..} repo = databasePath lmdbRoot repo </> "db"
 
 instance DatabaseOps (Database LMDB) where
-  close (Database db) = close db
-  predicateStats (Database db) = predicateStats db
-  store (Database db) = store db
-  retrieve (Database db) = retrieve db
-  commit (Database db) = commit db
-  addOwnership (Database db) = addOwnership db
-  optimize (Database db) = optimize db
-  computeOwnership (Database db) = computeOwnership db
-  storeOwnership (Database db) = storeOwnership db
-  getOwnership (Database db) = getOwnership db
-  getUnitId (Database db) = getUnitId db
-  getUnit (Database db) = getUnit db
-  addDefineOwnership (Database db) = addDefineOwnership db
-  computeDerivedOwnership (Database db) = computeDerivedOwnership db
-  cacheOwnership (Database db) = cacheOwnership db
-  prepareFactOwnerCache (Database db) = prepareFactOwnerCache db
-
-  backup (Database db) cfg scratch process =
+  close (Database db@(DB _ repo) lmdb) = do
+    close db
+    let path = containerPath lmdb repo; squash = path <.> "squashfs"
+    haveSquash <- doesFileExist squash
+    when haveSquash $ callProcess "umount" [path]
+      `catch` \(_ :: IOException) -> return ()
+
+  predicateStats (Database db _) = predicateStats db
+  store (Database db _) = store db
+  retrieve (Database db _) = retrieve db
+  commit (Database db _) = commit db
+  addOwnership (Database db _) = addOwnership db
+  optimize (Database db _) = optimize db
+  computeOwnership (Database db _) = computeOwnership db
+  storeOwnership (Database db _) = storeOwnership db
+  getOwnership (Database db _) = getOwnership db
+  getUnitId (Database db _) = getUnitId db
+  getUnit (Database db _) = getUnit db
+  addDefineOwnership (Database db _) = addDefineOwnership db
+  computeDerivedOwnership (Database db _) = computeDerivedOwnership db
+  cacheOwnership (Database db _) = cacheOwnership db
+  prepareFactOwnerCache (Database db _) = prepareFactOwnerCache db
+
+  backup (Database db _) cfg scratch process =
     backup db cfg scratch $ \path _ -> do
       withTempDirectory scratch "out" $ \tmpdir -> do
         let out = tmpdir </> "db.squashfs"

glean/db/Glean/Database/Storage/Memory.hs

@@ -79,7 +79,7 @@ instance Storage Memory where
   withScratchRoot _ f = withSystemTempDirectory "glean" f
 
   -- TODO
-  restore _ repo _ _ = dbError repo "unimplemented 'restore'"
+  restore _ _ repo _ _ = dbError repo "unimplemented 'restore'"
 
 instance CanLookup (Database Memory) where
   lookupName Database{..} = "memory:" <> repoToText dbRepo

glean/db/Glean/Database/Storage/RocksDB.hs

@@ -162,7 +162,7 @@ instance Storage RocksDB where
 
   withScratchRoot rocks f = f $ rocksRoot rocks </> ".scratch"
 
-  restore rocks repo scratch scratch_file =
+  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,

glean/test/lib/Glean/Database/Test.hs

@@ -20,6 +20,7 @@ module Glean.Database.Test
   , setDBVersion
   , setCompactOnCompletion
   , setMaxSetSize
+  , setLMDBNoUnpack
   , enableTcDebug
   , enableQueryDebug
   , enableRocksDBCache
@@ -113,6 +114,11 @@ setMaxSetSize i cfg = cfg
   { cfgServerConfig = cfgServerConfig cfg <&> \scfg -> scfg
       { ServerConfig.config_max_set_size_bytes = Just i } }
 
+setLMDBNoUnpack :: Setting
+setLMDBNoUnpack cfg = cfg
+  { cfgServerConfig = cfgServerConfig cfg <&> \scfg -> scfg
+      { ServerConfig.config_db_lmdb_restore_unpack = False } }
+
 enableTcDebug :: Setting
 enableTcDebug cfg = cfg
   { cfgDebug = (cfgDebug cfg) { tcDebug = True } }

glean/test/tests/BackupTest.hs

@@ -312,7 +312,8 @@ backends :: ([Setting] -> Test) -> Test
 backends fn =
   TestList [
     TestLabel lbl (fn settings)
-    | (lbl, settings) <- allStorage
+    | (lbl, settings) <- allStorage <>
+        [("lmdb.squashfs", [setLMDBStorage, setLMDBNoUnpack])]
     , lbl /= "memory" -- doesn't support ownership yet
   ]