Haskellでwebsocketサーバを作る

haskellでの、websocketのサーバの作り方です。
サーバとしてはwarpのみを使用しています。

確認したバージョンなどは以下。

* ghc-7.6.3
* warp-3.0.0.3
* wai-3.0.0.2
* websockets-0.8.2.5
* wai-websockets-3.0.0
* http-types-0.8.5
* OS: windows8.1, lubuntu-14.04

# とりあえず動かす

## サーバコード

解説っぽいものはほとんどソースコード中にコメントで記載しました。
websocketでクライアントから投げられたメッセージをそのまま返すようなアプリケーションです。
普段はmainは下から書いていくのですが、上から読めるように、上からmainを書いています。
あと、他にも初心者にわかりにくそうな書き方は少しだけ避けたつもり。

````haskell:simple_websocket_server.hs
{-# LANGUAGE OverloadedStrings #-}

module Main where

-- for warp
import qualified Network.Wai as Wai
import qualified Network.Wai.Handler.Warp as Warp
import qualified Network.HTTP.Types.Status as Status

-- for websocket
import qualified Network.Wai.Handler.WebSockets as WaiWS
import qualified Network.WebSockets as WS

-- Data.*
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as BSC

main :: IO ()
main = do
    let port = 3000 -- とりあえずポート3000番で。
    let setting = Warp.setPort port Warp.defaultSettings
    putStrLn $ "start server port=" ++ show port
    Warp.runSettings setting simpleWSWaiApp

-- websocketの場合はwebsocketの処理を行い、
-- そうでない場合はsimpleWaiAppで処理を行うアプリケーション
-- websocketかどうかはupgradeヘッダーがwebsocketとなっているかどうかで判断される。
simpleWSWaiApp :: Wai.Application
simpleWSWaiApp = WaiWS.websocketsOr WS.defaultConnectionOptions simpleWSServerApp simpleWaiApp

-- websocketじゃないリクエストを処理するアプリケーション
simpleWaiApp :: Wai.Application
simpleWaiApp req respond = do
    putStrLn . show . Wai.pathInfo $ req -- とりあえずリクエストのパスとか表示してみる。
    respond $ Wai.responseFile Status.status200 [] "index.html" Nothing

-- websocketのリクエストを処理するアプリケーション
simpleWSServerApp :: WS.ServerApp
simpleWSServerApp pdconn = do
    putStrLn "Websocket Request received"
    -- ここで、websocketのリクエストをpathなどに応じてacceptRequestするか、
    -- rejectRequestするかなどを判断・処理する。
    -- リクエストの情報はpendingRequestを使用してpdconnから取得できる。
    conn <- WS.acceptRequest pdconn
    -- リクエストをacceptしたらメッセージループに入る。
    -- リクエスト内容に応じて異なるメッセージループを使い分けることもできる。
    -- また、接続直後のメッセージのみ特別な処理をしたり、
    -- 接続直後に一方的にメッセージを送り付けるような処理をしたりする場合は、
    -- メッセージループに入る前に処理してしまうことも可能。
    wsMessageLoop conn

-- websocketのメッセージループ
wsMessageLoop :: WS.Connection -> IO ()
wsMessageLoop conn = do
    -- メッセージの受信方法にはreceiveとreceiveDataMessageとreceiveDataの3種類がある。
    -- * receiveはClose,Ping,PongなどのControlMessageも自前で処理したい場合に使用する。
    --   接続が切断されている場合は、例外（ConnectionClosed）がthrowされる。
    -- * receiveDataMessageはControlMessageは自動で処理させてしまい、
    --   DataMessageの場合のみその中身をByteStringとして処理したい場合に使用する。
    --   Closeメッセージを受信した場合は、
    --   自動処理内で例外（ConnectionClosed）がthrowされる。
    -- * receiveDataはデータがWS.TextかWS.Binaryかは気にせず中身を推論された型で受け取る。
    --   受け取るデータの型をWS.WebSocketsDataクラスのインスタンスにする必要がある。
    --   ByteStringやTextはデフォルトでWS.WebSocketsDataクラスのインスタンスになっている。
    msgData <- WS.receiveData conn
    -- メッセージハンドラでメッセージを処理。
    handlerResult <- handleMessage conn msgData
    -- メッセージハンドラの処理結果次第でループさせます。
    -- 抜けるときは単に再帰を止めるだけ。
    -- 例外で抜けることも可能なので必ずこの形式でなければいけないことは全くない。
    case handlerResult of
        Break -> WS.sendClose conn ("exit" :: BS.ByteString)
        Continue -> wsMessageLoop conn

-- メッセージハンドラの結果。ループ処理を続けるか否かを表す。
data HandlerResult = Continue
                   | Break

-- メッセージハンドラ
-- ここでは、とりあえず、ByteStringを取り出して表示したうえで、
-- クライアントに送り返してみる。
handleMessage :: WS.Connection -> BS.ByteString -> IO HandlerResult
handleMessage conn msgData = do
    BSC.putStrLn $ BS.append "Message received : " msgData
    -- 受信と同様、メッセージの送信にも複数の関数が用意されているので、
    -- （send, sendDataMessage, sendTextData, sendBinaryData）
    -- 必要に応じて使い分ける。
    -- ここでは受け取ったメッセージをそのまま送り返してみる。
    -- ただしBinaryで受け取ってもTextで送り返している点に注意。
    WS.sendTextData conn msgData
    -- その後、受信したメッセージが"exit"の場合だけループを抜ける。
    case msgData of
        "exit" -> do
          putStrLn "exit"
          return Break
        _      -> return Continue
````

以上がサーバのコードです。
動作確認用にwebsocketじゃないリクエストを送るととりあえずindex.htmlを返すようにしてあります。

ビルドと実行は以下。

````
ghc -O2 -Wall simple_websocket_server.hs
./simple_websocket_server
````

## 確認用クライアントコード

動作確認のため、index.htmlにwebsocketを使うコードを書きます。
何本かコネクションを張って、各コネクションごとに毎秒メッセージを投げる処理を数回繰り返して終了するような内容になっています。

````html:index.html
<!DOCTYPE html>
<html>
<title>websocket test</title>
<body>
websocket test

<script type=text/javascript>
console.log("websocket test")

function startWebSocket(name, count) {
  try {
    var ws = new WebSocket('ws://localhost:3000/'+count);
  } catch (err) {
    console.error(err);
  }

  function sendMessage() {
    ws.send('message : ' + name + ' : ' + count)
    count--;
  }

  ws.onopen = function () {
    console.log('onOpen', name);
    sendMessage();
  };

  ws.onerror = function (error) {
    console.log('onError', name, error);
  };

  ws.onmessage = function (e) {
    console.log('onMessage', name, e);
    if (count > 0) {
      setTimeout(sendMessage, 1000);
    } else {
      // exitを送ってサーバ側からクローズさせる。
      console.log('send exit', name);
      ws.send('exit');
      // clientからクローズする場合は上記をコメントアウトして以下を有効にする。
      //console.log('close websocket', name);
      //ws.close(1000, 'closed by client');
    }
  };

  ws.onclose = function () {
    console.log('onClose', name);
  };
  console.log(ws);
};

for (var i = 1; i <= 5; i++) {
  startWebSocket('test' + i, 10);
}

</script>

</body>
</html>
````

最後のstartWebsocketの呼び出し部分を変えれば並列度やメッセージを送る回数などを変更できます。
ただし、ブラウザによって同時に扱えるwebsocketの数などが異なるようですので注意してください。

# サーバに状態を持たせる

さて、上記でとりあえずの動作はするのですが、実用しようとするとサーバに状態を持たせたくなると思います。
そこで、全体として受け取ったメッセージの数をカウントして（サーバ側のコンソールに）表示するように修正してみます。
具体的にはMvarに状態を作成してこれを必要に応じて更新します。
（この程度の内容ならMVarで十分ですが、本格的なものならSTMとかの使用を考えた方がいいと思います）

あと、上記のサーバで実際に動かすと、Websocketをクライアント（ブラウザ）側でクローズした時にクライアント（ブラウザ）側でエラーが出ます。特に実害はないように見えますが、気持ち悪いので一緒に修正してしまいます。

````haskell:simple_websocket_server2.hs
{-# LANGUAGE OverloadedStrings #-}

module Main where

-- for warp
import qualified Network.Wai as Wai
import qualified Network.Wai.Handler.Warp as Warp
import qualified Network.HTTP.Types.Status as Status

-- for websocket
import qualified Network.Wai.Handler.WebSockets as WaiWS
import qualified Network.WebSockets as WS

-- Data.*
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as BSC

-- Control.*
import qualified Control.Exception as E
import qualified Control.Concurrent as CC

main :: IO ()
main = do
    let port = 3000 -- とりあえずポート3000番で。
    let setting = Warp.setPort port Warp.defaultSettings
    putStrLn $ "start server port=" ++ show port
    simpleWSWaiApp <- newSimpleWSWaiApp
    Warp.runSettings setting simpleWSWaiApp

-- websocketの場合はwebsocketの処理を行い、
-- そうでない場合はsimpleWaiAppで処理を行うアプリケーション
-- websocketかどうかはupgradeヘッダーがwebsocketとなっているかどうかで判断される。
newSimpleWSWaiApp :: IO Wai.Application
newSimpleWSWaiApp = do
    state <- newServerState
    let wsServer = simpleWSServerApp state
    return $ WaiWS.websocketsOr WS.defaultConnectionOptions wsServer simpleWaiApp

-- server state
type ServerState = CC.MVar Int

newServerState :: IO ServerState
newServerState = CC.newMVar 0

updateTotalCount :: ServerState -> IO Int
updateTotalCount state = CC.modifyMVar state $ \ c -> return (c+1, c+1)

-- websocketじゃないリクエストを処理するアプリケーション
simpleWaiApp :: Wai.Application
simpleWaiApp req respond = do
    putStrLn . show . Wai.pathInfo $ req -- とりあえずリクエストのパスとか表示してみる。
    respond $ Wai.responseFile Status.status200 [] "index.html" Nothing

-- websocketのリクエストを処理するアプリケーション
simpleWSServerApp :: ServerState -> WS.ServerApp
simpleWSServerApp state pdconn = do
    putStrLn "Websocket Request received"
    -- ここで、websocketのリクエストをpathなどに応じてacceptRequestするか、
    -- rejectRequestするかなどを判断・処理する。
    -- リクエストの情報はpendingRequestを使用してpdconnから取得できる。
    conn <- WS.acceptRequest pdconn
    -- リクエストをacceptしたらメッセージループに入る。
    -- リクエスト内容に応じて異なるメッセージループを使い分けることもできる。
    -- また、接続直後のメッセージのみ特別な処理をしたり、
    -- 接続直後に一方的にメッセージを送り付けるような処理をしたりする場合は、
    -- メッセージループに入る前に処理してしまうことも可能。
    wsMessageLoop state conn

-- websocketのメッセージループ
wsMessageLoop :: ServerState -> WS.Connection -> IO ()
wsMessageLoop state conn = do
    -- メッセージの受信方法にはreceiveとreceiveDataMessageとreceiveDataの3種類がある。
    -- * receiveはClose,Ping,PongなどのControlMessageも自前で処理したい場合に使用する。
    --   接続が切断されている場合は、例外（ConnectionClosed）がthrowされる。
    -- * receiveDataMessageはControlMessageは自動で処理させてしまい、
    --   DataMessageの場合のみその中身をByteStringとして処理したい場合に使用する。
    --   Closeメッセージを受信した場合は、
    --   自動処理内で例外（ConnectionClosed）がthrowされる。
    -- * receiveDataはデータがWS.TextかWS.Binaryかは気にせず中身を推論された型で受け取る。
    --   受け取るデータの型をWS.WebSocketsDataクラスのインスタンスにする必要がある。
    --   ByteStringやTextはデフォルトでWS.WebSocketsDataクラスのインスタンスになっている。
    msgData <- receiveData conn
    -- メッセージハンドラでメッセージを処理。
    handlerResult <- handleMessage state conn msgData
    -- メッセージハンドラの処理結果次第でループさせます。
    -- 抜けるときは単に再帰を止めるだけ。
    -- 例外で抜けることも可能なので必ずこの形式でなければいけないことは全くない。
    case handlerResult of
        Break -> WS.sendClose conn ("exit" :: BS.ByteString)
        Continue -> wsMessageLoop state conn

-- Closeの自動処理でブラウザ側で警告が出るため、自前で修正する。
-- 内容はWS.receiveDataとほぼ同じ。
receiveData :: WS.WebSocketsData a => WS.Connection -> IO a
receiveData conn = do
    dm <- receiveDataMessage conn
    case dm of
        WS.Text x   -> return (WS.fromLazyByteString x)
        WS.Binary x -> return (WS.fromLazyByteString x)

-- Closeの自動処理でブラウザ側で警告が出るため、自前で修正する。
-- 内容はWS.receiveDataMessageとだいたい同じ。
-- ついでにデバッグ情報を表示するように修正してみる。
receiveDataMessage :: WS.Connection -> IO WS.DataMessage
receiveDataMessage conn = do
    msg <- WS.receive conn
    case msg of
        WS.DataMessage am    -> return am
        WS.ControlMessage cm -> case cm of
            WS.Close bs   -> do
                putStrLn $ "Close message received : " ++ show bs
                WS.sendClose conn ("-- server received close message" :: BS.ByteString)
                E.throw WS.ConnectionClosed
            WS.Pong bs   -> do
                putStrLn $ "Pong message received : " ++ show bs
                -- WS.connectionOnPong (WS.connectionOptions conn)
                receiveDataMessage conn
            WS.Ping bs   -> do
                putStrLn $ "Ping message received : " ++ show bs
                WS.send conn (WS.ControlMessage (WS.Pong bs))
                receiveDataMessage conn

-- メッセージハンドラの結果。ループ処理を続けるか否かを表す。
data HandlerResult = Continue
                   | Break

-- メッセージハンドラ
-- ここでは、とりあえず、ByteStringを取り出して表示したうえで、
-- クライアントに送り返してみる。
handleMessage :: ServerState -> WS.Connection -> BS.ByteString -> IO HandlerResult
handleMessage state conn msgData = do
    count <- updateTotalCount state
    let str = BSC.pack $ "Message received : " ++ show count ++ " : "
    BSC.putStrLn $ BS.append str msgData
    -- 受信と同様、メッセージの送信にも複数の関数が用意されているので、
    -- （send, sendDataMessage, sendTextData, sendBinaryData）
    -- 必要に応じて使い分ける。
    -- ここでは受け取ったメッセージをそのまま送り返してみる。
    -- ただしBinaryで受け取ってもTextで送り返している点に注意。
    WS.sendTextData conn msgData
    -- その後、受信したメッセージが"exit"の場合だけループを抜ける。
    case msgData of
        "exit" -> do
          putStrLn "exit"
          return Break
        _      -> return Continue
````

ServerStateの中身を変えて、それを使う部分を変えれば状態が必要なだいたいのことはできるようになると思います。

# 感想

websocketsにはWebsocketのクライアントを作る機能もあるのですが、私の環境ではなぜか動きませんでした。
原因は未調査ですが、クライアント側からクローズした場合に警告が出ることなどから考えて、このライブラリには微妙に問題が残っているのかもしれません。