1{- Benchmark server based upon Jasper van der Jeugt's 'BenchmarkServer.lhs'
2   from blaze-html. Modified for network-2.3 by Simon Meier <iridcode@gmail.com>
3-}
4
5{-# LANGUAGE OverloadedStrings #-}
6module BenchmarkServer where
7
8import Prelude hiding (putStrLn)
9
10import Data.Char   (ord)
11import Data.Monoid
12import Data.ByteString.Char8 () -- IsString instance only
13import qualified Data.ByteString               as S
14import qualified Data.ByteString.Lazy          as L
15import qualified Data.ByteString.Lazy.Internal as L
16
17import Control.Concurrent (forkIO, putMVar, takeMVar, newEmptyMVar)
18import Control.Exception  (bracket)
19import Control.Monad
20
21import Network.Socket   (Socket, accept, sClose)
22import Network          (listenOn, PortID (PortNumber))
23import Network.Socket.ByteString      as S
24import Network.Socket.ByteString.Lazy as L
25
26import System (getArgs)
27
28import Blaze.ByteString.Builder
29import Blaze.ByteString.Builder.Internal (defaultBufferSize, defaultMinimalBufferSize)
30import Blaze.ByteString.Builder.Char.Utf8
31
32import Criterion.Main
33
34httpOkHeader :: S.ByteString
35httpOkHeader = S.concat
36    [ "HTTP/1.1 200 OK\r\n"
37    , "Content-Type: text/html; charset=UTF-8\r\n"
38    , "\r\n" ]
39
40response :: Int -> Builder
41response n =
42  fromByteString httpOkHeader `mappend`
43  fromString (take n $ cycle "hello λ-world! ")
44
45sendVectoredBuilderLBS :: Socket -> Builder -> IO ()
46sendVectoredBuilderLBS s = L.sendAll s . toLazyByteString
47{-# NOINLINE sendVectoredBuilderLBS #-}
48
49sendBuilderLBS :: Socket -> Builder -> IO ()
50sendBuilderLBS s =
51  -- mapM_ (S.sendAll s) . L.toChunks . toLazyByteString
52  L.foldrChunks (\c -> (S.sendAll s c >>)) (return ()). toLazyByteString
53{-# NOINLINE sendBuilderLBS #-}
54
55sendBuilderBSIO :: Socket -> Builder -> IO ()
56sendBuilderBSIO s = toByteStringIO $ S.sendAll s
57{-# NOINLINE sendBuilderBSIO #-}
58
59-- criterion benchmark determining the speed of response
60main2 = defaultMain
61    [ bench ("response " ++ show n) $ whnf
62        (L.length . toLazyByteString . response) n
63    ]
64  where
65    n :: Int
66    n = 1000000
67
68main :: IO ()
69main = do
70    [port, nChars] <- map read `liftM` getArgs
71    killSignal <- newEmptyMVar
72    bracket (listenOn . PortNumber . fromIntegral $ port) sClose
73        (\socket -> do
74            _ <- forkIO $ loop (putMVar killSignal ()) nChars socket
75            takeMVar killSignal)
76  where
77    loop killServer nChars socket = forever $ do
78        (s, _) <- accept socket
79        forkIO (respond s nChars)
80      where
81        respond s n = do
82            input <- S.recv s 1024
83            let requestUrl = (S.split (fromIntegral $ ord ' ') input) !! 1
84            case tail (S.split (fromIntegral $ ord '/') requestUrl) of
85                ["lbs"]     -> sendBuilderLBS s         $ response n
86                ["lbs-vec"] -> sendVectoredBuilderLBS s $ response n
87                ["bs-io"]   -> sendBuilderBSIO   s      $ response n
88                ["kill"]    -> notFound s >> killServer
89                _           -> notFound s
90            sClose s
91
92    notFound s = do
93        _ <- S.sendAll s $ "HTTP/1.1 404 Not Found\r\n"
94            `mappend` "Content-Type: text/html; charset=UTF-8\r\n"
95            `mappend` "\r\n"
96            `mappend` "<h1>Page not found</h1>"
97        return ()
98