content: {-
    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.
  -}
  
  {-# LANGUAGE DeriveGeneric #-}
  module Glean.Database.Backup
    ( backuper
    , Event(..)
    , backupDatabase
    -- for testing
    , newestByRepo
    , bestRestore
    ) where
  
  import Control.Applicative
  import Control.Concurrent.Async
  import Control.Exception
  import Control.Monad
  import Control.Trace (traceMsg)
  import qualified Data.HashMap.Strict as HashMap
  import Data.HashSet (HashSet)
  import qualified Data.HashSet as HashSet
  import Data.Hashable (Hashable)
  import Data.Int
  import Data.List
  import qualified Data.Map as Map
  import Data.Maybe
  import Data.Ord
  import Data.Text (Text)
  import qualified Data.Text as Text
  import Data.Typeable (Typeable)
  import GHC.Generics hiding (Meta)
  import System.Directory
  import System.FilePath
  import System.Timeout
  import Text.Printf
  import TextShow
  
  import Util.Control.Exception
  import Util.Graph
  import Util.IO (safeRemovePathForcibly)
  import Util.Log
  import Util.Logger
  import Util.STM
  
  import Glean.Database.Backup.Backend as Backend
  import Glean.Database.Backup.Locator
  import qualified Glean.Database.Catalog as Catalog
  import Glean.Database.Catalog.Filter
  import Glean.Database.CompletePredicates
  import qualified Glean.Database.Logger as Logger
  import Glean.Database.Meta
  import Glean.Database.Repo
  import qualified Glean.Database.Storage as Storage
  import Glean.Database.Open
  import Glean.Database.Trace
  import Glean.Database.Types
  import Glean.Database.Schema
  import Glean.Logger
  import Glean.RTS.Foreign.Ownership (getOwnershipStats, showOwnershipStats)
  import Glean.ServerConfig.Types (DatabaseBackupPolicy(..))
  import qualified Glean.ServerConfig.Types as ServerConfig
  import Glean.Internal.Types as Thrift
  import Glean.Types as Thrift hiding (Exception)
  import Glean.Util.Some
  import Glean.Util.Observed as Observed
  import Glean.Util.Trace
  
  data Event
    = BackupStarted Repo
    | BackupFinished Repo
    | BackupFailed Repo
    | RestoreStarted Repo
    | RestoreFinished Repo
    | RestoreFailed Repo
    | RestoreAborted Repo
    | FinalizeStarted Repo
    | FinalizeFinished Repo
    | FinalizeFailed Repo
    | Waiting
    deriving (Eq, Generic, Show, Typeable)
  
  instance Hashable Event
  
  type Todo = (Repo, IO Bool)
  
  -- | Set of Repos that shouldn't be backed up (because we already tried and
  -- failed). We currently only remove things from here when they get deleted from
  -- our DB list which matches previous behaviour. Eventually, we should retry
  -- with exponential back-off.
  type SinBin = HashSet Repo
  
  -- | In order to download a DB we need twice the space to untar it, plus a bit
  --   more for RocksDB to expand it
  dbSizeDownloadFactor :: Double
  dbSizeDownloadFactor = 2.1
  
  newtype BackupException = BackupException Text
    deriving (Typeable)
  
  instance Show BackupException where
    show (BackupException t) = Text.unpack t
  instance Exception BackupException where
  
  -- | Get the next thing to do. Prioritises restores over backups and prefers
  -- newer databases.
  getTodo :: Env -> SinBin -> STM Todo
  getTodo env@Env{..} sinbin = getFinalize <|> getRestore <|> getBackup
    where
      getFinalize = do
        dbs <- Catalog.list envCatalog [Local] $ do
          statusV .==. DatabaseStatus_Finalizing
          latest
        case dbs of
          [] -> retry
          Item{..} : _ -> return (itemRepo, doFinalize env itemRepo)
  
      getRestore = do
        restoring <- Catalog.list envCatalog [Restoring] $ do
          repoV `notInF` sinbin
          newestByRepo
        when (null restoring) retry
        avail <-  Catalog.list envCatalog [Local] $ do
          repoV `notInF` sinbin
          statusV .==. DatabaseStatus_Complete
          newestByRepo
        case bestRestore restoring avail of
          [] -> retry
          Item{..} : _ -> return (itemRepo, doRestore env itemRepo itemMeta)
  
      getBackup = do
        policy <- ServerConfig.config_backup <$> Observed.get envServerConfig
        let allowedRepoNames = databaseBackupPolicy_allowed policy
        dbs <- Catalog.list envCatalog [Local] $ do
          repoNameV `inF` allowedRepoNames
          statusV .==. DatabaseStatus_Complete
          backedUpV .==. False
          latest
        case dbs of
          [] -> retry
          Item{..} : _ -> do
            dsite <- getSite env (repo_name itemRepo)
            case dsite of
              Just (prefix, Some site) ->
                return (itemRepo, doBackup env itemRepo prefix site)
              Nothing ->
                return (itemRepo, throwIO $ BackupException $
                  "Could not backup " <> repo_name itemRepo <>
                  ". Could not resolve the site")
  
      latest = do
        repoV `notInF` sinbin
        sortF createdV Descending
        limitF 1
  
  newestByRepo :: Filter ()
  newestByRepo = groupF repoNameV $ do
    sortF createdV Descending
    limitF 1
  -- To find the next DB to restore, we
  --   - find the newest DB of each repo that needs restoring
  --   - restore the DB which is currently the most stale; that is
  --     where the newest, complete, local DB of the same repo is the oldest
  bestRestore :: [Item] -> [Item] -> [Item]
  bestRestore restoring avail = depsFirst (sortBy order restoring)
    where
      -- stalest first, and then most recent if none are stale
      order a b =
          comparing staleness a b <>
          comparing (metaCreated . itemMeta) b a -- NB. descending
  
      availMap = HashMap.fromList
        [ (repo_name (itemRepo item), item) | item <- avail ]
  
      -- creation time of the local DB if stale, or maxBound
      staleness backup =
        case HashMap.lookup (repo_name (itemRepo backup)) availMap of
          Nothing -> NoLocalDb
          Just local
            | tlocal > tbackup -> LocalDbNewer
            | otherwise -> LocalDbOlder tlocal
            where
              tlocal = timestamp local
              tbackup = timestamp backup
              repoHashTime = metaRepoHashTime . itemMeta
              creationTime = metaCreated . itemMeta
              timestamp item = fromMaybe (creationTime item) (repoHashTime item)
              -- ^ prefer the exact repoHash time if we have it, otherwise use the
              -- creationTime as the best guess we have
  
      -- ensure we're restoring dependencies of a stacked DB first, if
      -- they need restoring.
      depsFirst items = postorder items itemRepo itemDeps
        where
        itemDeps item = case metaDependencies (itemMeta item) of
          Nothing -> []
          Just (Dependencies_stacked Thrift.Stacked{..}) ->
            [Thrift.Repo stacked_name stacked_hash]
          Just (Dependencies_pruned Pruned{..}) -> [pruned_base]
  
  data Staleness
    = NoLocalDb
    | LocalDbOlder PosixEpochTime
    | LocalDbNewer
    deriving (Eq,Ord)
  
  hasExcludeProperty
    :: Repo
    -> DatabaseProperties
    -> ServerConfig.DatabaseRetentionPolicy
    -> Bool
  hasExcludeProperty repo props policy = do
    let retentionPolicy = ServerConfig.databaseRetentionPolicy_repos policy
    case Map.lookup (repo_name repo) retentionPolicy of
      Just ret -> do
        let exclude_props = ServerConfig.retention_excluded_properties ret
        any (hasProperty props) (HashMap.toList exclude_props)
      Nothing -> False
    where
      hasProperty props (name,val) = HashMap.lookup name props == Just val
  
  
  doBackup :: Site site => Env -> Repo -> Text -> site -> IO Bool
  doBackup env@Env{..} repo prefix site =
    backup `catch` \exc -> do
      atomically $ notify envListener $ BackupFailed repo
      say logError $ "failed: " ++ show (exc :: SomeException)
      rethrowAsync exc
      return False
    where
    say log s = log $ inRepo repo $ "backup: " ++ s
  
    backup = loggingAction (runLogRepo "backup" env repo) (const mempty) $ do
      ServerConfig.Config{..} <- Observed.get envServerConfig
      meta <- atomically $ Catalog.readMeta envCatalog repo
      let excluded =
            hasExcludeProperty repo (metaProperties meta) config_retention
      atomically $ notify envListener $ BackupStarted repo
      say logInfo "starting"
      withOpenDatabaseStorage env repo $ \_storage 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
  
      Backend.Data{..} <- withScratchDirectory envStorage repo $ \scratch ->
        Storage.backup odbHandle scratch $ \path Data{dataSize} -> do
          say logInfo "uploading"
          let policy = ServerConfig.databaseBackupPolicy_repos config_backup
              ttl = case ServerConfig.backup_delete_after_seconds <$>
                Map.lookup (repo_name repo) policy of
                  Just 0 -> Nothing
                  ttl -> fromIntegral <$> ttl
              metaWithBytes = meta {
                metaCompleteness = case metaCompleteness meta of
                  Complete DatabaseComplete{..} ->
                    Complete DatabaseComplete
                      {databaseComplete_bytes = Just (fromIntegral dataSize)
                      ,..}
                  _ -> metaCompleteness meta
          }
          Backend.backup site repo metaWithBytes ttl path
      let locator = toRepoLocator prefix site repo
      Logger.logDBStatistics
        env
        repo
        stats
        ownershipStats
        dataSize
        locator
        excluded
      say logInfo "finished"
      atomically $ do
        void $ Catalog.modifyMeta envCatalog repo $ \meta -> return meta
          { metaBackup = Just locator
          , metaCompleteness = case metaCompleteness meta of
            Complete DatabaseComplete{..} ->
              Complete DatabaseComplete
                {databaseComplete_bytes = Just (fromIntegral dataSize)
                ,..}
            other -> other
          }
        notify envListener $ BackupFinished repo
      return True
  
  backupDatabase :: Env -> Repo -> Text -> IO Bool
  backupDatabase env repo loc
    | Just (prefix, site) <- fromSiteLocator (envBackupBackends env) loc =
        doBackup env repo prefix site
    | otherwise = throwIO $
        Thrift.InvalidLocator $ "invalid locator '" <> loc <>  "'"
  
  doRestore :: Env -> Repo -> Meta -> IO Bool
  doRestore env@Env{..} repo meta
    | Just loc <- metaBackup meta
    , Complete DatabaseComplete{databaseComplete_bytes = size}
    <- metaCompleteness 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
        let maybeTimeout =
              case config_restore_timeout of
                Just seconds -> void . timeout (fromIntegral seconds * 1000000)
                Nothing -> id
        maybeTimeout $ restore site size `catch` handler
  
        -- 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
        return True
  
    | otherwise = do
        atomically $ notify envListener $ RestoreStarted repo
        say logError $ case (metaCompleteness meta, metaBackup meta) of
          (Complete{}, Just loc) -> "invalid location " ++ show loc
          (incomplete, Just{}) -> "database incomplete: " <> show incomplete
          _ -> "missing location"
        atomically $ notify envListener $ RestoreFailed repo
        return False
    where
    say log s = log $ inRepo repo $ "restore: " ++ s
  
    restore :: Site s => s -> Maybe Int64 -> IO ()
    restore site bytes = traceMsg envTracer (GleanTraceDownload repo) $ do
      atomically $ notify envListener $ RestoreStarted repo
      mbFreeBytes <- (Just <$> Storage.getFreeCapacity envStorage)
                    `catch` \(_ :: IOException) -> return Nothing
      case (mbFreeBytes, bytes)  of
        (Just freeBytes, Just size) -> do
          let neededBytes = ceiling $ dbSizeDownloadFactor * fromIntegral size
          when (freeBytes < neededBytes) $
            -- the catch-all exception handler will log and cancel the download
            throwIO $ ErrorCall $
              printf "Not enough disk space: %d needed, %d available"
                neededBytes freeBytes
        _ -> return ()
  
      withScratchDirectory envStorage repo $ \scratch -> do
      say logInfo "starting"
      say logInfo "downloading"
      let scratch_restore = scratch </> "restore"
          scratch_file = scratch </> "file"
      -- TODO: implement buffered downloads in Manifold client
      void $ traceMsg envTracer GleanTraceSiteRestore $
        Backend.restore site repo scratch_file
      say logInfo "restoring"
      createDirectoryIfMissing True scratch_restore
      traceMsg envTracer GleanTraceStorageRestore $
        Storage.restore envStorage 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
      failed <- atomically $ do
        failed <- Catalog.exists envCatalog [Restoring] repo
        when failed $ do
          Catalog.abortRestoring envCatalog repo
          notify envListener $ RestoreFailed repo
        return failed
      when failed $ do
        say logError $ "failed: " ++ show exc
        swallow $ Storage.safeRemoveForcibly envStorage repo
      rethrowAsync exc
  
  
  doFinalize :: Env -> Repo -> IO Bool
  doFinalize env@Env{..} repo =
    loggingAction (runLogRepo "finalize" env repo) (const mempty) $ do
      atomically $ notify envListener $ FinalizeStarted repo
  
      -- If the client didn't explicitly call completePredicates, we'll
      -- do that now. Do this *before* optimising/compacting, because
      -- that deletes all the temporary ownership information consumed
      -- by this pass.
      meta <- atomically $ Catalog.readMeta envCatalog repo
      when (not (metaAxiomComplete meta)) $ syncCompletePredicates env repo
  
      config <- Observed.get envServerConfig
      withOpenDatabase env repo $ \OpenDB{..} -> do
        Storage.prepareFactOwnerCache odbHandle
        maybeOwnership <- readTVarIO odbOwnership
        forM_ maybeOwnership $ \ownership -> do
          stats <- getOwnershipStats ownership
          logInfo $ "final ownership: " <> Text.unpack (showOwnershipStats stats)
        let compact = ServerConfig.config_compact_on_completion config
        say logInfo $ if compact then "optimizing(compacting)" else "optimizing"
        Storage.optimize odbHandle compact
        Storage.cacheOwnership odbHandle
  
      -- update and re-merge our internal representation of the schema
      schemaUpdated env (Just repo)
  
      time <- envGetCurrentTime
      atomically $ do
        void $ Catalog.modifyMeta envCatalog repo $ \meta -> return meta
          { metaCompleteness = Complete $ DatabaseComplete
            (utcTimeToPosixEpochTime time)
            Nothing -- the size gets updated after a backup
          }
        notify envListener $ FinalizeFinished repo
      say logInfo "finished"
      return True
  
    `catch` \exc -> do
      let
        interrupted
          | Just AsyncCancelled{} <- fromException exc = True
          | otherwise = False
  
      -- If we were interrupted, then keep the DB in the finalizing state
      -- and we'll try to finalize it again later.  If there was some other
      -- exception, we'll mark the DB broken.
      atomically $ do
        notify envListener $ FinalizeFailed repo
        when (not interrupted) $
          void $ Catalog.modifyMeta envCatalog repo $ \meta -> return meta
            { metaCompleteness = Broken $ DatabaseBroken "finalize" (showt exc) }
      say logError $
        (if interrupted then "interrupted" else "failed") <>
        ": " <> show (exc :: SomeException)
      rethrowAsync exc
      return False
    where
      say log s = log $ inRepo repo $ "finalize: " ++ s
  
  -- | Back up databases forever.
  backuper :: Env -> IO ()
  backuper env@Env{..} = loop mempty `catchAll` \exc -> do
    logError $ "unexpected error in backup thread: " ++ show exc
    backuper env
    where
      loop sinbin = do
        (repo, action) <- do
          r <- atomically $ Just <$> getTodo env sinbin <|> return Nothing
          case r of
            Just todo -> return todo
            Nothing -> do
              atomically $ notify envListener Waiting
              atomically $ getTodo env sinbin
        ok <- action `catchAll` \exc -> do
          -- NOTE: action shouldn't throw so this shouldn't ever run
          logError $ "unexpected error during backup: " ++ show exc
          return False
        sinbin' <- HashSet.fromList . map itemRepo <$> atomically
          (Catalog.list envCatalog [Local, Restoring] $
            repoV `inF` if ok then sinbin else HashSet.insert repo sinbin)
        loop sinbin'
  
  withScratchDirectory
    :: Storage.Storage s => s -> Repo -> (FilePath -> IO a) -> IO a
  withScratchDirectory storage repo action =
    Storage.withScratchRoot storage $ \root -> do
      let scratch = root </> databaseSubdir repo
      safeRemovePathForcibly scratch
      bracket_
        (createDirectoryIfMissing True scratch)
        (logExceptions id $ safeRemovePathForcibly scratch)
        (action scratch)
  
  rethrowAsync :: SomeException -> IO ()
  rethrowAsync exc
    | isAsyncException exc = throwIO exc
    | otherwise = return ()
  