[lmdb facebookincubator#2] Add support for multiple storage backends

Author: simonmar

glean/db/Glean/Database/Backup.hs

@@ -60,6 +60,7 @@ import Glean.Database.Open
 import Glean.Database.Trace
 import Glean.Database.Types
 import Glean.Database.Schema
+import Glean.Database.Storage
 import Glean.Logger
 import Glean.RTS.Foreign.Ownership (getOwnershipStats, showOwnershipStats)
 import Glean.ServerConfig.Types (DatabaseBackupPolicy(..))
@@ -241,16 +242,17 @@ doBackup env@Env{..} repo prefix site =
           hasExcludeProperty repo (metaProperties meta) config_retention
     atomically $ notify envListener $ BackupStarted repo
     say logInfo "starting"
-    withOpenDatabaseStorage env repo $ \_storage OpenDB{..} -> do
+    withOpenDatabase env repo $ \OpenDB{..} -> do
     say logInfo "packing"
     stats <- mapMaybe
       (\(pid,stats) -> (,stats) . predicateRef <$> lookupPid pid odbSchema)
       <$> Storage.predicateStats odbHandle
     ownershipStats <- do
       maybeOwnership <- readTVarIO odbOwnership
       mapM getOwnershipStats maybeOwnership
+    withStorageFor env repo meta $ \storage -> do
 
-    Backend.Data{..} <- withScratchDirectory envStorage repo $ \scratch ->
+    Backend.Data{..} <- withScratchDirectory storage repo $ \scratch ->
       Storage.backup odbHandle scratch $ \path Data{dataSize} -> do
         say logInfo "uploading"
         let policy = ServerConfig.databaseBackupPolicy_repos config_backup
@@ -313,7 +315,8 @@ doRestore env@Env{..} repo meta
             case config_restore_timeout of
               Just seconds -> void . timeout (fromIntegral seconds * 1000000)
               Nothing -> id
-      maybeTimeout $ restore site size `catch` handler
+      withStorageFor env repo meta $ \storage ->
+        maybeTimeout $ restore site storage 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
@@ -330,10 +333,10 @@ doRestore env@Env{..} repo meta
   where
   say log s = log $ inRepo repo $ "restore: " ++ s
 
-  restore :: Site s => s -> Maybe Int64 -> IO ()
-  restore site bytes = traceMsg envTracer (GleanTraceDownload repo) $ do
+  restore :: (Storage st, Site s) => s -> st -> Maybe Int64 -> IO ()
+  restore site storage bytes = traceMsg envTracer (GleanTraceDownload repo) $ do
     atomically $ notify envListener $ RestoreStarted repo
-    mbFreeBytes <- (Just <$> Storage.getFreeCapacity envStorage)
+    mbFreeBytes <- (Just <$> Storage.getFreeCapacity storage)
                   `catch` \(_ :: IOException) -> return Nothing
     case (mbFreeBytes, bytes)  of
       (Just freeBytes, Just size) -> do
@@ -345,7 +348,7 @@ doRestore env@Env{..} repo meta
               neededBytes freeBytes
       _ -> return ()
 
-    withScratchDirectory envStorage repo $ \scratch -> do
+    withScratchDirectory storage repo $ \scratch -> do
     say logInfo "starting"
     say logInfo "downloading"
     let scratch_restore = scratch </> "restore"
@@ -356,14 +359,15 @@ doRestore env@Env{..} repo meta
     say logInfo "restoring"
     createDirectoryIfMissing True scratch_restore
     traceMsg envTracer GleanTraceStorageRestore $
-      Storage.restore envStorage repo scratch_restore scratch_file
+      Storage.restore storage repo scratch_restore scratch_file
     say logInfo "adding"
     traceMsg envTracer GleanTraceFinishRestore $
       Catalog.finishRestoring envCatalog repo
     atomically $ notify envListener $ RestoreFinished repo
     say logInfo "finished"
 
-  handler exc = do
+  handler :: Storage s => s -> SomeException -> IO ()
+  handler storage exc = do
     failed <- atomically $ do
       failed <- Catalog.exists envCatalog [Restoring] repo
       when failed $ do
@@ -372,7 +376,7 @@ doRestore env@Env{..} repo meta
       return failed
     when failed $ do
       say logError $ "failed: " ++ show exc
-      swallow $ Storage.safeRemoveForcibly envStorage repo
+      swallow $ Storage.safeRemoveForcibly storage repo
     rethrowAsync exc
 
 

glean/db/Glean/Database/Close.hs

@@ -39,7 +39,7 @@ closeDatabases env@Env{..} = do
   dbs <- readTVarIO envActive
   forM_ (HashMap.keys dbs) $ closeDatabase env
 
-isIdle :: (TimePoint -> Bool) -> DB s -> OpenDB s -> STM Bool
+isIdle :: (TimePoint -> Bool) -> DB -> OpenDB -> STM Bool
 isIdle long_enough db odb = and <$> sequence
   [ (== 1) <$> readTVar (dbUsers db)  -- we are the only user
   , long_enough <$> readTVar (odbIdleSince odb)
@@ -52,7 +52,7 @@ isIdle long_enough db odb = and <$> sequence
   ]
 
 closeIf
-  :: (forall s . DB s -> DBState s -> STM (Maybe (OpenDB s)))
+  :: (DB -> DBState -> STM (Maybe OpenDB))
   -> Env
   -> Repo
   -> IO ()
@@ -132,7 +132,7 @@ exportOpenDBStats Env{..} = do
   forM_ (HashMap.toList repoOpenCounts) $ \(repoNm,count) -> do
     setCounter ("glean.db." <> Text.encodeUtf8 repoNm <> ".open") count
 
-closeOpenDB :: Storage.Storage s => Env -> OpenDB s -> IO ()
+closeOpenDB :: Env -> OpenDB -> IO ()
 closeOpenDB env OpenDB{..} = do
   case odbWriting of
     Just Writing{..} -> do

glean/db/Glean/Database/Config.hs

@@ -10,9 +10,11 @@
 module Glean.Database.Config (
   -- * DataStore
   DataStore(..),
+  StorageName,
   fileDataStore,
   tmpDataStore,
   memoryDataStore,
+  rocksdbName,
 
   -- * Config, and options parser
   options,
@@ -84,6 +86,7 @@ import qualified Glean.Database.Storage.Memory as Memory
 import qualified Glean.Database.Storage.RocksDB as RocksDB
 import Glean.Database.Trace
 import qualified Glean.Internal.Types as Internal
+import Glean.Internal.Types (StorageName(..))
 import Glean.DefaultConfigs
 import Glean.Logger.Database
 import Glean.Logger.Server
@@ -104,36 +107,51 @@ import Paths_glean
 #endif
 
 data DataStore = DataStore
-  { withDataStore
-      :: forall a. ServerConfig.Config
-      -> (forall c s. (Catalog.Store c, Storage s) => c -> s -> IO a)
-      -> IO a
+  { withStorage ::
+      -- setup the storage backend. Must be scoped, because it might
+      -- involve creating temporary resources.
+      forall a.
+      ServerConfig.Config ->
+      ((HashMap StorageName (Some Storage), Some Catalog.Store) -> IO a) ->
+      IO a
+  , defaultStorage :: StorageName
   , dataStoreTag :: String
   }
 
+rocksdbName, memoryName :: StorageName
+rocksdbName = StorageName "rocksdb"
+memoryName = StorageName "memory"
+
 fileDataStore :: FilePath -> DataStore
 fileDataStore path = DataStore
-  { withDataStore = \scfg f -> do
+  { withStorage = \scfg f -> do
       rocksdb <- RocksDB.newStorage path scfg
-      f (Catalog.fileCatalog path) rocksdb
-  , dataStoreTag = "rocksdb:" <> path
+      f (
+        HashMap.fromList
+          [ (rocksdbName, Some rocksdb)
+          ],
+        Some (Catalog.fileCatalog path)
+        )
+  , defaultStorage = rocksdbName
+  , dataStoreTag = "db:" <> path
   }
 
 tmpDataStore :: DataStore
 tmpDataStore = DataStore
-  { withDataStore = \scfg f -> withSystemTempDirectory "glean" $ \tmp -> do
+  { withStorage = \scfg f -> withSystemTempDirectory "glean" $ \tmp -> do
       logInfo $ "Storing temporary DBs in " <> tmp
-      rocksdb <- RocksDB.newStorage tmp scfg
-      f (Catalog.fileCatalog tmp) rocksdb
-  , dataStoreTag = "rocksdb:{TMP}"
+      withStorage (fileDataStore tmp) scfg f
+  , defaultStorage = rocksdbName
+  , dataStoreTag = dataStoreTag (fileDataStore "<tmp>")
   }
 
 memoryDataStore :: DataStore
 memoryDataStore = DataStore
-  { withDataStore = \_ f -> do
+  { withStorage = \_ f -> do
       cat <- Catalog.memoryCatalog
       mem <- Memory.newStorage
-      f cat mem
+      f (HashMap.fromList [(memoryName, Some mem)], Some cat)
+  , defaultStorage = memoryName
   , dataStoreTag = "memory"
   }
 
@@ -488,10 +506,12 @@ schemaLocationOption = option (eitherReader schemaLocationParser)
 options :: Parser Config
 options = do
   let
-    dbRoot = fileDataStore <$> strOption (
-      long "db-root" <>
-      metavar "DIR" <>
-      help "Directory containing databases")
+    dbRoot = do
+      path <- strOption (
+        long "db-root" <>
+        metavar "DIR" <>
+        help "Directory containing databases")
+      pure $ fileDataStore path
     dbTmp = tmpDataStore <$ flag' () (
       long "db-tmp" <>
       help "Store databases in a temporary directory (default)")

glean/db/Glean/Database/Create.hs

@@ -115,29 +115,28 @@ kickOffDatabase env@Env{..} kickOff@Thrift.KickOff{..}
       handle
         (\Catalog.EntryAlreadyExists{} ->
             return $ Thrift.KickOffResponse True) $
-        mask $ \unmask ->
+        mask $ \unmask -> do
         -- FIXME: There is a tiny race here where we might fail in a weird way
         -- if kick off a DB that is being deleted after it got removed from
         -- the Catalog but before it got removed from the storage. The entire
         -- concept of deleting DBs will change with the new metadata handling so
         -- it's not worth fixing at this point, especially since we aren't
         -- supposed to be kicking off DBs we've previously deleted.
+        let meta = newMeta envDefaultStorage version time
+                     (Incomplete def) allProps
+                     (lightDeps kickOff_dependencies')
         bracket_
-          (Catalog.create
-            envCatalog
-            kickOff_repo
-            (newMeta version time (Incomplete def) allProps
-              (lightDeps kickOff_dependencies')) $ do
-                modifyTVar' envActive $ HashMap.insert kickOff_repo db
-                writeTVar (dbState db) Opening
-                acquireDB db)
+          (Catalog.create envCatalog kickOff_repo meta $ do
+             modifyTVar' envActive $ HashMap.insert kickOff_repo db
+             writeTVar (dbState db) Opening
+             acquireDB db)
           (atomically $ releaseDB envCatalog envActive db) $
-          do
+            withStorageFor env kickOff_repo meta $ \storage -> do
             -- Open the new db in Create mode which will create the
             -- physical storage. This might fail - in that case, we
             -- mark the db as failed. NB. pass the full dependencies
             -- here, not lightDeps.
-            opener <- asyncOpenDB env envStorage db version mode
+            opener <- asyncOpenDB env storage db version mode
                 kickOff_dependencies'
               (do
                 logInfo $ inRepo kickOff_repo "created")

glean/db/Glean/Database/Data.hs

@@ -22,7 +22,7 @@ import Data.ByteString.Lazy (toStrict, fromStrict)
 import Thrift.Protocol.Compact
 
 import Glean.Database.Exception
-import Glean.Database.Storage (Storage, Database)
+import Glean.Database.Storage (DatabaseOps)
 import qualified Glean.Database.Storage as Storage
 import Glean.RTS.Foreign.Ownership
 import Glean.Types (Repo)
@@ -39,21 +39,21 @@ uNITS_KEY = "units"
 sLICES_KEY :: ByteString
 sLICES_KEY = "slices"
 
-storeSchema :: Storage s => Database s -> StoredSchema -> IO ()
+storeSchema :: DatabaseOps db => db -> StoredSchema -> IO ()
 storeSchema db = Storage.store db sCHEMA_KEY . serializeCompact
 
-retrieveSchema :: Storage s => Repo -> Database s -> IO (Maybe StoredSchema)
+retrieveSchema :: DatabaseOps db => Repo -> db -> IO (Maybe StoredSchema)
 retrieveSchema repo db = do
   value <- Storage.retrieve db sCHEMA_KEY
   case deserializeCompact <$> value of
     Just (Right info) -> return $ Just info
     Just (Left msg) -> dbError repo $ "invalid schema: " ++ msg
     Nothing -> return Nothing
 
-storeUnits :: Storage s => Database s -> [ByteString] -> IO ()
+storeUnits :: DatabaseOps db => db -> [ByteString] -> IO ()
 storeUnits db = Storage.store db uNITS_KEY . toStrict . encode
 
-retrieveUnits :: Storage s => Repo -> Database s -> IO (Maybe [ByteString])
+retrieveUnits :: DatabaseOps db => Repo -> db -> IO (Maybe [ByteString])
 retrieveUnits repo db = do
   value <- Storage.retrieve db uNITS_KEY
   case decodeOrFail . fromStrict <$> value of
@@ -64,12 +64,12 @@ retrieveUnits repo db = do
 -- Slices are each serialized using the RTS Slice::serialize(), and then
 -- the list of serialized slices :: [ByteString] is serialized with
 -- the Haskell Binary encoder.
-storeSlices :: Storage s => Database s -> [Slice] -> IO ()
+storeSlices :: DatabaseOps db => db -> [Slice] -> IO ()
 storeSlices db slices = do
   bytestrings <- mapM serializeSlice slices
   Storage.store db sLICES_KEY $ toStrict $ encode bytestrings
 
-retrieveSlices :: Storage s => Repo -> Database s -> IO (Maybe [Slice])
+retrieveSlices :: DatabaseOps db => Repo -> db -> IO (Maybe [Slice])
 retrieveSlices repo db = do
   value <- Storage.retrieve db sLICES_KEY
   case decodeOrFail . fromStrict <$> value of

glean/db/Glean/Database/Delete.hs

@@ -63,10 +63,9 @@ expireDatabase delay env@Env{..} repo = do
 
 -- | Database deletion thread
 removeDatabase
-  :: Storage.Storage s
-  => Env
+  :: Env
   -> Repo
-  -> TMVar (Maybe (DB s))
+  -> TMVar (Maybe DB)
   -> IO ()
 removeDatabase env@Env{..} repo todo = uninterruptibleMask_ $
   -- This runs under uninterruptibleMask_ because there is really nothing
@@ -94,12 +93,14 @@ removeDatabase env@Env{..} repo todo = uninterruptibleMask_ $
           Open odb -> closeOpenDB env odb
             `finally` atomically (writeTVar dbState Closed)
           _ -> return ()
-        Storage.delete envStorage repo
-        Catalog.delete envCatalog repo
-        Storage.safeRemoveForcibly envStorage repo
+        meta <- atomically $ Catalog.readMeta envCatalog repo
+        withStorageFor env repo meta $ \storage -> do
+          Storage.delete storage repo
+          Catalog.delete envCatalog repo
+          Storage.safeRemoveForcibly storage repo
         atomically $ modifyTVar envDerivations $
           HashMap.filterWithKey (\(repo',_) _ -> repo' /= repo)
-      logInfo $ inRepo repo "deleted"
+        logInfo $ inRepo repo "deleted"
 
 -- | Schedule a DB for deletion and return the 'Async' which can be used to
 -- obtain the result.