{-
  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.CompletePredicates
  ( completePredicates
  , syncCompletePredicates
  ) where

import Control.Concurrent.Async
import Control.Exception
import Control.Monad
import qualified Data.HashMap.Strict as HashMap
import qualified Data.Text as Text
import System.Timeout

import Util.Control.Exception
import Util.Defer
import Util.Log.Text
import Util.Logger
import Util.STM

import qualified Glean.Database.Catalog as Catalog
import Glean.Database.Open
import Glean.Database.Schema
import Glean.Database.Schema.Types
  ( lookupPredicateSourceRef
  , SchemaSelector(..)
  )
import qualified Glean.Database.Storage as Storage
import Glean.Database.Types
import Glean.Internal.Types as Thrift
import Glean.Logger
import Glean.RTS.Foreign.Ownership
import Glean.RTS.Types (Pid)
import Glean.Schema.Util (convertRef)
import Glean.Types as Thrift
import qualified Glean.Util.Warden as Warden
import Glean.Util.Mutex


-- | Client API: kick off the completion process for a DB asynchronously
completePredicates
  :: Env
  -> Repo
  -> Thrift.CompletePredicates
  -> IO CompletePredicatesResponse
completePredicates env repo complete =
  case complete of
    Thrift.CompletePredicates_axiom _ ->
      completeAxiomPredicates env repo
    Thrift.CompletePredicates_derived (Thrift.CompleteDerivedPredicate pred) ->
      completeDerivedPredicate env repo pred

-- For internal use: actually perform completion for a DB
syncCompletePredicates :: Env -> Repo -> IO ()
syncCompletePredicates env repo =
  -- log this, because the completePredicates request itself is async
  -- and doesn't record the total time spent.
  loggingAction
    (runLogRepo "completePredicates(server)" env repo) (const mempty) $ do
  maybeBase <- repoParent env repo
  let withBase repo f =
        readDatabase env repo $ \_ lookup -> f (Just lookup)
  maybe ($ Nothing) withBase maybeBase $ \base -> do
  withOpenDatabase env repo $ \OpenDB{..} -> do
    own <- Storage.computeOwnership odbHandle base
      (schemaInventory odbSchema)
    withWriteLock odbWriting $ \lock ->
      Storage.storeOwnership odbHandle lock own
    maybeOwnership <- readTVarIO odbOwnership
    forM_ maybeOwnership $ \ownership -> do
      stats <- getOwnershipStats ownership
      logInfo $ "ownership propagation complete: " <> showOwnershipStats stats

completeAxiomPredicates :: Env -> Repo -> IO CompletePredicatesResponse
completeAxiomPredicates env@Env{..} repo = do
  let
    doCompletion = -- we are masked in here
      (`finally` deregister) $ do
      r <- tryAll $ syncCompletePredicates env repo
      case r of
        Left (e :: SomeException) -> do
          setBroken repo envCatalog "completePredicates" e
          throwIO e
        Right{} -> setComplete

    deregister =
      atomically $ modifyTVar envCompleting $ HashMap.delete repo

    setComplete = void $ atomically $
      Catalog.modifyMeta envCatalog repo $ \meta ->
        return meta { metaAxiomComplete = True }

    isInProgress = do
      completing <- now $ readTVar envCompleting
      return (HashMap.lookup repo completing)

    storeComputation async =
      modifyTVar envCompleting (HashMap.insert repo async)

  scheduleCompletion
    env repo SkipIfComplete doCompletion isInProgress storeComputation

-- | Propagate ownership information for an externally derived predicate.
syncCompleteDerivedPredicate :: Env -> Repo -> Pid -> IO ()
syncCompleteDerivedPredicate env repo pid =
  withOpenDatabase env repo $ \OpenDB{..} -> do
  maybeOwnership <- readTVarIO odbOwnership
  forM_ maybeOwnership $ \ownership -> do
    maybeBase <- repoParent env repo
    let withBase repo f = readDatabase env repo $ \_ lookup -> f (Just lookup)
    maybe ($ Nothing) withBase maybeBase $ \base ->
      withWriteLock odbWriting $ \lock -> do
        computed <- Storage.computeDerivedOwnership
          odbHandle lock ownership base pid
        Storage.storeOwnership odbHandle lock computed

withWriteLock
  :: Maybe Writing
  -> (forall w . Storage.WriteLock w -> IO b)
  -> IO b
withWriteLock Nothing f = f (undefined :: Storage.WriteLock ())
  -- if there's no write lock, we must be in the finalization
  -- phase and the DB has already been marked read-only. We have
  -- exclusive write access at this point so it's safe to
  -- continue.
withWriteLock (Just writing) f =
  withMutexSafe (wrLock writing) f

-- | Kick off completion of externally derived predicate asynchronously
completeDerivedPredicate
  :: Env
  -> Repo
  -> PredicateRef
  -> IO CompletePredicatesResponse
completeDerivedPredicate env@Env{..} repo pred = do
  details <- withOpenDatabase env repo $ \odb ->
    predicateDetails (odbSchema odb) pred
  completing <- readTVarIO envCompletingDerived
  let
      doCompletion = do -- we are masked in here
        r <- tryAll $
          syncCompleteDerivedPredicate env repo (predicatePid details)
        case r of
          Left (e :: SomeException) -> do
            setBroken repo envCatalog "completeDerivedPredicate" e
            throwIO e
          Right{} -> return ()

      derivations = HashMap.lookupDefault mempty repo completing

      predId = predicateId details

      isInProgress = do
        return (HashMap.lookup predId derivations)

      storeComputation async =
        modifyTVar envCompletingDerived $
          HashMap.insert repo $
          HashMap.insert predId async derivations

  scheduleCompletion
    env repo FailIfNotComplete doCompletion isInProgress storeComputation
  where
  predicateDetails schema pred =
    case lookupPredicateSourceRef (convertRef pred) LatestSchema schema of
      Right details -> return details
      Left err ->
        throwIO $ Thrift.Exception $ "completeDerivedPredicate: " <> err

data OnAxiomComplete
  = SkipIfComplete
  | FailIfNotComplete

scheduleCompletion
  :: Env
  -> Repo
  -> OnAxiomComplete
  -> IO ()
  -> Defer IO STM (Maybe (Async ()))
  -> (Async () -> STM ())
  -> IO CompletePredicatesResponse
scheduleCompletion Env{..} repo onAxiomComplete
  doCompletion isInProgress storeComputation = do
    mask_ $ do
    -- speculatively spawn a thread to do the completion, we'll cancel
    -- this if we don't need it. This is so that we can atomically
    -- start the job and update the Env state together.
    tmvar <- newEmptyTMVarIO
    async <- Warden.spawn envWarden $
      atomically (takeTMVar tmvar) >> doCompletion
    join $ immediately $ do
      meta <- now $ Catalog.readMeta envCatalog repo
      if
        | SkipIfComplete <- onAxiomComplete, metaAxiomComplete meta -> do
          later $ cancel async -- already done
          return (return CompletePredicatesResponse{})
        | FailIfNotComplete <- onAxiomComplete, not $ metaAxiomComplete meta ->
          now $ throwSTM $ Thrift.Exception "DB is not complete."
        | Broken b <- metaCompleteness meta -> do
          later $ cancel async
          now $ throwSTM $ Exception $ databaseBroken_reason b
        | envReadOnly -> do
          later $ cancel async
          now $ throwSTM $ Exception "DB is read-only"
        | otherwise -> do
          inProgress <- isInProgress
          case inProgress of
            Just existingAsync -> do  -- in progress
              later $ cancel async
              return $ waitFor existingAsync CompletePredicatesResponse{}
            Nothing -> now $ do  -- start async completion
              void $ tryPutTMVar tmvar ()
              storeComputation async
              return $ waitFor async CompletePredicatesResponse{}

-- If a synchronous exception is thrown during completion, this is
-- not recoverable, so mark the DB as broken. The exception is also
-- propagated to the caller.
setBroken :: Repo -> Catalog.Catalog -> Text.Text -> SomeException -> IO ()
setBroken repo catalog context e = void $ atomically $
  Catalog.modifyMeta catalog repo $ \meta ->
    return meta {
      metaCompleteness = Broken
        (DatabaseBroken context (Text.pack (show e)))
      }

waitFor :: Async () -> b -> IO b
waitFor async result = do
  r <- timeout 1000000 $ wait async
  case r of
    Nothing -> throwIO $ Retry 1
    Just _ -> return result