2 files changed, 44 insertions, 22 deletions

diff --git a/src/Network/BitTorrent/Client.hs b/src/Network/BitTorrent/Client.hs
index 349283c3..905d3384 100644
--- a/src/Network/BitTorrent/Client.hs
+++ b/src/Network/BitTorrent/Client.hs
@@ -39,10 +39,10 @@ import Control.Exception
 import Control.Concurrent
 import Control.Monad.Logger
 import Control.Monad.Trans
+import Control.Monad.Trans.Resource
 
 import Data.Default
 import Data.HashMap.Strict as HM
-import Data.Maybe
 import Data.Text
 import Network
 
@@ -86,38 +86,52 @@ exchangeOptions pid Options {..} = Exchange.Options
 connHandler :: MVar (HashMap InfoHash Handle) -> Exchange.Handler
 connHandler _tmap = undefined
 
-newClient :: Options -> LogFun -> IO Client
-newClient opts @ Options {..} logger = do
-  pid  <- genPeerId
-  tmap <- newMVar HM.empty
-  tmgr <- Tracker.newManager def (PeerInfo pid Nothing optPort)
-  emgr <- Exchange.newManager (exchangeOptions pid opts) (connHandler tmap)
-  node <- do
-    node <- startNode defaultHandlers def optNodeAddr logger
-    runDHT node $ bootstrap (maybeToList optBootNode)
-    return node
+initClient :: Options -> LogFun -> ResIO Client
+initClient opts @ Options {..} logFun = do
+  pid  <- liftIO genPeerId
+  tmap <- liftIO $ newMVar HM.empty
+
+  let peerInfo = PeerInfo pid Nothing optPort
+  let mkTracker = Tracker.newManager def peerInfo
+  (_, tmgr) <- allocate mkTracker Tracker.closeManager
+
+  let mkEx = Exchange.newManager (exchangeOptions pid opts) (connHandler tmap)
+  (_, emgr) <- allocate mkEx Exchange.closeManager
+
+  let mkNode = startNode defaultHandlers def optNodeAddr logFun
+  (_, node) <- allocate mkNode stopNode
+
+  resourceMap <- getInternalState
   return Client
     { clientPeerId       = pid
     , clientListenerPort = optPort
     , allowedExtensions  = toCaps optExtensions
+    , clientResources    = resourceMap
     , trackerManager     = tmgr
     , exchangeManager    = emgr
     , clientNode         = node
     , clientTorrents     = tmap
-    , clientLogger       = logger
+    , clientLogger       = logFun
     }
 
+newClient :: Options -> LogFun -> IO Client
+newClient opts logFun = do
+  s <- createInternalState
+  runInternalState (initClient opts logFun) s
+    `onException` closeInternalState s
+
 closeClient :: Client -> IO ()
-closeClient Client {..} = do
-  Exchange.closeManager exchangeManager
-  Tracker.closeManager  trackerManager
-  DHT.stopNode clientNode
-  return ()
---  closeNode clientNode
+closeClient Client {..} = closeInternalState clientResources
 
 withClient :: Options -> LogFun -> (Client -> IO a) -> IO a
 withClient opts lf action = bracket (newClient opts lf) closeClient action
 
+-- do not perform IO in 'initClient', do it in the 'boot'
+--boot :: BitTorrent ()
+--boot = do
+--  Options {..} <- asks options
+--  liftDHT $ bootstrap (maybeToList optBootNode)
+
 -- | Run bittorrent client with default options and log to @stderr@.
 --
 --   For testing purposes only.
diff --git a/src/Network/BitTorrent/Client/Types.hs b/src/Network/BitTorrent/Client/Types.hs
index 63971518..c2f79460 100644
--- a/src/Network/BitTorrent/Client/Types.hs
+++ b/src/Network/BitTorrent/Client/Types.hs
@@ -13,6 +13,7 @@ module Network.BitTorrent.Client.Types
        , MonadBitTorrent (..)
        ) where
 
+import Control.Applicative
 import Control.Concurrent
 import Control.Monad.Logger
 import Control.Monad.Reader
@@ -40,6 +41,7 @@ data Client = Client
   { clientPeerId       :: !PeerId
   , clientListenerPort :: !PortNumber
   , allowedExtensions  :: !Caps
+  , clientResources    :: !InternalState
   , trackerManager     :: !Tracker.Manager
   , exchangeManager    :: !Exchange.Manager
   , clientNode         :: !(Node IPv4)
@@ -68,8 +70,10 @@ externalAddr Client {..} = PeerAddr
 -----------------------------------------------------------------------}
 
 newtype BitTorrent a = BitTorrent
-  { unBitTorrent :: ReaderT Client (ResourceT IO) a
-  } deriving (Functor, Monad, MonadIO)
+  { unBitTorrent :: ReaderT Client IO a
+  } deriving ( Functor, Applicative, Monad
+             , MonadIO, MonadThrow, MonadUnsafeIO
+             )
 
 class MonadBitTorrent m where
   liftBT :: BitTorrent a -> m a
@@ -77,6 +81,11 @@ class MonadBitTorrent m where
 instance MonadBitTorrent BitTorrent where
   liftBT = id
 
+instance MonadResource BitTorrent where
+  liftResourceT m = BitTorrent $ do
+    s <- asks clientResources
+    liftIO $ runInternalState m s
+
 instance MonadDHT BitTorrent where
   liftDHT action = BitTorrent $ do
     node <- asks clientNode
@@ -88,8 +97,7 @@ instance MonadLogger BitTorrent where
     liftIO $ logger loc src lvl (toLogStr msg)
 
 runBitTorrent :: Client -> BitTorrent a -> IO a
-runBitTorrent client action = runResourceT $
-  runReaderT (unBitTorrent action) client
+runBitTorrent client action = runReaderT (unBitTorrent action) client
 {-# INLINE runBitTorrent #-}
 
 getClient :: BitTorrent Client