content: {-# LANGUAGE TupleSections #-}
  {-| Ad-hoc rate limiting for the JQScheduler based on reason trails.
  
  -}
  
  {-
  
  Copyright (C) 2014 Google Inc.
  All rights reserved.
  
  Redistribution and use in source and binary forms, with or without
  modification, are permitted provided that the following conditions are
  met:
  
  1. Redistributions of source code must retain the above copyright notice,
  this list of conditions and the following disclaimer.
  
  2. Redistributions in binary form must reproduce the above copyright
  notice, this list of conditions and the following disclaimer in the
  documentation and/or other materials provided with the distribution.
  
  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
  IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
  TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
  PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
  CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
  EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
  PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
  PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
  LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
  NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  
  -}
  
  module Ganeti.JQScheduler.ReasonRateLimiting
    ( reasonRateLimit
    -- * For testing only
    , parseReasonRateLimit
    , countMapFromJob
    , slotMapFromJobs
    ) where
  
  import Control.Monad.Fail (MonadFail)
  import Data.List
  import Data.Maybe
  import qualified Data.Map as Map
  
  import Ganeti.Lens hiding (chosen)
  import Ganeti.JQScheduler.Types
  import Ganeti.JQueue (QueuedJob(..))
  import Ganeti.JQueue.Lens
  import Ganeti.OpCodes.Lens
  import Ganeti.SlotMap
  import Ganeti.Utils
  
  
  
  -- | Ad-hoc rate limiting buckets are identified by the /combination/
  -- `REASONSTRING:n`, so "mybucket:3" and "mybucket:4" are /different/ buckets.
  type AdHocReasonKey = String
  
  
  -- | Parses an ad-hoc rate limit from a reason trail, as defined under
  -- "Ad-Hoc Rate Limiting" in `doc/design-optables.rst`.
  --
  -- The parse succeeds only on reasons of form `rate-limit:n:REASONSTRING`
  -- where `n` is a positive integer and `REASONSTRING` is an arbitrary
  -- string (may include spaces).
  parseReasonRateLimit :: (MonadFail m) => String -> m (String, Int)
  parseReasonRateLimit reason = case sepSplit ':' reason of
    "rate-limit":nStr:rest
      | Just n <- readMaybe nStr
      , n > 0 -> return (intercalate ":" (nStr:rest), n)
    _ -> fail $ "'" ++ reason ++ "' is not a valid ad-hoc rate limit reason"
  
  
  -- | Computes the bucket slots required by a job, also extracting how many
  -- slots are available from the reason rate limits in the job reason trails.
  --
  -- A job can have multiple `OpCode`s, and the `ReasonTrail`s
  -- can be different for each `OpCode`. The `OpCode`s of a job are
  -- run sequentially, so a job can only take 1 slot.
  -- Thus a job takes part in a set of buckets, requiring 1 slot in
  -- each of them.
  labelCountMapFromJob :: QueuedJob -> CountMap (String, Int)
  labelCountMapFromJob job =
    let reasonsStrings =
          job ^.. qjOpsL . traverse . qoInputL . validOpCodeL
                  . metaParamsL . opReasonL . traverse . _2
  
        buckets = ordNub . mapMaybe parseReasonRateLimit $ reasonsStrings
  
    -- Buckets are already unique from `ordNub`.
    in Map.fromList $ map (, 1) buckets
  
  
  -- | Computes the bucket slots required by a job.
  countMapFromJob :: QueuedJob -> CountMap AdHocReasonKey
  countMapFromJob = Map.mapKeys (\(str, n) -> str ++ ":" ++ show n)
                      . labelCountMapFromJob
  
  
  -- | Map of how many slots are in use for a given bucket, for a list of jobs.
  -- The slot limits are taken from the ad-hoc reason rate limiting strings.
  slotMapFromJobs :: [QueuedJob] -> SlotMap AdHocReasonKey
  slotMapFromJobs jobs =
    Map.mapKeys (\(str, n) -> str ++ ":" ++ show n)
      . Map.mapWithKey (\(_str, limit) occup -> Slot occup limit)
      . Map.unionsWith (+) . map labelCountMapFromJob
      $ jobs
  
  
  -- | Like `slotMapFromJobs`, but setting all occupation counts to 0.
  -- Useful to find what the bucket limits of a set of jobs are.
  unoccupiedSlotMapFromJobs :: [QueuedJob] -> SlotMap AdHocReasonKey
  unoccupiedSlotMapFromJobs = Map.map (\s -> s{ slotOccupied = 0 })
                                . slotMapFromJobs
  
  
  -- | Implements ad-hoc rate limiting using the reason trail as specified
  -- in `doc/design-optables.rst`.
  --
  -- Reasons of form `rate-limit:n:REASONSTRING` define buckets that limit
  -- how many jobs with that reason can be running at the same time to
  -- a positive integer n of available slots.
  --
  -- The used buckets map is currently not cached across `selectJobsToRun`
  -- invocations because the number of running jobs is typically small
  -- (< 100).
  reasonRateLimit :: Queue -> [JobWithStat] -> [JobWithStat]
  reasonRateLimit queue jobs =
    let -- For the purpose of rate limiting, manipulated jobs count as running.
        running    = map jJob $ qRunning queue ++ qManipulated queue
        candidates = map jJob jobs
  
        -- Reason rate limiting slot map of the jobs running in the queue.
        -- All jobs determine the reason buckets, but only running jobs count
        -- to the initial limits.
        initSlotMap = unoccupiedSlotMapFromJobs (running ++ candidates)
                      `occupySlots`
                      toCountMap (slotMapFromJobs running)
  
        -- A job can be run (fits) if all buckets it takes part in have
        -- a free slot. If yes, accept the job and update the slotMap.
        -- Note: If the slotMap is overfull in some slots, but the job
        -- doesn't take part in any of those, it is to be accepted.
        accumFittingJobs slotMap job =
          let jobBuckets = countMapFromJob (jJob job)
          in if slotMap `hasSlotsFor` jobBuckets
            then (slotMap `occupySlots` jobBuckets, Just job) -- job fits
            else (slotMap, Nothing)                           -- job doesn't fit
  
    in catMaybes . snd . mapAccumL accumFittingJobs initSlotMap $ jobs
  