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.
  -}
  
  module Glean.Database.Retention (
      dbIndex, DbIndex(..),
      retentionChanges,
      RetentionChanges(..),
      repoRetention,
      hasAllProperties,
    ) where
  
  import Control.Monad
  import Control.Monad.Extra
  import qualified Data.HashMap.Strict as HashMap
  import Data.List
  import Data.List.Extra (nubOrdOn)
  import Data.List.NonEmpty (NonEmpty(..))
  import qualified Data.List.NonEmpty as NonEmpty
  import Data.Maybe
  import Data.Ord
  import qualified Data.Map as Map
  import qualified Data.Set as Set
  import Data.Text (Text)
  import Data.Time
  
  import Util.List
  
  import Glean.Database.Catalog.Filter
  import Glean.Database.Meta
  import Glean.Database.Restore (restorable)
  import qualified Glean.ServerConfig.Types as ServerConfig
  import Glean.Types
  import Glean.Util.ShardManager
  import Glean.Util.TransitiveClosure
  
  data RetentionChanges shard = RetentionChanges
    { retentionLocal :: [(Item,shard)]
    , retentionDelete :: [Item]
    , retentionRestore :: [Item]
    , retentionElsewhere :: [Item]
    , localMissingDependencies :: [Item]
    , allMissingDependencies :: [Item]
    }
  
  -- | Compute the changes to the current set of local DBs based on the
  -- set of known DBs (DbIndex), the retention/restore policy, the local
  -- shard, and an oracle to determine whether a DB is available elsewhere.
  
  retentionChanges
    :: forall m shard . (Monad m, Show shard, Ord shard)
    => ServerConfig.DatabaseRetentionPolicy
      -- ^ Retention policy
    -> ServerConfig.DatabaseRestorePolicy
      -- ^ Restore policy
    -> UTCTime
      -- ^ Current time
    -> DbIndex
      -- ^ All DBs
    -> (Item -> m Bool)
      -- ^ Is a DB available on another server?
    -> (BaseOfStack -> Repo -> shard)
      -- ^ Compute shard for DB
    -> Set.Set shard
      -- ^ Shards to retain in this node
    -> m (RetentionChanges shard)
  
  retentionChanges
      retentionPolicy restorePolicy
      now
      index
      isAvailable
      dbToShard myShards = do
  
    keep <- computeRetentionSet
      retentionPolicy restorePolicy now isAvailable index
  
    let
      itemToShard :: Item -> Maybe shard
      itemToShard item = do
        stack <- repoStack index item
        return $
          dbToShard (BaseOfStack $ last $ itemRepo item : stack) (itemRepo item)
  
      keepAnnotatedWithShard =
        [ (item, guard (shard `Set.member` myShards) >> pure shard)
        | item <- keep
        -- We get Nothing when a dependency is missing, which should never happen.
        -- If it does a 'logWarning' will be emitted by 'checkDependencies' below
        , Just shard <- [itemToShard item]
        ]
  
      keepInThisNode =
        [ (item, shard)
        | (item, Just shard) <- keepAnnotatedWithShard ]
  
      delete =
        [ item
        | item@(Item repo Local _ _) <- Map.elems (byRepo index)
        , repo `notElem` map (itemRepo . fst) keepInThisNode ]
  
      fetch =
        [ item
        | (item@Item{..},_) <- keepInThisNode
        , itemLocality == Cloud ]
  
      elsewhere =
        [ item
          -- Nothing means it is not in any of the shards assigned to this node
        | (item, Nothing) <- keepAnnotatedWithShard
        ]
  
      missingDeps dbs =
        nubOrdOn itemRepo [db | db <- dbs, Nothing <- dependencies index db]
  
    return RetentionChanges {
      retentionLocal = keepInThisNode,
      retentionDelete = delete,
      retentionRestore = fetch,
      retentionElsewhere = elsewhere,
      localMissingDependencies = missingDeps (map fst keepInThisNode),
      allMissingDependencies = missingDeps keep
    }
  
  repoStack :: DbIndex -> Item -> Maybe [Repo]
  repoStack index@DbIndex{..} Item{..} =
    case metaDependencies itemMeta of
      Just (Dependencies_stacked Stacked{..}) -> do
        let base = Repo stacked_name stacked_hash
        baseItem <- Map.lookup base byRepo
        rest <- repoStack index baseItem
        return (base : rest)
  
      Just (Dependencies_pruned Pruned{..}) -> do
        baseItem <- Map.lookup pruned_base byRepo
        rest <- repoStack index baseItem
        return (pruned_base : rest)
      Nothing -> return []
  
  -- | Information about the set of DBs that we have
  data DbIndex = DbIndex
    { byRepo :: Map.Map Repo Item
      -- ^ Items indexed by Repo
    , byRepoName :: [(Text, NonEmpty Item)]
      -- ^ Items grouped by Repo name
    , dependencies :: Item -> [Maybe Item]
      -- ^ DB dependencies, Nothing indicates that there is a dependency
      -- but it is missing in the catalog.
    }
  
  dbIndex :: [Item] -> DbIndex
  dbIndex items = DbIndex{..}
    where
      byRepo =
        Map.fromListWith pickFromDuplicates
          $ [(itemRepo item, item) | item <- items]
      pickFromDuplicates x y = -- Pick local one in a conflict
        minimumBy (\x y -> compare (itemLocality x) (itemLocality y)) [x, y]
  
      byRepoName =
        HashMap.toList $ HashMap.fromListWith (<>)
          [ (repo_name $ itemRepo item, item :| []) | item <- items ]
  
      dependencies = stacked . metaDependencies . itemMeta
  
      stacked (Just (Dependencies_stacked Stacked{..})) =
        [ Repo stacked_name stacked_hash `Map.lookup` byRepo ]
      stacked (Just (Dependencies_pruned update)) =
        [ pruned_base update `Map.lookup` byRepo ]
      stacked Nothing = []
  
  -- | The final set of DBs we want usable on disk.
  -- This is the set of 'keepRoots' DBs extended with all the stacked dependencies
  computeRetentionSet
    :: forall m . Monad m
    => ServerConfig.DatabaseRetentionPolicy
    -> ServerConfig.DatabaseRestorePolicy
    -> UTCTime
    -> (Item -> m Bool)
    -> DbIndex
    -> m [Item]
  computeRetentionSet config_retention config_restore
      time isAvailableM dbIndex@DbIndex{..} =
    transitiveClosureBy itemRepo (catMaybes . depsRestored) <$>
      concatMapM allRetention byRepoName
    where
      -- Add transitive dependencies of a DB to the retention set only
      -- if the DB is local or will be restored.
      depsRestored :: Item -> [Maybe Item]
      depsRestored item@Item{..}
        | itemLocality == Local
          || restorable config_restore itemRepo = dependencies item
        | otherwise = []
  
      allRetention :: (Text, NonEmpty Item) -> m [Item]
      allRetention (repo, dbs) = do
        let policies = repoRetention config_retention repo
        uniqBy (comparing itemRepo) . concat <$>
          mapM (\pol ->
            dbRetentionForRepo pol time isAvailableM dbs dbIndex) policies
  
  
  -- | The target set of DBs we want usable on the disk. This is a set of
  -- DBs that satisfies the policy.
  dbRetentionForRepo
    :: Monad m
    => ServerConfig.Retention
    -> UTCTime
    -> (Item -> m Bool)
    -> NonEmpty Item
    -> DbIndex
    -> m [Item]
  dbRetentionForRepo ServerConfig.Retention{..} t isAvailableM dbs dbIndex = do
    let
      -- retention policy parameters
      retainAtLeast' = fromIntegral $ fromMaybe 0 retention_retain_at_least
      retainAtMost' = fmap fromIntegral retention_retain_at_most
      -- enforce invariant: atLeast <= atMost
      retainAtLeast = min retainAtLeast' (fromMaybe maxBound retainAtMost')
      retainAtMost  = max retainAtLeast' <$> retainAtMost'
      deleteIfOlder = fmap fromIntegral retention_delete_if_older
      deleteIncompleteIfOlder =
        fmap fromIntegral retention_delete_incomplete_if_older
  
      f &&& g = \x -> f x && g x
      f ||| g = \x -> f x || g x
      f &&&> g = \x -> if not(f x) then return False else g x
      f |||> g = \x -> if f x then return True else g x
  
      ifSet (Just a) f = f a
      ifSet Nothing _ = const False
  
      -- predicates
      isLocal Item{..} = itemLocality == Local
      isComplete Item{..} =
        completenessStatus itemMeta == DatabaseStatus_Complete
      isOlderThan secs Item{..} = dbAge t itemMeta >= secs
      isAvailable = isLocal |||> isAvailableM
      hasDependencies = not . missingDependencies dbIndex
  
      -- all DBs with the required/excluded properties, sorted by most recent first
      sorted =
        sortOn (Down . dbTime . itemMeta) $
        filter (hasAllProperties retention_required_properties) $
        filter (not . hasAnyProperties retention_excluded_properties) $
        NonEmpty.toList dbs
  
      -- whether to delete a DB according to the deletion policy
      delete =
        ifSet deleteIfOlder isOlderThan |||
        (ifSet deleteIncompleteIfOlder $ \secs ->
          (not . isComplete) &&& isOlderThan secs)
  
      -- implements retain_at_least
      atLeast =
        take retainAtLeast $
        filter (isComplete &&& hasDependencies) sorted
  
    -- Ensure we have retain_at_least DBs available somewhere.  This
    -- prevents us deleting an older DB before the newer one has been
    -- restored. The isAvailable check ensures that we don't keep an old
    -- DB on this server when the newer one is on another server.
    atLeastAvailable <- takeFilterM
      retainAtLeast
      (isComplete &&& hasDependencies &&&> isAvailable)
      -- bound the search since isAvailable is expensive
      -- this matters only for tier bootstraps where all DBs are unavailable
      (take (retainAtLeast*10) sorted)
  
      -- delete DBs according to the deletion policy, and keep retain_at_most
    let atMost = maybe id take retainAtMost (filter (not . delete) sorted)
  
    return $ uniqBy (comparing itemRepo) $
      atLeast ++ atLeastAvailable ++ atMost
  
  missingDependencies :: DbIndex -> Item -> Bool
  missingDependencies dbIndex item = any isNothing (dependencies dbIndex item)
  
  hasProperty :: Item -> (Text, Text) -> Bool
  hasProperty Item{..} (name,val) =
    HashMap.lookup name (metaProperties itemMeta) == Just val
  
  hasAllProperties :: HashMap.HashMap Text Text -> Item -> Bool
  hasAllProperties req item = all (hasProperty item) (HashMap.toList req)
  
  hasAnyProperties :: HashMap.HashMap Text Text -> Item -> Bool
  hasAnyProperties exclude_props item =
    any (hasProperty item) (HashMap.toList exclude_props)
  
  repoRetention
    :: ServerConfig.DatabaseRetentionPolicy
    -> Text
    -> NonEmpty ServerConfig.Retention
  repoRetention ServerConfig.DatabaseRetentionPolicy{..} repoNm =
    case NonEmpty.nonEmpty (old_retention <> new_retention) of
      Nothing -> databaseRetentionPolicy_default_retention :| []
      Just some -> some
    where
    old_retention =
      maybeToList $
        Map.lookup
          repoNm
          databaseRetentionPolicy_repos
    new_retention =
      Map.findWithDefault
        []
        repoNm
        databaseRetentionPolicy_by_repo
  
  -- | Take the first n items that satisfy the predicate
  takeFilterM :: (Monad m) => Int -> (a -> m Bool) -> [a] -> m [a]
  takeFilterM n pred = loop [] n where
    loop acc _ [] = return (reverse acc)
    loop acc 0 _ = return (reverse acc)
    loop acc n (x:xx) = do
      accept <- pred x
      if accept
        then loop (x:acc) (n-1) xx
        else loop acc n xx
  