content: {-
    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.
  -}
  
  -- |
  -- This layer provides synchronous writing on top of the lower-level
  -- asynchronous write API, by doing polling and waiting as necessary.
  --
  module Glean.Write.SendBatch
    ( sendBatch
    , sendBatchAsync
    , sendJsonBatch
    , sendJsonBatchAsync
    , sendBatchDescriptor
    , sendBatchDescriptorAsync
    , waitBatch
    ) where
  
  import Control.Exception
  import Data.Default
  import qualified Data.Text as Text
  import System.Time.Extra
  
  import Glean.Backend.Types
  import qualified Glean.Types as Thrift
  
  sendBatch
    :: Backend be
    => be
    -> Thrift.Repo
    -> Thrift.Batch
    -> IO Thrift.Subst
  sendBatch backend repo batch = do
    handle <- sendBatchAsync backend repo batch
    waitBatch backend handle
  
  sendBatchAsync
    :: Backend be
    => be
    -> Thrift.Repo
    -> Thrift.Batch
    -> IO Thrift.Handle
  sendBatchAsync backend repo batch = do
    r <- enqueueBatch backend $ Thrift.ComputedBatch repo True batch
    case r of
      Thrift.SendResponse_handle h -> return h
      Thrift.SendResponse_retry (Thrift.BatchRetry r) ->
        retry r $ sendBatchAsync backend repo batch
  
  sendJsonBatchAsync
    :: Backend be
    => be
    -> Thrift.Repo
    -> [Thrift.JsonFactBatch]
    -> Maybe Thrift.SendJsonBatchOptions
    -> IO Thrift.Handle
  sendJsonBatchAsync backend repo batches opts = do
    r <- try $ enqueueJsonBatch backend repo
      Thrift.SendJsonBatch
        { Thrift.sendJsonBatch_batches = batches
        , Thrift.sendJsonBatch_options = opts
        , Thrift.sendJsonBatch_remember = True }
    case r of
      Right Thrift.SendJsonBatchResponse{..} ->
        return sendJsonBatchResponse_handle
      Left Thrift.Retry{..} ->
        retry retry_seconds $
          sendJsonBatchAsync backend repo batches opts
  
  sendJsonBatch
    :: Backend be
    => be
    -> Thrift.Repo
    -> [Thrift.JsonFactBatch]
    -> Maybe Thrift.SendJsonBatchOptions
    -> IO Thrift.Subst
  sendJsonBatch backend repo batches opts = do
    handle <- sendJsonBatchAsync backend repo batches opts
    -- Cope with talking to an old server where this API was synchronous
    -- and didn't return a handle.
    if Text.null handle
      then return def
      else waitBatch backend handle
  
  sendBatchDescriptor
    :: Backend be
    => be
    -> Thrift.Repo
    -> Thrift.BatchDescriptor
    -> IO Thrift.Subst
  sendBatchDescriptor backend repo descriptor = do
    handle <- sendBatchDescriptorAsync backend repo descriptor
    if Text.null handle
      then return def
      else waitBatch backend handle
  
  sendBatchDescriptorAsync
    :: Backend be
    => be
    -> Thrift.Repo
    -> Thrift.BatchDescriptor
    -> IO Thrift.Handle
  sendBatchDescriptorAsync backend repo descriptor = do
    let batch = Thrift.EnqueueBatch_descriptor descriptor
    r <- try $ enqueueBatchDescriptor backend repo batch
      Thrift.EnqueueBatchWaitPolicy_Remember
    case r of
      Right Thrift.EnqueueBatchResponse{..} ->
        return enqueueBatchResponse_handle
      Left Thrift.Retry{..} ->
        retry retry_seconds $
          sendBatchDescriptorAsync backend repo descriptor
  
  retry :: Double -> IO a -> IO a
  retry secs action = do
    sleep $ max secs 1
    action
  
  waitBatch
    :: Backend be
    => be
    -> Thrift.Handle
    -> IO Thrift.Subst
  waitBatch backend handle = do
    e <- try $ pollBatch backend handle
    case e of
      Left Thrift.Retry{..} ->
        retry retry_seconds $ waitBatch backend handle
      Right r -> case r of
        Thrift.FinishResponse_subst subst -> return subst
        Thrift.FinishResponse_retry (Thrift.BatchRetry r) ->
          retry r $ waitBatch backend handle
  