{-
  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 DatabaseJanitorTest
  ( main
  , dbConfig
  , serverConfig
  , makeFakeDB
  , withTest
  ) where

import Control.Concurrent (threadDelay)
import Util.STM
import Control.Exception
import Control.Monad
import Data.Aeson
import qualified Data.ByteString.Char8 as BS
import qualified Data.ByteString.Lazy as LB
import Data.Default
import qualified Data.HashMap.Strict as HashMap
import qualified Data.HashSet as HashSet
import Data.Int (Int64)
import Data.IORef
import Data.List
import qualified Data.Map as Map
import Data.Maybe
import qualified Data.Set as Set
import Data.Text (Text)
import qualified Data.Text as Text
import Data.Time.Clock
import GHC.Stack (HasCallStack)
import System.Directory
import System.FilePath
import System.IO.Temp
import System.Time.Extra (sleep)
import System.Timeout
import Test.HUnit

import ServiceData.GlobalStats (getCounters)
import TestRunner
import Util.EventBase
import Util.TimeSec

import qualified Glean.Database.Catalog as Catalog
import Glean.Database.Config
import Glean.Database.Data
import Glean.Database.Delete
import Glean.Database.Env
import Glean.Database.Janitor
import Glean.Database.Meta
import Glean.Database.Repo
import qualified Glean.Database.Storage as Storage
import qualified Glean.Database.Storage.RocksDB as RocksDB
import Glean.Database.Close
import Glean.Database.List
import Glean.Database.Open (isDatabaseClosed, withOpenDatabase)
import Glean.Database.Types
import Glean.Database.Schema
import Glean.Database.Schema.Types
import Glean.Init
import Glean.RTS.Types (lowestFid)
import Glean.ServerConfig.Types as ServerTypes
import Glean.Internal.Types
import Glean.Types as Thrift
import Glean.Util.ConfigProvider
import Glean.Util.ShardManager
import Glean.Util.ThriftSource as ThriftSource
import Util.Time (seconds)
import Glean.Database.Backup.Backend
import Glean.Database.Backup.Mock


withTest
  :: (FilePath -> DbSchema -> IO ())
  -> (FilePath -> DbSchema -> IO ())
  -> (EventBaseDataplane -> NullConfigProvider -> FilePath -> FilePath
       -> IO ())
  -> IO ()
withTest setup setupBackup action =
  withEventBaseDataplane $ \evb ->
  withConfigProvider defaultConfigOptions $ \cfgAPI ->
  withSystemTempDirectory "glean-dbtest" $ \dbdir -> do
  withSystemTempDirectory "glean-dbtest-backup" $ \backupdir -> do
    schema <- parseSchemaDir schemaSourceDir
    schema <- newDbSchema Nothing schema LatestSchema readWriteContent def
    setup dbdir schema
    setupBackup backupdir schema
    action evb cfgAPI dbdir backupdir

broken :: UTCTime -> Completeness
broken _ = Broken (DatabaseBroken "index" "TESTING")

complete :: Int64 -> UTCTime -> Completeness
complete size =
  Complete . (`DatabaseComplete` Just size) . utcTimeToPosixEpochTime

repo0001 :: Repo
repo0001 = Repo "test" "0001"

setupBasicDBs :: FilePath -> DbSchema -> IO ()
setupBasicDBs dbdir schema = do
  now <- getCurrentTime
  let age t = addUTCTime (negate (fromIntegral (timeSpanInSeconds t))) now
  -- populate a dir with various DBs
  makeFakeDB schema dbdir repo0001 (age (days 0)) (complete 1)
    (props [("bool","yes")])
  makeFakeDB schema dbdir (Repo "test" "0002") (age (days 2)) broken id
  makeFakeDB schema dbdir (Repo "test" "0003") (age (days 3)) (complete 3) $
    stacked (Stacked "test" "0004" Nothing)
  makeFakeDB schema dbdir (Repo "test" "0004") (age (days 0)) (complete 4) $
    sourceTime (age (days 4)) .  -- source 4 days ago, created 0 days ago
    stacked (Stacked "test" "0005" Nothing)
  makeFakeDB schema dbdir (Repo "test" "0005") (age (days 5)) (complete 5) $
    stacked (Stacked "test2" "0006" Nothing) . props [("bool","no")]
  makeFakeDB schema dbdir (Repo "test2" "0006") (age (days 6)) (complete 6) id
  makeFakeDB schema dbdir (Repo "test3" "0007") (age (days 0)) (complete 1)
    (props [("exclude_me","true")])


setupBasicCloudDBs :: FilePath -> DbSchema -> IO ()
setupBasicCloudDBs backupDir schema = do
  now <- getCurrentTime
  let age t = addUTCTime (negate (fromIntegral (timeSpanInSeconds t))) now
  makeFakeCloudDB schema backupDir (Repo "test" "0008")
    (age(days 8)) (complete 8) id
  makeFakeCloudDB schema backupDir (Repo "test2" "0009")
    (age(days 7)) (complete 9) id
  makeFakeCloudDB schema backupDir (Repo "test2" "0010")
    (age(days 6)) (complete 9) id
  makeFakeCloudDB schema backupDir (Repo "test2" "0011")
    (age(days 5)) (complete 9) id
  makeFakeCloudDB schema backupDir (Repo "test2" "0012")
    (age(days 4)) (complete 9) id
  makeFakeCloudDB schema backupDir (Repo "test2" "0013")
    (age(days 3)) (complete 9) id
  makeFakeCloudDB schema backupDir (Repo "test2" "0014")
    (age(days 2)) (complete 9) id
  makeFakeCloudDB schema backupDir (Repo "test" "0015")
    (age(days 7)) (complete 7) (stacked (Stacked "test2" "0013" Nothing))

setupBasicCloudDBsWithMissingDependencyOne :: FilePath -> DbSchema -> IO ()
setupBasicCloudDBsWithMissingDependencyOne backupDir schema = do
  setupBasicCloudDBs backupDir schema
  dbtime <- getCurrentTime
  makeFakeCloudDB schema backupDir (Repo "test" "0016") dbtime
    (complete 8) (stacked (Stacked "missing" "0001" Nothing))

withFakeDBs
  :: (EventBaseDataplane -> NullConfigProvider -> FilePath -> FilePath
       -> IO ())
  -> IO ()
withFakeDBs action = withTest setupBasicDBs (\_ _ -> pure ()) action

withFakeCloudDBs
  :: (EventBaseDataplane -> NullConfigProvider -> FilePath -> FilePath
       -> IO ())
  -> IO ()
withFakeCloudDBs = withTest (\_ _ -> pure ()) setupBasicCloudDBs

stacked :: Stacked -> Meta -> Meta
stacked st meta = meta { metaDependencies = Just (Thrift.Dependencies_stacked st) }

props :: [(Text, Text)] -> Meta -> Meta
props list meta = meta { metaProperties = HashMap.fromList list }

sourceTime :: UTCTime -> Meta -> Meta
sourceTime t meta = meta { metaRepoHashTime = Just (utcTimeToPosixEpochTime t) }

makeFakeDB
  :: DbSchema
  -> FilePath
  -> Repo
  -> UTCTime
  -> (UTCTime -> Completeness)
  -> (Meta -> Meta)
  -> IO ()
makeFakeDB schema root repo dbtime completeness opts = do
  let
    meta = opts $ Meta
      { metaVersion = Storage.currentVersion
      , metaCreated = utcTimeToPosixEpochTime dbtime
      , metaRepoHashTime = Nothing
      , metaCompleteness = completeness dbtime
      , metaBackup = Nothing
      , metaProperties = HashMap.empty
      , metaDependencies = Nothing
      , metaCompletePredicates = mempty
      , metaAxiomComplete = False
      }
  let repoPath = databasePath root repo
  createDirectoryIfMissing True repoPath
  storage <- RocksDB.newStorage root def
  bracket
    (Storage.open
      storage
      repo
      (Storage.Create lowestFid Nothing Storage.UseDefaultSchema)
      Storage.currentVersion)
    Storage.close
    (\hdl -> storeSchema hdl $ toStoredSchema schema)
  LB.writeFile (repoPath </> "meta") (encode meta)

makeFakeCloudDB
  :: DbSchema
  -> FilePath
  -> Repo
  -> UTCTime
  -> (UTCTime -> Completeness)
  -> (Meta -> Meta)
  -> IO ()
makeFakeCloudDB schema backupDir repo dbtime completeness opts = do
  let repoPath = databasePath backupDir repo
  createDirectoryIfMissing True repoPath
  storage <- RocksDB.newStorage backupDir def
  bracket
    (Storage.open
      storage
      repo
      (Storage.Create lowestFid Nothing Storage.UseDefaultSchema)
      Storage.currentVersion)
    Storage.close
    (\hdl -> do
      storeSchema hdl $ toStoredSchema schema
      tmpDir <- getCanonicalTemporaryDirectory
      withTempDirectory tmpDir "scratch" $ \scratch ->
        Storage.backup hdl scratch $ \file _data ->
          void $ backup (mockSite backupDir) repo props Nothing file
    )
  where
    props = opts $ Meta
        { metaVersion = Storage.currentVersion
        , metaCreated = utcTimeToPosixEpochTime dbtime
        , metaRepoHashTime = Nothing
        , metaCompleteness = completeness dbtime
        , metaBackup = Nothing
        , metaProperties = HashMap.empty
        , metaDependencies = Nothing
        , metaCompletePredicates = mempty
        , metaAxiomComplete = False
        }

dbConfig :: FilePath -> ServerTypes.Config -> Glean.Database.Config.Config
dbConfig dbdir serverConfig = def
  { cfgDataStore = fileDataStore dbdir
  , cfgSchemaLocation = Just schemaLocationFiles
  , cfgServerConfig = ThriftSource.value serverConfig
  , cfgReadOnly = False
      -- If we set cfgReadOnly=True, then backing up a DB will fail
      -- because backup first opens the DB, and opening a DB tries to
      -- load the schema from the DB.  If we use read/write mode, then
      -- opening will commit the schema instead of failing.
  , cfgMockWrites = False
  , cfgListener = mempty
  }

serverConfig :: FilePath -> ServerTypes.Config
serverConfig backupdir = def
  { config_backup = def
    { databaseBackupPolicy_allowed = mempty
    , databaseBackupPolicy_location = "mock:" <> Text.pack backupdir
    }
  , config_restore = def
    { databaseRestorePolicy_enabled = False
    }
  , config_janitor_period = Nothing -- no auto janitor, we'll run it manually
  , config_db_rocksdb_cache_mb = 0
  }

listDBs :: Env -> IO [Thrift.Database]
listDBs env = filter hereDBs <$> listAllDBs env
  where
    hereDBs Database{..} = database_status /= DatabaseStatus_Available

listHereDBs :: Env -> IO [Database]
listHereDBs = listDBs

listAllDBs :: Env -> IO [Database]
listAllDBs env = listDatabasesResult_databases <$> listDatabases env def

waitDel :: Env -> IO ()
waitDel env = atomically $ do
  done <- HashMap.null <$> readTVar (envDeleting env)
  when (not done) retry

deleteOldDBsTest :: Test
deleteOldDBsTest = TestCase $ withFakeDBs $ \evb cfgAPI dbdir backupdir -> do
  let cfg = dbConfig dbdir $ (serverConfig backupdir)
        { config_retention = def
          { databaseRetentionPolicy_default_retention = def
            { retention_delete_if_older =
                Just $ fromIntegral $ timeSpanInSeconds $ days 1 }
          }
        }
  withDatabases evb cfg cfgAPI $ \env -> do

  dbs <- listDBs env
  assertEqual "before"
    [ "0001", "0002", "0003", "0004", "0005", "0006", "0007"]
    (sort (map (repo_hash . database_repo) dbs))

  runDatabaseJanitor env
  waitDel env
  dbs <- listDBs env
  assertEqual "after"
    [ "0001", "0004", "0005", "0006","0007"]
      -- we don't delete 0004 because even though its source is old,
      -- it was created recently. 0004 depends on 0005, and 0005
      -- depends on 0006
    (sort (map (repo_hash . database_repo) dbs))

  dbdirs1 <- listDirectory (dbdir </> "test")
  dbdirs2 <- listDirectory (dbdir </> "test2")
  assertEqual "directories deleted"
    [ "0001", "0004", "0005", "0006" ]
    (dbdirs1 ++ dbdirs2)

  runDatabaseJanitor env
  waitDel env
  dbs <- listDBs env
  assertEqual "after-repeat"
    [ "0001", "0004", "0005", "0006","0007" ]
    (sort (map (repo_hash . database_repo) dbs))

deleteIncompleteDBsTest :: Test
deleteIncompleteDBsTest = TestCase $
  withFakeDBs $ \evb cfgAPI dbdir backupdir -> do
  let cfg = dbConfig dbdir $ (serverConfig backupdir)
        { config_retention = def
          { databaseRetentionPolicy_default_retention = def
            { retention_delete_if_older =
                Just $ fromIntegral $ timeSpanInSeconds $ days 10
            , retention_delete_incomplete_if_older =
                Just $ fromIntegral $ timeSpanInSeconds $ days 1 }
          }
        }
  withDatabases evb cfg cfgAPI $ \env -> do

  dbs <- listDBs env
  assertEqual "before"
    [ "0001", "0002", "0003", "0004", "0005", "0006","0007"]
    (sort (map (repo_hash . database_repo) dbs))

  runDatabaseJanitor env
  waitDel env
  dbs <- listDBs env
  print (map database_repo dbs)
  assertEqual "after"
    [ "0001", "0003", "0004", "0005", "0006", "0007"]
    (sort (map (repo_hash . database_repo) dbs))

  runDatabaseJanitor env
  waitDel env
  dbs <- listDBs env
  assertEqual "after-repeat"
    [ "0001", "0003", "0004", "0005", "0006","0007"]
    (sort $ map (repo_hash . database_repo) dbs)


retainAtMostTest :: Test
retainAtMostTest = TestCase $ withFakeDBs $ \evb cfgAPI dbdir backupdir -> do
  let cfg = dbConfig dbdir $ (serverConfig backupdir)
        { config_retention = def
          { databaseRetentionPolicy_default_retention = def
            { retention_retain_at_most = Just 2 }
          }
        }
  withDatabases evb cfg cfgAPI $ \env -> do

  runDatabaseJanitor env
  waitDel env
  dbs <- listDBs env
  let repos = map database_repo dbs
  assertEqual "after"
    [ "0001", "0002", "0006", "0007"]
    (sort $ map repo_hash repos)

  runDatabaseJanitor env
  waitDel env
  dbs <- listDBs env
  assertEqual "after-repeat"
    [ "0001", "0002", "0006", "0007"]
    (sort $ map (repo_hash . database_repo) dbs)


retainAtLeastTest :: Test
retainAtLeastTest = TestCase $ withFakeDBs $ \evb cfgAPI dbdir backupdir -> do
  let cfg = dbConfig dbdir $ (serverConfig backupdir)
        { config_retention = def
          { databaseRetentionPolicy_default_retention = def
            { retention_delete_if_older =
                Just $ fromIntegral $ timeSpanInSeconds $ days 1
            , retention_retain_at_least = Just 2 }
          }
        }
  withDatabases evb cfg cfgAPI $ \env -> do

  runDatabaseJanitor env
  waitDel env
  dbs <- listDBs env
  let repos = map database_repo dbs
  assertEqual "after"
    [ "0001", "0003", "0004", "0005", "0006", "0007"]
    (sort $ map repo_hash repos)
    -- should drop 0002, because its metaIndexing=Failed
    -- should retain 0003 even though it is older than delete_if_older,
    -- due to retain_at_least.

  runDatabaseJanitor env
  waitDel env
  dbs <- listDBs env
  let repos = map database_repo dbs
  assertEqual "after-repeat"
    [ "0001", "0003", "0004", "0005", "0006","0007"]
    (sort $ map repo_hash repos)

excludedPropsTest :: Test
excludedPropsTest = TestCase $ withFakeDBs $ \evb cfgAPI dbdir backupdir -> do
  let cfg = dbConfig dbdir $ (serverConfig backupdir)
        { config_retention = def
          { databaseRetentionPolicy_default_retention = def
            { retention_excluded_properties =
                HashMap.fromList [("exclude_me","true")]
            }
          }
        }
  withDatabases evb cfg cfgAPI $ \env -> do
  runDatabaseJanitor env
  dbs <- listDBs env
  let repos = map database_repo dbs
  assertEqual "after"
    [ "0001", "0002", "0003", "0004", "0005", "0006"]
    (sort $ map repo_hash repos)

requiredPropsTest :: Test
requiredPropsTest = TestCase $ withFakeDBs $ \evb cfgAPI dbdir backupdir -> do
  let cfg = dbConfig dbdir $ (serverConfig backupdir)
        { config_retention = def
          { databaseRetentionPolicy_default_retention = def
            { retention_required_properties =
                HashMap.fromList [("bool","yes")]
            }
          }
        }
  withDatabases evb cfg cfgAPI $ \env -> do
  runDatabaseJanitor env
  dbs <- listDBs env
  let repos = map database_repo dbs
  assertEqual "after" [ "0001" ] (map repo_hash repos)

multiRetentionTest :: Test
multiRetentionTest = TestCase $ withFakeDBs $ \evb cfgAPI dbdir backupdir -> do
  let cfg = dbConfig dbdir $ (serverConfig backupdir)
        { config_retention = def
          { databaseRetentionPolicy_by_repo =
            Map.fromList
              [ ("test",
                [ def
                  { retention_required_properties =
                      HashMap.fromList [("bool","yes")]
                  , retention_retain_at_least = Just 1
                  }
                , def
                  { retention_required_properties =
                      HashMap.fromList [("bool","no")]
                  , retention_retain_at_least = Just 1
                  , retention_keep_open = True
                  }
                ])
              ]
            }
          }
  withDatabases evb cfg cfgAPI $ \env -> do
  runDatabaseJanitor env
  dbs <- listDBs env
  let repos = sort $ map (repo_hash . database_repo) dbs
  assertEqual "after" [ "0001", "0005", "0006", "0007" ] repos

  -- 0005 should be open due to keep_open=True in the retention policy
  closed <- atomically $ isDatabaseClosed env (Repo "test" "0005")
  assertBool "open" (not closed)

-- | If we want to restore only one type of DB with a retention
-- policy, check that we don't restore additional instances of that DB
-- that are dependencies of other (non-restored) DBs.
retentionRestoreDepsTest :: Test
retentionRestoreDepsTest = TestCase $
  withFakeCloudDBs $ \evb cfgAPI dbdir backupdir -> do
    let cfg = dbConfig dbdir $ (serverConfig backupdir)
          { config_retention = def
            { databaseRetentionPolicy_default_retention = def
              { retention_retain_at_most = Just 1
              }
            },
            config_restore = def
            { databaseRestorePolicy_enabled = False
            , databaseRestorePolicy_override = Set.fromList ["test2"]
            }
          }
    withDatabases evb cfg cfgAPI $ \env -> do
    runDatabaseJanitor env
    dbs <- listDBs env
    let repos = map database_repo dbs
    assertEqual "after" [ "0014" ] (map repo_hash repos)

backupRestoreTest :: Test
backupRestoreTest = TestCase $ withFakeDBs $ \evb cfgAPI dbdir backupdir -> do
  let cfg = dbConfig dbdir $ (serverConfig backupdir)
        { config_backup = (config_backup (serverConfig backupdir))
          { databaseBackupPolicy_allowed =
              HashSet.fromList ["test", "test2", "test3"] }
        }
  withDatabases evb cfg cfgAPI $ \env -> do
    runDatabaseJanitor env
    atomically $ do
      dbs <- Catalog.getLocalDatabases $ envCatalog env
      let should_wait Thrift.GetDatabaseResult
              {getDatabaseResult_database = Thrift.Database{..}} =
            database_status == DatabaseStatus_Complete
            && isNothing database_location
      when (any should_wait dbs) retry

  backups <- listDirectory backupdir
  let withProps filename = [filename, filename ++ ".props"]
  let expected = concatMap withProps
        [ "test.0001", "test.0003", "test.0004", "test.0005" , "test2.0006", "test3.0007"]
  assertEqual "after"
    expected
    (sort backups)
  -- should not backup failed db test.0002

  let cfg = dbConfig dbdir $ (serverConfig backupdir)
        { config_restore = (config_restore (serverConfig backupdir))
          { databaseRestorePolicy_enabled = True }
        , config_retention = def
          { databaseRetentionPolicy_default_retention = def
            { retention_delete_if_older =
                Just $ fromIntegral $ timeSpanInSeconds $ days 1
            , retention_retain_at_least = Just 2 }
          }
        }

  -- remove all the DBs, so we have to restore them:
  removeDirectoryRecursive dbdir
  createDirectoryIfMissing True (dbdir </> "test")
  createDirectoryIfMissing True (dbdir </> "test2")
  createDirectoryIfMissing True (dbdir </> "test3")
  withDatabases evb cfg cfgAPI $ \env -> do
  runDatabaseJanitor env  -- this should kick off the restore

  let waitForRestore = do
        dbs <- listDBs env
        let
          available =
            [ db | db <- dbs
            , database_status db == Thrift.DatabaseStatus_Complete ]
        when (length dbs < 2 || length available < 2) waitForRestore

  r <- timeout (60*1000000) $ waitForRestore
  assertBool "timeout" $ isJust r

  -- we should have restored both 0001 and 0003, even though 0003 is
  -- older than delete_if_older, because retain_at_least = 2
  dbs <- listDBs env
  let
    repos = map database_repo dbs
    db1 = Repo "test" "0001"
    db3 = Repo "test" "0003"


  assertEqual "after"
    [ "0001", "0003", "0004", "0005", "0006", "0007" ]
    (sort $ map repo_hash repos)

  areClosed <- atomically
    $ liftM2 (&&) (isDatabaseClosed env db1) (isDatabaseClosed env db3)
  assertBool "dbs available but closed after restore" areClosed


openNewestTest :: Test
openNewestTest = TestCase $ withFakeDBs $ \evb cfgAPI dbdir backupdir -> do
  let cfg = dbConfig dbdir $ serverConfig backupdir

  withDatabases evb cfg cfgAPI $ \env -> do

  let
    newestDb = Repo "test" "0001"
    oldestDb = Repo "test" "0005"

  newestClosed <- atomically $ isDatabaseClosed env newestDb
  oldestClosed <- atomically $ isDatabaseClosed env oldestDb

  assertBool "newest closed before" newestClosed
  assertBool "oldest closed before" oldestClosed

  runDatabaseJanitor env

  newestOpen <- atomically $ not <$> isDatabaseClosed env newestDb
  oldestStillClosed <- atomically $ isDatabaseClosed env oldestDb

  assertBool "newest open after" newestOpen
  assertBool "oldest closed after" oldestStillClosed

  -- delete repo001, repo002 is the newest but is broken, so we
  -- should open repo003.
  deleteDatabase env newestDb
  runDatabaseJanitor env

  let nextBroken = Repo "test" "0002"
  let nextNotBroken = Repo "test" "0003"
  dontOpenBroken <- atomically $ isDatabaseClosed env nextBroken
  openNotBroken <- atomically $ not <$> isDatabaseClosed env nextNotBroken

  assertBool "don't open broken" dontOpenBroken
  assertBool "open not broken" openNotBroken

closeIdleDBsTest :: Test
closeIdleDBsTest = TestCase $ withFakeDBs $ \evb cfgAPI dbdir backupdir -> do
  let cfg = dbConfig dbdir $ serverConfig backupdir

  withDatabases evb cfg cfgAPI $ \env -> do

  let
    normalDb = Repo "test" "0001"
    blackListedDb = Repo "test" "0003"

  withOpenDatabase env normalDb (\_ -> return ())
  withOpenDatabase env blackListedDb (\_ -> return ())

  normalDbOpen <- atomically $ not <$> isDatabaseClosed env normalDb
  blackListedDbOpen <- atomically $ not <$> isDatabaseClosed env blackListedDb

  assertBool "regular DB open before" normalDbOpen
  assertBool "blacklisted DB open before" blackListedDbOpen

  -- Close all available DBs unless blacklisted
  closeIdleDatabases env (seconds 0) [blackListedDb]

  normalDbClosed <- atomically $ isDatabaseClosed env normalDb
  blackListedDbStillOpen <- atomically
    $ not <$> isDatabaseClosed env blackListedDb

  assertBool "regular DB closed after" normalDbClosed
  assertBool "blacklisted DB open after" blackListedDbStillOpen

shardingTest :: Test
shardingTest = TestCase $ withFakeDBs $ \evb cfgAPI dbdir backupdir -> do
  myShards <- newIORef ["0001"] -- initial shard assignment
  let cfg = (dbConfig dbdir (serverConfig backupdir))
        {cfgShardManager = \_ _ k -> k $ SomeShardManager $
          shardByRepoHash (Just <$> readIORef myShards)}
  withDatabases evb cfg cfgAPI $ \env -> do
    runDatabaseJanitor env
    waitDel env
    dbs <- listHereDBs env
    assertEqual "initial shard assignment"
      ["0001"]
      (sort $ map (repo_hash . database_repo) dbs)

    -- update the shard assignment and verify
    writeIORef myShards ["nil"]
    runDatabaseJanitor env
    waitDel env
    dbs <- listHereDBs env
    assertEqual "shard assignment: removed 0003"
      []
      (sort $ map (repo_hash . database_repo) dbs)

shardingStacksTest :: Test
shardingStacksTest = TestCase $ withFakeDBs $ \evb cfgAPI dbdir backupdir -> do
  myShards <- newIORef ["0006"] -- initial shard assignment
  let cfg = (dbConfig dbdir (serverConfig backupdir))
        {cfgShardManager = \_ _ k -> k $ SomeShardManager $
          shardByBaseOfStackRepoHash (Just <$> readIORef myShards)}
  withDatabases evb cfg cfgAPI $ \env -> do
    runDatabaseJanitor env
    waitDel env
    dbs <- listHereDBs env
    assertEqual "all dbs in the stack belong to the shard"
      ["0003", "0004", "0005", "0006"]
      (sort $ map (repo_hash . database_repo) dbs)

shardingFallbackTest :: Test
shardingFallbackTest = TestCase $ withFakeDBs $ \evb cfgAPI dbdir backupdir -> do
  let cfg = (dbConfig dbdir (serverConfig backupdir))
        {cfgShardManager = \_ _ k ->
          k $ SomeShardManager $ shardByRepoHash (pure Nothing)}
  withDatabases evb cfg cfgAPI $ \env -> do
    runDatabaseJanitor env
    waitDel env
    dbs <- listHereDBs env
    assertEqual "falls back to no sharding"
      ["0001", "0002", "0003", "0004", "0005", "0006", "0007"]
      (sort $ map (repo_hash . database_repo) dbs)

shardingByRepoNameTest :: Test
shardingByRepoNameTest = TestCase $ withFakeDBs $ \evb cfgAPI dbdir backupdir -> do
  let cfg = (dbConfig dbdir (serverConfig backupdir))
        {cfgShardManager = \_ _ k -> k $ SomeShardManager $
          shardByRepo (pure $ Just ["test2"])}
  withDatabases evb cfg cfgAPI $ \env -> do
    runDatabaseJanitor env
    waitDel env
    dbs <- listHereDBs env
    assertEqual "only test2 dbs belong to the shard"
      ["0003", "0004", "0005", "0006"]
      (sort $ map (repo_hash . database_repo) dbs)


elsewhereTest :: Test
elsewhereTest = TestCase $ withFakeDBs $ \evb cfgAPI dbdir backupdir -> do
  let myShards = pure $ Just ["0001", "0003"]
  let cfg = (dbConfig dbdir (serverConfig backupdir)
        { config_retention = def
          { databaseRetentionPolicy_default_retention = def
            { retention_delete_if_older =
                Just $ fromIntegral $ timeSpanInSeconds $ days 10
            , retention_retain_at_least = Just 10 }
          }
        })
        {cfgShardManager = \_ _ k -> k $ SomeShardManager $ shardByRepoHash myShards}
  withDatabases evb cfg cfgAPI $ \env -> do

    dbs <- listHereDBs env
    assertEqual "before"
      [ "0001", "0002", "0003", "0004", "0005", "0006", "0007"]
      (sort (map (repo_hash . database_repo) dbs))

    runDatabaseJanitor env
    waitDel env
    dbs <- listAllDBs env

    assertEqual "after: dbs available with the retention policy"
      [ "0001", "0002", "0003", "0004", "0005", "0006", "0007"]
      (sort $ map (repo_hash . database_repo) dbs)

makeFilterDBsWithInvariant :: ([Repo] -> IO [Repo]) -> IO ([Repo] -> IO [Repo])
makeFilterDBsWithInvariant pred = do
  seenAvailabeDBsRef <- newIORef mempty
  return $ \dbs -> do
    available <- pred dbs
    let availableSet = HashSet.fromList available
    beenAvailableBefore <- readIORef seenAvailabeDBsRef
    assertEqual "called cfgFilterAvailableDBs twice on the same DB"
      mempty (HashSet.intersection availableSet beenAvailableBefore)
    writeIORef seenAvailabeDBsRef (availableSet <> beenAvailableBefore)
    return available

elsewhereNotYetAvailableTest :: Test
elsewhereNotYetAvailableTest =
  TestCase $ withFakeCloudDBs $ \evb cfgAPI dbdir backupdir -> do
    let myShards = pure $ Just []
    filterDBs <- makeFilterDBsWithInvariant $
      return . filter ((`elem` ["0008","0009"]) . repo_hash)
    let cfg = (dbConfig dbdir (serverConfig backupdir)
          { config_retention = def
            { databaseRetentionPolicy_default_retention = def
              { retention_delete_if_older =
                  Just $ fromIntegral $ timeSpanInSeconds $ days 10
              , retention_retain_at_least = Just 2
              , retention_retain_at_most = Just 4
              }
            },
            config_restore = def {
              databaseRestorePolicy_enabled = True
            }
          })
          {cfgShardManager = \_ _ k ->
              k $ SomeShardManager $ shardByRepoHash myShards
          ,cfgFilterAvailableDBs = filterDBs
          }
    withDatabases evb cfg cfgAPI $ \env -> do
      runDatabaseJanitor env
      dbs <- listAllDBs env

      assertEqual
        "at least 2 dbs actually available + at most 4 more not yet available"
        ["0008", "0009", "0011", "0012", "0013", "0014", "0015"]
        (sort $ map (repo_hash . database_repo) dbs)

shardUnexpireTest :: Test
shardUnexpireTest = TestCase $ withFakeDBs $ \evb cfgAPI dbdir backupdir -> do
  myShards <- newIORef ["nil"] -- initial incomplete shard assignment
  let cfg = (dbConfig dbdir (serverConfig backupdir)
        { config_retention = def
          { databaseRetentionPolicy_default_retention = def
            {retention_expire_delay = Just 30}
          }
        })
        {cfgShardManager = \_ _ k -> k $ SomeShardManager $
          shardByRepoHash (Just <$> readIORef myShards)}
  withDatabases evb cfg cfgAPI $ \env -> do
    runDatabaseJanitor env
    dbs <- listDBs env

    expiring <- atomically $
      mapM (Catalog.readExpiring (envCatalog env) . database_repo) dbs
    assertBool "All expiring" (all isJust expiring)

    -- update the shard assignment and verify
    writeIORef myShards [repo_hash repo0001]
    runDatabaseJanitor env

    expiring <- atomically $ Catalog.readExpiring (envCatalog env) repo0001
    assertBool "0001 not expiring now" (isNothing expiring)

expireTest :: Test
expireTest = TestCase $ withFakeDBs $ \evb cfgAPI dbdir backupdir -> do
  let emptyShardAssignment = shardByRepoHash $ pure $ Just ["nil"]
  let cfg = (dbConfig dbdir (serverConfig backupdir)
        { config_retention = def
          { databaseRetentionPolicy_default_retention = def
            {retention_expire_delay = Just 1}
          }
        , config_janitor_period = Nothing
        })
        {cfgShardManager = \_ _ k -> k $ SomeShardManager emptyShardAssignment}
  withDatabases evb cfg cfgAPI $ \env -> do
    runDatabaseJanitor env
    dbs <- listDBs env

    expiring <- atomically $
      mapM (Catalog.readExpiring (envCatalog env) . database_repo) dbs
    assertBool "All expiring" (all isJust expiring)

    -- run the Janitor before the expire delay to exercise all code paths
    runDatabaseJanitor env
    sleep 1
    -- run the Janitor twice after the expire delay to exercise all code paths
    runDatabaseJanitor env
    runDatabaseJanitor env
    -- check that the DBs were really deleted
    res <- timeout 10000000 -- 10s
                   (waitDel env)
    waitingDeletion <- readTVarIO (envDeleting env)
    assertBool ("timeout: " <> show (HashMap.keys waitingDeletion)) (isJust res)
    assertEqual "envDeleting" [] (HashMap.keys waitingDeletion)
    localDBs <- listHereDBs env
    assertEqual "All deleted" [] (map database_repo localDBs)

ageCountersTestEx
  :: HasCallStack
  => SomeShardManager
  -> (Env -> [BS.ByteString] -> IO ())
  -> IO ()
ageCountersTestEx shardManager k =
  withTest setupBasicDBs setupBasicCloudDBs $ \evb cfgAPI dbdir backupdir -> do
    let cfg = (dbConfig dbdir (serverConfig backupdir)
          { config_restore = def {
              databaseRestorePolicy_enabled = True
            }, config_retention = def
          }) {cfgShardManager = \_ _ k -> k shardManager}
    withDatabases evb cfg cfgAPI $ \env -> do
      runDatabaseJanitor env
      waitDel env
      sideEffects <- runDatabaseJanitorPureish env
      let countersToPublish =
            [ c
            | PublishCounter c _ <- sideEffects
            , ".age" `BS.isSuffixOf` c || ".span" `BS.isSuffixOf` c
            ]
      k env $ sort countersToPublish

ageCountersCompleteTest :: Test
ageCountersCompleteTest = TestCase $ ageCountersTestEx
  (SomeShardManager $ shardByRepoHash (pure $ Just ["0001"]))
  $ \_ -> assertEqual
    "Should publish age counters for all newest DBs restored locally"
    ["glean.db.test.age", "glean.db.test.span"]

ageCountersOnlyNewestTest :: Test
ageCountersOnlyNewestTest = TestCase $ ageCountersTestEx
  (SomeShardManager $ shardByRepoHash (pure $ Just []))
  $ \_ -> assertEqual
      "Should not publish age for newest DBs not restored locally"
      []

ageCountersClearTest:: Test
ageCountersClearTest = TestCase $ do
  shardAssignmentRef <- newIORef ["0001"]
  ageCountersTestEx
    (SomeShardManager $ shardByRepoHash (Just <$> readIORef shardAssignmentRef))
    $ \env _ -> do
    writeIORef shardAssignmentRef []
    -- Run the Janitor to pick up the new assignment and delete unassigned DBs
    runDatabaseJanitor env
    waitDel env
    -- Run the Janitor again after DBs are deleted to clear the counters
    runDatabaseJanitor env
    counters <- getCounters
    assertBool "glean.db.test.age cleared"
      $ not (HashMap.member "glean.db.test.age" counters)

stuckTest :: Test
stuckTest = TestCase $ withFakeCloudDBs $ \evb cfgAPI dbdir backupdir -> do
  let cfg = dbConfig dbdir $ (serverConfig backupdir)
        { config_janitor_period = Just 0
        , config_restore = def {
              databaseRestorePolicy_enabled = True
            }
        , config_retention = def {
            databaseRetentionPolicy_default_retention = def {
            retention_retain_at_least = Just 2
          }
         }
        }
      cfg' = cfg {
         cfgFilterAvailableDBs  = \_ -> do
           uninterruptibleMask_ $ threadDelay maxBound
           return []
      }
  withDatabases evb cfg' cfgAPI $ \env -> do
    (_, janitorResult) <- atomically $ do
      res <- readTVar $ envDatabaseJanitor env
      maybe retry return res
    assertBool (show janitorResult) $ case janitorResult of
      JanitorStuck -> True
      _ -> False

retentionCheckMissingDepsTest :: Test
retentionCheckMissingDepsTest = TestCase $
  withTest (\_ _ -> pure ()) setupBasicCloudDBsWithMissingDependencyOne
     $ \evb cfgAPI dbdir backupdir -> do
    let cfg = dbConfig dbdir $ (serverConfig backupdir)
          { config_retention = def
            { databaseRetentionPolicy_default_retention = def
              { retention_retain_at_least = Just 1,
                retention_retain_at_most = Just 1
              }
            },
            config_restore = def
            { databaseRestorePolicy_enabled = True
            }
          }

    withDatabases evb cfg cfgAPI $ \env -> do
      runDatabaseJanitor env
      dbs <- listDBs env
      let repos = map database_repo dbs
      --- 0014 is the newest test2 db
      --- 0016 is the newest test db, but it's missing dependency
      --- so restoring next available test db 0015 and it's dependency 0013
      assertEqual "after" [ "0013", "0014","0015"] (map repo_hash repos)

main :: IO ()
main = withUnitTest $ testRunner $ TestList
  [ TestLabel "deleteOldDBs" deleteOldDBsTest
  , TestLabel "deleteIncompleteDBs" deleteIncompleteDBsTest
  , TestLabel "retainAtMost" retainAtMostTest
  , TestLabel "retainAtLeast" retainAtLeastTest
  , TestLabel "backupRestore" backupRestoreTest
  , TestLabel "openNewestDB" openNewestTest
  , TestLabel "closeIdleDBs" closeIdleDBsTest
  , TestLabel "expireDelay" expireTest
  , TestLabel "sharding" shardingTest
  , TestLabel "shardingStacks" shardingStacksTest
  , TestLabel "shardingFallback" shardingFallbackTest
  , TestLabel "shardingByRepoName" shardingByRepoNameTest
  , TestLabel "shardingExpiring" shardUnexpireTest
  , TestLabel "availableElsewhere" elsewhereTest
  , TestLabel "notAvailableElsewhere" elsewhereNotYetAvailableTest
  , TestLabel "ageCountersForAllNewestDBs" ageCountersCompleteTest
  , TestLabel "ageCountersForOnlyNewestDBs" ageCountersOnlyNewestTest
  , TestLabel "ageCountersClear" ageCountersClearTest
  , TestLabel "stuck" stuckTest
  , TestLabel "requiredPropsTest" requiredPropsTest
  , TestLabel "retentionRestoreDepsTest" retentionRestoreDepsTest
  , TestLabel "multiRetentionTest" multiRetentionTest
  , TestLabel "retentionCheckMissingDepsTest" retentionCheckMissingDepsTest
  , TestLabel "excludedPropsTest" excludedPropsTest
  ]