{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE RankNTypes #-}

module Network.HTTP2.H2.OutBodyIface (
    StreamTerminated (..),
    withOutBodyIface,
) where

import Control.Concurrent.STM
import Control.Exception
import Network.HTTP.Semantics
import Network.HTTP.Semantics.IO
import Network.HTTP2.H2.Context
import Network.HTTP2.H2.Sync
import Network.HTTP2.H2.Types

----------------------------------------------------------------

data StreamTerminated
    = StreamPushedFinal
    | StreamCancelled
    | StreamOutOfScope
    | StreamRemoteReset ClosedCode
    deriving (Int -> StreamTerminated -> ShowS
[StreamTerminated] -> ShowS
StreamTerminated -> String
(Int -> StreamTerminated -> ShowS)
-> (StreamTerminated -> String)
-> ([StreamTerminated] -> ShowS)
-> Show StreamTerminated
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> StreamTerminated -> ShowS
showsPrec :: Int -> StreamTerminated -> ShowS
$cshow :: StreamTerminated -> String
show :: StreamTerminated -> String
$cshowList :: [StreamTerminated] -> ShowS
showList :: [StreamTerminated] -> ShowS
Show)
    deriving anyclass (Show StreamTerminated
Typeable StreamTerminated
(Typeable StreamTerminated, Show StreamTerminated) =>
(StreamTerminated -> SomeException)
-> (SomeException -> Maybe StreamTerminated)
-> (StreamTerminated -> String)
-> Exception StreamTerminated
SomeException -> Maybe StreamTerminated
StreamTerminated -> String
StreamTerminated -> SomeException
forall e.
(Typeable e, Show e) =>
(e -> SomeException)
-> (SomeException -> Maybe e) -> (e -> String) -> Exception e
$ctoException :: StreamTerminated -> SomeException
toException :: StreamTerminated -> SomeException
$cfromException :: SomeException -> Maybe StreamTerminated
fromException :: SomeException -> Maybe StreamTerminated
$cdisplayException :: StreamTerminated -> String
displayException :: StreamTerminated -> String
Exception)

----------------------------------------------------------------

withOutBodyIface
    :: Context
    -> Stream
    -> TBQueue StreamingChunk
    -> (forall a. IO a -> IO a)
    -> (OutBodyIface -> IO r)
    -> IO r
withOutBodyIface :: forall r.
Context
-> Stream
-> TBQueue StreamingChunk
-> (forall a. IO a -> IO a)
-> (OutBodyIface -> IO r)
-> IO r
withOutBodyIface ctx :: Context
ctx@Context{TQueue Output
outputQ :: TQueue Output
outputQ :: Context -> TQueue Output
outputQ} Stream
strm TBQueue StreamingChunk
tbq forall a. IO a -> IO a
unmask OutBodyIface -> IO r
k = do
    TVar (Maybe StreamTerminated)
terminated <- Maybe StreamTerminated -> IO (TVar (Maybe StreamTerminated))
forall a. a -> IO (TVar a)
newTVarIO Maybe StreamTerminated
forall a. Maybe a
Nothing
    let checkNotTerminated :: STM ()
        checkNotTerminated :: STM ()
checkNotTerminated = do
            Maybe StreamTerminated
mTerminated <- TVar (Maybe StreamTerminated) -> STM (Maybe StreamTerminated)
forall a. TVar a -> STM a
readTVar TVar (Maybe StreamTerminated)
terminated
            STM ()
-> (StreamTerminated -> STM ()) -> Maybe StreamTerminated -> STM ()
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (() -> STM ()
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return ()) StreamTerminated -> STM ()
forall e a. Exception e => e -> STM a
throwSTM Maybe StreamTerminated
mTerminated

        -- Check if the peer is still listening for messages
        --
        -- It is important to call 'checkNotClosed' prior to enqueuing stream
        -- chunks to ensure that 'writeTBQueue' will not block indefinitely
        -- (because nothing is consuming elements from the queue anymore).
        --
        -- Assumes 'checkNotTerminated'.
        checkNotClosed :: STM ()
        checkNotClosed :: STM ()
checkNotClosed = do
            Maybe ClosedCode
mClosed <- STM (Maybe ClosedCode)
getIsClosed
            case Maybe ClosedCode
mClosed of
                Just ClosedCode
code ->
                    -- When the stream is closed, but /we/ did not close it (or
                    -- 'checkNotTerminated' would have thrown an exception), it
                    -- must mean that our peer send us a RST_STREAM, indicating
                    -- that they do not want to receive any further messages.
                    StreamTerminated -> STM ()
forall e a. Exception e => e -> STM a
throwSTM (StreamTerminated -> STM ()) -> StreamTerminated -> STM ()
forall a b. (a -> b) -> a -> b
$ ClosedCode -> StreamTerminated
StreamRemoteReset ClosedCode
code
                Maybe ClosedCode
_otherwise ->
                    () -> STM ()
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return ()

        getIsClosed :: STM (Maybe ClosedCode)
        getIsClosed :: STM (Maybe ClosedCode)
getIsClosed = do
            StreamState
st <- TVar StreamState -> STM StreamState
forall a. TVar a -> STM a
readTVar (Stream -> TVar StreamState
streamState Stream
strm)
            case StreamState
st of
                Closed ClosedCode
code -> Maybe ClosedCode -> STM (Maybe ClosedCode)
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe ClosedCode -> STM (Maybe ClosedCode))
-> Maybe ClosedCode -> STM (Maybe ClosedCode)
forall a b. (a -> b) -> a -> b
$ ClosedCode -> Maybe ClosedCode
forall a. a -> Maybe a
Just ClosedCode
code
                StreamState
_otherwise -> Maybe ClosedCode -> STM (Maybe ClosedCode)
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe ClosedCode
forall a. Maybe a
Nothing

        cancelAfterFinish :: Maybe SomeException -> STM ()
        cancelAfterFinish :: Maybe SomeException -> STM ()
cancelAfterFinish Maybe SomeException
mErr =
            TQueue Output -> Output -> STM ()
forall a. TQueue a -> a -> STM ()
writeTQueue TQueue Output
outputQ (Output -> STM ()) -> Output -> STM ()
forall a b. (a -> b) -> a -> b
$ Context
-> Stream -> Maybe (TBQueue StreamingChunk) -> OutputType -> Output
makeOutputIO Context
ctx Stream
strm Maybe (TBQueue StreamingChunk)
forall a. Maybe a
Nothing (Maybe SomeException -> OutputType
OReset Maybe SomeException
mErr)

        iface :: OutBodyIface
        iface :: OutBodyIface
iface =
            OutBodyIface
                { outBodyUnmask :: forall a. IO a -> IO a
outBodyUnmask = IO x -> IO x
forall a. IO a -> IO a
unmask
                , outBodyPush :: Builder -> IO ()
outBodyPush = \Builder
b -> STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
                    STM ()
checkNotTerminated
                    STM ()
checkNotClosed
                    TBQueue StreamingChunk -> StreamingChunk -> STM ()
forall a. TBQueue a -> a -> STM ()
writeTBQueue TBQueue StreamingChunk
tbq (StreamingChunk -> STM ()) -> StreamingChunk -> STM ()
forall a b. (a -> b) -> a -> b
$ Builder -> IsEndOfStream -> StreamingChunk
StreamingBuilder Builder
b IsEndOfStream
NotEndOfStream
                , outBodyPushFinal :: Builder -> IO ()
outBodyPushFinal = \Builder
b -> STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
                    STM ()
checkNotTerminated
                    STM ()
checkNotClosed
                    TVar (Maybe StreamTerminated) -> Maybe StreamTerminated -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (Maybe StreamTerminated)
terminated (StreamTerminated -> Maybe StreamTerminated
forall a. a -> Maybe a
Just StreamTerminated
StreamPushedFinal)
                    TBQueue StreamingChunk -> StreamingChunk -> STM ()
forall a. TBQueue a -> a -> STM ()
writeTBQueue TBQueue StreamingChunk
tbq (StreamingChunk -> STM ()) -> StreamingChunk -> STM ()
forall a b. (a -> b) -> a -> b
$ Builder -> IsEndOfStream -> StreamingChunk
StreamingBuilder Builder
b (Maybe (IO ()) -> IsEndOfStream
EndOfStream Maybe (IO ())
forall a. Maybe a
Nothing)
                    TBQueue StreamingChunk -> StreamingChunk -> STM ()
forall a. TBQueue a -> a -> STM ()
writeTBQueue TBQueue StreamingChunk
tbq (StreamingChunk -> STM ()) -> StreamingChunk -> STM ()
forall a b. (a -> b) -> a -> b
$ Maybe (IO ()) -> StreamingChunk
StreamingFinished Maybe (IO ())
forall a. Maybe a
Nothing
                , outBodyFlush :: IO ()
outBodyFlush = STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
                    STM ()
checkNotTerminated
                    STM ()
checkNotClosed
                    TBQueue StreamingChunk -> StreamingChunk -> STM ()
forall a. TBQueue a -> a -> STM ()
writeTBQueue TBQueue StreamingChunk
tbq StreamingChunk
StreamingFlush
                , outBodyCancel :: Maybe SomeException -> IO ()
outBodyCancel = \Maybe SomeException
mErr -> STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
                    Maybe StreamTerminated
mTerminated <- TVar (Maybe StreamTerminated) -> STM (Maybe StreamTerminated)
forall a. TVar a -> STM a
readTVar TVar (Maybe StreamTerminated)
terminated
                    Maybe ClosedCode
mClosed <- STM (Maybe ClosedCode)
getIsClosed
                    case (Maybe ClosedCode
mClosed, Maybe StreamTerminated
mTerminated) of
                        (Maybe ClosedCode
Nothing, Maybe StreamTerminated
Nothing) -> do
                            TVar (Maybe StreamTerminated) -> Maybe StreamTerminated -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (Maybe StreamTerminated)
terminated (StreamTerminated -> Maybe StreamTerminated
forall a. a -> Maybe a
Just StreamTerminated
StreamCancelled)
                            TBQueue StreamingChunk -> StreamingChunk -> STM ()
forall a. TBQueue a -> a -> STM ()
writeTBQueue TBQueue StreamingChunk
tbq (StreamingChunk -> STM ()) -> StreamingChunk -> STM ()
forall a b. (a -> b) -> a -> b
$ Maybe SomeException -> StreamingChunk
StreamingCancelled Maybe SomeException
mErr
                        (Maybe ClosedCode
Nothing, Just StreamTerminated
StreamCancelled) ->
                            -- Already cancelled
                            () -> STM ()
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
                        (Maybe ClosedCode
Nothing, Just StreamTerminated
_) -> do
                            -- We finished streaming (that is, sending messages to the peer),
                            -- but we must still be able to cancel the stream entirely
                            -- (that is, tell the peer that we no longer want to /receive/ messages: RST_STREAM)
                            TVar (Maybe StreamTerminated) -> Maybe StreamTerminated -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (Maybe StreamTerminated)
terminated (StreamTerminated -> Maybe StreamTerminated
forall a. a -> Maybe a
Just StreamTerminated
StreamCancelled)
                            Maybe SomeException -> STM ()
cancelAfterFinish Maybe SomeException
mErr
                        (Just ClosedCode
_code, Maybe StreamTerminated
_) ->
                            -- Peer already closed
                            () -> STM ()
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
                }

        -- The body is done with: it returned, or it threw.  One that returned
        -- has sent all it had, and END_STREAM goes out.  One that threw has
        -- not: ending the stream all the same delivered a truncated body to
        -- the peer as if it were complete.  So it is cancelled instead, and
        -- the peer gets RST_STREAM (INTERNAL_ERROR).
        finished :: Maybe SomeException -> IO ()
        finished :: Maybe SomeException -> IO ()
finished Maybe SomeException
mErr = STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
            Maybe StreamTerminated
mTerminated <- TVar (Maybe StreamTerminated) -> STM (Maybe StreamTerminated)
forall a. TVar a -> STM a
readTVar TVar (Maybe StreamTerminated)
terminated
            Maybe ClosedCode
mClosed <- STM (Maybe ClosedCode)
getIsClosed
            case (Maybe ClosedCode
mClosed, Maybe StreamTerminated
mTerminated) of
                (Maybe ClosedCode
Nothing, Maybe StreamTerminated
Nothing) -> case Maybe SomeException
mErr of
                    Maybe SomeException
Nothing -> do
                        TVar (Maybe StreamTerminated) -> Maybe StreamTerminated -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (Maybe StreamTerminated)
terminated (StreamTerminated -> Maybe StreamTerminated
forall a. a -> Maybe a
Just StreamTerminated
StreamOutOfScope)
                        TBQueue StreamingChunk -> StreamingChunk -> STM ()
forall a. TBQueue a -> a -> STM ()
writeTBQueue TBQueue StreamingChunk
tbq (StreamingChunk -> STM ()) -> StreamingChunk -> STM ()
forall a b. (a -> b) -> a -> b
$ Maybe (IO ()) -> StreamingChunk
StreamingFinished Maybe (IO ())
forall a. Maybe a
Nothing
                    Just SomeException
_ -> do
                        TVar (Maybe StreamTerminated) -> Maybe StreamTerminated -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (Maybe StreamTerminated)
terminated (StreamTerminated -> Maybe StreamTerminated
forall a. a -> Maybe a
Just StreamTerminated
StreamCancelled)
                        TBQueue StreamingChunk -> StreamingChunk -> STM ()
forall a. TBQueue a -> a -> STM ()
writeTBQueue TBQueue StreamingChunk
tbq (StreamingChunk -> STM ()) -> StreamingChunk -> STM ()
forall a b. (a -> b) -> a -> b
$ Maybe SomeException -> StreamingChunk
StreamingCancelled Maybe SomeException
mErr
                (Maybe ClosedCode
Nothing, Just StreamTerminated
_) ->
                    -- We already terminated
                    () -> STM ()
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
                (Just ClosedCode
_code, Maybe StreamTerminated
_) ->
                    -- Peer already closed
                    () -> STM ()
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return ()

    -- Like 'finally', telling the two apart.
    ((forall a. IO a -> IO a) -> IO r) -> IO r
forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
mask (((forall a. IO a -> IO a) -> IO r) -> IO r)
-> ((forall a. IO a -> IO a) -> IO r) -> IO r
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
restore -> do
        r
r <- IO r -> IO r
forall a. IO a -> IO a
restore (OutBodyIface -> IO r
k OutBodyIface
iface) IO r -> (SomeException -> IO r) -> IO r
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` \SomeException
e -> Maybe SomeException -> IO ()
finished (SomeException -> Maybe SomeException
forall a. a -> Maybe a
Just SomeException
e) IO () -> IO r -> IO r
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> SomeException -> IO r
forall e a. Exception e => e -> IO a
throwIO SomeException
e
        Maybe SomeException -> IO ()
finished Maybe SomeException
forall a. Maybe a
Nothing
        r -> IO r
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return r
r