ConduitベースのTCPクライアント Ver 0.2

こちらの記事で書いたTCPクライアントを改良したので掲載。アプリケーションから見ればTCPサーバとの通信はConduit ByteString IO ByteStringだと気づいたので、そうなるように隠蔽してみました。

TCP Client Ver 0.2

TCPClient.hs

{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE OverloadedStrings #-}  -- converts String to ByteString implicitly

module Main where

import ClassyPrelude
import ClassyPrelude.Conduit
import qualified Data.Conduit.Network as DCN
import Data.ByteString.Char8 (ByteString)

import Data.Conduit.Binary.Utils (snocC)
import Data.Conduit.Network.Utils (initTCPClient)

main :: IO ()
main = do
  (tid, requestC) <- initTCPClient settings
  app requestC

app :: Conduit ByteString IO ByteString -> IO ()
app requestC = forever $ stdinC $$ requestC =$= snocC '\n' =$= stdoutC  -- consume 1 responce and resume

settings :: DCN.ClientSettings
settings = DCN.clientSettings 69 "127.0.0.1"

Client側は入力ソースがTCP ServerとStdinの二箇所あるため、一問一答形式にするためにconnect-and-resumeを用いる。initTCPClientでTCP Serverと送受信するConduit ByteString IO ByteStringを返す。

Data/Conduit/Network/Utils.hs

{-# LANGUAGE NoImplicitPrelude #-}

module Data.Conduit.Network.Utils (
  initTCPClient
) where

import ClassyPrelude
import ClassyPrelude.Conduit
import qualified Data.Conduit.List as DCL
import qualified Data.Conduit.Network as DCN
import Data.ByteString.Char8 (ByteString)
import qualified Control.Concurrent as CC
import qualified Control.Concurrent.MVar as MV

import Data.Conduit.Utils (mvSource, mvSink)
import Data.Conduit.Binary.Utils (trimC)

initTCPClient :: DCN.ClientSettings -> IO (CC.ThreadId, Conduit ByteString IO ByteString)
initTCPClient settings = do
  sendMV <- MV.newEmptyMVar
  recvMV <- MV.newEmptyMVar
  tid <- CC.forkIO $ tcpClient settings sendMV recvMV
  return (tid, request sendMV recvMV)
{-# INLINE initTCPClient #-}

request :: MV.MVar ByteString -> MV.MVar ByteString -> Conduit ByteString IO ByteString
request sendMV recvMV = mapMC $ \bs -> MV.putMVar sendMV bs >> MV.takeMVar recvMV

tcpClient :: DCN.ClientSettings -> MV.MVar ByteString -> MV.MVar ByteString -> IO ()
tcpClient settings sendMV recvMV = DCN.runTCPClient settings $ \ad ->
  tcpLoop (newResumableSource (DCN.appSource ad), DCN.appSink ad) (mvSource sendMV, mvSink recvMV)
{-# INLINE tcpClient #-}

tcpLoop :: Monad m => (ResumableSource m ByteString, Sink ByteString m ()) -> (Source m ByteString, Sink ByteString m ()) -> m ()
tcpLoop (serverP, serverC) (appP, appC) = do
  appP $$ trimC =$= DCL.isolate 1 =$= serverC    -- consume 1 command
  (serverP', _) <- serverP $$++ DCL.isolate 1 =$= appC  -- consume 1 responce and resume
  tcpLoop (serverP', serverC) (appP, appC)

Data/Conduit/Utils.hs

{-# LANGUAGE NoImplicitPrelude #-}

module Data.Conduit.Utils (
  , mvSource
  , mvSink
  ) where

import ClassyPrelude
import ClassyPrelude.Conduit
import qualified Control.Concurrent.MVar as MV

mvSource :: MVar a -> Source IO a
mvSource mv = lift (MV.takeMVar mv) >>= yield >> mvSource mv
{-# INLINEABLE mvSource #-}

mvSink :: MVar a -> Sink a IO ()
mvSink mv = awaitForever $ liftIO . putMVar mv
{-# INLINEABLE mvSink #-}

Data/Conduit/Binary/Utils.hs

{-# LANGUAGE NoImplicitPrelude #-}

module Data.Conduit.Binary.Utils (
    snocC
  , trimC
  ) where

import ClassyPrelude
import ClassyPrelude.Conduit
import qualified Data.Char as DC
import Data.ByteString.Char8 (ByteString)
import qualified Data.ByteString.Char8 as DBC

-- add a char at the last of ByteString
snocC :: Monad m => Char -> Conduit ByteString m ByteString
snocC c = mapC $ flip DBC.snoc c
{-# INLINEABLE snocC #-}

-- trim left and right spaces of ByteString (for removing \n in the end of the commands from stdin)
trimC :: Monad m => Conduit ByteString m ByteString
trimC = mapC trim
  where
    trim :: ByteString -> ByteString
    trim = dropWhileEnd DC.isSpace . DBC.dropWhile DC.isSpace
    dropWhileEnd :: (Char -> Bool) -> ByteString -> ByteString
    dropWhileEnd f = fst . DBC.spanEnd f
{-# INLINEABLE trimC #-}