5 files changed, 303 insertions, 185 deletions

diff --git a/krpc.cabal b/krpc.cabal
index aa081a54..bccdd6c3 100644
--- a/krpc.cabal
+++ b/krpc.cabal
@@ -37,7 +37,10 @@ library
                      , RecordWildCards
   hs-source-dirs:      src
   exposed-modules:     Network.KRPC
-                     , Network.KRPC.Message
+                       Network.KRPC.Message
+                       Network.KRPC.Method
+                       Network.KRPC.Manager
+
   build-depends:       base          == 4.*
                      , bytestring   >= 0.10
                      , lifted-base   >= 0.1.1
@@ -45,6 +48,7 @@ library
                      , monad-control >= 0.3
                      , bencoding    >= 0.4.3
                      , network      >= 2.3
+                     , containers
 
   if impl(ghc < 7.6)
     build-depends:     ghc-prim
diff --git a/src/Network/KRPC.hs b/src/Network/KRPC.hs
index a96d8da9..09d1c5b2 100644
--- a/src/Network/KRPC.hs
+++ b/src/Network/KRPC.hs
@@ -97,190 +97,23 @@
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE FunctionalDependencies #-}
 module Network.KRPC
-       ( KRPC (..)
+       ( -- * Methods
+         Method
+       , KRPC (..)
 
-         -- * Exception
-       , KError (..)
-
-         -- * Method
-       , Method(..)
+         -- * RPC
+       , handler
+       , query
 
-         -- * Client
-       , call
+         -- * Manager
+       , MonadKRPC (..)
+       , newManager
+--       , closeManager
 
-         -- * Server
-       , MethodHandler
-       , handler
-       , server
+         -- * Exceptions
+       , KError (..)
        ) where
 
-import Control.Applicative
-import Control.Exception.Lifted as Lifted
-import Control.Monad
-import Control.Monad.Trans.Control
-import Control.Monad.IO.Class
-import Data.BEncode as BE
-import Data.ByteString.Char8 as BC
-import Data.ByteString.Lazy as BL
-import Data.List as L
-import Data.Monoid
-import Data.String
-import Data.Typeable
-import Network
-import Network.Socket
-import Network.Socket.ByteString as BS
-
 import Network.KRPC.Message
-
-
-class (BEncode req, BEncode resp) => KRPC req resp | req -> resp where
-  method :: Method req resp
-
--- | Method datatype used to describe name, parameters and return
---   values of procedure. Client use a method to /invoke/, server
---   /implements/ the method to make the actual work.
---
---   We use the following fantom types to ensure type-safiety:
---
---     * param: Type of method parameters. Ordinary Tuple type used
---     to specify more than one parameter, so for example @Method
---     (Int, Int) result@ will take two arguments.
---
---     * result: Type of return value of the method. Similarly,
---     tuple used to specify more than one return value, so for
---     exsample @Method (Foo, Bar) (Bar, Foo)@ will take two arguments
---     and return two values.
---
-newtype Method param result = Method MethodName
-  deriving (Eq, Ord, IsString, BEncode)
-
-instance (Typeable a, Typeable b) => Show (Method a b) where
-  showsPrec _ = showsMethod
-
-showsMethod :: forall a. forall b. Typeable a => Typeable b
-            => Method a b -> ShowS
-showsMethod (Method name) =
-    shows name <>
-    showString " :: " <>
-    shows paramsTy <>
-    showString " -> " <>
-    shows valuesTy
-  where
-    impossible = error "KRPC.showsMethod: impossible"
-    paramsTy = typeOf (impossible :: a)
-    valuesTy = typeOf (impossible :: b)
-
-{-----------------------------------------------------------------------
---  Client
------------------------------------------------------------------------}
-
-sendMessage :: BEncode msg => msg -> SockAddr -> Socket -> IO ()
-sendMessage msg addr sock = sendManyTo sock (BL.toChunks (encode msg)) addr
-{-# INLINE sendMessage #-}
-
-maxMsgSize :: Int
---maxMsgSize = 512 -- release: size of payload of one udp packet
-maxMsgSize = 64 * 1024 -- bench: max UDP MTU
-{-# INLINE maxMsgSize #-}
-
-recvResponse :: Socket -> IO (Either KError KResponse)
-recvResponse sock = do
-  (raw, _) <- BS.recvFrom sock maxMsgSize
-  return $ case decode raw of
-    Right resp -> Right resp
-    Left decE -> Left $ case decode raw of
-      Right kerror -> kerror
-      _ -> KError ProtocolError (BC.pack decE) undefined
-
-withRemote :: (MonadBaseControl IO m, MonadIO m) => (Socket -> m a) -> m a
-withRemote = bracket (liftIO (socket AF_INET6 Datagram defaultProtocol))
-                     (liftIO .  sClose)
-{-# SPECIALIZE withRemote :: (Socket -> IO a) -> IO a #-}
-
-getResult :: BEncode result => Socket -> IO result
-getResult sock = do
-  KResponse {..} <- either throw return =<< recvResponse sock
-  case fromBEncode respVals of
-    Left msg -> throw $ KError ProtocolError (BC.pack msg) respId
-    Right r  -> return r
-
--- | Makes remote procedure call. Throws RPCException on any error
--- occurred.
-call :: forall req resp host.
-        (MonadBaseControl IO host, MonadIO host, KRPC req resp)
-     => SockAddr -> req -> host resp
-call addr arg = liftIO $ withRemote $ \sock -> do
-     sendMessage (KQuery (toBEncode arg) name undefined) addr sock
-     getResult sock
-  where
-    Method name = method :: Method req resp
-
-{-----------------------------------------------------------------------
---  Server
------------------------------------------------------------------------}
-
-type HandlerBody remote = SockAddr -> KQuery -> remote (Either KError KResponse)
-
--- | Procedure signature and implementation binded up.
-type MethodHandler remote = (MethodName, HandlerBody remote)
-
--- | Similar to '==>@' but additionally pass caller address.
-handler  :: forall (remote :: * -> *) (req :: *) (resp :: *).
-           (KRPC req resp, Monad remote)
-        => (SockAddr -> req -> remote resp) -> MethodHandler remote
-handler body = (name, newbody)
-  where
-    Method name = method :: Method req resp
-
-    {-# INLINE newbody #-}
-    newbody addr KQuery {..} =
-      case fromBEncode queryArgs of
-        Left  e -> return $ Left $ KError ProtocolError (BC.pack e) queryId
-        Right a -> do
-          r <- body addr a
-          return $ Right $ KResponse (toBEncode r) queryId
-
-sockAddrFamily :: SockAddr -> Family
-sockAddrFamily (SockAddrInet  _ _    ) = AF_INET
-sockAddrFamily (SockAddrInet6 _ _ _ _) = AF_INET6
-sockAddrFamily (SockAddrUnix  _      ) = AF_UNIX
-
--- | Run server using a given port. Method invocation should be done manually.
-remoteServer :: (MonadBaseControl IO remote, MonadIO remote)
-             => SockAddr -- ^ Port number to listen.
-             -> (SockAddr -> KQuery -> remote (Either KError KResponse))
-             -> remote ()
-remoteServer servAddr action = bracket (liftIO bindServ) (liftIO . sClose) loop
-  where
-    bindServ = do
-        let family = sockAddrFamily servAddr
-        sock <- socket family Datagram defaultProtocol
-        when (family == AF_INET6) $ do
-          setSocketOption sock IPv6Only 0
-        bindSocket sock servAddr
-        return sock
-
-    loop sock = forever $ do
-        (bs, addr) <- liftIO $ BS.recvFrom sock maxMsgSize
-        reply <- handleMsg bs addr
-        liftIO $ sendMessage reply addr sock
-      where
-        handleMsg bs addr = case decode bs of
-          Right query -> (either toBEncode toBEncode <$> action addr query)
-              `Lifted.catch` (return . toBEncode . (`serverError` undefined ))
-          Left decodeE   -> return $ toBEncode $
-             KError ProtocolError (BC.pack decodeE) undefined
-
--- | Run RPC server on specified port by using list of handlers.
---   Server will dispatch procedure specified by callee, but note that
---   it will not create new thread for each connection.
---
-server :: (MonadBaseControl IO remote, MonadIO remote)
-       => SockAddr                -- ^ Port used to accept incoming connections.
-       -> [MethodHandler remote]  -- ^ Method table.
-       -> remote ()
-server servAddr handlers = do
-  remoteServer servAddr $ \addr q @ KQuery {..} -> do
-    case L.lookup  queryMethod handlers of
-      Nothing -> return $ Left $ KError MethodUnknown queryMethod queryId
-      Just  m -> m addr q
+import Network.KRPC.Method
+import Network.KRPC.Manager
\ No newline at end of file
diff --git a/src/Network/KRPC/Manager.hs b/src/Network/KRPC/Manager.hs
new file mode 100644
index 00000000..9aa1bea7
--- /dev/null
+++ b/src/Network/KRPC/Manager.hs
@@ -0,0 +1,179 @@
+{-# LANGUAGE FlexibleContexts    #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+module Network.KRPC.Manager
+       ( MonadKRPC (..)
+       , newManager
+       , query
+       , handler
+       ) where
+
+import Control.Applicative
+import Control.Arrow
+import Control.Concurrent
+--import Control.Exception hiding (Handler)
+import Control.Exception.Lifted as Lifted hiding (Handler)
+import Control.Monad
+import Control.Monad.Trans.Control
+import Control.Monad.IO.Class
+import Data.BEncode as BE
+import Data.ByteString.Char8 as BC
+import Data.ByteString.Lazy as BL
+import Data.IORef
+import Data.List as L
+import Data.Map as M
+import Network.KRPC.Message
+import Network.KRPC.Method
+import Network.Socket
+import Network.Socket.ByteString as BS
+
+
+type KResult = Either KError KResponse
+
+type TransactionCounter = IORef Int
+type CallId             = (TransactionId, SockAddr)
+type CallRes            = MVar KResult
+type PendingCalls       = IORef (Map CallId CallRes)
+
+type HandlerBody m = SockAddr -> BValue -> m (BE.Result BValue)
+type Handler     m = (MethodName, HandlerBody m)
+
+data Manager m = Manager
+  { sock               :: !Socket
+  , transactionCounter :: {-# UNPACK #-} !TransactionCounter
+  , pendingCalls       :: {-# UNPACK #-} !PendingCalls
+  , handlers           :: [Handler m]
+  }
+
+class (MonadBaseControl IO m, MonadIO m) => MonadKRPC m where
+  getManager :: m (Manager a)
+
+sockAddrFamily :: SockAddr -> Family
+sockAddrFamily (SockAddrInet  _ _    ) = AF_INET
+sockAddrFamily (SockAddrInet6 _ _ _ _) = AF_INET6
+sockAddrFamily (SockAddrUnix  _      ) = AF_UNIX
+
+seedTransaction :: Int
+seedTransaction = 0
+
+newManager :: SockAddr -> IO (Manager a)
+newManager servAddr = do
+    sock  <- bindServ
+    tran  <- newIORef seedTransaction
+    calls <- newIORef M.empty
+    return $ Manager sock tran calls []
+  where
+    bindServ = do
+      let family = sockAddrFamily servAddr
+      sock <- socket family Datagram defaultProtocol
+      when (family == AF_INET6) $ do
+        setSocketOption sock IPv6Only 0
+      bindSocket sock servAddr
+      return sock
+
+sendMessage :: MonadIO m => BEncode a => Socket -> SockAddr -> a -> m ()
+sendMessage sock addr a =
+  liftIO $ sendManyTo sock (BL.toChunks (BE.encode a)) addr
+
+{-----------------------------------------------------------------------
+--  Client
+-----------------------------------------------------------------------}
+
+genTransactionId :: TransactionCounter -> IO TransactionId
+genTransactionId ref = do
+  cur <- atomicModifyIORef' ref $ \ cur -> (succ cur, cur)
+  return $ BC.pack (show cur)
+
+registerQuery :: CallId -> PendingCalls -> IO CallRes
+registerQuery cid ref = do
+  ares <- newEmptyMVar
+  atomicModifyIORef' ref $ \ m -> (M.insert cid ares m, ())
+  return ares
+
+unregisterQuery :: CallId -> PendingCalls -> IO ()
+unregisterQuery cid ref = do
+  atomicModifyIORef' ref $ \ m -> (M.delete cid m, ())
+
+queryResponse :: BEncode a => CallRes -> IO a
+queryResponse ares = do
+  res  <- readMVar ares
+  case res of
+    Left  e -> throwIO e
+    Right (KResponse {..}) ->
+      case fromBEncode respVals of
+        Left e  -> throwIO (KError ProtocolError (BC.pack e) respId)
+        Right a -> return a
+
+query :: forall m a b. (MonadKRPC m, KRPC a b) => SockAddr -> a -> m b
+query addr params = do
+  Manager {..} <- getManager
+  liftIO $ do
+    tid <- genTransactionId transactionCounter
+    let Method name = method :: Method a b
+    let q = KQuery (toBEncode params) name tid
+    ares <- registerQuery (tid, addr) pendingCalls
+    sendMessage sock addr q
+      `onException` unregisterQuery (tid, addr) pendingCalls
+    queryResponse ares
+
+{-----------------------------------------------------------------------
+--  Handlers
+-----------------------------------------------------------------------}
+
+handler :: forall m a b. (KRPC a b, MonadKRPC m)
+        => (SockAddr -> a -> m b) -> Handler m
+handler body = (name, wrapper)
+  where
+    Method name = method :: Method a b
+    wrapper addr args =
+      case fromBEncode args of
+        Left  e -> return $ Left e
+        Right a -> (Right . toBEncode) <$> body addr a
+
+runHandler :: MonadKRPC m => HandlerBody m -> SockAddr -> KQuery -> m KResult
+runHandler handler addr KQuery {..} = wrapper `Lifted.catch` failback
+  where
+    wrapper = ((`decodeError` queryId) +++ (`KResponse` queryId))
+           <$> handler addr queryArgs
+    failback e = return $ Left $ serverError e queryId
+
+dispatchHandler :: MonadKRPC m => KQuery -> SockAddr -> m KResult
+dispatchHandler q @ KQuery {..} addr = do
+  Manager {..} <- getManager
+  case L.lookup queryMethod handlers of
+    Nothing      -> return $ Left $ unknownMethod queryMethod queryId
+    Just handler -> runHandler handler addr q
+
+{-----------------------------------------------------------------------
+--  Listener
+-----------------------------------------------------------------------}
+
+handleQuery :: MonadKRPC m => KQuery -> SockAddr -> m ()
+handleQuery q addr = do
+  Manager {..} <- getManager
+  res <- dispatchHandler q addr
+  sendMessage sock addr $ either toBEncode toBEncode res
+
+handleResponse :: MonadKRPC m => KResult -> SockAddr -> m ()
+handleResponse result addr = do
+  Manager {..} <- getManager
+  mcall <- undefined (addr, respId) pendingCalls
+  case mcall of
+    Nothing   -> return ()
+    Just ares -> liftIO $ putMVar ares result
+
+handleMessage :: MonadKRPC m => KMessage -> SockAddr -> m ()
+handleMessage (Q q) = handleQuery    q
+handleMessage (R r) = handleResponse (Right r)
+handleMessage (E e) = handleResponse (Left  e)
+
+maxMsgSize :: Int
+maxMsgSize = 64 * 1024
+
+listener :: MonadKRPC m => m ()
+listener = do
+  Manager {..} <- getManager
+  forever $ do
+    (bs, addr) <- liftIO $ BS.recvFrom sock maxMsgSize
+    case BE.decode bs of
+      Left  e -> liftIO $ sendMessage sock addr  $ unknownMessage e
+      Right m -> handleMessage m addr
diff --git a/src/Network/KRPC/Message.hs b/src/Network/KRPC/Message.hs
index a70c2ea9..3bbfb1db 100644
--- a/src/Network/KRPC/Message.hs
+++ b/src/Network/KRPC/Message.hs
@@ -19,10 +19,17 @@
 {-# LANGUAGE DefaultSignatures      #-}
 {-# LANGUAGE DeriveDataTypeable     #-}
 module Network.KRPC.Message
-       ( -- * Error
-         ErrorCode (..)
+       ( -- * Transaction
+         TransactionId
+       , unknownTransaction
+
+         -- * Error
+       , ErrorCode (..)
        , KError(..)
        , serverError
+       , decodeError
+       , unknownMethod
+       , unknownMessage
 
          -- * Query
        , KQuery(..)
@@ -30,6 +37,9 @@ module Network.KRPC.Message
 
          -- * Response
        , KResponse(..)
+
+         -- * Message
+       , KMessage (..)
        ) where
 
 import Control.Applicative
@@ -47,6 +57,9 @@ import Data.Typeable
 -- are enough as they cover 2^16 outstanding queries.
 type TransactionId = ByteString
 
+unknownTransaction :: TransactionId
+unknownTransaction = ""
+
 {-----------------------------------------------------------------------
 -- Error messages
 -----------------------------------------------------------------------}
@@ -120,6 +133,15 @@ instance Exception KError
 serverError :: SomeException -> TransactionId -> KError
 serverError e = KError ServerError (BC.pack (show e))
 
+decodeError :: String -> TransactionId -> KError
+decodeError msg = KError ProtocolError (BC.pack msg)
+
+unknownMethod :: MethodName -> TransactionId -> KError
+unknownMethod = KError MethodUnknown
+
+unknownMessage :: String -> KError
+unknownMessage msg = KError ProtocolError (BC.pack msg) ""
+
 {-----------------------------------------------------------------------
 -- Query messages
 -----------------------------------------------------------------------}
@@ -183,4 +205,23 @@ instance BEncode KResponse where
   fromBEncode = fromDict $ do
     lookAhead $ match "y" (BString "r")
     KResponse <$>! "r" <*>! "t"
-  {-# INLINE fromBEncode #-}
\ No newline at end of file
+  {-# INLINE fromBEncode #-}
+
+{-----------------------------------------------------------------------
+-- Summed messages
+-----------------------------------------------------------------------}
+
+data KMessage
+  = Q KQuery
+  | R KResponse
+  | E KError
+
+instance BEncode KMessage where
+  toBEncode (Q q) = toBEncode q
+  toBEncode (R r) = toBEncode r
+  toBEncode (E e) = toBEncode e
+
+  fromBEncode b =
+        Q <$> fromBEncode b
+    <|> R <$> fromBEncode b
+    <|> E <$> fromBEncode b
diff --git a/src/Network/KRPC/Method.hs b/src/Network/KRPC/Method.hs
new file mode 100644
index 00000000..54aa8ef0
--- /dev/null
+++ b/src/Network/KRPC/Method.hs
@@ -0,0 +1,61 @@
+{-# LANGUAGE RankNTypes                 #-}
+{-# LANGUAGE MultiParamTypeClasses      #-}
+{-# LANGUAGE FunctionalDependencies     #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+module Network.KRPC.Method
+       ( Method (..)
+       , KRPC (..)
+       ) where
+
+import Data.BEncode (BEncode)
+import Data.Monoid
+import Data.String
+import Data.Typeable
+import Network.KRPC.Message
+
+
+-- | Method datatype used to describe name, parameters and return
+--   values of procedure. Client use a method to /invoke/, server
+--   /implements/ the method to make the actual work.
+--
+--   We use the following fantom types to ensure type-safiety:
+--
+--     * param: Type of method parameters. Ordinary Tuple type used
+--     to specify more than one parameter, so for example @Method
+--     (Int, Int) result@ will take two arguments.
+--
+--     * result: Type of return value of the method. Similarly,
+--     tuple used to specify more than one return value, so for
+--     exsample @Method (Foo, Bar) (Bar, Foo)@ will take two arguments
+--     and return two values.
+--
+newtype Method param result = Method MethodName
+  deriving (Eq, Ord, IsString, BEncode)
+
+instance (Typeable a, Typeable b) => Show (Method a b) where
+  showsPrec _ = showsMethod
+
+showsMethod :: forall a. forall b. Typeable a => Typeable b
+            => Method a b -> ShowS
+showsMethod (Method name) =
+    shows name <>
+    showString " :: " <>
+    shows paramsTy <>
+    showString " -> " <>
+    shows valuesTy
+  where
+    impossible = error "KRPC.showsMethod: impossible"
+    paramsTy = typeOf (impossible :: a)
+    valuesTy = typeOf (impossible :: b)
+
+-- | Example:
+--   @
+--   data Ping = Ping Text deriving BEncode
+--   data Pong = Pong Text deriving BEncode
+--
+--   instance KRPC Ping Pong where
+--     method = "ping"
+--   @
+class (BEncode req, BEncode resp) => KRPC req resp | req -> resp where
+  method :: Method req resp