{-# LANGUAGE BangPatterns #-}
{-| Implementation of lock allocation.

-}

{-

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.Locking.Allocation
  ( LockAllocation
  , emptyAllocation
  , OwnerState(..)
  , lockOwners
  , listLocks
  , listAllLocks
  , listAllLocksOwners
  , holdsLock
  , LockRequest(..)
  , requestExclusive
  , requestShared
  , requestRelease
  , updateLocks
  , freeLocks
  ) where

import Control.Applicative (liftA2, (<$>), (<*>), pure)
import Control.Arrow (second, (***))
import Control.Monad
import Data.Foldable (for_, find)
import Data.List (foldl')
import qualified Data.Map as M
import Data.Maybe (fromMaybe)
import qualified Data.Set as S
import qualified Text.JSON as J

import Ganeti.BasicTypes
import Ganeti.JSON (toArray)
import Ganeti.Locking.Types

{-

This module is parametric in the type of locks and lock owners.
While we only state minimal requirements for the types, we will
consistently use the type variable 'a' for the type of locks and
the variable 'b' for the type of the lock owners throughout this
module.

-}

-- | Data type describing the way a lock can be owned.
data OwnerState = OwnShared | OwnExclusive deriving (Ord, Eq, Show)

-- | Type describing indirect ownership on a lock. We keep the set
-- of all (lock, owner)-pairs for locks that are implied in the given
-- lock, annotated with the type of ownership (shared or exclusive).
type IndirectOwners a b = M.Map (a, b) OwnerState

-- | The state of a lock that is taken. Besides the state of the lock
-- itself, we also keep track of all other lock allocation that affect
-- the given lock by means of implication.
data AllocationState a b = Exclusive b (IndirectOwners a b)
                         | Shared (S.Set b) (IndirectOwners a b)
                         deriving (Eq, Show)

-- | Compute the set of indirect owners from the information about
-- indirect ownership.
indirectOwners :: (Ord a, Ord b) => M.Map (a, b) OwnerState -> S.Set b
indirectOwners = S.map snd . M.keysSet

-- | Compute the (zero or one-elment) set of exclusive indirect owners.
indirectExclusives :: (Ord a, Ord b) => M.Map (a, b) OwnerState -> S.Set b
indirectExclusives = indirectOwners . M.filter (== OwnExclusive)

{-| Representation of a Lock allocation

To keep queries for locks efficient, we keep two
associations, with the invariant that they fit
together: the association from locks to their
allocation state, and the association from an
owner to the set of locks owned. As we do not
export the constructor, the problem of keeping
this invariant reduces to only exporting functions
that keep the invariant.

-}

data LockAllocation a b =
  LockAllocation { laLocks :: M.Map a (AllocationState a b)
                 , laOwned :: M.Map b (M.Map a OwnerState)
                 }
  deriving (Eq, Show)

-- | A state with all locks being free.
emptyAllocation :: (Ord a, Ord b) => LockAllocation a b
emptyAllocation =
  LockAllocation { laLocks = M.empty
                 , laOwned = M.empty
                 }

-- | Obtain the list of all owners holding at least a single lock.
lockOwners :: Ord b => LockAllocation a b -> [b]
lockOwners = M.keys . laOwned

-- | Obtain the locks held by a given owner. The locks are reported
-- as a map from the owned locks to the form of ownership (OwnShared
-- or OwnExclusive).
listLocks :: Ord b => b -> LockAllocation a b -> M.Map a OwnerState
listLocks owner = fromMaybe M.empty . M.lookup owner . laOwned

-- | List all locks currently (directly or indirectly) owned by someone.
listAllLocks :: Ord b => LockAllocation a b -> [a]
listAllLocks = M.keys . laLocks

-- | Map an AllocationState to a list of pairs of owners and type of
-- ownership, showing the direct owners only.
toOwnersList :: AllocationState a b -> [(b, OwnerState)]
toOwnersList (Exclusive owner _) = [(owner, OwnExclusive)]
toOwnersList (Shared owners _) = map (flip (,) OwnShared) . S.elems $ owners

-- | List all locks currently (directly of indirectly) in use together
-- with the direct owners.
listAllLocksOwners :: LockAllocation a b -> [(a, [(b, OwnerState)])]
listAllLocksOwners = M.toList . M.map toOwnersList . laLocks

-- | Returns 'True' if the given owner holds the given lock at the given
-- ownership level or higher. This means that querying for a shared lock
-- returns 'True' of the owner holds the lock in shared or exlusive mode.
holdsLock :: (Ord a, Ord b)
          => b -> a -> OwnerState -> LockAllocation a b -> Bool
holdsLock owner lock state = (>= Just state) . M.lookup lock . listLocks owner

-- | Data Type describing a change request on a single lock.
data LockRequest a = LockRequest { lockAffected :: a
                                 , lockRequestType :: Maybe OwnerState
                                 }
                     deriving (Eq, Show, Ord)

instance J.JSON a => J.JSON (LockRequest a) where
  showJSON (LockRequest a Nothing) = J.showJSON (a, "release")
  showJSON (LockRequest a (Just OwnShared)) = J.showJSON (a, "shared")
  showJSON (LockRequest a (Just OwnExclusive)) = J.showJSON (a, "exclusive")
  readJSON (J.JSArray [a, J.JSString tp]) =
    case J.fromJSString tp of
      "release" -> LockRequest <$> J.readJSON a <*> pure Nothing
      "shared" -> LockRequest <$> J.readJSON a <*> pure (Just OwnShared)
      "exclusive" -> LockRequest <$> J.readJSON a <*> pure (Just OwnExclusive)
      _ -> J.Error $ "malformed request type: " ++ J.fromJSString tp
  readJSON x = J.Error $ "malformed lock request: " ++ show x

-- | Lock request for an exclusive lock.
requestExclusive :: a -> LockRequest a
requestExclusive lock = LockRequest { lockAffected = lock
                                    , lockRequestType = Just OwnExclusive }

-- | Lock request for a shared lock.
requestShared :: a -> LockRequest a
requestShared lock = LockRequest { lockAffected = lock
                                 , lockRequestType = Just OwnShared }

-- | Request to release a lock.
requestRelease :: a -> LockRequest a
requestRelease lock = LockRequest { lockAffected = lock
                                  , lockRequestType = Nothing }

-- | Update the Allocation state of a lock according to a given
-- function.
updateAllocState :: (Ord a, Ord b)
                  => (Maybe (AllocationState a b) -> AllocationState a b)
                  -> LockAllocation a b -> a -> LockAllocation a b
updateAllocState f state lock =
  let !locks' = M.alter (find (/= Shared S.empty M.empty) . Just . f)
                        lock (laLocks state)
  in state { laLocks = locks' }

-- | Internal function to update the state according to a single
-- lock request, assuming all prerequisites are met.
updateLock :: (Ord a, Ord b)
           => b
           -> LockAllocation a b -> LockRequest a -> LockAllocation a b
updateLock owner state (LockRequest lock (Just OwnExclusive)) =
  let locks = laLocks state
      lockstate' = case M.lookup lock locks of
        Just (Exclusive _ i) -> Exclusive owner i
        Just (Shared _ i) -> Exclusive owner i
        Nothing -> Exclusive owner M.empty
      !locks' = M.insert lock lockstate' locks
      ownersLocks' = M.insert lock OwnExclusive $ listLocks owner state
      !owned' = M.insert owner ownersLocks' $ laOwned state
  in state { laLocks = locks', laOwned = owned' }
updateLock owner state (LockRequest lock (Just OwnShared)) =
  let ownersLocks' = M.insert lock OwnShared $ listLocks owner state
      !owned' = M.insert owner ownersLocks' $ laOwned state
      locks = laLocks state
      lockState' = case M.lookup lock locks of
        Just (Exclusive _ i) -> Shared (S.singleton owner) i
        Just (Shared s i) -> Shared (S.insert owner s) i
        _ -> Shared (S.singleton owner) M.empty
      !locks' = M.insert lock lockState' locks
  in state { laLocks = locks', laOwned = owned' }
updateLock owner state (LockRequest lock Nothing) =
  let ownersLocks' = M.delete lock $ listLocks owner state
      owned = laOwned state
      owned' = if M.null ownersLocks'
                 then M.delete owner owned
                 else M.insert owner ownersLocks' owned
      update (Just (Exclusive x i)) = if x == owner
                                        then Shared S.empty i
                                        else Exclusive x i
      update (Just (Shared s i)) = Shared (S.delete owner s) i
      update Nothing = Shared S.empty M.empty
  in updateAllocState update (state { laOwned = owned' }) lock

-- | Update the set of indirect ownerships of a lock by the given function.
updateIndirectSet :: (Ord a, Ord b)
                  => (IndirectOwners a b -> IndirectOwners a b)
                  -> LockAllocation a b -> a -> LockAllocation a b
updateIndirectSet f =
  let update (Just (Exclusive x i)) = Exclusive x (f i)
      update (Just (Shared s i)) = Shared s (f i)
      update Nothing = Shared S.empty (f M.empty)
  in updateAllocState update

-- | Update all indirect onwerships of a given lock.
updateIndirects :: (Lock a, Ord b)
                => b
                -> LockAllocation a b -> LockRequest a -> LockAllocation a b
updateIndirects owner state req =
  let lock = lockAffected req
      fn = case lockRequestType req of
             Nothing -> M.delete (lock, owner)
             Just tp -> M.insert (lock, owner) tp
  in foldl' (updateIndirectSet fn) state $ lockImplications lock

-- | Update the locks of an owner according to the given request. Return
-- the pair of the new state and the result of the operation, which is the
-- the set of owners on which the operation was blocked on. so an empty set is
-- success, and the state is updated if, and only if, the returned set is emtpy.
-- In that way, it can be used in atomicModifyIORef.
updateLocks :: (Lock a, Ord b)
            => b
            -> [LockRequest a]
            -> LockAllocation a b -> (LockAllocation a b, Result (S.Set b))
updateLocks owner reqs state = genericResult ((,) state . Bad) (second Ok) $ do
  unless ((==) (length reqs) . S.size . S.fromList $ map lockAffected reqs)
    . runListHead (return ())
                  (fail . (++) "Inconsitent requests for lock " . show) $ do
      r <- reqs
      r' <- reqs
      guard $ r /= r'
      guard $ lockAffected r == lockAffected r'
      return $ lockAffected r
  let current = listLocks owner state
  unless (M.null current) $ do
    let (highest, _) = M.findMax current
        notHolding = not
                     . any (uncurry (==) . ((M.lookup `flip` current) *** Just))
        orderViolation l = fail $ "Order violation: requesting " ++ show l
                                   ++ " while holding " ++ show highest
    for_ reqs $ \req -> case req of
      LockRequest lock (Just OwnExclusive)
        | lock < highest && notHolding ((,) <$> lock : lockImplications lock
                                            <*> [OwnExclusive])
        -> orderViolation lock
      LockRequest lock (Just OwnShared)
        | lock < highest && notHolding ((,) <$> lock : lockImplications lock
                                            <*> [OwnExclusive, OwnShared])
        -> orderViolation lock
      _ -> Ok ()
  let sharedsHeld = M.keysSet $ M.filter (== OwnShared) current
      exclusivesRequested = map lockAffected
                            . filter ((== Just OwnExclusive) . lockRequestType)
                            $ reqs
  runListHead (return ()) fail $ do
    x <- exclusivesRequested
    i <- lockImplications x
    guard $ S.member i sharedsHeld
    return $ "Order violation: requesting exclusively " ++ show x
              ++ " while holding a shared lock on the group lock " ++ show i
              ++ " it belongs to."
  let blockedOn (LockRequest  _ Nothing) = S.empty
      blockedOn (LockRequest lock (Just OwnExclusive)) =
        case M.lookup lock (laLocks state) of
          Just (Exclusive x i) ->
            S.singleton x `S.union` indirectOwners i
          Just (Shared xs i) ->
            xs `S.union` indirectOwners i
          _ -> S.empty
      blockedOn (LockRequest lock (Just OwnShared)) =
        case M.lookup lock (laLocks state) of
          Just (Exclusive x i) ->
            S.singleton x `S.union` indirectExclusives i
          Just (Shared _ i) -> indirectExclusives i
          _ -> S.empty
  let indirectBlocked Nothing _ = S.empty
      indirectBlocked (Just OwnShared) lock =
        case M.lookup lock (laLocks state) of
          Just (Exclusive x _) -> S.singleton x
          _ -> S.empty
      indirectBlocked (Just OwnExclusive) lock =
        case M.lookup lock (laLocks state) of
          Just (Exclusive x _) -> S.singleton x
          Just (Shared xs _) -> xs
          _ -> S.empty
  let direct = S.unions $ map blockedOn reqs
      indirect = reqs >>= \req ->
        map (indirectBlocked (lockRequestType req))
          . lockImplications $ lockAffected req
  let blocked = S.delete owner . S.unions $ direct:indirect
  let state' = foldl' (updateLock owner) state reqs
      state'' = foldl' (updateIndirects owner) state' reqs
  return (if S.null blocked then state'' else state, blocked)

-- | Manipluate all locks of the owner with a given property.
manipulateLocksPredicate :: (Lock a, Ord b)
                         => (a -> LockRequest a)
                         -> (a -> Bool)
                         -> b -> LockAllocation a b -> LockAllocation a b
manipulateLocksPredicate req prop owner state =
  fst . flip (updateLocks owner) state . map req
    . filter prop
    . M.keys
    $ listLocks owner state

-- | Compute the state after an owner releases all its locks that
-- satisfy a certain property.
freeLocksPredicate :: (Lock a, Ord b)
                   => (a -> Bool)
                   -> LockAllocation a b -> b -> LockAllocation a b
freeLocksPredicate prop = flip $ manipulateLocksPredicate requestRelease prop

-- | Compute the state after an onwer releases all its locks.
freeLocks :: (Lock a, Ord b) => LockAllocation a b -> b -> LockAllocation a b
freeLocks = freeLocksPredicate (const True)

{-| Serializaiton of Lock Allocations

To serialize a lock allocation, we only remember which owner holds
which locks at which level (shared or exclusive). From this information,
everything else can be reconstructed, simply using updateLocks.
-}

instance J.JSON OwnerState where
  showJSON OwnShared = J.showJSON "shared"
  showJSON OwnExclusive = J.showJSON "exclusive"
  readJSON (J.JSString x) = let s = J.fromJSString x
                            in case s of
                              "shared" -> J.Ok OwnShared
                              "exclusive" -> J.Ok OwnExclusive
                              _ -> J.Error $ "Unknown owner type " ++ s
  readJSON _ = J.Error "Owner type not encoded as a string"

-- | Read a lock-ownerstate pair from JSON.
readLockOwnerstate :: (J.JSON a) => J.JSValue -> J.Result (a, OwnerState)
readLockOwnerstate (J.JSArray [x, y]) = liftA2 (,) (J.readJSON x) (J.readJSON y)
readLockOwnerstate x = fail $ "lock-ownerstate pairs are encoded as arrays"
                              ++ " of length 2, but found " ++ show x

-- | Read an owner-lock pair from JSON.
readOwnerLock :: (J.JSON a, J.JSON b)
              => J.JSValue -> J.Result (b, [(a, OwnerState)])
readOwnerLock (J.JSArray [x, J.JSArray ys]) =
  liftA2 (,) (J.readJSON x) (mapM readLockOwnerstate ys)
readOwnerLock x = fail $ "Expected pair of owner and list of owned locks,"
                         ++ " but found " ++ show x

-- | Transform a lock-ownerstate pair into a LockRequest.
toRequest :: (a, OwnerState) -> LockRequest a
toRequest (a, OwnExclusive) = requestExclusive a
toRequest (a, OwnShared) = requestShared a

-- | Obtain a LockAllocation from a given owner-locks list.
-- The obtained allocation is the one obtained if the respective owners
-- requested their locks sequentially.
allocationFromOwners :: (Lock a, Ord b, Show b)
                     => [(b, [(a, OwnerState)])]
                     -> J.Result (LockAllocation a b)
allocationFromOwners =
  let allocateOneOwner s (o, req) = do
        let (s', result) = updateLocks o (map toRequest req) s
        when (result /= Ok S.empty) . fail
          . (++) ("Inconsistent lock status for " ++ show o ++ ": ")
          $ case result of
            Bad err -> err
            Ok blocked -> "blocked on " ++ show (S.toList blocked)
        return s'
  in foldM allocateOneOwner emptyAllocation

instance (Lock a, J.JSON a, Ord b, J.JSON b, Show b)
           => J.JSON (LockAllocation a b) where
  showJSON = J.showJSON . M.toList . M.map M.toList . laOwned
  readJSON x = do
    xs <- toArray x
    owned <- mapM readOwnerLock xs
    allocationFromOwners owned