StreamingSpec.hs

{-# LANGUAGE OverloadedStrings #-}

module StreamingSpec where

import           Control.Concurrent
import           Control.Concurrent.STM
import           Control.Exception
import qualified Control.Monad
import           Data.ByteString
    ( ByteString
    , append
    , empty
    , hGetSome
    , hPut
    , head
    , isPrefixOf
    , pack
    , tail
    )
import qualified Data.ByteString.Lazy      as LBS
import           Data.Char                 (chr)
import           Data.Word8                (isUpper)
import           GHC.IO.Handle
    ( BufferMode (NoBuffering)
    , Handle
    , hClose
    , hSetBuffering
    )
import           GHC.IO.IOMode
import           Lib.Logger
import           Network.ConnectionAPI     (ReaderAPI (..), WriterAPI (..))
import           Network.Socket            hiding (Debug, close)
import           Parser.API
    ( ParseStep (DropPrefix, Emit, NeedMore)
    , ParserAPI (ParserAPI)
    )
import           Pipeline.Broadcasting     (broadcastingApi)
import           Pipeline.Broadcasting.API (BroadcastingAPI (BroadcastingAPI))
import           Pipeline.Streaming        (streamingApi)
import           Pipeline.Streaming.API    (StreamingAPI (StreamingAPI))
import           Prelude                   hiding
    ( head
    , last
    , length
    , null
    , replicate
    , tail
    , take
    )
import           Queue.API                 (QueueItem (..), close, enqueue)
import           Queue.TransactionalQueue  (newQueue)
import           Test.Hspec
import qualified Types.Pub                 as Pub

defaultLogger' :: IO LoggerConfig
defaultLogger' = do
  lock <- newTMVarIO ()
  pure $ LoggerConfig Debug (putStrLn . renderLogEntry) lock

bufferLimit :: Int
bufferLimit = 4096

newtype CaptureWriter = CaptureWriter (TVar [ByteString])

captureWriterApi :: WriterAPI CaptureWriter
captureWriterApi =
  WriterAPI
    { writeData = \(CaptureWriter output) bytes -> do
        atomically $ modifyTVar' output (<> [bytes])
        pure (Right ())
    , writeDataLazy = \writer bytes ->
        writeData captureWriterApi writer (LBS.toStrict bytes)
    , closeWriter = \_ -> pure ()
    , openWriter = \_ -> pure ()
    }

handleReaderApi :: ReaderAPI Handle
handleReaderApi =
  ReaderAPI
    { readData = \h n -> do
        result <- try (hGetSome h n) :: IO (Either SomeException ByteString)
        case result of
          Left err    -> return $ Left (show err)
          Right bytes -> return $ Right bytes
    , closeReader = hClose
    , openReader = \_ -> pure ()
    }

spec :: Spec
spec = do
  describe "Streaming" $ do
      it "reads from source and writes to sink" $ do
        (server, client) <- makeSocketPair
        result <- newTVarIO "" :: IO (TVar ByteString)
        let sink curr = atomically $ modifyTVar' result (`append` curr)
        dl <- defaultLogger'
        ctx <- newLogContext
        let StreamingAPI runStreaming = streamingApi
        (forkIO . runWithLogger dl ctx) (runStreaming bufferLimit handleReaderApi client singleByteParser sink :: AppM ())
        hPut server "Hello, World"
        atomically $ assertTVarWithRetry result "Hello, World"
        hClose server
        hClose client
      it "continues reading from the handle after reaching end" $ do
        (server, client) <- makeSocketPair
        result <- newTVarIO "" :: IO (TVar ByteString)
        let sink curr = atomically $ modifyTVar' result (`append` curr)
        dl <- defaultLogger'
        ctx <- newLogContext
        let StreamingAPI runStreaming = streamingApi
        (forkIO . runWithLogger dl ctx) (runStreaming bufferLimit handleReaderApi client singleByteParser sink :: AppM ())
        hPut server "Hello, World"
        atomically $ assertTVarWithRetry result "Hello, World"
        hPut server "Hello, again"
        atomically $ assertTVarWithRetry result "Hello, WorldHello, again"
        hClose server
        hClose client
      it "applies the parser healing" $ do
        (server, client) <- makeSocketPair
        result <- newTVarIO "" :: IO (TVar ByteString)
        let sink curr = atomically $ modifyTVar' result (`append` curr)
        dl <- defaultLogger'
        ctx <- newLogContext
        let StreamingAPI runStreaming = streamingApi
        (forkIO . runWithLogger dl ctx) (runStreaming bufferLimit handleReaderApi client excludingUpperParser sink :: AppM ())
        hPut server "HELLO WORLD hello, world"
        atomically $ assertTVarWithRetry result "  hello, world"
        hClose server
        hClose client
      it "waits for more data" $ do
        (server, client) <- makeSocketPair
        result <- newTVarIO "" :: IO (TVar ByteString)
        let sink curr = atomically $ modifyTVar' result (`append` curr)
        dl <- defaultLogger'
        ctx <- newLogContext
        let StreamingAPI runStreaming = streamingApi
        (forkIO . runWithLogger dl ctx) (runStreaming bufferLimit handleReaderApi client explicitWordParser sink :: AppM ())
        hPut server "part1"
        threadDelay 100000
        atomically $ ensureTVarIsEmpty result
        hPut server "part2"
        atomically $ assertTVarWithRetry result "part1part2"
        hClose server
        hClose client
  describe "Broadcasting" $ do
      it "drops messages larger than the buffer limit" $ do
        dl <- defaultLogger'
        ctx <- newLogContext
        q <- newQueue
        output <- newTVarIO []
        let pub = Pub.Pub "FOO" Nothing Nothing (Just "0123456789")
        let BroadcastingAPI runBroadcasting = broadcastingApi
        enqueue q (QueueItem pub) `shouldReturn` Right ()
        close q
        runWithLogger dl ctx (runBroadcasting 5 q captureWriterApi (CaptureWriter output) :: AppM ())
        readTVarIO output `shouldReturn` []

ensureTVarIsEmpty :: TVar ByteString -> STM ()
ensureTVarIsEmpty tvar = do
  content <- readTVar tvar
  Control.Monad.when (content /= empty) retry

assertTVarWithRetry :: Eq a => TVar a -> a -> STM ()
assertTVarWithRetry tvar expected = do
  actual <- readTVar tvar
  Control.Monad.unless (actual == expected) retry

singleByteParser :: ParserAPI ByteString
singleByteParser = ParserAPI (\bs -> Emit (pack [head bs]) (tail bs))

excludingUpperParser :: ParserAPI ByteString
excludingUpperParser =
  ParserAPI $ \bs ->
    case isUpper (head bs) of
      False ->
        Emit (pack [head bs]) (tail bs)
      True ->
        DropPrefix 1 [chr . fromIntegral . head $ bs]

explicitWordParser :: ParserAPI ByteString
explicitWordParser = ParserAPI parseWord
  where
    parseWord bs
      | bs == "part1part2" =
          Emit bs empty
      | bs `isPrefixOf` "part1part2" =
          NeedMore
      | otherwise =
          DropPrefix 1 [chr . fromIntegral . head $ bs]

makeSocketPair :: IO (Handle, Handle)
makeSocketPair = do
  serverSock <- socket AF_INET Stream defaultProtocol
  setSocketOption serverSock ReuseAddr 1
  bind serverSock (SockAddrInet 0 (tupleToHostAddress (127,0,0,1)))
  listen serverSock 1
  SockAddrInet port _ <- getSocketName serverSock

  -- initiate client connection
  clientSock <- socket AF_INET Stream defaultProtocol
  connect clientSock (SockAddrInet port (tupleToHostAddress (127,0,0,1)))

  -- accept connection from server side
  (serverConn, _) <- accept serverSock

  -- convert both sides to handles
  serverHandle <- socketToHandle serverConn ReadWriteMode
  clientHandle <- socketToHandle clientSock ReadWriteMode
  hSetBuffering serverHandle NoBuffering
  hSetBuffering clientHandle NoBuffering

  return (serverHandle, clientHandle)