Client.hs

{-# LANGUAGE RoleAnnotations #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -Wno-orphans #-}

module Test.Marionette.Client where

import Control.Exception (AssertionFailed (AssertionFailed))
import Control.Monad (void, when)
import Control.Monad.Catch
    ( Exception (..)
    , MonadCatch
    , MonadMask
    , MonadThrow
    , catch
    , catchAll
    , throwM
    )
import Control.Monad.Error.Class (MonadError (..))
import Control.Monad.IO.Class (MonadIO (liftIO))
import Control.Monad.Reader (ReaderT (runReaderT))
import Control.Monad.Reader qualified as Reader
import Control.Monad.Reader.Class (MonadReader)
import Data.Aeson (AesonException (AesonException), FromJSON)
import Data.Aeson qualified as Aeson
import Data.Aeson.Types qualified as Aeson
import Data.Binary (Binary)
import Data.Binary qualified as Binary
import Data.Binary.Get qualified as Binary
import Data.ByteString (ByteString)
import Data.ByteString.Builder.Extra qualified as ByteString
import Data.IntMap.Strict qualified as IntMap
import Data.Maybe (isNothing)
import GHC.Stack (HasCallStack)
import Network.Simple.TCP
    ( HostName
    , ServiceName
    , SockAddr
    , Socket
    , closeSock
    , connectSock
    , sendLazy
    )
import Network.Socket.ByteString (recv)
import System.Timeout (timeout)
import Test.Marionette.Class (Marionette (..))
import Test.Marionette.Protocol
import UnliftIO
    ( MonadUnliftIO
    , TMVar
    , TQueue
    , async
    , bracket
    , isEmptyTMVar
    , link
    , readTMVar
    )
import UnliftIO.Concurrent (forkIO)
import UnliftIO.Retry (constantDelay, limitRetriesByCumulativeDelay, recoverAll)
import UnliftIO.STM
    ( atomically
    , modifyTVar'
    , newEmptyTMVarIO
    , newTQueueIO
    , newTVarIO
    , putTMVar
    , readTQueue
    , stateTVar
    , writeTQueue
    )
import Prelude hiding (log)

data SocketClosed = SocketClosed
    deriving stock (Show)
    deriving anyclass (Exception)

newtype DecodeError = DecodeError String
    deriving stock (Show)
    deriving anyclass (Exception)

incoming
    :: forall m a
     . (MonadUnliftIO m, MonadThrow m, Binary a)
    => Socket
    -> m (TQueue a)
incoming socket = do
    q <- newTQueueIO
    link =<< async (go (atomically . writeTQueue q) (newDecoder ""))
    pure q
  where
    newDecoder :: ByteString -> Binary.Decoder a
    newDecoder bs = Binary.runGetIncremental Binary.get `Binary.pushChunk` bs
    go _ (Binary.Fail _ _ err) = throwM . DecodeError $ err
    go f (Binary.Done rest _ a) = do
        f a
        go f $ newDecoder rest
    go f dec =
        go f . (dec `Binary.pushChunk`)
            =<< liftIO
                (recv socket ByteString.defaultChunkSize `catchAll` \_ -> throwM SocketClosed)

connect
    :: (MonadUnliftIO m)
    => HostName
    -> ServiceName
    -> ((Socket, SockAddr) -> m a)
    -> m a
connect host port =
    bracket
        ( recoverAll (limitRetriesByCumulativeDelay 5_000_000 $ constantDelay 50_000)
            . const
            $ connectSock host port
        )
        (closeSock . fst)

newtype MarionetteTimeout = MarionetteTimeout MarionetteMessage
    deriving stock (Show)
    deriving anyclass (Exception)

newtype UnexpectedResult = UnexpectedResult Result
    deriving stock (Show)
    deriving anyclass (Exception)

data CommandWithCallback = CommandWithCallback Command (TMVar Result)

-- | A monad transformer that speaks the Marionette wire protocol over a TCP
-- socket. Run it with 'runMarionetteT'.
newtype MarionetteT m a = MarionetteT (ReaderT (TQueue CommandWithCallback) m a)
    deriving newtype
        ( Functor
        , Applicative
        , Monad
        , MonadThrow
        , MonadCatch
        , MonadMask
        , MonadIO
        , MonadUnliftIO
        , MonadReader (TQueue CommandWithCallback)
        )

instance (MonadUnliftIO m, MonadThrow m, MonadCatch m) => Marionette (MarionetteT m) where
    sendCommand :: (HasCallStack, FromJSON a) => Command -> MarionetteT m a
    sendCommand command = do
        q <- Reader.ask
        result <- newEmptyTMVarIO
        atomically . writeTQueue q $ CommandWithCallback command result
        either throwM pure =<< parseResult =<< atomically (readTMVar result)
      where
        parseResult :: (FromJSON a) => Result -> MarionetteT m (Either Error a)
        parseResult =
            either (pure . Left) $
                either (throwM . AesonException) (pure . Right)
                    . Aeson.parseEither Aeson.parseJSON

instance (MonadThrow m, MonadCatch m) => MonadError Error (MarionetteT m) where
    throwError = throwM
    catchError = catch

-- | Connect to a Marionette server listening on @localhost:2828@ (started
-- by launching Firefox with the @--marionette@ flag), and run the given
-- action against it.
runMarionetteT
    :: forall m a
     . (MonadUnliftIO m, MonadMask m)
    => MarionetteT m a
    -> m a
runMarionetteT action = do
    sendQueue :: TQueue CommandWithCallback <- newTQueueIO
    pendingCommands <- newTVarIO mempty
    done <- newEmptyTMVarIO
    let repeatUntilDone :: forall m'. (MonadIO m') => m' () -> m' ()
        repeatUntilDone a =
            atomically (isEmptyTMVar done) >>= \case
                False -> pure ()
                True -> a >> repeatUntilDone a
        handleIncoming :: MarionetteMessage -> MarionetteT m ()
        handleIncoming message =
            decodeMarionetteM message >>= \Message{..} ->
                atomically
                    ( stateTVar pendingCommands $
                        IntMap.updateLookupWithKey (\_ _ -> Nothing) messageId
                    )
                    >>= \case
                        Nothing -> throwM . UnexpectedResult $ messageContent
                        Just result -> atomically $ putTMVar result messageContent
        handleCommand :: Socket -> Int -> CommandWithCallback -> MarionetteT m ()
        handleCommand socket messageId (CommandWithCallback messageContent callback) = do
            let message = MarionetteMessage . Aeson.encode $ Message{..}
            liftIO . sendLazy socket . Binary.encode $ message
            forkIO do
                result <- liftIO . timeout 5_000_000 . atomically . readTMVar $ callback
                when (isNothing result) . throwM . MarionetteTimeout $ message
            atomically . modifyTVar' pendingCommands $
                IntMap.insert messageId callback
        runSocket :: MarionetteT m ()
        runSocket =
            connect "localhost" "2828" \(socket, _) -> do
                incomingQueue <- incoming socket
                void . decodeMarionetteM @_ @Greeting =<< atomically (readTQueue incomingQueue)
                let send messageId =
                        atomically
                            ( isEmptyTMVar done >>= \case
                                True -> Just <$> readTQueue sendQueue
                                False -> pure Nothing
                            )
                            >>= \case
                                Nothing -> pure ()
                                Just command -> do
                                    handleCommand socket messageId command
                                    send $ messageId + 1
                forkIO $ send 1
                forkIO . repeatUntilDone $ handleIncoming =<< atomically (readTQueue incomingQueue)
                atomically $ readTMVar done
    flip runReaderT sendQueue . (\(MarionetteT r) -> r) $ do
        forkIO runSocket
        a <- action
        atomically $ putTMVar done ()
        pure a
  where
    decodeMarionetteM :: forall m' a'. (MonadThrow m', FromJSON a') => MarionetteMessage -> m' a'
    decodeMarionetteM = either (throwM . AssertionFailed) pure . decodeMarionette

    decodeMarionette :: forall a'. (FromJSON a') => MarionetteMessage -> Either String a'
    decodeMarionette (MarionetteMessage lbs) = Aeson.eitherDecode lbs