{-
  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.
-}

-- | Server-side write queue implementation.  Currently this is just a
-- layer that makes every write operation asynchronous with support
-- for polling the status of a previous write.
--
-- Rationale: why are writes async?
--
-- The alternative is to put writes into a separate Thrift queue with
-- a lower priority, to avoid writes from blocking other
-- requests. This would work, but under heavy write load we would
-- still have a lot of concurrent open connections to the server,
-- which is not ideal.
--
-- Longer term we might want to completely separate the write queue,
-- e.g. using Scribe, so making the write API async is a step in the
-- right direction.
--
-- ToDo:
--   - make it a bounded queue
--   - add logging for the size of the queue etc.
--   - decode JSON eagerly but do writes async, so that we can
--     parallelise the JSON writing. Might run into issues with the
--     decoded JSON taking more space and adding GC overhead.

module Glean.Database.Writes
  ( enqueueBatch
  , enqueueJsonBatch
  , enqueueBatchDescriptor
  , enqueueCheckpoint
  , pollBatch
  , reapWrites
  , writerThread
  , deleteWriteQueues
  , batchOwnedSize
  , batchDependenciesSize
  ) where

import Control.Concurrent
import Control.Exception
import Control.Monad
import Control.Monad.Extra (whenM)
import Control.Trace (traceMsg)
import qualified Data.ByteString as ByteString
import Data.Default
import Data.Either
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HashMap
import qualified Data.Map as Map
import qualified Data.Monoid as Monoid
import Data.Text as Text (Text)
import qualified Data.Text.Encoding as Text
import qualified Data.UUID as UUID
import qualified Data.UUID.V4 as UUID
import qualified Data.Vector.Storable as VS
import Foreign.Storable
import System.Clock
import System.Timeout
import Text.Printf

import ServiceData.GlobalStats
import ServiceData.Types
import Util.Control.Exception
import Util.Defer
import Util.Log
import Util.STM

import Glean.Database.BatchLocation as BatchLocation
import Glean.Database.Exception
import Glean.Database.Open
import Glean.Database.Schema.Types
import Glean.Database.Trace
import Glean.Database.Write.Batch
import Glean.Database.Types
import qualified Glean.RTS.Foreign.Subst as Subst
import Glean.RTS.Foreign.Ownership (DefineOwnership)
import qualified Glean.ServerConfig.Types as ServerConfig
import Glean.Types as Thrift hiding (Database)
import Glean.Write.JSON
import Glean.Util.Metric
import Glean.Util.Observed as Observed
import Util.Time

{-
Write Queue Monitoring
----------------------

glean.db.write.failed.avg.60
  Total size of write batches that failed (kB)

glean.db.write.succeeded.avg.60
  Total size of write batches that succeeded (kB)

glean.db.write.queue.avg.60
  Total size of write queues (kB)

glean.db.write.queue.<repo>.avg.60
  Size of write queues for <repo> (kB)

glean.db.write.enqueued.avg.60
  Total size of write batches enqueued (kB)

glean.db.write.rejected.avg.60
  Total size of write batches rejected due to queue limit (kB)
-}

-- | Create threads to process the write queues
writerThread :: Env -> WriteQueues -> IO ()
writerThread env WriteQueues{..} = mask $ \restore ->
  forever $ handler restore $
    void $ tryBracket
      dequeue
      done
      execute
 where
  done (Just (WriteJob{..}, repo, WriteQueue{..})) result = do
    latency <- do
      writeEnd <- getTime Monotonic
      return (writeEnd - pointStart writeStart)
    if isLeft result then
      addStatValueType "glean.db.write.failed" (writeSize `div` k) Sum
    else
      addStatValueType "glean.db.write.succeeded" (writeSize `div` k) Sum
    immediately $ do
      now $ writeTVar writeQueueLatency latency
      now $ modifyTVar' writeQueueActive (subtract 1)
      queueCount <- now $ updateTVar writeQueueCount (subtract 1)
      qSize <- now $ updateTVar writeQueueSize (subtract writeSize)
      size <- now $ updateTVar writeQueuesSize (subtract writeSize)
      later $ reportQueueSizes repo queueCount qSize size (Just latency)
    -- don't put the MVar until we have updated writeQueueActive, otherwise
    -- there is a race condition with finish database which will fail because
    -- active /= 0.
    void $ tryPutMVar writeDone result
  done _ _ = return ()

  handler restore action = do
    r <- tryAll (restore action)
    case r of
      Left e -> do logError (show e); return ()
      Right () -> return ()

  dequeue = atomically $ dequeueLoop (pure ())

  -- Round-robin processing of write queues.
  --  * If we take a job from a queue, and it has other jobs left,
  --    we put it to the back of writeQueues
  --  * If we hit a checkpoint that's not ready, we keep looking
  --    and put the queue with the checkpoint to the back of
  --    writeQueues
  --  * But if there are only queues with checkpoints that aren't
  --    ready, we will block and not busy-wait.
  --
  -- This is an O(n) transaction which could be bad, but in
  -- practice the number of queues with checkpoints will be small.
  dequeueLoop requeueCheckpoints = do
    (repo, queue@WriteQueue{..}) <- readTQueue writeQueues
    maybeJob <- tryReadTQueue writeQueue
    case maybeJob of
      Nothing -> do requeueCheckpoints; return Nothing
      Just job -> do
        let
          requeue = do
            requeueCheckpoints
            whenM (not <$> isEmptyTQueue writeQueue) $
              writeTQueue writeQueues (repo, queue)
        case job of
          WriteJob{} -> do
            requeue
            modifyTVar' writeQueueActive (+1)
            return (Just (job, repo, queue))
          WriteCheckpoint{} -> do
            active <- readTVar writeQueueActive
            if active == 0  -- we get to execute the action
              then do requeue; return (Just (job, repo, queue))
              else do unGetTQueue writeQueue job; dequeueLoop requeue

  execute (Just (WriteJob{..}, repo, _)) = do
    writeContent <- writeContentIO
    writeDatabase env repo writeContent writeStart
  execute (Just (WriteCheckpoint io, _, _)) = do io; return Subst.empty
  execute _ = return Subst.empty

-- | Check and update the in-memory write queue size
checkMemoryAvailable
  :: Env
  -> ServerConfig.Config
  -> Int -- ^ requested size
  -> STM Bool
checkMemoryAvailable Env{..} ServerConfig.Config{..} size = do
  let WriteQueues{..} = envWriteQueues
  pending <- readTVar writeQueuesSize
  let !newSize = pending + size
  if roundUp mb newSize <= fromIntegral config_db_write_queue_limit_mb
    then do
      writeTVar writeQueuesSize newSize
      return True
    else return False

-- | Update stats and throw a Retry exception
rejectWrite :: Repo -> Int -> Double -> IO retry
rejectWrite repo size elapsed = do
  let clamped = min 1000.0 $ max 1.0 elapsed
      repoB = Text.encodeUtf8 (repo_name repo)
  addStatValueType "glean.db.write.rejected" (size `div` k) Sum
  addStatValueType ("glean.db.write.retry_ms." <> repoB)
    (round (elapsed*1000)) Avg
  throwIO (Retry clamped)

-- | Add a write job to the queue, or throw 'Retry' if no memory available
enqueueWrite
  :: Env
  -> Repo
  -> Int
  -> Maybe SchemaId
  -> Bool
  -> IO WriteContent
  -> IO (MVar (Either SomeException Subst.Subst))
enqueueWrite env@Env{..} repo size optSchemaId checkQueueSize writeContent = do
  start <- beginTick 1
  config <- Observed.get envServerConfig
  mvar <- newEmptyMVar
  withWritableDatabase env repo $ \(queue@WriteQueue{..}, odbSchema) -> do

  -- check the schema ID in the batch matches the DB
  case optSchemaId of
    Just schemaId | ServerConfig.config_check_write_schema_id config ->
      when (not (Map.member schemaId (schemaEnvs odbSchema))) $
        dbError repo $ printf
          "schema ID in batch (%s) does not match schema ID of DB (%s)"
          (unSchemaId schemaId)
          (show (Map.keys (schemaEnvs odbSchema)))
    _ -> return ()

  let WriteQueues{..} = envWriteQueues
      enqueueIt = do
        pending <- now $ readTVar writeQueuesSize
        let !newSize = pending + size
        now $ do
          addToWriteQueue
            repo
            queue
            envWriteQueues
            WriteJob
              { writeSize = size
              , writeContentIO = writeContent
              , writeDone = mvar
              , writeStart = start }
        queueCount <- now $ updateTVar writeQueueCount (+1)
        queueSize <- now $ updateTVar writeQueueSize (+ size)
        later $ do
          addStatValueType "glean.db.write.enqueued" (size `div` k) Sum
          reportQueueSizes repo queueCount queueSize newSize Nothing
  immediately $ do
    check <- if checkQueueSize
      then now $ checkMemoryAvailable env config size
      else return True
    if check then enqueueIt else do
      latency <- now $ readTVar writeQueueLatency
      let elapsed = fromIntegral (toNanoSecs latency) / 1000000000.0
      later $ rejectWrite repo size elapsed
  return mvar

-- | Add a checkpoint to the write queue, which will be performed when
-- all previous writes have completed.
enqueueCheckpoint
  :: Env
  -> Repo
  -> IO ()
  -> IO ()
enqueueCheckpoint env repo io = withWritableDatabase env repo $ \(queue, _) ->
  atomically $ void $
    addToWriteQueue repo queue (envWriteQueues env) (WriteCheckpoint io)

addToWriteQueue
  :: Repo
  -> WriteQueue
  -> WriteQueues
  -> WriteJob
  -> STM ()
addToWriteQueue repo queue@WriteQueue{..} WriteQueues{..} job = do
  wasEmpty <- isEmptyTQueue writeQueue
  writeTQueue writeQueue job
  -- if this repo previously had no writes, add it to the round-robin
  when wasEmpty $ writeTQueue writeQueues (repo, queue)

reportQueueSizes :: Repo -> Int -> Int -> Int -> Maybe TimeSpec -> IO ()
reportQueueSizes repo repoQueueCount repoQueueSize totalQueueSize mLatency = do
  let repoB = Text.encodeUtf8 (repo_name repo)
  addStatValueType ("glean.db.write.queue_count." <> repoB) repoQueueCount Avg
  addStatValueType ("glean.db.write.queue." <> repoB) (repoQueueSize `div` k)
    Avg
  addStatValueType "glean.db.write.queue" (totalQueueSize `div` k) Avg
  forM_ mLatency $ \ latency -> do
    let elapsedMilliSeconds = fromInteger (toNanoSecs latency `div` 1000000)
    when (elapsedMilliSeconds > 0) $
      addStatValueType ("glean.db.write.queue_ms." <> repoB)
        elapsedMilliSeconds Avg

enqueueBatch :: Env -> ComputedBatch -> Maybe DefineOwnership -> IO SendResponse
enqueueBatch env ComputedBatch{..} ownership = do
  let size = ByteString.length (batch_facts computedBatch_batch)
  traceMsg (envTracer env)
    (GleanTraceEnqueue computedBatch_repo EnqueueBatch size) $ do
  -- NOTE: we use UUIDs here rather than, say, consecutive
  -- numbers because we want to avoid conflicts when the
  -- server restarts/crashes
  handle <- UUID.toText <$> UUID.nextRandom

  let size = batchSize computedBatch_batch
      optSchemaId = batch_schema_id computedBatch_batch
  r <- try $ enqueueWrite env computedBatch_repo size optSchemaId True $ pure $
        (writeContentFromBatch computedBatch_batch) {
          writeOwnership= ownership
        }
  case r of
    -- ToDo: make sendBatch use Retry exceptions instead of results too
    Left (Retry n) ->
      return $ Thrift.SendResponse_retry (BatchRetry n)
    Right write -> do
     when computedBatch_remember $ rememberWrite env handle write
     return $ Thrift.SendResponse_handle handle

batchSize :: Thrift.Batch -> Int
batchSize Thrift.Batch{..} =
  ByteString.length batch_facts +
  batchOwnedSize batch_owned +
  batchDependenciesSize batch_dependencies

batchOwnedSize :: HashMap UnitName ListOfIds -> Int
batchOwnedSize = Monoid.getSum . foldMap (Monoid.Sum . storableSize)

batchDependenciesSize :: HashMap Id [FactDependencies] -> Int
batchDependenciesSize deps =
  Monoid.getSum (foldMap (Monoid.Sum . depsSize) deps)
  where
  depsSize deps = sum
    [ storableSize f + storableSize d
    | FactDependencies f d <- deps
    ]

storableSize :: forall a . VS.Storable a => VS.Vector a -> Int
storableSize = (sizeOf (undefined :: a) *) . snd . VS.unsafeToForeignPtr0

enqueueJsonBatch
  :: Env
  -> Repo
  -> Thrift.SendJsonBatch
  -> IO Thrift.SendJsonBatchResponse
enqueueJsonBatch env repo batch = do
  let
    jsonFactBatchSize JsonFactBatch{..} =
      sum (map ByteString.length jsonFactBatch_facts) +
      maybe 0 ByteString.length jsonFactBatch_unit
    size = sum (map jsonFactBatchSize (sendJsonBatch_batches batch))
  traceMsg (envTracer env) (GleanTraceEnqueue repo EnqueueJsonBatch size) $ do
  handle <- UUID.toText <$> UUID.nextRandom
  let optSchemaId =
        sendJsonBatch_options batch >>= sendJsonBatchOptions_schema_id
  write <- enqueueWrite env repo size optSchemaId True $
    writeJsonBatch env repo batch
  when (sendJsonBatch_remember batch) $ rememberWrite env handle write
  return $ def { sendJsonBatchResponse_handle = handle }

enqueueBatchDescriptor
  :: Env
  -> Repo
  -> EnqueueBatch
  -> EnqueueBatchWaitPolicy
  -> IO Thrift.EnqueueBatchResponse
enqueueBatchDescriptor env repo enqueueBatch waitPolicy = do
  traceMsg (envTracer env)
    (GleanTraceEnqueue repo EnqueueBatchDescriptor 0) $ do
  handle <- UUID.toText <$> UUID.nextRandom
  descriptor <- case enqueueBatch of
    Thrift.EnqueueBatch_descriptor descriptor -> return descriptor
    Thrift.EnqueueBatch_EMPTY -> throwIO $ Thrift.Exception "empty batch"
  write <- enqueueWrite env repo 0 Nothing False $
    writeContentFromBatch <$> downloadBatchFromLocation env descriptor
  when (waitPolicy == Thrift.EnqueueBatchWaitPolicy_Remember)
    $ rememberWrite env handle write
  return $ def { enqueueBatchResponse_handle = handle }

pollBatch :: Env -> Handle -> IO FinishResponse
pollBatch env@Env{..} handle = do
  r <- HashMap.lookup handle <$> readTVarIO envWrites
  case r of
    Just write -> do
      -- for tiny writes that will complete in a few ms, we would like
      -- to wait synchronously. Otherwise we'll return a Retry to the
      -- caller which will wait 1s before polling again. In particular
      -- all those one-second delays make tests run slowly.
      s <- timeout 100000 $ readMVar (writeWait write)
      case s of
        Just x -> do
          atomically $ void $ updateTVar envWrites $ HashMap.delete handle
          case x of
            Right subst ->
              return $ Thrift.FinishResponse_subst $ Subst.serialize subst
            Left exc -> case fromException exc of
              Just Retry{..} -> return $ Thrift.FinishResponse_retry $
                Thrift.BatchRetry retry_seconds
              Nothing -> throwIO exc
        Nothing -> do
          timeout <- getWriteTimeout env
          atomically
            $ void
            $ updateTVar envWrites
            $ HashMap.adjust
                (\w -> w { writeTimeout = timeout })
                handle
          return $ Thrift.FinishResponse_retry $ Thrift.BatchRetry 0
    Nothing -> throwIO Thrift.UnknownBatchHandle

rememberWrite
  :: Env
  -> Thrift.Handle
  -> MVar (Either SomeException Subst.Subst)
  -> IO ()
rememberWrite env@Env{..} handle write = do
  timeout <- getWriteTimeout env
  atomically
    $ modifyTVar' envWrites
    $ HashMap.insert handle Write
        { writeWait = write
        , writeTimeout = timeout
        }

getWriteTimeout :: Env -> IO TimePoint
getWriteTimeout Env{..} = do
  ServerConfig.Config{..} <- Observed.get envServerConfig
  now <- getTimePoint
  return $ addToTimePoint now $ seconds $ fromIntegral config_db_writes_keep

-- | Periodically remove write handles that have timed out. NB: The writes
-- themselves will still be executed but the substitutions they produce can
-- no longer be queried.
reapWrites :: Env -> TVar (HashMap Text Write) -> IO ()
reapWrites Env{..} writes = forever $ do
  ServerConfig.Config{..} <- Observed.get envServerConfig
  threadDelay $ fromIntegral config_db_writes_reap * 1000000
  now <- getTimePoint
  atomically $ modifyTVar' writes $ HashMap.filter $ \x -> writeTimeout x > now

deleteWriteQueues :: Env -> OpenDB s -> STM ()
deleteWriteQueues env OpenDB{odbWriting = Just Writing{..}} = do
  let !WriteQueue{..} = wrQueue
  size <- readTVar writeQueueSize
  modifyTVar' (writeQueuesSize (envWriteQueues env)) (subtract size)
  void $ flushTQueue writeQueue
  -- This WriteQueue might still be on the writeQueues, but it will
  -- get removed by the next write thread to encounter it.
deleteWriteQueues _ _ = return ()

downloadBatchFromLocation
  :: Env
  -> Thrift.BatchDescriptor
  -> IO Thrift.Batch
downloadBatchFromLocation Env{..} batchDescriptor =
  let
    batchFormat = batchDescriptor_format batchDescriptor
    locationString = batchDescriptor_location batchDescriptor
  in
    BatchLocation.fromString envBatchLocationParser locationString batchFormat


k :: Int
k = 1024

mb :: Int
mb = 1024*1024

roundUp :: Int -> Int -> Int
roundUp unit x = (x + unit-1) `div` unit