2 files changed, 9 insertions, 6 deletions

diff --git a/src/Network/BitTorrent/Exchange.hs b/src/Network/BitTorrent/Exchange.hs
index 5be9ae73..a0e7df8c 100644
--- a/src/Network/BitTorrent/Exchange.hs
+++ b/src/Network/BitTorrent/Exchange.hs
@@ -171,10 +171,7 @@ runP2P se addr p2p = waitVacancy se $ runSession se addr p2p
 -- some other loop. Note that this function may block while waiting
 -- for a vacant place: use forkIO and runP2P instead.
 spawnP2P :: SwarmSession -> PeerAddr -> P2P () -> IO ThreadId
-spawnP2P se addr p2p = do
-  enterSwarm se
-  forkIO $ do
-    runSession se addr p2p `finally` leaveSwarm se
+spawnP2P se addr p2p = forkThrottle se $ runSession se addr p2p
 
 -- TODO unify this all using PeerConnection
 {-
diff --git a/src/Network/BitTorrent/Internal.lhs b/src/Network/BitTorrent/Internal.lhs
index 7edcc977..5f6ad458 100644
--- a/src/Network/BitTorrent/Internal.lhs
+++ b/src/Network/BitTorrent/Internal.lhs
@@ -56,9 +56,8 @@
 >        , newSeeder
 >        , getClientBitfield
 >
->        , enterSwarm
->        , leaveSwarm
 >        , waitVacancy
+>        , forkThrottle
 >
 >        , pieceLength
 >
@@ -616,6 +615,13 @@ Peer sessions throttling
 >   bracket (enterSwarm se) (const (leaveSwarm se))
 >                   . const
 
+> forkThrottle :: SwarmSession -> IO () -> IO ThreadId
+> forkThrottle se action = do
+>   enterSwarm se
+>   (forkIO $ do
+>     action `finally` leaveSwarm se)
+>        `onException` leaveSwarm se
+
 Peer sessions
 ------------------------------------------------------------------------