-- Disabling to stop warnings on HasCallStack
{-# LANGUAGE DeepSubsumption #-}
{-# OPTIONS_GHC -Wno-redundant-constraints #-}

-- This file is part of the Wire Server implementation.
--
-- Copyright (C) 2022 Wire Swiss GmbH <opensource@wire.com>
--
-- This program is free software: you can redistribute it and/or modify it under
-- the terms of the GNU Affero General Public License as published by the Free
-- Software Foundation, either version 3 of the License, or (at your option) any
-- later version.
--
-- This program is distributed in the hope that it will be useful, but WITHOUT
-- ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
-- FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
-- details.
--
-- You should have received a copy of the GNU Affero General Public License along
-- with this program. If not, see <https://www.gnu.org/licenses/>.

module Gundeck.ThreadBudget.Internal where

import Control.Exception.Safe (catchAny)
import Control.Lens
import Control.Monad.Catch (MonadCatch)
import Data.HashMap.Strict qualified as HM
import Data.Set qualified as Set
import Data.Time
import Data.UUID (UUID, toText)
import Data.UUID.V4 (nextRandom)
import Gundeck.Options
import Imports
import Prometheus (MonadMonitor)
import Prometheus qualified as Prom
import System.Logger.Class qualified as LC
import UnliftIO.Async
import UnliftIO.Exception (finally)

data ThreadBudgetState = ThreadBudgetState
  { threadBudgetLimits :: MaxConcurrentNativePushes,
    _threadBudgetRunning :: IORef BudgetMap
  }
  deriving (Generic)

-- | Store all handles for cleanup in 'watchThreadBudgetState'.
data BudgetMap = BudgetMap
  { bspent :: Int,
    bmap :: HashMap UUID (Int, Maybe (Async ()))
  }
  deriving (Eq, Generic)

-- | Instead of taking the pre-computed total budget spent of the 'BudgetMap' (O(1)), this
-- counts all the threads that are successfully running (dropping the ones that are just about
-- to try to grab a token).
--
-- WARNING: takes O(n)), use with care.  See 'bench_BudgetSpent''.
budgetSpent :: ThreadBudgetState -> IO Int
budgetSpent (ThreadBudgetState _ running) = budgetSpent' <$> readIORef running

budgetSpent' :: BudgetMap -> Int
budgetSpent' = sum . fmap fst . filter (isJust . snd) . HM.elems . bmap

cancelAllThreads :: ThreadBudgetState -> IO ()
cancelAllThreads (ThreadBudgetState _ ref) =
  readIORef ref
    >>= mapM_ cancel . mapMaybe snd . HM.elems . bmap

mkThreadBudgetState :: (HasCallStack) => MaxConcurrentNativePushes -> IO ThreadBudgetState
mkThreadBudgetState limits = ThreadBudgetState limits <$> newIORef (BudgetMap 0 HM.empty)

-- | Allocate the resources for a new action to be called (but don't call the action yet).
allocate ::
  IORef BudgetMap -> UUID -> Int -> (MonadIO m) => m Int
allocate ref key newspent =
  atomicModifyIORef' ref $
    \(BudgetMap spent hm) ->
      ( BudgetMap (spent + newspent) (HM.insert key (newspent, Nothing) hm),
        spent
      )

-- | Register an already-allocated action with its 'Async'.
register ::
  IORef BudgetMap -> UUID -> Async () -> (MonadIO m) => m Int
register ref key handle =
  atomicModifyIORef' ref $
    \(BudgetMap spent hm) ->
      ( BudgetMap spent (HM.adjust (_2 ?~ handle) key hm),
        spent
      )

-- | Remove an registered and/or allocated action from a 'BudgetMap'.
unregister ::
  IORef BudgetMap -> UUID -> (MonadIO m) => m ()
unregister ref key =
  atomicModifyIORef' ref $
    \bhm@(BudgetMap spent hm) ->
      case HM.lookup key hm of
        Just (newspent, _) -> (BudgetMap (spent - newspent) (HM.delete key hm), ())
        Nothing -> (bhm, ())

-- | If there is budget available, execute the action synchronously; otherwise, log a warning
-- and return immediately.  Make sure the budget state is updated accordingly both when
-- starting and ending the execution.
--
-- The hard limit in the 'ThreadBudgetState' argument is guaranteed to be an upper bound for
-- the number of concurrently spent budget tokens; surpassing the soft limit will trigger a
-- warning, but still execute the action.  One action can use up any integer number of budget
-- tokens, including 0 and negative.
--
-- The action is called in an 'Async', but 'runWithBudget' waits for it to finish so it can
-- update the budget.
runWithBudget ::
  forall m.
  (LC.MonadLogger m, MonadUnliftIO m, MonadMonitor m) =>
  ThreadBudgetState ->
  Int ->
  m () ->
  m ()
runWithBudget tbs spent = runWithBudget' tbs spent ()

-- | More flexible variant of 'runWithBudget' that allows the action to return a value.  With
-- a default in case of budget exhaustion.
runWithBudget' ::
  forall m a.
  (MonadIO m, LC.MonadLogger m, MonadUnliftIO m, MonadMonitor m) =>
  ThreadBudgetState ->
  Int ->
  a ->
  m a ->
  m a
runWithBudget' (ThreadBudgetState limits ref) spent fallback action = do
  key <- liftIO nextRandom
  (`finally` unregister ref key) $ do
    oldsize <- allocate ref key spent
    let softLimitBreached = maybe False (oldsize >=) (limits ^. soft)
        hardLimitBreached = maybe False (oldsize >=) (limits ^. hard)
    warnNoBudget softLimitBreached hardLimitBreached oldsize
    if maybe True (oldsize <) (limits ^. hard)
      then go key oldsize
      else pure fallback
  where
    go :: UUID -> Int -> m a
    go key oldsize = do
      LC.debug $
        "key"
          LC..= toText key
          LC.~~ "spent"
          LC..= oldsize
          LC.~~ LC.msg (LC.val "runWithBudget: go")
      handle <- async action
      _ <- register ref key (void handle)
      wait handle
    -- iff soft and/or hard limit are breached, log a warning-level message.
    warnNoBudget :: Bool -> Bool -> Int -> m ()
    warnNoBudget False False _ = pure ()
    warnNoBudget soft' hard' oldsize = do
      let limit :: ByteString = if hard' then "hard" else "soft"
          counter =
            if hard'
              then threadBudgetHardLimitBreachedCounter
              else threadBudgetSoftLimitBreachedCounter
      Prom.incCounter counter
      LC.warn $
        "spent"
          LC..= show oldsize
          LC.~~ "soft-breach"
          LC..= soft'
          LC.~~ "hard-breach"
          LC..= hard'
          LC.~~ LC.msg ("runWithBudget: " <> limit <> " limit reached")

-- | Fork a thread that checks with the given frequency if any async handles stored in the
-- state are stale (ie., have terminated with or without exception, but not been removed).  If
-- that happens, log a warning.
--
-- 'runWithBudget' should keep track of the state itself; 'watchThreadBudgetState' is solely a
-- safety precaution to see if there aren't any corner cases we missed.
--
-- Also, issue some metrics.
watchThreadBudgetState ::
  forall m.
  (MonadIO m, LC.MonadLogger m, MonadCatch m, MonadMonitor m) =>
  ThreadBudgetState ->
  NominalDiffTime ->
  m ()
watchThreadBudgetState (ThreadBudgetState limits ref) freq = safeForever $ do
  recordMetrics limits ref
  removeStaleHandles ref
  threadDelayNominalDiffTime freq

recordMetrics ::
  forall m.
  (MonadIO m, MonadMonitor m) =>
  MaxConcurrentNativePushes ->
  IORef BudgetMap ->
  m ()
recordMetrics limits ref = do
  (BudgetMap spent _) <- readIORef ref
  Prom.setGauge threadBudgetAllocatedGauge (fromIntegral spent)
  forM_ (limits ^. hard) $ \lim ->
    Prom.setGauge threadBudgetHardLimitGauge (fromIntegral lim)
  forM_ (limits ^. soft) $ \lim ->
    Prom.setGauge threadBudgetSoftLimitGauge (fromIntegral lim)

{-# NOINLINE threadBudgetAllocatedGauge #-}
threadBudgetAllocatedGauge :: Prom.Gauge
threadBudgetAllocatedGauge =
  Prom.unsafeRegister $
    Prom.gauge
      Prom.Info
        { Prom.metricName = "net_nativepush_thread_budget_allocated",
          Prom.metricHelp = "Number of allocated threads for native pushes"
        }

{-# NOINLINE threadBudgetHardLimitGauge #-}
threadBudgetHardLimitGauge :: Prom.Gauge
threadBudgetHardLimitGauge =
  Prom.unsafeRegister $
    Prom.gauge
      Prom.Info
        { Prom.metricName = "net_nativepush_thread_budget_hard_limit",
          Prom.metricHelp = "Hard limit for threads for native pushes"
        }

{-# NOINLINE threadBudgetSoftLimitGauge #-}
threadBudgetSoftLimitGauge :: Prom.Gauge
threadBudgetSoftLimitGauge =
  Prom.unsafeRegister $
    Prom.gauge
      Prom.Info
        { Prom.metricName = "net_nativepush_thread_budget_soft_limit",
          Prom.metricHelp = "Soft limit for threads for native pushes"
        }

{-# NOINLINE threadBudgetHardLimitBreachedCounter #-}
threadBudgetHardLimitBreachedCounter :: Prom.Counter
threadBudgetHardLimitBreachedCounter =
  Prom.unsafeRegister $
    Prom.counter
      Prom.Info
        { Prom.metricName = "net_nativepush_thread_budget_hard_limit_breached",
          Prom.metricHelp = "Number of times hard limit for threads for native pushes was breached"
        }

{-# NOINLINE threadBudgetSoftLimitBreachedCounter #-}
threadBudgetSoftLimitBreachedCounter :: Prom.Counter
threadBudgetSoftLimitBreachedCounter =
  Prom.unsafeRegister $
    Prom.counter
      Prom.Info
        { Prom.metricName = "net_nativepush_thread_budget_soft_limit_breached",
          Prom.metricHelp = "Number of times soft limit for threads for native pushes was breached"
        }

threadDelayNominalDiffTime :: NominalDiffTime -> (MonadIO m) => m ()
threadDelayNominalDiffTime = threadDelay . round . (* 1000000) . toRational

staleTolerance :: NominalDiffTime
staleTolerance = 3

-- | Get all handles for asyncs that have terminated, but not been removed from the state.  Do
-- that again after 'staleTolerance' to make sure that we don't catch any handles that would
-- have been removed during the 'runWithBudget' cleanup, but we were faster.  The intersection
-- between the two rounds constitutes the legitimately stale handles: warn about them, and
-- then remove them from the budgetmap.
removeStaleHandles ::
  forall m.
  (MonadIO m, LC.MonadLogger m) =>
  IORef BudgetMap ->
  m ()
removeStaleHandles ref = do
  round1 <- getStaleHandles
  threadDelayNominalDiffTime staleTolerance
  round2 <- getStaleHandles
  let staleHandles = Set.intersection round1 round2
  unless (null staleHandles) $ do
    warnStaleHandles (Set.size staleHandles) =<< readIORef ref
    forM_ staleHandles $ \key -> do
      (mapM_ waitCatch . (snd =<<)) . HM.lookup key . bmap =<< readIORef ref
      unregister ref key
  isSanitary <- (\bm -> bspent bm == budgetSpent' bm) <$> readIORef ref
  unless isSanitary . LC.warn . LC.msg . LC.val $
    "watchThreadBudgetState: total overall thread budget diverged from async weights (repaired)."
  where
    getStaleHandles :: m (Set UUID)
    getStaleHandles =
      Set.fromList . mconcat <$> do
        handles <- HM.toList . bmap <$> readIORef ref
        forM handles $ \case
          (_, (_, Nothing)) -> do
            pure []
          (key, (_, Just handle)) -> do
            status <- poll handle
            pure [key | isJust status]
    warnStaleHandles :: Int -> BudgetMap -> m ()
    warnStaleHandles num (BudgetMap spent _) =
      LC.warn $
        "spent"
          LC..= show spent
          LC.~~ LC.msg ("watchThreadBudgetState: removed " <> show num <> " stale handles.")

safeForever ::
  forall m.
  (MonadIO m, LC.MonadLogger m, MonadCatch m) =>
  m () ->
  m ()
safeForever action =
  forever $
    action `catchAny` \exc -> do
      LC.err $ "error" LC..= displayException exc LC.~~ LC.msg (LC.val "watchThreadBudgetState: crashed; retrying")
      threadDelay 60000000 -- pause to keep worst-case noise in logs manageable
