1 files changed, 75 insertions, 0 deletions

diff --git a/bittorrent/bench/Main.hs b/bittorrent/bench/Main.hs
new file mode 100644
index 00000000..f04485ab
--- /dev/null
+++ b/bittorrent/bench/Main.hs
@@ -0,0 +1,75 @@
+{-# LANGUAGE OverloadedStrings     #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
+{-# OPTIONS -fno-warn-orphans #-}
+module Main (main) where
+
+import Control.DeepSeq
+import Network
+import Control.Monad
+import Control.Monad.Logger
+import Control.Monad.Reader
+import Criterion.Main
+import Data.ByteString as BS
+import Network.DatagramServer
+
+
+import Network.BitTorrent.Exchange.Protocol as BT
+import Data.Torrent.Block as BT
+import Data.Torrent.Bitfield as BT
+
+instance KRPC ByteString ByteString where
+  method = "echo"
+
+instance MonadLogger IO where
+  monadLoggerLog _ _ _ _ = return ()
+
+
+instance NFData PortNumber where
+  rnf = rnf . (fromIntegral :: PortNumber -> Int)
+
+instance NFData BlockIx where
+  rnf (BlockIx a b c) = a `deepseq` b `deepseq` rnf c
+
+instance NFData Block where
+  rnf (Block a b c) = a `deepseq` b `deepseq` rnf c
+
+instance NFData Bitfield
+
+instance NFData Message where
+  rnf (Have i)     = rnf i
+  rnf (Bitfield b) = rnf b
+  rnf (Request  b) = rnf b
+  rnf (Piece    b) = rnf b
+  rnf (Cancel   b) = rnf b
+  rnf (Port     i) = rnf i
+  rnf _ = ()  -- other fields are forced by pattern matching
+
+{-
+encodeMessages :: [Message] -> ByteString
+encodeMessages xs = runPut (mapM_ put xs)
+
+decodeMessages :: ByteString -> Either String [Message]
+decodeMessages = runGet (many get)
+-}
+
+echo :: Handler IO
+echo = handler $ \ _ bs -> return (bs :: ByteString)
+
+addr :: SockAddr
+addr = SockAddrInet 6000 (256 * 256 * 256 + 127)
+
+-- main :: IO ()
+-- main = defaultMain []
+main :: IO ()
+main = withManager def addr [echo] $ \ m -> (`runReaderT` m) $ do
+    listen
+    liftIO $ defaultMain (benchmarks m)
+  where
+    sizes        = [10, 100, 1000, 10000, 16 * 1024]
+    repetitions  = [1, 10, 100, 1000]
+    benchmarks m = [mkbench m r s | r <- repetitions, s <- sizes]
+      where
+        mkbench action r n =
+          bench (show r ++ "times" ++ "/" ++ show n ++ "bytes") $ nfIO $
+            replicateM r $
+              runReaderT (query addr (BS.replicate n 0)) action