{-
  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.Write.SendQueue
  ( SendQueueSettings(..)
  , SendQueueEvent(..)
  , SendQueue
  , Callback
  , withSendQueue
  , writeSendQueue
  , writeSendQueueJson
  , writeSendQueueDescriptor
  ) where

import qualified Control.Concurrent.Async as Async
import Control.Exception
import Control.Monad
import qualified Data.ByteString as BS
import Data.Default
import System.Time.Extra (sleep)
import Text.Printf (printf)

import Util.Control.Exception (tryBracket)
import Util.Log.String
import Util.STM
import Util.Time

import Glean.Backend.Types (Backend)
import Glean.Backend.Retry
import qualified Glean.Backend.Types as Backend
import qualified Glean.Types as Thrift
import Glean.Write.SendBatch
import Glean.Util.RetryChannelException


-- | An event that happens in a 'SendQueue' and can be logged
data SendQueueEvent
  = -- | Started sending a batch
    SendQueueSending Thrift.Batch

    -- | Started sending a JSON batch
  | SendQueueSendingJson [Thrift.JsonFactBatch]

    -- | Started sending a batch descriptor
  | SendQueueSendingDescriptor Thrift.BatchDescriptor

    -- | Finished sending a batch
  | SendQueueSent
      Int -- batch size
      DiffTimePoints -- time taken

    -- | Finished sending all batches and shutting down
  | SendQueueFinished

    -- | Sending failed
  | SendQueueFailed SomeException


data QueueStatus
  = Open
  | Closed
  | Failed SomeException

-- | When the send operation fails or is completed this is called once.
type Callback = Either SomeException Thrift.Subst -> STM ()

-- | The GleanService can return a Handle for an async operation, and the
-- client can later poll to determine the actual result/retry/error.
data Wait = Wait
  { waitHandle :: !Thrift.Handle  -- ^ handle from GleanService
  , waitStart :: !TimePoint  -- ^ When this 'Wait' was constructed
  , waitCallback :: !Callback  -- ^ use-once callback with result.
  , waitOriginalBatch :: !Batch  -- ^ the batch, so we can resend it
  }

data Batch
  = BinaryBatch Thrift.Batch
  | JsonBatch !Int {- size -} [Thrift.JsonFactBatch]
  | BatchDescriptor Thrift.BatchDescriptor

batchSize :: Batch -> Int
batchSize (BinaryBatch bin) = BS.length $ Thrift.batch_facts bin
batchSize (JsonBatch sz _) = sz
batchSize (BatchDescriptor _) = 0

data SendQueue = SendQueue
  { -- | Batches we haven't sent yet
    sqOutQueue :: TQueue (Batch, Callback)

    -- | Batches we've sent and are now waiting for the server to commit
  , sqWaitQueue :: TQueue Wait

    -- | Memory used by batches in sqOutQueue and sqWaitQueue
  , sqMemory :: TVar Int

    -- | Number of batches in flight (in sqOutQueue and sqWaitQueue)
  , sqCount :: TVar Int

    -- | Limit on sqMemory
  , sqMaxMemory :: Int

    -- | Limit on sqCount
  , sqMaxCount :: Int

    -- | Queue status
  , sqStatus :: TVar QueueStatus
  }

data WriteQueueClosed = WriteQueueClosed deriving(Show)
instance Exception WriteQueueClosed

newSendQueue :: Int -> Int -> IO SendQueue
newSendQueue max_mem max_count = SendQueue
  <$> newTQueueIO
  <*> newTQueueIO
  <*> newTVarIO 0
  <*> newTVarIO 0
  <*> pure max_mem
  <*> pure max_count
  <*> newTVarIO Open

releaseBatch :: SendQueue -> Int -> STM ()
releaseBatch sq size = do
  modifyTVar' (sqCount sq) $ \n -> n - 1
  modifyTVar' (sqMemory sq) $ \n -> n - size

acquireBatch :: SendQueue -> Int -> STM ()
acquireBatch sq size = do
  mem <- readTVar $ sqMemory sq
  count <- readTVar $ sqCount sq
  when (mem > sqMaxMemory sq || count > sqMaxCount sq) retry
  writeTVar (sqMemory sq) $! mem + size
  writeTVar (sqCount sq) $! count + 1

writeSendQueue :: SendQueue -> Thrift.Batch -> Callback -> STM ()
writeSendQueue sq batch callback =
  writeSendQueue_ sq (BinaryBatch batch) callback

writeSendQueueJson :: SendQueue -> [Thrift.JsonFactBatch] -> Callback -> STM ()
writeSendQueueJson sq json callback =
  writeSendQueue_ sq (JsonBatch size json) callback
  where
  -- size is approximate for JSON
  size = sum (map jsonFactBatchSize json)
  jsonFactBatchSize Thrift.JsonFactBatch{..} =
    sum (map BS.length jsonFactBatch_facts) +
    maybe 0 BS.length jsonFactBatch_unit

writeSendQueueDescriptor
  :: SendQueue
  -> Thrift.BatchDescriptor
  -> Callback
  -> STM ()
writeSendQueueDescriptor sq descriptor callback =
  writeSendQueue_ sq (BatchDescriptor descriptor) callback

writeSendQueue_ :: SendQueue -> Batch -> Callback -> STM ()
writeSendQueue_ sq batch callback = do
  status <- readTVar $ sqStatus sq
  case status of
    Open -> return ()
    Closed -> throwSTM WriteQueueClosed
    Failed exc -> throwSTM exc
  let !size = batchSize batch
  acquireBatch sq size
  writeTQueue (sqOutQueue sq) (batch, callback)

closeSendQueue :: SendQueue -> STM ()
closeSendQueue sq = do
  status <- readTVar $ sqStatus sq
  case status of
    Open -> writeTVar (sqStatus sq) Closed
    Closed -> return ()
    Failed exc -> throwSTM exc

failSendQueue :: SendQueue -> SomeException -> STM ()
failSendQueue sq exc = do
  writeTVar (sqStatus sq) $ Failed exc
  xs <- map waitCallback <$> flushTQueue (sqWaitQueue sq)
  ys <- map snd <$> flushTQueue (sqOutQueue sq)
  mapM_ ($ Left exc) (xs ++ ys)

readSendQueue :: SendQueue -> STM (Maybe (Batch, Callback))
readSendQueue sq = do
  status <- readTVar $ sqStatus sq
  case status of
    Open -> get
    Closed -> do
      -- Don't exit the sendFromQueue until all batches are safely
      -- written. We might have to resend batches if the server
      -- forgets them.
      count <- readTVar $ sqCount sq
      if count /= 0 then get else return Nothing
    Failed exc -> throwSTM exc
  where
    get = do
      (batch, callback) <- readTQueue $ sqOutQueue sq
      return $ Just (batch, callback)

readWaitQueue :: SendQueue -> STM (Maybe Wait)
readWaitQueue sq = do
  status <- readTVar $ sqStatus sq
  case status of
    Open -> get
    Closed -> do
      -- NOTE: There can still be outstanding batches even if both queues are
      -- empty - sendFromQueue might be writing one.
      count <- readTVar $ sqCount sq
      if count /= 0 then get else return Nothing
    Failed exc -> throwSTM exc
  where
    get = Just <$> readTQueue (sqWaitQueue sq)

pollFromWaitQueue
  :: Backend.Backend e => e -> SendQueueSettings -> SendQueue -> IO ()
pollFromWaitQueue backend settings sq = do
  cont <- tryBracket
    (atomically $ readWaitQueue sq)
    (\r ex -> case (ex,r) of
      (Left exc, Just wait) -> do
        logWarning $ "Thread killed: " <> show exc
        atomically $ do
          waitCallback wait $ Left exc
          failSendQueue sq exc
        sendQueueLog settings $ SendQueueFailed exc
      _ -> return ())
    $ \r ->
    case r of
      Just Wait{..} -> do
        result <- try $ waitBatch backend waitHandle
        let !size = batchSize waitOriginalBatch
        case result of
          Left e@Thrift.UnknownBatchHandle{} -> do
            logWarning $ "Server forgot batch; resending (" <> show e <> ")"
            atomically $ do
              -- don't writeSendQueue; retries should work even if the
              -- queue is in the Closed state and shutting down
              writeTQueue (sqOutQueue sq) (waitOriginalBatch, waitCallback)
              return True
          Right subst -> do
            atomically $ do
              waitCallback $ Right subst
              releaseBatch sq size
            elapsed <- getElapsedTime waitStart
            sendQueueLog settings $ SendQueueSent size elapsed
            return True

      Nothing -> return False

  when cont $ pollFromWaitQueue backend settings sq

sendFromQueue
  :: Backend.Backend e
  => e
  -> Thrift.Repo
  -> SendQueueSettings
  -> SendQueue
  -> IO ()
sendFromQueue backend repo settings sq = do
  cont <- tryBracket
    (atomically $ readSendQueue sq)
    (\r ex -> case (ex,r) of
      (Left exc, Just (_, callback)) -> do
        logWarning $ "Thread killed: " <> show exc
        atomically $ do
          callback $ Left exc
          failSendQueue sq exc
        sendQueueLog settings $ SendQueueFailed exc
      _ -> return ())
    $ \r ->
    case r of
      Just (batch, callback) -> do
        sendQueueLog settings $ case batch of
          BinaryBatch bin -> SendQueueSending bin
          JsonBatch _ json -> SendQueueSendingJson json
          BatchDescriptor descriptor -> SendQueueSendingDescriptor descriptor
        start <- getTimePoint
        handle <- case batch of
          BinaryBatch bin -> sendBatchAsync backend repo bin
          JsonBatch _ json -> sendJsonBatchAsync backend repo json Nothing
          BatchDescriptor descriptor
            -> sendBatchDescriptorAsync backend repo descriptor
        atomically $ writeTQueue (sqWaitQueue sq) Wait
          { waitHandle = handle
          , waitStart = start
          , waitCallback = callback
          , waitOriginalBatch = batch
          }
        return True

      Nothing -> return False
  when cont $ sendFromQueue backend repo settings sq

-- | Settings for a 'SendQueue'
data SendQueueSettings = SendQueueSettings
  { -- | Max memory that the send queue should occupy (soft limit)
    sendQueueMaxMemory :: !Int

    -- | Max number of batches that the queue should hold
  , sendQueueMaxBatches :: !Int

    -- | Number of threads to use for sending
  , sendQueueThreads :: !Int

    -- | A function called whenever a 'SendQueueEvent' occurs
  , sendQueueLog :: SendQueueEvent -> IO ()

    -- | How to retry failures
  , sendQueueRetry :: RetryPolicy

    -- | Optionally log stats every 10s
  , sendQueueLogStats :: Bool
  }

instance Show SendQueueSettings where
  show _ = "SendQueueSettings{}"

instance Default SendQueueSettings where
  def = SendQueueSettings
    { sendQueueMaxMemory = 1000000000
    , sendQueueMaxBatches = 10000
    , sendQueueThreads = 1
    , sendQueueLog = const $ return ()
    , sendQueueRetry = defaultRetryPolicy
    , sendQueueLogStats = False
    }

withSendQueue
  :: Backend e
  => e
  -> Thrift.Repo
  -> SendQueueSettings
  -> (SendQueue -> IO a)
  -> IO a
withSendQueue backend repo settings@SendQueueSettings{..} action =
  mask $ \restore -> do
    q <- newSendQueue sendQueueMaxMemory sendQueueMaxBatches
    let retryBackend = backendRetryWrites backend sendQueueRetry
    withStats sendQueueLogStats q $ do
      snd <$> Async.concurrently
        (Async.mapConcurrently_ id
          $ restore (pollFromWaitQueue retryBackend settings q)
          : replicate
              sendQueueThreads
              (restore (sendFromQueue retryBackend repo settings q)))
        (restore (action q)
          `finally` atomically (closeSendQueue q))

withStats :: Bool -> SendQueue -> IO a -> IO a
withStats on SendQueue{..} act = do
  verbose <- vlogIsOn 1
  if on || verbose
    then Async.withAsync (forever $ logQueueStats >> sleep 10) $ \_ -> act
    else act
  where
  logQueueStats = do
    count <- readTVarIO sqCount
    memory <- readTVarIO sqMemory
    logInfo $ printf "count:%d/%d memory:%d/%d"
      count sqMaxCount memory sqMaxMemory