{-
  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.Write.Batch
  ( syncWriteDatabase
  , syncWriteContentDatabase
  , writeDatabase
  , writeContentFromBatch
  ) where

import Control.Exception
import Control.Monad.Extra
import Control.Trace (traceMsg)
import qualified Data.ByteString as BS
import qualified Data.HashMap.Strict as HashMap
import Data.HashMap.Strict (HashMap)
import Data.Default
import Data.Int (Int64)
import Data.IORef
import Data.Maybe
import Data.Word

import Util.Control.Exception
import Util.STM

import qualified Glean.Database.Catalog as Catalog
import Glean.Database.Open
import Glean.Database.Exception
import Glean.Database.Repo
import qualified Glean.Database.Storage as Storage
import Glean.Database.Schema
import Glean.Database.Trace
import Glean.Database.Types
import Glean.FFI
import Glean.Internal.Types as Thrift
import Glean.RTS.Foreign.Define (trustRefs)
import qualified Glean.RTS.Foreign.FactSet as FactSet
import Glean.RTS.Foreign.Lookup (Lookup, withCanLookup)
import qualified Glean.RTS.Foreign.Lookup as Lookup
import qualified Glean.RTS.Foreign.LookupCache as LookupCache
import Glean.RTS.Foreign.Ownership as Ownership
import Glean.RTS.Foreign.Stacked (stacked)
import Glean.RTS.Foreign.Subst (Subst)
import qualified Glean.RTS.Foreign.Subst as Subst
import Glean.RTS.Types (Pid(..), Fid(..))
import Glean.Types (Repo)
import qualified Glean.Types as Thrift
import Glean.Util.Metric
import Glean.Util.Mutex
import Glean.Util.Vector
import qualified Glean.Write.Stats as Stats

writeContentFromBatch :: Thrift.Batch -> WriteContent
writeContentFromBatch writeBatch = WriteContent
  { writeBatch = writeBatch
  , writeOwnership = Nothing
  }

syncWriteDatabase
  :: Env
  -> Repo
  -> Thrift.Batch
  -> IO Subst
syncWriteDatabase env repo batch =
  syncWriteContentDatabase env repo (writeContentFromBatch batch)

syncWriteContentDatabase :: Env -> Repo -> WriteContent -> IO Subst
syncWriteContentDatabase env repo content = do
  tick <- beginTick 1
  writeDatabase env repo content tick

makeDefineOwnership
  :: Env
  -> Repo
  -> Fid
  -> HashMap Int64 [Thrift.FactDependencies]
  -> IO (Maybe DefineOwnership)
makeDefineOwnership env repo nextId deps
  | HashMap.null deps = return Nothing
  | otherwise = do
  readDatabase env repo $ \odb lookup -> do
  maybeOwnership <- readTVarIO (odbOwnership odb)
  forM maybeOwnership $ \ownership -> do
    define <- Ownership.newDefineOwnership ownership nextId
    forM_ (HashMap.toList deps) $ \(pid, ownerMap) ->
      Ownership.addDerivedOwners lookup define (Pid pid) ownerMap
    return define

checkWritable :: Repo -> OpenDB s -> IO Writing
checkWritable repo OpenDB{..} =
  case odbWriting of
    Nothing -> dbError repo "can't write to a read only database"
    Just writing -> return writing

-- | We can only add fact ownership *before* @glean complete@, otherwise
-- it won't get propagated. Similarly, we can only add facts with
-- explicit dependencies *after* @glean complete@.
checkComplete :: Env -> Repo -> Thrift.Batch -> IO ()
checkComplete env repo Thrift.Batch{..} = do
  when (not $ HashMap.null batch_owned) $ do
    meta <- atomically $ Catalog.readMeta (envCatalog env) repo
    when (metaAxiomComplete meta) $
      throwIO $ Thrift.Exception
        "Attempting to write facts with ownership after 'glean complete'"
  when (not $ HashMap.null batch_dependencies) $ do
    -- TODO: enforce that not only the axiom predicates be complete
    -- but that every predicate depended upon be complete.
    meta <- atomically $ Catalog.readMeta (envCatalog env) repo
    when (not (metaAxiomComplete meta)) $
      throwIO $ Thrift.Exception
        "Attempting to write facts with dependencies before 'glean complete'"

-- | Write a batch of facts to the database, returning a substitution
writeDatabase
  :: Env
  -> Repo
  -> WriteContent
  -> Point
  -> IO Subst
writeDatabase env repo WriteContent{..} latency =
  readDatabase env repo $ \odb lookup -> do
    writing <- checkWritable repo odb
    checkComplete env repo writeBatch
    Stats.bump (envStats env) Stats.mutatorLatency =<< endTick latency
    let !size = batchSize writeBatch

    tick env repo WriteTraceInput Stats.mutatorInput size $ do
      -- If nobody is writing to the DB just write the batch directly.
      --
      -- TODO: What if someone is already deduplicating another batch? Should we
      -- not write in that case?
      r <- tryWithMutexSafe (wrLock writing) $ \lock ->
        reallyWriteBatch
          env repo odb lock lookup writing size False writeBatch writeOwnership
      case r of
        Just cont -> cont
        Nothing ->
          -- Somebody is already writing to the DB - deduplicate the batch
          deDupBatch env repo odb lookup writing size writeBatch writeOwnership

reallyWriteBatch
  :: Storage.Storage s
  => Env
  -> Repo
  -> OpenDB s
  -> Storage.WriteLock w
  -> Lookup
  -> Writing
  -> Word64  -- ^ original size of the batch
  -> Bool  -- ^ has the batch already been de-duped?
  -> Thrift.Batch
  -> Maybe DefineOwnership
  -> IO (IO Subst)
reallyWriteBatch env repo OpenDB{..} lock lookup writing original_size deduped
    batch@Thrift.Batch{..} maybeOwn = do
  let !real_size = batchSize batch
  Stats.tick (envStats env) Stats.mutatorThroughput original_size
    $ Stats.tick (envStats env)
        (if deduped
          then Stats.mutatorDedupedThroughput
          else Stats.mutatorDupThroughput) real_size
    $ do

    next_id <- readIORef (wrNextId writing)

    (facts, subst) <-
      logExceptions (\s -> inRepo repo $ "rename error: " ++ s) $
        tick env repo WriteTraceRename Stats.renameThroughput real_size $ do
          let withBase f = do
                r <- readTVarIO (wrCommit writing)
                case r of
                  Nothing -> f lookup
                  Just (_, facts) -> withCanLookup (stacked lookup facts) f
          withBase $ \base -> do
            withLookupCache writing base $ \cache -> do
              FactSet.renameFacts
                (schemaInventory odbSchema)
                cache
                next_id
                batch
                def { trustRefs = deduped }

    logExceptions (\s -> inRepo repo $ "commit error: " ++ s) $ do
      updateLookupCacheStats env

      let
        commitOwnership = do
          let apply v = unsafeCoerceVector <$>
                Subst.unsafeSubstIntervalsAndRelease subst
                  (unsafeCoerceVector v)
          owned <- mapM apply batch_owned
          Storage.addOwnership odbHandle lock owned
          deps <- mapM (substDependencies subst) batch_dependencies
          derivedOwners <-
            if | Just owners <- maybeOwn -> do
                  Ownership.substDefineOwnership owners subst
                  return $ Just owners
               | not $ HashMap.null deps ->
                  makeDefineOwnership env repo next_id deps
               | otherwise -> return Nothing
          forM_ derivedOwners $ \ownBatch ->
            Storage.addDefineOwnership odbHandle lock ownBatch

        doCommit =
          tick env repo WriteTraceCommit
            Stats.commitThroughput real_size $ do
              Storage.commit odbHandle facts

      -- we're going to perform the actual commit outside the write
      -- lock, so that the next batch can start renaming while we're
      -- committing. But the next rename will need the current batch
      -- to de-duplicate against, so we save it in wrCommit. There can
      -- only be one commit happening at a time and there's no way to
      -- queue more than one pending commit currently.
      commit <-
        if envMockWrites env
          then do release facts; return (return ())
          else do
            commitOwnership -- must be done under the write lock
            atomically $ do
               m <- readTVar (wrCommit writing)
               case m of
                 Just{} -> retry
                 Nothing -> writeTVar (wrCommit writing) (Just (next_id, facts))
            return $
              doCommit
                `finally`
              do atomically $ writeTVar (wrCommit writing) Nothing
                 withMutexSafe (wrLock writing) $ const $ release facts

      new_next_id <- Lookup.firstFreeId facts
      atomicWriteIORef (wrNextId writing) new_next_id
      return (commit >> return subst)
        -- commit takes place outside the write lock

deDupBatch
  :: Storage.Storage s
  => Env
  -> Repo
  -> OpenDB s
  -> Lookup
  -> Writing
  -> Word64
  -> Thrift.Batch
  -> Maybe DefineOwnership
  -> IO Subst
deDupBatch env repo odb lookup writing original_size
    batch@Thrift.Batch{..} maybeOwn =
  logExceptions (\s -> inRepo repo $ "dedup error: " ++ s) $ do
    next_id <- do
      r <- readTVarIO (wrCommit writing)
      -- if there is a commit in progress, we use the ID boundary
      -- before that commit for the snapshot.
      case r of
        Just (commit_id, _) -> return commit_id
        _otherwise -> readIORef $ wrNextId writing
    (maybe_deduped_batch, dsubst) <- bracket
      (
        withLookupCache writing lookup $ \cache ->
        -- We need a snapshot here because we don't want lookups to
        -- return fact ids which conflict with ids in the renamed batch.
        Lookup.withSnapshot cache next_id $ \snapshot ->
          FactSet.renameFacts
            (schemaInventory (odbSchema odb))
            snapshot
            next_id
            batch
            def
      )
      (\(deduped_facts, _) -> release deduped_facts)
        -- release the FactSet now that we're done with it, don't wait for
        -- the GC to free it.
      (\(deduped_facts, dsubst) -> do
        factCount <- FactSet.factCount deduped_facts
        if factCount == 0 && isNothing maybeOwn then
          return (Nothing, dsubst)
        else do
          deduped_batch <- FactSet.serialize deduped_facts
          return (Just deduped_batch, dsubst)
      )

    case maybe_deduped_batch of
      Nothing -> return dsubst
      Just deduped_batch -> do
        let apply v = unsafeCoerceVector <$>
              Subst.unsafeSubstIntervalsAndRelease dsubst
                (unsafeCoerceVector v)
        is <- mapM apply batch_owned
        deps <- mapM (substDependencies dsubst) batch_dependencies
        forM_ maybeOwn $ \ownBatch ->
          Ownership.substDefineOwnership ownBatch dsubst
        -- And now write it do the DB, deduplicating again
        cont <- withMutexSafe (wrLock writing) $ \lock ->
          reallyWriteBatch env repo odb lock lookup writing original_size True
            deduped_batch
              { Thrift.batch_owned = is
              , Thrift.batch_dependencies = deps
              }
            maybeOwn
        wsubst <- cont
        return $ dsubst <> wsubst

withLookupCache
  :: Writing
  -> Lookup.Lookup
  -> (Lookup.Lookup -> IO a)
  -> IO a
withLookupCache Writing{..} lookup f = do
  LookupCache.withCache lookup wrLookupCache LookupCache.FIFO f

substDependencies
 :: Subst
 -> [Thrift.FactDependencies]
 -> IO [Thrift.FactDependencies]
substDependencies subst dmap = mapM substFD dmap
  where
  substFD (Thrift.FactDependencies facts deps) = do
    Thrift.FactDependencies <$> apply facts <*> apply deps
  apply v =
    unsafeCoerceVector <$>
      Subst.substVector subst (unsafeCoerceVector v)

batchSize :: Thrift.Batch -> Word64
batchSize = fromIntegral . BS.length . Thrift.batch_facts

tick
  :: Env -> Repo -> WriteTraceEvent -> Stats.Bump Tick -> Word64 -> IO a -> IO a
tick env repo event ticker value cont =
  traceMsg (envTracer env) (GleanTraceWrite repo event value)
  $ Stats.tick (envStats env) ticker value cont