module Resolve.DNS.Channel where
import qualified Resolve.Types as T
import Resolve.DNS.Types
import qualified Resolve.DNS.Encode as E
import qualified Resolve.DNS.Decode as D
import Data.Word
import Data.Hashable
import Data.Either
import Data.Maybe
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BSL
import Control.Monad.STM
import Control.Concurrent
import Control.Exception
import Control.Monad
import Control.Monad.Trans.Class
import Control.Monad.Trans.Except
import Control.Monad.Trans.Reader
import Control.Concurrent.STM.TVar
import Control.Concurrent.STM.TMVar
import Data.Typeable
import System.Log.Logger
import qualified STMContainers.Map as M
nameM = "Resolve.DNS.Channel"
type Send = ByteString -> IO ()
type Recv = IO ByteString
data EncodeError = EncodeError String
deriving (Typeable, Show)
data Dead = Dead
deriving (Typeable, Show)
instance Exception EncodeError where
toException = dnsExceptionToException
fromException = dnsExceptionFromException
instance Exception Dead where
toException = dnsExceptionToException
fromException = dnsExceptionFromException
data Config = Config { send :: Send
, recv :: Recv
, nick :: String
}
data Resolver = Resolver { tid :: ThreadId
, book :: M.Map Word16 (TMVar Message)
, config :: Config
, dead :: TVar Bool
}
new :: Config -> IO (T.Resolver Message Message)
new c = do
b <- M.newIO
si <- newTVarIO False
bracketOnError
(forkIOWithUnmask $ \unmask -> unmask $ finally (forever $ runExceptT $ do
let nameF = nameM ++ ".recv"
bs <- lift $ recv c
m <- ExceptT $ return $ D.decodeMessage bs
lift $ debugM nameF $ (nick c) ++ " -> recvd " ++ (show m)
let ident' = (ident $ header $ m)
r <- lift $ atomically $ lookupAndDelete b ident'
case r of
Nothing -> throwE "ID not in book"
Just mvar -> lift $ atomically $ tryPutTMVar mvar m )
(do
debugM nameM $ (nick c) ++ " died"
atomically $ writeTVar si True))
killThread
(\t -> do
let chan = Resolver { tid = t
, book = b
, config = c
, dead = si
}
debugM nameM $ (nick c) ++ " created"
return $ T.Resolver { T.delete = delete chan
, T.resolve = resolve chan
}
)
delete r = killThread (tid r)
resolve r a = do
let nameF = nameM ++ ".resolve"
mvar <- newEmptyTMVarIO
bracketOnError
(atomically $ allocate (book r) (ident $ header $ a) mvar)
(\ident_ -> atomically $ lookupAndDelete (book r) ident_)
(\ident_ -> do
let a_ = a { header = (header a) { ident = ident_ }}
bs <- case E.encode E.message a_ of
Left s -> throw $ EncodeError s
Right b -> return $ BSL.toStrict b
forkIO $ catch
(do
debugM nameF $ (nick $ config $ r) ++ " <- sending " ++ (show a_)
(send $ config $ r) bs
debugM nameF $ (nick $ config $ r) ++ " <- sent ")
(\e -> debugM nameF $ "exception when sending: " ++ show (e :: SomeException))
x <- atomically $ do
a <- readTVar $ dead r
if a then tryTakeTMVar mvar
else
Just <$> takeTMVar mvar
case x of
Nothing -> throwIO Dead
Just x -> return x
)
allocate :: (Eq i, Hashable i, Num i) => M.Map i a -> i -> a -> STM i
allocate b i a = let loop i = do
m <- do
r <- M.lookup i b
case r of
Nothing -> do
M.insert a i b
return True
Just _ -> return False
if m then return i
else loop (i + 1)
in loop i
lookupAndDelete :: (Eq i, Hashable i) => M.Map i a -> i -> STM (Maybe a)
lookupAndDelete b i = do
mvar <- M.lookup i b
case mvar of
Nothing -> return Nothing
Just mvar -> do
M.delete i b
return $ Just mvar