From 79e7fa8060b82871eab64f448521ca20928e7cc8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Domen=20Ko=C5=BEar?= <domen@dev.si>
Date: Thu, 21 Dec 2023 13:33:30 +0000
Subject: [PATCH] Add benchmark for roundtrip connections

---
 .github/workflows/ci.yml  |  3 +++
 benchmarks/connections.hs | 44 +++++++++++++++++++++++++++++++++++++++
 websockets.cabal          | 14 +++++++++++++
 3 files changed, 61 insertions(+)
 create mode 100644 benchmarks/connections.hs

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 30198d2..1c32b89 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -71,6 +71,9 @@ jobs:
       - name: Build documentation
         run: cabal haddock all
 
+      - name: Benchmark
+        run: cabal bench all
+
       - name: Install virtualenv
         if: matrix.os == 'ubuntu-latest'
         run: | 
diff --git a/benchmarks/connections.hs b/benchmarks/connections.hs
new file mode 100644
index 0000000..94b570e
--- /dev/null
+++ b/benchmarks/connections.hs
@@ -0,0 +1,44 @@
+{-# LANGUAGE OverloadedStrings #-}
+module Main where
+
+import Control.Monad (forever, unless)
+import Data.ByteString (ByteString)
+import qualified Network.WebSockets as WS
+import qualified Network.WebSockets.Client as WC
+import Criterion.Main
+import Control.Concurrent.Async (mapConcurrently_, race_, async, wait, Async)
+
+
+server :: WS.ServerApp
+server pending = do
+    conn <- WS.acceptRequest pending
+    msg  <- WS.receiveData conn
+    WS.sendBinaryData conn (msg :: ByteString)
+
+client :: WC.ClientApp ()
+client conn = do
+    -- Send and receive a message back
+    let msg = "Hello, world!" :: ByteString
+    WS.sendBinaryData conn msg
+    msg' <- WS.receiveData conn
+    unless (msg == msg') $ error "Message mismatch"
+
+    WS.sendClose conn ("Bye!" :: ByteString)
+
+run :: Int -> Async () -> IO ()
+run n server = race_ runServer runClient
+    where 
+    runClient = mapConcurrently_ (\_ -> WC.runClient "127.0.0.1" 8089 "/" client) [1..n]
+    runServer = wait server
+
+main :: IO ()
+main = do
+    server <- async $ WS.runServer "127.0.0.1" 8089 server
+    defaultMain [
+        bgroup "connections"
+            [ bench "100" $ nfIO $ run 100 server
+            , bench "1000" $ nfIO $ run 1000 server
+            , bench "10000" $ nfIO $ run 10000 server
+            , bench "100000" $ nfIO $ run 100000 server
+            ]
+        ]
\ No newline at end of file
diff --git a/websockets.cabal b/websockets.cabal
index 35e9f02..4021fe5 100644
--- a/websockets.cabal
+++ b/websockets.cabal
@@ -248,6 +248,20 @@ Benchmark bench-mask
     text              >= 0.10   && < 2.2,
     entropy           >= 0.2.1  && < 0.5
 
+Benchmark bench-connections
+  type:             exitcode-stdio-1.0
+  main-is:          connections.hs
+  Hs-source-dirs:   benchmarks
+  ghc-options:      -threaded -O2 -rtsopts "-with-rtsopts=-N"
+  Default-language: Haskell2010
+  Build-depends:
+    async,
+    base,
+    bytestring,
+    criterion,
+    async,
+    websockets
+
 Executable websockets-website
   If !flag(Website)
     Buildable: False