HaskellでのWebSocketsサーバの効率改善

前回の記事で、HTMLのCanvas要素とWebSocketsの組み合わせで画像転送と表示の実験を行いました。想定よりもHaskellで作ったWebSocketsサーバの動作が重いため、少し調査して改善しました。

import Control.Exception (finally)
import Control.Monad (forM_, forever)

import qualified Network.WebSockets as WS
import qualified Data.ByteString.Char8 as C
import qualified Data.ByteString.Lazy as L
import qualified Data.ByteString.Internal as BI

import Foreign.Ptr (plusPtr, Ptr)
import Foreign.Storable (poke)
import Data.Word (Word8)

main :: IO ()
main = WS.runServer "127.0.0.1" 9160 application

-- BI.create :: Int -> (Ptr Word8 -> IO ()) -> IO ByteString
fill :: Int32 -> Ptr Word8 -> IO ()
fill seed ptr = 
    go (0, 0) ptr
      where
        go (j, i) ptr = do
            let val = fromIntegral $ seed + i + j
            poke ptr val
            if i == 1919 then
                if j == 1079 then
                    return ()
                else
                    go (j+1, 0) $ ptr `plusPtr` 1
            else
                go (j, i+1) $ ptr `plusPtr` 1

application :: WS.ServerApp
application pending = do
    conn <- WS.acceptRequest pending
    WS.forkPingThread conn 30
    forever $ do
        bytes <- WS.receiveData conn
        let seed:_ = L.unpack bytes
        let seed2 = fromIntegral (toInteger seed)
        buf <- BI.create (1920*1080) (fill seed2)
        WS.sendBinaryData conn buf

ByteString.Internalのcreate関数を使って、転送するバッファの生成と初期化を同時に行います。 fill関数には、createで確保される領域へのポインタが渡されるので、poke関数を使って当該領域に1つずつデータを書き込みます。 go関数で、1画面分ループします。1ループごとにポインタを1進める(plusPtr)ことで、次のWord8にアクセスできます。

このプログラムを試してみると、FirefoxはCPUを60%程度消費する一方、サーバ側は15%程度で済んでいます。

クライアントからリクエストがあるごとに毎回create関数を呼び出しているので、あまり効率は良くないと思いますが、それでも初期のサーバ律速であった状況からは大幅に改善しました。

もっとも、このようなプログラムにしてしまうと、中身はCで記述するのと大差ないので、あえてHaskellで書く意義はかなり失われている感があります。