content: {-# LANGUAGE TemplateHaskell, Rank2Types, FlexibleContexts #-}
  
  {-| Pure functions for manipulating reservations of temporary objects
  
  NOTE: Reservations aren't released specifically, they're just all
  released at the end of a job. This could be improved in the future.
  
  -}
  
  {-
  
  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.WConfd.TempRes
    ( TempRes
    , mkTempRes
    , TempResState(..)
    , emptyTempResState
    , NodeUUID
    , InstanceUUID
    , DiskUUID
    , NetworkUUID
    , DRBDMinor
    , DRBDMap
    , trsDRBDL
    , computeDRBDMap
    , computeDRBDMap'
    , allocateDRBDMinor
    , releaseDRBDMinors
    , MAC
    , generateMAC
    , reserveMAC
    , generateDRBDSecret
    , reserveLV
    , IPv4ResAction(..)
    , IPv4Reservation(..)
    , reserveIp
    , releaseIp
    , generateIp
    , commitReleaseIp
    , commitReservedIps
    , listReservedIps
    , dropAllReservations
    , isReserved
    , reserve
    , dropReservationsFor
    , reserved
    ) where
  
  import Control.Lens.At
  import Control.Monad.Except
  import Control.Monad.State
  import Control.Monad.Trans.Maybe
  import Control.Monad (liftM, when, unless, mfilter, forM, forM_, (>=>))
  import qualified Data.ByteString as BS
  import qualified Data.ByteString.UTF8 as UTF8
  import qualified Data.Foldable as F
  import Data.Maybe
  import Data.Map (Map)
  import qualified Data.Map as M
  import Data.Monoid (Monoid(mempty, mappend), (<>))
  import qualified Data.Semigroup as Sem
  import qualified Data.Set as S
  import System.Random
  import qualified Text.JSON as J
  
  import Ganeti.BasicTypes
  import Ganeti.Config
  import qualified Ganeti.Constants as C
  import Ganeti.Errors
  import qualified Ganeti.JSON as J
  import Ganeti.Lens
  import qualified Ganeti.Network as N
  import Ganeti.Locking.Locks (ClientId)
  import Ganeti.Logging
  import Ganeti.Objects
  import Ganeti.THH
  import Ganeti.Objects.Lens (configNetworksL)
  import Ganeti.Utils
  import Ganeti.Utils.Monad
  import Ganeti.Utils.Random
  import qualified Ganeti.Utils.MultiMap as MM
  
  -- * The main reservation state
  
  -- ** Aliases to make types more meaningful:
  
  type NodeUUID = BS.ByteString
  
  type InstanceUUID = BS.ByteString
  
  type DiskUUID = BS.ByteString
  
  type NetworkUUID = BS.ByteString
  
  type DRBDMinor = Int
  
  -- | A map of the usage of DRBD minors
  type DRBDMap = Map NodeUUID (Map DRBDMinor DiskUUID)
  
  -- | A map of the usage of DRBD minors with possible duplicates
  type DRBDMap' = Map NodeUUID (Map DRBDMinor [DiskUUID])
  
  -- * The state data structure
  
  -- | Types of IPv4 reservation actions.
  data IPv4ResAction = IPv4Reserve | IPv4Release
    deriving (Eq, Ord, Show, Bounded, Enum)
  
  instance J.JSON IPv4ResAction where
    showJSON IPv4Reserve = J.JSString . J.toJSString $ C.reserveAction
    showJSON IPv4Release = J.JSString . J.toJSString $ C.releaseAction
    readJSON = J.readEitherString
               >=> \s -> case () of
                      _ | s == C.reserveAction -> return IPv4Reserve
                        | s == C.releaseAction -> return IPv4Release
                        | otherwise -> fail $ "Invalid IP reservation action: "
                                             ++ s
  
  -- | The values stored in the IPv4 reservation table.
  data IPv4Reservation = IPv4Res
    { ipv4ResAction :: IPv4ResAction
    , ipv4ResNetwork :: NetworkUUID
    , ipv4ResAddr :: Ip4Address
    } deriving (Eq, Ord, Show)
  
  instance J.JSON IPv4Reservation where
    -- Notice that addr and net are in a different order, to be compatible
    -- with the original Python representation (while it's used).
    showJSON (IPv4Res a net addr) = J.showJSON (a, addr, net)
    readJSON = fmap (\(a, addr, net) -> IPv4Res a net addr) . J.readJSON
  
  -- | A polymorphic data structure for managing temporary resources assigned
  -- to jobs.
  newtype TempRes j a = TempRes { getTempRes :: MM.MultiMap j a }
    deriving (Eq, Ord, Show)
  
  instance (Ord j, Ord a) => Sem.Semigroup (TempRes j a) where
    (TempRes x) <> (TempRes y) = TempRes $ x <> y
  
  instance (Ord j, Ord a) => Monoid (TempRes j a) where
    mempty = TempRes mempty
    mappend = (Sem.<>)
  
  instance (J.JSON j, Ord j, J.JSON a, Ord a) => J.JSON (TempRes j a) where
    showJSON = J.showJSON . getTempRes
    readJSON = liftM TempRes . J.readJSON
  
  -- | Create a temporary reservations from a given multi-map.
  mkTempRes :: MM.MultiMap j a -> TempRes j a
  mkTempRes = TempRes
  
  -- | The state of the temporary reservations
  $(buildObject "TempResState" "trs"
    [ simpleField "dRBD"             [t| DRBDMap |]
    , simpleField "mACs"             [t| TempRes ClientId MAC |]
    , simpleField "dRBDSecrets"      [t| TempRes ClientId DRBDSecret |]
    , simpleField "lVs"              [t| TempRes ClientId LogicalVolume |]
    , simpleField "iPv4s"            [t| TempRes ClientId IPv4Reservation |]
    ])
  
  emptyTempResState :: TempResState
  emptyTempResState = TempResState M.empty mempty mempty mempty mempty
  
  $(makeCustomLenses ''TempResState)
  
  -- ** Utility functions
  
  -- | Issues a reservation error.
  resError :: (MonadError GanetiException m) => String -> m a
  resError = throwError . ReservationError
  
  -- | Converts 'GenericError' into a 'ReservationError'.
  toResError :: (MonadError GanetiException m) => m a -> m a
  toResError = flip catchError (throwError . f)
    where
      f (GenericError msg) = ReservationError msg
      f e = e
  
  -- | Filter values from the nested map and remove any nested maps
  -- that become empty.
  filterNested :: (Ord a, Ord b)
               => (c -> Bool) -> Map a (Map b c) -> Map a (Map b c)
  filterNested p = M.filter (not . M.null) . fmap (M.filter p)
  
  -- | Converts a lens that works on maybe values into a lens that works
  -- on regular ones. A missing value on the input is replaced by
  -- 'mempty'.
  -- The output is is @Just something@ iff @something /= mempty@.
  maybeLens :: (Monoid a, Monoid b, Eq b)
            => Lens s t (Maybe a) (Maybe b) -> Lens s t a b
  maybeLens l f = l (fmap (mfilter (/= mempty) . Just) . f . fromMaybe mempty)
  
  -- * DRBD functions
  
  -- | Compute the map of used DRBD minor/nodes, including possible
  -- duplicates.
  -- An error is returned if the configuration isn't consistent
  -- (for example if a referenced disk is missing etc.).
  computeDRBDMap' :: (MonadError GanetiException m)
                  => ConfigData -> TempResState -> m DRBDMap'
  computeDRBDMap' cfg trs =
      flip execStateT (fmap (fmap (: [])) (trsDRBD trs))
      $ F.forM_ (configDisks cfg) addMinors
    where
      -- | Creates a lens for modifying the list of instances
      nodeMinor :: NodeUUID -> DRBDMinor -> Lens' DRBDMap' [DiskUUID]
      nodeMinor node minor = maybeLens (at node) . maybeLens (at minor)
      -- | Adds minors of a disk within the state monad
      addMinors disk = do
        let minors = getDrbdMinorsForDisk disk
        forM_ minors $ \(minor, node) ->
          nodeMinor (UTF8.fromString node) minor %=
              (UTF8.fromString (uuidOf disk) :)
  
  -- | Compute the map of used DRBD minor/nodes.
  -- Report any duplicate entries as an error.
  --
  -- Unlike 'computeDRBDMap'', includes entries for all nodes, even if empty.
  computeDRBDMap :: (MonadError GanetiException m)
                 => ConfigData -> TempResState -> m DRBDMap
  computeDRBDMap cfg trs = do
    m <- computeDRBDMap' cfg trs
    let dups = filterNested ((>= 2) . length) m
    unless (M.null dups) . resError
      $ "Duplicate DRBD ports detected: " ++ show (M.toList $ fmap M.toList dups)
    return $ fmap (M.mapMaybe listToMaybe) m
             `M.union` (fmap (const mempty) . J.fromContainer . configNodes $ cfg)
  
  -- Allocate a drbd minor.
  --
  -- The free minor will be automatically computed from the existing devices.
  -- A node can not be given multiple times.
  -- The result is the list of minors, in the same order as the passed nodes.
  allocateDRBDMinor :: (MonadError GanetiException m, MonadState TempResState m)
                    => ConfigData -> DiskUUID -> [NodeUUID]
                    -> m [DRBDMinor]
  allocateDRBDMinor cfg disk nodes = do
    unless (nodes == ordNub nodes) . resError
      $ "Duplicate nodes detected in list '" ++ show nodes ++ "'"
    dMap <- computeDRBDMap' cfg =<< get
    let usedMap = fmap M.keysSet dMap
    let alloc :: S.Set DRBDMinor -> Map DRBDMinor DiskUUID
              -> (DRBDMinor, Map DRBDMinor DiskUUID)
        alloc used m = let k = findFirst 0 (M.keysSet m `S.union` used)
                       in (k, M.insert k disk m)
    forM nodes $ \node -> trsDRBDL . maybeLens (at node)
                          %%= alloc (M.findWithDefault mempty node usedMap)
  
  -- Release temporary drbd minors allocated for a given disk using
  -- 'allocateDRBDMinor'.
  releaseDRBDMinors :: (MonadState TempResState m) => DiskUUID -> m ()
  releaseDRBDMinors disk = trsDRBDL %= filterNested (/= disk)
  
  -- * Other temporary resources
  
  -- | Tests if a given value is reserved for a given job.
  isReserved :: (Ord a, Ord j) => a -> TempRes j a -> Bool
  isReserved x = MM.elem x . getTempRes
  
  -- | Tries to reserve a given value for a given job.
  reserve :: (MonadError GanetiException m, Show a, Ord a, Ord j)
          => j -> a -> TempRes j a -> m (TempRes j a)
  reserve jobid x tr = do
    when (isReserved x tr) . resError $ "Duplicate reservation for resource '"
                                         ++ show x ++ "'"
    return . TempRes . MM.insert jobid x $ getTempRes tr
  
  dropReservationsFor :: (Ord a, Ord j) => j -> TempRes j a -> TempRes j a
  dropReservationsFor jobid = TempRes . MM.deleteAll jobid . getTempRes
  
  reservedFor :: (Ord a, Ord j) => j -> TempRes j a -> S.Set a
  reservedFor jobid = MM.lookup jobid . getTempRes
  
  reserved :: (Ord a, Ord j) => TempRes j a -> S.Set a
  reserved = MM.values . getTempRes
  
  -- | Computes the set of all reserved resources and passes it to
  -- the given function.
  -- This allows it to avoid resources that are already in use.
  withReserved :: (MonadError GanetiException m, Show a, Ord a, Ord j)
               => j -> (S.Set a -> m a) -> TempRes j a -> m (a, TempRes j a)
  withReserved jobid genfn tr = do
    x <- genfn (reserved tr)
    (,) x `liftM` reserve jobid x tr
  
  -- | Repeatedly tries to run a given monadic function until it succeeds
  -- and the returned value is free to reserve.
  -- If such a value is found, it's reserved and returned.
  -- Otherwise fails with an error.
  generate :: (MonadError GanetiException m, Show a, Ord a, Ord j)
           => j -> S.Set a -> m (Maybe a) -> TempRes j a -> m (a, TempRes j a)
  generate jobid existing genfn = withReserved jobid f
    where
      retries = 64 :: Int
      f res = do
        let vals = res `S.union` existing
        xOpt <- retryMaybeN retries
                            (\_ -> mfilter (`S.notMember` vals) (MaybeT genfn))
        maybe (resError "Not able generate new resource")
                         -- TODO: (last tried: " ++ %s)" % new_resource
              return xOpt
  
  -- | A variant of 'generate' for randomized computations.
  generateRand
    :: (MonadError GanetiException m, Show a, Ord a, Ord j, RandomGen g)
    => g -> j -> S.Set a -> (g -> (Maybe a, g)) -> TempRes j a
    -> m (a, TempRes j a)
  generateRand rgen jobid existing genfn tr =
    evalStateT (generate jobid existing (state genfn) tr) rgen
  
  -- | Embeds a stateful computation in a stateful monad.
  stateM :: (MonadState s m) => (s -> m (a, s)) -> m a
  stateM f = get >>= f >>= \(x, s) -> liftM (const x) (put s)
  
  -- | Embeds a state-modifying computation in a stateful monad.
  modifyM :: (MonadState s m) => (s -> m s) -> m ()
  modifyM f = get >>= f >>= put
  
  -- ** Functions common to all reservations
  
  -- | Removes all resources reserved by a given job.
  --
  -- If a new reservation resource type is added, it must be added here as well.
  dropAllReservations :: ClientId -> State TempResState ()
  dropAllReservations jobId = modify $
      (trsMACsL %~ dropReservationsFor jobId)
    . (trsDRBDSecretsL %~ dropReservationsFor jobId)
    . (trsLVsL %~ dropReservationsFor jobId)
    . (trsIPv4sL %~ dropReservationsFor jobId)
  
  -- | Looks up a network by its UUID.
  lookupNetwork :: (MonadError GanetiException m)
                => ConfigData -> NetworkUUID -> m Network
  lookupNetwork cd netId =
    J.lookupContainer (resError $ "Network '" ++ show netId ++ "' not found")
                      netId (configNetworks cd)
  
  -- ** IDs
  
  -- ** MAC addresses
  
  -- Randomly generate a MAC for an instance.
  -- Checks that the generated MAC isn't used by another instance.
  --
  -- Note that we only consume, but not return the state of a random number
  -- generator. This is because the whole operation needs to be pure (for atomic
  -- 'IORef' updates) and therefore we can't use 'getStdRandom'. Therefore the
  -- approach we take is to instead use 'newStdGen' and discard the split
  -- generator afterwards.
  generateMAC
    :: (RandomGen g, MonadError GanetiException m, Functor m)
    => g -> ClientId -> Maybe NetworkUUID -> ConfigData
    -> StateT TempResState m MAC
  generateMAC rgen jobId netId cd = do
    net <- case netId of
      Just n -> Just <$> lookupNetwork cd n
      Nothing -> return Nothing
    let prefix = fromMaybe (clusterMacPrefix . configCluster $ cd)
                           (networkMacPrefix =<< net)
    let existing = S.fromList $ getAllMACs cd
    StateT
      $ traverseOf2 trsMACsL
          (generateRand rgen jobId existing
                        (over _1 Just . generateOneMAC prefix))
  
  -- Reserves a MAC for an instance in the list of temporary reservations.
  reserveMAC
    :: (MonadError GanetiException m, MonadState TempResState m, Functor m)
    => ClientId -> MAC -> ConfigData -> m ()
  reserveMAC jobId mac cd = do
    let existing = S.fromList $ getAllMACs cd
    when (S.member mac existing)
      $ resError "MAC already in use"
    modifyM $ traverseOf trsMACsL (reserve jobId mac)
  
  -- ** DRBD secrets
  
  generateDRBDSecret
    :: (RandomGen g, MonadError GanetiException m, Functor m)
    => g -> ClientId -> ConfigData -> StateT TempResState m DRBDSecret
  generateDRBDSecret rgen jobId cd = do
    let existing = S.fromList $ getAllDrbdSecrets cd
    StateT $ traverseOf2 trsDRBDSecretsL
             (generateRand rgen jobId existing
                           (over _1 Just . generateSecret C.drbdSecretLength))
  
  -- ** LVs
  
  reserveLV
    :: (MonadError GanetiException m, MonadState TempResState m, Functor m)
    => ClientId -> LogicalVolume -> ConfigData -> m ()
  reserveLV jobId lv cd = do
    existing <- toError $ getAllLVs cd
    when (S.member lv existing)
      $ resError "MAC already in use"
    modifyM $ traverseOf trsLVsL (reserve jobId lv)
  
  -- ** IPv4 addresses
  
  -- | Lists all IPv4 addresses reserved for a given network.
  usedIPv4Addrs :: NetworkUUID -> S.Set IPv4Reservation -> S.Set Ip4Address
  usedIPv4Addrs netuuid =
    S.map ipv4ResAddr . S.filter ((== netuuid) . ipv4ResNetwork)
  
  -- | Reserve a given IPv4 address for use by an instance.
  reserveIp
    :: (MonadError GanetiException m, MonadState TempResState m, Functor m)
    => ClientId -> NetworkUUID -> Ip4Address
    -> Bool -- ^ whether to check externally reserved IPs
    -> ConfigData -> m ()
  reserveIp jobId netuuid addr checkExt cd = toResError $ do
      net <- lookupNetwork cd netuuid
      isres <- N.isReserved N.PoolInstances addr net
      when isres . resError $ "IP address already in use"
      when checkExt $ do
        isextres <- N.isReserved N.PoolExt addr net
        when isextres . resError $ "IP is externally reserved"
      let action = IPv4Res IPv4Reserve netuuid addr
      modifyM $ traverseOf trsIPv4sL (reserve jobId action)
  
  -- | Give a specific IP address back to an IP pool.
  -- The IP address is returned to the IP pool designated by network id
  -- and marked as reserved.
  releaseIp
    :: (MonadError GanetiException m, MonadState TempResState m, Functor m)
    => ClientId -> NetworkUUID -> Ip4Address -> m ()
  releaseIp jobId netuuid addr =
    let action = IPv4Res { ipv4ResAction = IPv4Release
                         , ipv4ResNetwork = netuuid
                         , ipv4ResAddr = addr }
    in modifyM $ traverseOf trsIPv4sL (reserve jobId action)
  
  -- Find a free IPv4 address for an instance and reserve it.
  generateIp
    :: (MonadError GanetiException m, MonadState TempResState m, Functor m)
    => ClientId -> NetworkUUID -> ConfigData -> m Ip4Address
  generateIp jobId netuuid cd = toResError $ do
      net <- lookupNetwork cd netuuid
      let f res = do
            let ips = usedIPv4Addrs netuuid res
            addr <- N.findFree (`S.notMember` ips) net
            maybe (resError "Cannot generate IP. Network is full")
                  (return . IPv4Res IPv4Reserve netuuid) addr
      liftM ipv4ResAddr . stateM $ traverseOf2 trsIPv4sL (withReserved jobId f)
  
  -- | Commit a reserved/released IP address to an IP pool.
  -- The IP address is taken from the network's IP pool and marked as
  -- reserved/free for instances.
  commitIp
    :: (MonadError GanetiException m, Functor m)
    => IPv4Reservation -> ConfigData -> m ConfigData
  commitIp (IPv4Res actType netuuid addr) cd = toResError $ do
      let call = case actType of
                  IPv4Reserve -> N.reserve
                  IPv4Release -> N.release
          f Nothing = resError $ "Network '" ++ show netuuid ++ "' not found"
          f (Just net) = Just `liftM` call N.PoolInstances addr net
      traverseOf (configNetworksL . J.alterContainerL netuuid) f cd
  
  -- | Immediately release an IP address, without using the reservations pool.
  commitReleaseIp
    :: (MonadError GanetiException m, Functor m)
    => NetworkUUID -> Ip4Address -> ConfigData -> m ConfigData
  commitReleaseIp netuuid addr =
    commitIp (IPv4Res IPv4Release netuuid addr)
  
  -- | Commit all reserved/released IP address to an IP pool.
  -- The IP addresses are taken from the network's IP pool and marked as
  -- reserved/free for instances.
  --
  -- Note that the reservations are kept, they are supposed to be cleaned
  -- when a job finishes.
  commitReservedIps
    :: (MonadError GanetiException m, Functor m, MonadLog m)
    => ClientId -> TempResState -> ConfigData -> m ConfigData
  commitReservedIps jobId tr cd = do
    let res = reservedFor jobId (trsIPv4s tr)
    logDebug $ "Commiting reservations: " ++ show res
    F.foldrM commitIp cd res
  
  listReservedIps :: ClientId -> TempResState -> S.Set IPv4Reservation
  listReservedIps jobid = reservedFor jobid . trsIPv4s
  