{-
  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.Finish
  ( finishDatabase
  , unfinishDatabase
  , finalizeDatabase
  , finalizeWait
  ) where

import Control.Monad.Catch
import Control.Monad
import Util.STM
import Data.Default
import qualified Data.HashMap.Strict as HashMap

import qualified Glean.Database.Catalog as Catalog
import Glean.Database.Meta
import Glean.Types as Thrift
import Glean.Database.Types
import Glean.Internal.Types as Thrift

import qualified Data.HashSet as HashSet
import qualified Glean.ServerConfig.Types as ServerConfig
import Glean.Util.Observed as Observed

-- | Tell the server that the database is complete
-- starting the finalization process.
-- Throws an exception if the database is not incomplete
-- or there are pending writes.
finishDatabase :: Env -> Repo -> IO Thrift.FinishDatabaseResponse
finishDatabase Env{..} repo  = do
  atomically $ do
    void $ Catalog.modifyMeta envCatalog repo $ \oldmeta ->
      case completenessStatus oldmeta of
        DatabaseStatus_Incomplete ->
          return oldmeta { metaCompleteness = Finalizing def }
        _ -> throwM $ DatabaseNotIncomplete $ completenessStatus oldmeta
    makeReadOnly
    return Thrift.FinishDatabaseResponse{}
  where
    -- When a DB is complete, make it read-only to prevent further
    -- writes. It is an error to call finishDatabase
    -- if there are outstanding writes in the queue.
    makeReadOnly = do
      mdb <- HashMap.lookup repo <$> readTVar envActive
      forM_ mdb $ \db -> do
        st <- readTVar (dbState db)
        case st of
          Open odb@OpenDB { odbWriting = Just Writing{..} } -> do
            -- NB. check the active counter as well as the queue,
            -- because this will tell us if there are writes currently
            -- in progress.
            active <- readTVar (writeQueueActive wrQueue)
            empty <- isEmptyTQueue (writeQueue wrQueue)
            -- If there are outstanding writes then the client is
            -- either broken or is intentionally trying to complete
            -- the DB early. But we can't complete the DB with
            -- outstanding writes, so we'll ask the client to retry
            -- the request later.
            when (active /= 0 || not empty) $
              throwM $ Exception
                "finishDatabase called but there are queued writes"
            writeTVar (dbState db) $ Open odb { odbWriting = Nothing }
          _ -> return ()


-- | Change a database's state from Complete to Incomplete.
-- WARNING! This is for testing only, and should
-- never be used on a production database.
unfinishDatabase :: Env -> Repo -> IO ()
unfinishDatabase Env{..} repo  = do
  backupPolicy <- ServerConfig.config_backup <$> Observed.get envServerConfig
  let isBackupAllowed = repo_name repo `HashSet.member`
        ServerConfig.databaseBackupPolicy_allowed backupPolicy
  if isBackupAllowed
    then do
      throwM $ Thrift.Exception
        "The backup is enabled for this Repo so we cannot unfinish it"
    else atomically $ do
      void $ Catalog.modifyMeta envCatalog repo $ \oldmeta ->
        case Thrift.metaCompleteness oldmeta of
          Thrift.Incomplete{} -> return oldmeta
          Thrift.Complete{} -> return oldmeta
            { Thrift.metaCompleteness = Thrift.Incomplete def }
          Thrift.Finalizing{} -> return oldmeta
            { Thrift.metaCompleteness = Thrift.Incomplete def }
          someState -> throwM $ Thrift.Exception
            (  "Cannot unfinish a database in state: "
            <> showCompleteness someState)

-- | Poll for finalization of a database
--   Throws Exception if the database is broken or incomplete.
finalizeDatabase :: Env -> Repo -> IO FinalizeResponse
finalizeDatabase env repo = do
  atomically $ do
    meta <- Catalog.readMeta (envCatalog env) repo
    case metaCompleteness meta of
      Finalizing{} -> throwM $ Retry 1.0
      Incomplete{} -> throwM $ Exception "incomplete database"
      Broken b -> throwM $ Exception $ databaseBroken_reason b
      _ -> return ()
  return FinalizeResponse{}

-- | Wait for finalization of a database. Use for local DBs only.
--   Throws Exception if the database is broken or incomplete.
finalizeWait :: Env -> Repo -> IO ()
finalizeWait env repo =
  atomically $ do
    meta <- Catalog.readMeta (envCatalog env) repo
    case metaCompleteness meta of
      Finalizing{} -> retry
      Incomplete{} -> throwM $ Exception "incomplete database"
      Broken b -> throwM $ Exception $ databaseBroken_reason b
      _ -> return ()