Ticket #4885: udpcount.hs

{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE ScopedTypeVariables #-}

import Control.Concurrent (forkIO, threadDelay)
import Control.Concurrent.MVar
import Control.Exception (handle, IOException)
import Control.Monad
import Data.Char (ord)
import Data.Ratio
import Data.Time.Clock.POSIX (getPOSIXTime)
import Foreign.Marshal.Alloc (mallocBytes)
import Network.Socket
import System.Console.CmdArgs
import System.IO
import System.IO.Unsafe

import qualified Data.ByteString as B
import qualified Data.ByteString.Unsafe as U

--import qualified Data.ByteString as B

buffSize = 65000 :: Int

data UdpCount = UdpCount { port :: Int
                         } deriving (Show, Data, Typeable)

udpCount = cmdArgsMode $ UdpCount
  { port = def &= argPos 0 &= typ "PORT"
  } &= summary "udp v0.1 (c) Itai Zukerman 2011"
  
counts = unsafePerformIO $ newMVar (0, 0) :: (MVar (Int, Int))
timestamp = unsafePerformIO $ newMVar B.empty :: (MVar B.ByteString)

space = B.pack [32]
eol = B.pack [10]

main = do args <- cmdArgsRun udpCount
          free <- newEmptyMVar
          full <- newEmptyMVar
          forkIO updateTime
          replicateM_ 4 $ forkIO $ workerLoop free full
          forkIO $ recvLoop free full $ port args
          forever $ do threadDelay (10^6)
                       c <- readMVar counts
                       hPutStrLn stderr $ show c

updateTime = do t <- getPOSIXTime
                let b = t2b t
                swapMVar timestamp b
                threadDelay (10^3)
                updateTime
  where
    t2b t = let i = toRational t
                ms = (numerator i * 1000) `div` (denominator i)
            in B.pack $ map (fromIntegral . ord) (show ms)

recvLoop free full port = do sock <- socket AF_INET Datagram 0
                             bindSocket sock $ SockAddrInet (fromIntegral port) iNADDR_ANY
                             forever $ do buff <- takeMVar free
                                          (size, _) <- retryOnFail $ recvBufFrom sock buff buffSize
                                          putMVar full (size, buff)

retryOnFail action = handle (\(_ :: IOException) -> retryOnFail action) action

workerLoop free full = do buff <- mallocBytes buffSize
                          putMVar free buff
                          loop
  where 
    loop = do 
      (size, buff') <- takeMVar full
      --bs <- U.unsafePackCStringLen (buff', size)
      bs <- B.packCStringLen (buff', size)
      ts <- readMVar timestamp
      let line = B.concat [ts, space, bs, eol]
      --B.hPutStr stdout line
      (n, t) <- takeMVar counts
      putMVar counts $! (n+1, t+size)
      putMVar free buff'
      loop