{-
  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.Janitor
  ( runDatabaseJanitor
  , runDatabaseJanitorPureish
  , JanitorSideEffect(..)
  -- for testing
  , mergeLocalAndRemote
  , dbIndex
  , DbIndex(..)
  ) where

import Control.Concurrent.Stream (stream)
import Control.Concurrent (getNumCapabilities)
import Control.Exception.Safe
import Control.Monad.Extra
import Control.Monad.IO.Class (liftIO)
import Control.Monad.Trans (lift)
import Control.Monad.Trans.State.Strict (runStateT)
import qualified Control.Monad.Trans.State.Strict as State
import Control.Monad.Trans.Writer.CPS
import Data.ByteString (ByteString)
import Data.Foldable as Foldable
import qualified Data.HashMap.Strict as HashMap
import qualified Data.HashSet as HashSet
import Data.List (sortOn)
import qualified Data.List.NonEmpty as NonEmpty
import qualified Data.Map as Map
import Data.Maybe
import Data.Ord
import qualified Data.Set as Set
import qualified Data.Text.Encoding as Text
import Data.Time

import ServiceData.GlobalStats as Stats
import Util.Control.Exception
import Util.List
import Util.Log
import Util.Logger
import Util.STM
import Util.Time
import Util.TimeSec

import qualified Glean.Database.Catalog as Catalog
import Glean.Database.Catalog.Filter
import Glean.Database.Close
import Glean.Database.Delete
import Glean.Database.List
import Glean.Database.Meta
import Glean.Database.Repo
import Glean.Database.Restore
import Glean.Database.Open
import Glean.Database.Retention
import Glean.Database.Types
import Glean.Internal.Types
import Glean.Logger
import Glean.Repo.Text
import qualified Glean.ServerConfig.Types as ServerConfig
import Glean.Types hiding (Database)
import qualified Glean.Types as Thrift
import Glean.Util.Observed as Observed
import Glean.Util.ShardManager
    ( ShardManager(getAssignedShards, computeShardMapping),
      SomeShardManager(SomeShardManager), countersForShardSizes, noSharding )
import Glean.Util.TransitiveClosure

{- |
The database janitor has the following functions:
  - ensure the newest db for each repo is open to speed-up queries to it.
  - close databases that haven't been used in a while.
  - kick-off downloads of dbs available in remote backups.
  - delete dbs we don't need any longer.
  - publish counters of local db states.
-}
runDatabaseJanitor :: Env -> IO ()
runDatabaseJanitor env = do
  runDatabaseJanitorPureish env >>= executeJanitorSideEffects env

executeJanitorSideEffects :: Env -> [JanitorSideEffect] -> IO ()
executeJanitorSideEffects env sideEffects =
  do
  forM_ sideEffects $ \case
    PublishCounter n v -> void $ Stats.setCounter n v
    _ -> pure ()
  numCaps <- getNumCapabilities
  stream numCaps (forM_ sideEffects) $ \case
    PreOpenDB repo ->
      void $ tryAll $ logExceptions (inRepo repo) $
        withOpenDatabaseStack env repo (\_ -> return ())
    _ -> pure ()

  -- Record the published counters and clear stale ones
  let countersPublished =
        HashSet.fromList [n | PublishCounter n _ <- sideEffects]
  lastPublishedCounters <- atomically $ do
    last <- readTVar (envDatabaseJanitorPublishedCounters env)
    writeTVar (envDatabaseJanitorPublishedCounters env) countersPublished
    return last
  forM_ (HashSet.difference lastPublishedCounters countersPublished)
    clearCounter


data JanitorSideEffect
  = PublishCounter !ByteString !Int
  | PreOpenDB !Repo
  deriving (Eq, Show)

publishCounter :: ByteString -> Int -> WriterT [JanitorSideEffect] IO ()
publishCounter name value = tell [PublishCounter name value]

preOpenDB :: Monad m => Repo -> WriterT [JanitorSideEffect] m ()
preOpenDB repo = tell [PreOpenDB repo]

-- WIP making the Janitor more testable by returning the list of side effects
runDatabaseJanitorPureish :: Env -> IO [JanitorSideEffect]
runDatabaseJanitorPureish env@Env{envShardManager = SomeShardManager sm} = do
  maybeShards <- getAssignedShards sm
  case maybeShards of
    Just shards ->
        runWithShards env (Set.fromList shards) sm
    Nothing -> do
      Just myShards <- getAssignedShards noSharding
      runWithShards env (Set.fromList myShards) noSharding

runWithShards
  :: (Show shard, Ord shard)
  => Env
  -> Set.Set shard
  -> ShardManager shard
  -> IO [JanitorSideEffect]
runWithShards env myShards sm = do
  loggingAction (runLogCmd "janitor" env) (const mempty) $ do
  logInfo $ "running database janitor with shard assignment "
    <> show (toList myShards)

  ServerConfig.Config{..} <- Observed.get (envServerConfig env)

  let
    !ServerConfig.DatabaseRetentionPolicy{} = config_retention
    !ServerConfig.DatabaseClosePolicy{..} = config_close

  fetchBackupsResult <- fetchBackups env

  now <- envGetCurrentTime env

  backups <- case fetchBackupsResult of
    ReusedPreviousBackups{reusedBackups = (dbs,_utc)} -> return dbs
    FetchedNewBackups{newBackups} -> return newBackups
    FetchFailure{..} -> do
      logError $ "couldn't list restorable databases: " <> show fetchError

      case previousGoodBackups of
        Nothing -> throwIO $ JanitorFetchBackupsFailure fetchError
        Just (dbs, utc) -> do
          let previousAgeInSeconds = fromIntegral (timeDiffInSeconds now utc)
              tooOld = previousAgeInSeconds > config_max_remote_db_list_age
          when tooOld $ throwIO $ JanitorFetchBackupsFailure fetchError

          return dbs

  localAndRestoring <- atomically $
    Catalog.list (envCatalog env) [Local,Restoring] everythingF

  dbToShard <- computeShardMapping sm

  cachedAvailableDBs <- readTVarIO (envCachedAvailableDBs env)
  let
    isAvailable db
      | HashSet.member db cachedAvailableDBs = return True
      | otherwise = not . null <$> envFilterAvailableDBs env [db]

    itemAvailable Item{itemRepo} = do
      st <- State.get
      case HashMap.lookup itemRepo st of
        Just isAvailable -> return isAvailable
        Nothing -> do
          isAvailable <- lift $ isAvailable itemRepo
          State.put $! HashMap.insert itemRepo isAvailable st
          return isAvailable

    allDBsByAge :: [Item]
    allDBsByAge = mergeLocalAndRemote backups localAndRestoring

    index@DbIndex{..} = dbIndex allDBsByAge

  (RetentionChanges{..}, cachedAvailable) <- (`runStateT` mempty) $
    retentionChanges
      config_retention
      config_restore
      now
      index
      itemAvailable
      dbToShard
      myShards

  atomically $ writeTVar
    (envCachedAvailableDBs env)
    (HashMap.keysSet $ HashMap.filter id cachedAvailable)

  let
    checkDependencies dbs msg1 msg2 =
      unless (null dbs) $
        logWarning $ msg1 <> ": " <> show dbs <> "\n" <> msg2

  checkDependencies allMissingDependencies
    "dependencies missing in Catalog"
    ("This suggests a dependency db has been deleted from the cloud, " <>
     "or a failure to enumerate the cloud catalog")

  checkDependencies localMissingDependencies
    "dependencies not downloaded"
    ("This probably means a bug in the db->shard mapping, " <>
     "or a failure to enumerate the cloud catalog")

  forM_ retentionLocal $ \(Item{itemRepo, itemLocality}, _) ->
    when (itemLocality == Local) $
      atomically $ Catalog.unsetExpiring (envCatalog env) itemRepo

  forM_ retentionDelete $ \Item{itemRepo = repo} -> do
    let
      ServerConfig.Retention{..} =
        NonEmpty.head $ repoRetention config_retention $ Thrift.repo_name repo
    expireDatabase (fromIntegral <$> retention_expire_delay) env repo
      `catch` \UnknownDatabase{} ->
        -- DBs are deleted asynchronously and this code is not transactional,
        -- so it's possible that the DB is not in the Catalog at this point
        return ()

  restores <- fmap catMaybes $ forM retentionRestore $ \Item{..} ->
    ifRestoreRepo env Nothing itemRepo $ do
      logInfo $ "Restoring: " ++ showRepo itemRepo ++
        " ("  ++ showNominalDiffTime (dbAge now itemMeta) ++ " old)"
      return $ Just $ Catalog.startRestoring (envCatalog env) itemRepo itemMeta
  -- register all the restoring DBs together in a single transaction,
  -- so that the backup thread can't jump in early and pick one
  atomically $ sequence_ restores

  atomically $ Catalog.resetElsewhere (envCatalog env) retentionElsewhere

  deleting <- readTVarIO (envDeleting env)

  -- Open the most recent local, complete DB for each repo in order
  -- to avoid lag spikes. This will be the DB that clients will get
  -- by default unless they specify a particular DB instance.
  let
      toOpen = uniqBy (comparing itemRepo) $
        [ item
        | (_, dbsByAge) <- byRepoName
        , item : _ <- [filter shouldOpen (NonEmpty.toList dbsByAge)]
        ] <>
        [ item
        | (repoNm, dbs) <- byRepoName
        , ServerConfig.Retention{retention_keep_open = True, ..} <-
            NonEmpty.toList (repoRetention config_retention repoNm)
        , item <- NonEmpty.toList dbs
        , shouldOpen item
        , hasAllProperties retention_required_properties item
        ]

      isComplete Complete{} = True
      isComplete _ = False

      shouldOpen item =
        isComplete (metaCompleteness (itemMeta item)) &&
        itemLocality item == Local &&
        not (HashMap.member (itemRepo item) deleting)

      closeDeps = transitiveClosureBy itemRepo (catMaybes . dependencies)

  -- close any DBs that are idle, avoiding the set of DBs that we will
  -- be proactively keeping open (and their dependencies).
  closeIdleDatabases env
     (seconds $ fromIntegral databaseClosePolicy_close_after)
     (map itemRepo $ closeDeps toOpen)

  execWriterT $ do
    forM_ toOpen $ \item ->
      whenM (liftIO $ atomically $ isDatabaseClosed env $ itemRepo item) $
        preOpenDB (itemRepo item)

    publishCounter "glean.db.remote.oldness" $ case fetchBackupsResult of
      ReusedPreviousBackups{reusedBackups = (_,utc)} ->
        timeDiffInSeconds now utc
      FetchFailure{previousGoodBackups = Just (_, utc)} ->
        timeDiffInSeconds now utc
      FetchFailure{previousGoodBackups = Nothing} -> 0
      FetchedNewBackups{} -> 0

    forM_ byRepoName $ \(repoNm, dbsByAge) -> do
      let prefix = "glean.db." <> Text.encodeUtf8 repoNm
      let repoKeep =
            [ item
            | (item,_) <- retentionLocal
            , repoNm == Thrift.repo_name (itemRepo item) ]

      -- upsert counters
      publishCounter (prefix <> ".all") $ length repoKeep
      publishCounter (prefix <> ".available") $ length $ filter
        (\Item{..} ->
          itemLocality == Local
          && completenessStatus itemMeta == Thrift.DatabaseStatus_Complete)
        repoKeep
      publishCounter (prefix <> ".restoring") $
        length restores +
        length (NonEmpty.filter (\Item{..} -> itemLocality == Restoring)
          dbsByAge)
      publishCounter (prefix <> ".indexing") $ length $ NonEmpty.filter
        (\Item{..} ->
          itemLocality == Local
          && completenessStatus itemMeta == Thrift.DatabaseStatus_Incomplete)
        dbsByAge
      publishCounter (prefix <> ".backups") $ length $ NonEmpty.filter
        (\Item{..} -> itemLocality == Cloud)
        dbsByAge

      -- Report the age of the newest local DB. We only want to report
      -- the age for DBs that clients can query, hence the locality
      -- filter. The global DB age will be calculated by taking the
      -- minimum age reported by all the servers.
      case [ item | item <- NonEmpty.toList dbsByAge,
                itemLocality item == Local ] of
        [] -> return ()
        db:_ -> do
          let
              meta = itemMeta db
              dbStart = dbTime meta
              ageFrom t0 = timeSpanInSeconds $
                fromUTCTime now `timeDiff` posixEpochTimeToTime t0
          -- .age is the age of the data, .span is the age of the DB
          publishCounter (prefix <> ".age") (ageFrom dbStart)
          publishCounter (prefix <> ".span") (ageFrom (metaCreated meta))
          publishCounter (prefix <> ".newest")
            (fromIntegral (unPosixEpochTime dbStart))

    -- Report shard stats for dynamic sharding assignment
    mapM_ (\(n,v) -> publishCounter (Text.encodeUtf8 n) v) $
      countersForShardSizes sm $
      Map.fromListWith (+) $
      [ (shard, bytes)
      | (item, shard) <- retentionLocal
      , Complete DatabaseComplete{databaseComplete_bytes = Just bytes} <-
          [metaCompleteness (itemMeta item)]
      ] ++
      [ (shard, 0)
      | shard <- toList myShards
      , shard `notElem` Set.fromList (map snd retentionLocal)]


mergeLocalAndRemote :: [(Repo, Meta)] -> [Item] -> [Item]
mergeLocalAndRemote backups localAndRestoring =
  sortOn (dbTime . itemMeta) $
      localAndRestoring ++
        [ Item repo Cloud meta ItemMissing
          -- DBs we could restore
        | (repo, meta) <- backups
        , repo `notElem` map itemRepo localAndRestoring  ]

data FetchBackups
  = ReusedPreviousBackups {
    reusedBackups :: ([(Repo, Meta)], UTCTime)
    }
  | FetchedNewBackups {
    newBackups :: [(Repo, Meta)]
    }
  | FetchFailure {
    previousGoodBackups :: Maybe ([(Repo, Meta)], UTCTime),
    fetchError :: SomeException
  }

-- Fetches backups only if they haven't been fetched recently
fetchBackups :: Env -> IO FetchBackups
fetchBackups env = do
  ServerConfig.Config{..} <- Observed.get (envServerConfig env)
  let syncPeriodSeconds = fromIntegral config_backup_list_sync_period
  maybeLastFetch <- readTVarIO (envCachedRestorableDBs env)
  now <- envGetCurrentTime env
  case maybeLastFetch of
    Nothing -> fetch now Nothing
    Just (lastFetch, dbs)
      | now `timeDiffInSeconds` lastFetch > syncPeriodSeconds ->
        fetch now (Just (dbs, lastFetch))
      | otherwise ->
        return $ ReusedPreviousBackups (dbs, lastFetch)
  where
    fetch now previousGood = do
      logInfo "fetching restorable databases list"
      dbs <- tryAll $
        loggingAction (runLogCmd "listRestorable" env) (const mempty) $
          concatMap HashMap.toList <$>
            forRestoreSitesM env mempty listRestorable

      when (null dbs) $
        logWarning "found no restorable databases"

      for_ dbs $ \dbs ->
        atomically $ writeTVar (envCachedRestorableDBs env) (Just (now,dbs))

      return $ case dbs of
        Right dbs -> FetchedNewBackups dbs
        Left error -> FetchFailure previousGood error

timeDiffInSeconds :: UTCTime -> UTCTime -> Int
timeDiffInSeconds t1 t2 =
  timeSpanInSeconds $ fromUTCTime t1 `timeDiff` fromUTCTime t2