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.
  -}
  
  {- |
    Provides a configurable wrapper for retrying Thrift requests that may fail
    with ChannelException.
  -}
  module Glean.Util.RetryChannelException
    ( RetryPolicy(..)
    , defaultRetryPolicy
    , retryChannelExceptions
    ) where
  
  import Control.Exception
  import Text.Printf
  import System.Random
  import System.Time.Extra
  import Util.Log
  
  import Thrift.Channel
  
  data RetryPolicy = RetryPolicy
    { maxRetries :: Int
    , minRetryDelay :: Seconds
    , maxRetryDelay :: Seconds
    , retryJitter :: Double -- ^ fraction of delay
    , onError :: ChannelException -> Int -> Int -> Maybe Double -> IO ()
    }
  
  defaultRetryPolicy :: RetryPolicy
  defaultRetryPolicy = RetryPolicy
    { maxRetries = 5
    , minRetryDelay = 1
    , maxRetryDelay = 30
    , retryJitter = 0.2  -- 20%
    , onError = logRetry
    }
  
  logRetry :: ChannelException -> Int -> Int -> Maybe Double -> IO ()
  logRetry ex n maxRetries maybeDelay = do
    let msg = printf "failed (%s) (try %d/%d)" (show ex) n (maxRetries+1)
    logWarning $ msg <> case maybeDelay of
      Nothing -> ""
      Just delay -> printf ", retry after %.2fs" delay
  
  retryChannelExceptions :: RetryPolicy -> IO a -> IO a
  retryChannelExceptions r@RetryPolicy{..} act = go 1
    where
    go n = do -- attempt number n
      e <- try act
      case e of
        Right x -> return x
        Left ex@ChannelException{}
          | n > maxRetries -> do
            onError ex n maxRetries Nothing
            throwIO ex
          | otherwise -> do
            delay <- retryDuration r n
            onError ex n maxRetries (Just delay)
            sleep delay
            go (n+1)
  
  retryDuration :: RetryPolicy -> Int -> IO Double
  retryDuration RetryPolicy{..} n = do
    randomJitter <- randomRIO (-retryJitter, retryJitter)
    let dur = exp randomJitter * minRetryDelay * 2^(n-1)
    return $! max minRetryDelay $ min maxRetryDelay dur
  