{-# LANGUAGE Trustworthy, GADTs, ParallelListComp, Arrows, ImplicitParams, ScopedTypeVariables, FlexibleInstances #-}

-- | Automatic deadlock avoidance.
module Control.CUtils.Deadlock (AbstractLock(..), BoxedAbstractLock(..), LockSafetyException(LockTakenTwice, LockNotDeclared), acquire, release, test) where

import Control.Monad
import Control.Monad.Loops
import Control.Concurrent
import Control.Concurrent.STM
import Control.Concurrent.MVar
import qualified Control.Concurrent.STM.Map as MS
import Data.List (inits, tails, elemIndex, deleteBy)
import Data.Maybe
import Data.Function (on)
import Data.Typeable
import Data.BellmanFord
import Data.Hashable
import qualified Data.HashTable.IO as H
import System.IO.Unsafe
import Unsafe.Coerce
import Control.CUtils.Conc
import Control.Exception
import GHC.Conc

class AbstractLock l where
	lock :: l -> IO ()
	unlock :: l -> IO ()

instance AbstractLock (MVar ()) where
	lock m = takeMVar m
	unlock m = void (tryPutMVar m ())

-- | An automatic deadlock avoidance algorithm prevents deadlocks from occurring,
--   by dynamically controlling the way threads use locks, to prevent the appearance
--   of cycles in the ownership graph.
--
--   Another way of stating the intuition is: the evolution of multiple locks
--   can be visualized as a multi-dimensional space, with the regions in which
--   locks are held coloured black. This algorithm attempts to cover up
--   the concave spaces in these regions, in which the evolution can get stuck.
--
--   It is important to ensure that the management of the ownership graph is
--   fast. There are probably a lot of ways to do this; the one used here is
--   to use a concurrent hash table; and only process nodes that are in the
--   same connected component as the node in question. This at least gives
--   a system a performance related to the size of its interacting components.

reachable :: (Hashable t, Eq t, Show t) => MS.Map t [(u, [t])] -> t -> STM (H.BasicHashTable t [t])
reachable mp from = unsafeIOToSTM H.new>>= \tabu->recurse [from] tabu>>return tabu where
	recurse ls tabu =mapM_
		(\x->do
		may <- unsafeIOToSTM(H.lookup tabu x)
		unless(isJust may)$do
			ls' <- liftM(maybe [] (concatMap snd))$MS.lookup x mp
			unsafeIOToSTM$H.insert tabu x ls'
			recurse ls' tabu)
		ls

data BoxedAbstractLock where
	BoxedAbstractLock :: (AbstractLock l, Eq l, Hashable l, Typeable l) => l -> BoxedAbstractLock

instance Eq BoxedAbstractLock where
	BoxedAbstractLock l == BoxedAbstractLock l2 = maybe False (l==) (cast l2)
instance AbstractLock BoxedAbstractLock where
	lock (BoxedAbstractLock l) = lock l
	unlock (BoxedAbstractLock l) = unlock l
instance Show BoxedAbstractLock where
	showsPrec _ (BoxedAbstractLock _) = ("BoxedAbstractLock _"++)
instance Hashable BoxedAbstractLock where
	hashWithSalt x (BoxedAbstractLock l) = hashWithSalt x l

-- For each thread, we need to track what resources it currently
-- holds, and for each resource, the resources it may
-- potentially acquire while holding that resource.

resource :: MS.Map BoxedAbstractLock [(ThreadId,[BoxedAbstractLock])]
{-# NOINLINE resource #-}
resource = unsafePerformIO$atomically MS.empty

resourceThreads :: MS.Map ThreadId [(BoxedAbstractLock,[BoxedAbstractLock])]
{-# NOINLINE resourceThreads #-}
resourceThreads = unsafePerformIO$atomically MS.empty

instance Show (MVar t) where
	show x = show (unsafeCoerce x :: Int)
instance Hashable (MVar t) where
	hashWithSalt x m = hashWithSalt x(unsafeCoerce m :: Int)

-- If there is a HOLD-ACQUIRE cycle, find a lock which guards against this cycle.
-- I.e. by the time the lock becomes free, the cycle may no longer exist.
hazard m = reachable resource m>>=unsafeIOToSTM.H.toList>>= \ls->return$!cycles[((x,x2),())|(x, ls')<-ls,x2<-ls'] m

data LockSafetyException = LockTakenTwice -- Lock taken twice (thread deadlocks with itself)
	| LockNotDeclared -- Lock ordering improperly declared
	| Abort !BoxedAbstractLock -- This exception never propagates out of library code.
	deriving (Show, Typeable)

instance Exception LockSafetyException

insert x y ((x1, _):xs) | x == x1 = (x, y) : xs
insert x y (pr:xs) = ((:) $! pr) $! insert x y xs
insert x y [] = [(x, y)]

updateResource thd l possiblyAcq = do
	may <- MS.lookup thd resourceThreads
	let isDoubleTake = maybe False (\ls -> l `elem` map fst ls) may
	when isDoubleTake(throwSTM LockTakenTwice)
	let isPreDeclared = maybe True (\ls -> null ls || l `elem` concatMap snd ls) may
	unless isPreDeclared(throwSTM LockNotDeclared)

	-- Add this lock to held locks.
	may2 <- MS.lookup l resource
	let rtInsert = insert l possiblyAcq$maybe [] id may
	(MS.insert thd $! rtInsert) resourceThreads
	let rInsert = insert thd possiblyAcq$maybe [] id may2 
	(MS.insert l $! rInsert) resource

	-- Have to see if acquiring this lock risks encountering a cycle
	-- in the ownership graph.
	lock <- hazard l
	when(isJust lock)$throwSTM$!Abort$fromJust lock

-- | Acquire a lock. In order to implement deadlock avoidance, the function 'acquire'
-- requires that all locks a thread may take while holding the given lock are annotated
-- in parameter 'possiblyAcq'.
--
-- While programs with locks have rare deadlock conditions, they have common lock
-- use patterns in-thread. If we throw an exception whenever lock use patterns are not properly
-- declared, all lock use patterns will show up in testing and can be annotated.
--
-- Complexity overview:
--
--   Let N = the number of locks.
--
--   Let K = the size of the largest directed component of the ownership graph.
--
--   Let C = the number of capabilities.
--
--   'acquire' runs in O(K log N + K^3/C) time [provided C <= K].
acquire :: BoxedAbstractLock -> [BoxedAbstractLock] -> IO () 
acquire l possiblyAcq = resourceThreads `seq` resource `seq` whileM_
	(catch
	(do
	thd <- myThreadId
	atomically(updateResource thd l possiblyAcq)
	return False)
	(\ex -> case ex of
		-- Waits on the lock. This has the effect of denying service
		-- to this thread until the hazard has passed.
		Abort l' -> do{lock l';unlock l';return True}
		_ -> throwIO ex))
	(return ())
	>>lock l

-- | Release a lock so acquired.
--
--   'release' runs in O(log N) time.
release :: BoxedAbstractLock -> IO ()
release l = resourceThreads `seq` resource `seq` do
	thd <- myThreadId
	atomically$do
		may <- MS.lookup thd resourceThreads
		maybe
			(return())
			(\ls -> MS.insert thd(deleteBy((==)`on`fst) (l, []) ls) resourceThreads)
			may
		may2 <- MS.lookup l resource
		maybe
			(return())
			(\ls -> MS.insert l(deleteBy((==)`on`fst) (thd, []) ls) resource)
			may2
	unlock l

-- The dining philosophers problem
philosopher n lfFork rtFork delay printL = whileM_
	(return True)
	(do
	threadDelay (1000*delay)
	modifyMVar_ printL$ \_->print$"Take left fork (" ++ show n ++ ")"
	acquire lfFork [rtFork]
	modifyMVar_ printL$ \_->print$"Take right fork (" ++ show n ++ ")"
	acquire rtFork []
	threadDelay (1000*delay)
	release lfFork
	release rtFork)

-- | Test: The dining philosophers problem
test = do
	fork1 <- liftM BoxedAbstractLock (newMVar ())
	fork2 <- liftM BoxedAbstractLock (newMVar ())
	fork3 <- liftM BoxedAbstractLock (newMVar ())
	printL <- newMVar ()
	forkIO$philosopher 1 fork1 fork2 500 printL
	forkIO$philosopher 2 fork2 fork3 90 printL
	philosopher 3 fork3 fork1 1000 printL