Merge branch 'dht-only' of 10.0.0.137:bittorrent into dht-only

2 files changed, 83 insertions, 0 deletions

diff --git a/src/Control/Concurrent/Async/Lifted/Instrument.hs b/src/Control/Concurrent/Async/Lifted/Instrument.hs
new file mode 100644
index 00000000..eab0fadc
--- /dev/null
+++ b/src/Control/Concurrent/Async/Lifted/Instrument.hs
@@ -0,0 +1,5 @@
+module Control.Concurrent.Async.Lifted.Instrument
+    ( module Control.Concurrent.Async.Lifted
+    ) where
+
+import Control.Concurrent.Async.Lifted
diff --git a/src/Control/Concurrent/Lifted/Instrument.hs b/src/Control/Concurrent/Lifted/Instrument.hs
new file mode 100644
index 00000000..9ec5deef
--- /dev/null
+++ b/src/Control/Concurrent/Lifted/Instrument.hs
@@ -0,0 +1,78 @@
+{-# LANGUAGE FlexibleContexts #-}
+module Control.Concurrent.Lifted.Instrument
+    ( module Control.Concurrent.Lifted
+    , forkIO
+    , fork
+    , labelThread
+    , threadsInformation
+    , PerThread(..)
+    ) where
+
+import qualified Control.Concurrent.Lifted as Raw
+import Control.Concurrent.Lifted hiding (fork)
+import Control.Monad.Trans.Control
+import System.IO.Unsafe
+import System.Mem.Weak
+import qualified Data.Map as Map
+import qualified Data.IntMap as IntMap
+import Control.Exception.Lifted
+import Control.Monad.Base
+-- import Control.Monad.IO.Class
+import qualified GHC.Conc as GHC
+import Data.Time.Clock
+
+data PerThread = PerThread
+ { -- wkid :: Weak ThreadId
+   lbl :: String
+ , startTime :: UTCTime
+ }
+ deriving (Eq,Ord,Show) -- ,Data,Generic)
+
+data GlobalState = GlobalState
+ { threads :: Map.Map ThreadId PerThread -- IntMap.IntMap PerThread
+ -- , uniqSource :: Int
+ }
+
+globals :: MVar GlobalState
+globals = unsafePerformIO $ newMVar $ GlobalState
+    { threads = Map.empty
+    -- , uniqSource = 0
+    }
+{-# NOINLINE globals #-}
+
+
+forkIO :: IO () -> IO ThreadId
+forkIO = fork
+
+fork :: MonadBaseControl IO m => m () -> m ThreadId
+fork action = do
+    t <- Raw.fork $ do
+        -- wkid <- myThreadId >>= liftBase . mkWeakThreadId
+        -- tid <- newUniq
+        tid <- myThreadId
+        tm <- liftBase getCurrentTime
+        bracket_ (modifyThreads $ Map.insert tid (PerThread "" tm))
+                 (modifyThreads $ Map.delete tid)
+                 action
+    return t
+
+labelThread :: ThreadId -> String -> IO ()
+labelThread tid s = do
+    GHC.labelThread tid s
+    putStrLn $ "labelThread "++s++" "++show tid
+    modifyThreads $ Map.adjust (\pt -> pt { lbl = s }) tid
+
+threadsInformation :: IO [PerThread]
+threadsInformation = do
+    m <- threads <$> readMVar globals
+    return $ Map.elems m
+
+-- newUniq :: MonadBaseControl IO m => m Int
+-- newUniq = do
+--     modifyMVar globals (\st -> return (st { uniqSource = succ (uniqSource st) }, uniqSource st))
+ 
+modifyThreads :: MonadBase IO m => (Map.Map ThreadId PerThread -> Map.Map ThreadId PerThread) -> m ()
+modifyThreads f = do
+    g <- takeMVar globals
+    let f' st = st { threads = f (threads st) }
+    putMVar globals (f' g)