content: -- | Channel
  --
  -- You should not have to import this module directly; instead import
  -- "Network.GRPC.Util.Session".
  module Network.GRPC.Util.Session.Channel (
      -- * Main definition
      Channel(..)
    , initChannel
      -- ** Flow state
    , FlowState(..)
    , RegularFlowState(..)
    , initFlowStateRegular
      -- * Working with an open channel
    , getInboundHeaders
    , send
    , recvBoth
    , recvEither
    , RecvFinal(..)
    , RecvAfterFinal(..)
    , SendAfterFinal(..)
      -- * Closing
    , waitForOutbound
    , close
    , ChannelDiscarded(..)
    , ChannelAborted(..)
      -- * Support for half-closing
    , InboundResult
    , AllowHalfClosed(..)
    , linkOutboundToInbound
      -- * Constructing channels
    , sendMessageLoop
    , recvMessageLoop
    , outboundTrailersMaker
    ) where
  
  import Control.Concurrent.STM
  import Control.DeepSeq (NFData, force)
  import Control.Exception
  import Control.Monad
  import Control.Monad.Catch (ExitCase(..))
  import Data.Bifunctor
  import Data.ByteString.Builder (Builder)
  import Data.ByteString.Lazy qualified as BS.Lazy
  import GHC.Stack
  
  -- Doesn't really matter if we import from .Client or .Server
  import Network.HTTP2.Client qualified as HTTP2 (
      TrailersMaker
    , NextTrailersMaker(..)
    )
  
  import Network.GRPC.Common.StreamElem (StreamElem(..))
  import Network.GRPC.Common.StreamElem qualified as StreamElem
  import Network.GRPC.Spec.Util.Parser (Parser)
  import Network.GRPC.Spec.Util.Parser qualified as Parser
  import Network.GRPC.Util.HTTP2.Stream
  import Network.GRPC.Util.RedundantConstraint
  import Network.GRPC.Util.Session.API
  import Network.GRPC.Util.Thread
  
  {-------------------------------------------------------------------------------
    Definitions
  
    The fields of 'Channel' are its /implementation/, not its interface. It is
    kept opaque in the top-level @.Peer@ module.
  
    Implementation note: it is tempting to try and define 'Channel' purely in
    terms of bytestrings, and deal with serialization and deserialization to and
    from messages in a higher layer. However, this does not work:
  
    - For deserialization, if we make chunks of messages available in the 'TMVar',
      then if multiple threads are reading from that one 'TMVar', one thread might
      get the first chunk of a message and another thread the second.
    - Similarly, for serialization, if multiple threads are trying to write to the
      'TMVar', we might get interleaving of fragments of messages.
  
    Thus, to ensure thread safety, we work at the level of messages, not bytes.
  -------------------------------------------------------------------------------}
  
  -- | Bidirectional open channel on a node to a peer node
  --
  -- The node might be a client (and its peer a server), or the node might be
  -- a server (and its peer a client); the main purpose of this abstraction
  -- is precisely to abstract over that difference.
  --
  -- Each channel is constructed for a /single/ session (request/response).
  data Channel sess = Channel {
        -- | Thread state of the thread receiving messages from the peer
        channelInbound :: TVar (ThreadState (FlowState (Inbound sess)))
  
        -- | Thread state of the thread sending messages to the peer
      , channelOutbound :: TVar (ThreadState (FlowState (Outbound sess)))
  
        -- | Have we sent the final message?
        --
        -- The sole purpose of this 'TVar' is catching user mistakes: if there is
        -- another 'send' after the final message, we can throw an exception,
        -- rather than the message simply being lost or blockng indefinitely.
      , channelSentFinal :: TVar (Maybe CallStack)
  
        -- | Have we received the final message?
        --
        -- This is used to improve the user experience; see 'channelSentFinal'.
        -- It is also used when checking if a call should be considered
        -- \"cancelled\"; see 'withRPC'.
      , channelRecvFinal :: TVar (RecvFinal (Inbound sess))
      }
  
  data RecvFinal flow =
      -- | We have not yet delivered the final message to the client
      RecvNotFinal
  
      -- | We delivered the final message, but not yet the trailers
    | RecvWithoutTrailers (Trailers flow)
  
      -- | We delivered the final message and the trailers
    | RecvFinal CallStack
  
  -- | Data flow state
  data FlowState flow =
      FlowStateRegular (RegularFlowState flow)
    | FlowStateNoMessages (NoMessages flow)
  
  -- | Regular (streaming) flow state
  data RegularFlowState flow = RegularFlowState {
        -- | Headers
        --
        -- On the client side, the outbound headers are specified when the request
        -- is made ('callRequestMetadata'), and the inbound headers are recorded
        -- once the responds starts to come in; clients can block-and-wait for
        -- these headers ('getInboundHeaders').
        --
        -- On the server side, the inbound headers are recorded when the request
        -- comes in, and the outbound headers are specified
        -- ('setResponseMetadata') before the response is initiated
        -- ('initiateResponse'/'sendTrailersOnly').
        flowHeaders :: Headers flow
  
        -- | Messages
        --
        -- This TMVar is written to for incoming messages ('recvMessageLoop') and
        -- read from for outgoing messages ('sendMessageLoop'). It acts as a
        -- one-place buffer, providing backpressure in both directions.
      , flowMsg :: TMVar (StreamElem (Trailers flow) (Message flow))
  
        -- | Trailers
        --
        -- Unlike 'flowMsg', which is /written/ to in 'recvMessageLoop' and /read/
        -- from in 'sendMessageLoop', both loops /set/ 'flowTerminated', once,
        -- just before they terminate.
        --
        -- * For 'sendMessageLoop', this means that the last message has been
        --   written (that is, the last call to 'writeChunk' has happened).
        --   This has two consequences:
        --
        --   1. @http2@ can now construct the trailers ('outboundTrailersMaker')
        --   2. Higher layers can wait on 'flowTerminated' to be /sure/ that the
        --      last message has been written.
        --
        -- * For 'recvMessageLoop', this means that the trailers have been
        --   received from the peer. Higher layers can use this to check for, or
        --   block-and-wait, to receive those trailers.
        --
        -- == Relation to 'channelSentFinal'/'channelRecvFinal'
        --
        -- 'flowTerminated' is set at different times than 'channelSentFinal' and
        -- 'channelRecvFinal' are:
        --
        -- * 'channelSentFinal' is set on the last call to 'send', but /before/
        --   the message is processed by 'sendMessageLoop'.
        -- * 'channelRecvFinal', dually, is set on the last call to 'recv, which
        --   must (necessarily) happen /before/ that message is actually made
        --   available by 'recvMessageLoop'.
        --
        -- /Their/ sole purpose is to catch user errors, not capture data flow.
      , flowTerminated :: TMVar (Trailers flow)
      }
  
  -- | 'Show' instance is useful in combination with @stm-debug@ only
  deriving instance (
      Show (Headers flow)
    , Show (TMVar (StreamElem (Trailers flow) (Message flow)))
    , Show (TMVar (Trailers flow))
    ) => Show (RegularFlowState flow)
  
  {-------------------------------------------------------------------------------
    Initialization
  -------------------------------------------------------------------------------}
  
  initChannel :: HasCallStack => IO (Channel sess)
  initChannel = do
      channelInbound   <- newThreadState
      channelOutbound  <- newThreadState
      channelSentFinal <- newTVarIO Nothing
      channelRecvFinal <- newTVarIO RecvNotFinal
      return Channel{
          channelInbound
        , channelOutbound
        , channelSentFinal
        , channelRecvFinal
        }
  
  initFlowStateRegular :: Headers flow -> IO (RegularFlowState flow)
  initFlowStateRegular flowHeaders = do
     flowMsg        <- newEmptyTMVarIO
     flowTerminated <- newEmptyTMVarIO
     return RegularFlowState {
         flowHeaders
       , flowMsg
       , flowTerminated
       }
  
  {-------------------------------------------------------------------------------
    Working with an open channel
  -------------------------------------------------------------------------------}
  
  -- | The inbound headers
  --
  -- Will block if the inbound headers have not yet been received.
  getInboundHeaders ::
       Channel sess
    -> IO (Either (NoMessages (Inbound sess)) (Headers (Inbound sess)))
  getInboundHeaders Channel{channelInbound} =
      withThreadInterface channelInbound (return . aux)
    where
      aux :: forall flow.
           FlowState flow
        -> Either (NoMessages flow) (Headers flow)
      aux (FlowStateRegular    regular)  = Right $ flowHeaders regular
      aux (FlowStateNoMessages trailers) = Left trailers
  
  -- | Send a message to the node's peer
  --
  -- It is a bug to call 'send' again after the final message (that is, a message
  -- which 'StreamElem.whenDefinitelyFinal' considers to be final). Doing so will
  -- result in a 'SendAfterFinal' exception.
  send :: forall sess.
       (HasCallStack, NFData (Message (Outbound sess)))
    => Channel sess
    -> StreamElem (Trailers (Outbound sess)) (Message (Outbound sess))
    -> IO ()
  send Channel{channelOutbound, channelSentFinal} = \msg -> do
      msg' <- evaluate $ force <$> msg
      withThreadInterface channelOutbound $ aux msg'
    where
      aux ::
           StreamElem (Trailers (Outbound sess)) (Message (Outbound sess))
        -> FlowState (Outbound sess)
        -> STM ()
      aux msg st = do
          -- By checking that we haven't sent the final message yet, we know that
          -- this call to 'putMVar' will not block indefinitely: the thread that
          -- sends messages to the peer will get to it eventually (unless it dies,
          -- in which case the thread status will change and the call to
          -- 'getThreadInterface' will be retried).
          sentFinal <- readTVar channelSentFinal
          case sentFinal of
            Just cs -> throwSTM $ SendAfterFinal cs
            Nothing -> return ()
          case st of
            FlowStateRegular regular -> do
              StreamElem.whenDefinitelyFinal msg $ \_trailers ->
                writeTVar channelSentFinal $ Just callStack
  
              putTMVar (flowMsg regular) msg
            FlowStateNoMessages _ ->
              -- For outgoing messages, the caller decides to use Trailers-Only,
              -- so if they then subsequently call 'send', we throw an exception.
              -- This is different for /inbound/ messages; see 'recv', below.
              throwSTM $ SendButTrailersOnly
  
  -- | Receive a message from the node's peer
  --
  -- If the sender indicates that the message is final /when/ they send it, by
  -- sending the HTTP trailers in the same frame, then we will return the message
  -- and the trailers together. It is a bug to call 'recvBoth' again after this;
  -- doing so will result in a 'RecvAfterFinal' exception.
  recvBoth :: forall sess.
       HasCallStack
    => Channel sess
    -> IO ( Either
              (NoMessages (Inbound sess))
              (StreamElem (Trailers (Inbound sess)) (Message (Inbound sess)))
          )
  recvBoth =
      recv'
        StreamElem
        NoMoreElems
        ((,Nothing) . uncurry FinalElem)
  
  -- | Variant on 'recvBoth' where trailers are always returned separately
  --
  -- Unlike in 'recvBoth', even if the sender indicates that the final message is
  -- final when they send it, we will store these trailers internally and return
  -- only that final message. The trailers are then returned on the /next/ call to
  -- 'recvEither'. Call 'recvEither' again /after/ receiving the trailers is a
  -- bug; doing so will result in a 'RecvAfterFinal' exception.
  recvEither ::
       HasCallStack
    => Channel sess
    -> IO ( Either
              (NoMessages (Inbound sess))
              (Either (Trailers (Inbound sess)) (Message (Inbound sess)))
          )
  recvEither =
      recv'
        Right
        Left
        (bimap Right Just)
  
  -- | Internal generalization of 'recvBoth' and 'recvEither'
  recv' :: forall sess b.
       HasCallStack
    => (Message    (Inbound sess) -> b)  -- ^ Message without trailers
    -> (Trailers   (Inbound sess) -> b)  -- ^ Trailers without (final) message
    -> (    (Message (Inbound sess), Trailers (Inbound sess))
         -> (b, Maybe (Trailers (Inbound sess)))
       )
       -- ^ Message with trailers
       --
       -- In addition to the result, should also return the trailers to keep for
       -- the next call to 'recv'' (if any).
    -> Channel sess
    -> IO (Either (NoMessages (Inbound sess)) b)
  recv' messageWithoutTrailers
        trailersWithoutMessage
        messageWithTrailers
        Channel{channelInbound, channelRecvFinal} =
      withThreadInterface channelInbound aux
    where
      aux ::
           FlowState (Inbound sess)
        -> STM (Either (NoMessages (Inbound sess)) b)
      aux st = do
          -- By checking that we haven't received the final message yet, we know
          -- that this call to 'takeTMVar' will not block indefinitely: the thread
          -- that receives messages from the peer will get to it eventually
          -- (unless it dies, in which case the thread status will change and the
          -- call to 'getThreadInterface' will be retried).
          readFinal <- readTVar channelRecvFinal
          case readFinal of
            RecvNotFinal ->
              case st of
                FlowStateRegular regular -> Right <$> do
                  streamElem <- takeTMVar (flowMsg regular)
                  -- We update 'channelRecvFinal' in the same tx as the read, to
                  -- atomically change "there is a value" to "all values read".
                  case streamElem of
                    StreamElem msg ->
                      return $ messageWithoutTrailers msg
                    FinalElem msg trailers -> do
                      let (b, mTrailers) = messageWithTrailers (msg, trailers)
                      writeTVar channelRecvFinal $
                        maybe (RecvFinal callStack) RecvWithoutTrailers mTrailers
                      return $ b
                    NoMoreElems trailers -> do
                      writeTVar channelRecvFinal $ RecvFinal callStack
                      return $ trailersWithoutMessage trailers
                FlowStateNoMessages trailers -> do
                  writeTVar channelRecvFinal $ RecvFinal callStack
                  return $ Left trailers
            RecvWithoutTrailers trailers -> do
              writeTVar channelRecvFinal $ RecvFinal callStack
              return $ Right $ trailersWithoutMessage trailers
            RecvFinal cs ->
              throwSTM $ RecvAfterFinal cs
  
  -- | Thrown by 'send'
  --
  -- The 'CallStack' is the callstack of the final call to 'send'.
  --
  -- See 'send' for additional discussion.
  data SendAfterFinal =
      -- | Call to 'send' after the final message was sent
      SendAfterFinal CallStack
  
      -- | Call to 'send', but we are in the Trailers-Only case
    | SendButTrailersOnly
    deriving stock (Show)
    deriving anyclass (Exception)
  
  -- | Thrown by 'recv'
  --
  -- The 'CallStack' is the callstack of the final call to 'recv'.
  --
  -- See 'recv' for additional discussion.
  data RecvAfterFinal =
       -- | Call to 'recv' after the final message was already received
       RecvAfterFinal CallStack
    deriving stock (Show)
    deriving anyclass (Exception)
  
  {-------------------------------------------------------------------------------
    Closing
  -------------------------------------------------------------------------------}
  
  -- | Wait for the outbound thread to terminate
  --
  -- See 'close' for discussion.
  waitForOutbound :: Channel sess -> IO ()
  waitForOutbound Channel{channelOutbound} = atomically $
      waitForNormalThreadTermination channelOutbound
  
  -- | Close the channel
  --
  -- Before a channel can be closed, you should 'send' the final outbound message
  -- and then 'waitForOutbound' until all outbound messages have been processed.
  -- Not doing so is considered a bug (it is not possible to do this implicitly,
  -- because the final call to 'send' involves a choice of trailers, and calling
  -- 'waitForOutbound' /without/ a final close to 'send' will result in deadlock).
  -- Typically code will also process all /incoming/ messages, but doing so is of
  -- course not mandatory.
  --
  -- Calling 'close' will kill the outbound thread ('sendMessageLoop'), /if/ it is
  -- still running. If the thread was terminated with an exception, this could
  -- mean one of two things:
  --
  -- 1. The connection to the peer was lost
  -- 2. Proper procedure for outbound messages was not followed (see above)
  --
  -- In the case of (2) this is bug in the caller, and so 'close' will return an
  -- exception. In the case of (1), however, very likely an exception will
  -- /already/ have been thrown when a communication attempt was made, and 'close'
  -- will return 'Nothing'. This matches the design philosophy in @grapesy@ that
  -- exceptions are thrown \"lazily\" rather than \"strictly\".
  close ::
       HasCallStack
    => Channel sess
    -> ExitCase a    -- ^ The reason why the channel is being closed
    -> IO (Maybe SomeException)
  close Channel{channelOutbound} reason = do
      -- We leave the inbound thread running. Although the channel is closed,
      -- there might still be unprocessed messages in the queue. The inbound
      -- thread will terminate once it reaches the end of the queue.
      outbound <- cancelThread channelOutbound channelClosed
      case outbound of
        AlreadyTerminated _ ->
          return $ Nothing
        AlreadyAborted _err ->
          -- Connection to the peer was lost prior to closing
          return $ Nothing
        Cancelled ->
          -- Proper procedure for outbound messages was not followed
          return $ Just channelClosed
    where
      channelClosed :: SomeException
      channelClosed =
          case reason of
            ExitCaseSuccess _   -> toException $ ChannelDiscarded callStack
            ExitCaseAbort       -> toException $ ChannelAborted   callStack
            ExitCaseException e -> e
  
  -- | Channel was closed because it was discarded
  --
  -- This typically corresponds to leaving the scope of 'runHandler' or
  -- 'withRPC' (without throwing an exception).
  data ChannelDiscarded = ChannelDiscarded CallStack
    deriving stock (Show)
    deriving anyclass (Exception)
  
  -- | Channel was closed for an unknown reason
  --
  -- This will only be used in monad stacks that have error mechanisms other
  -- than exceptions.
  data ChannelAborted = ChannelAborted CallStack
    deriving stock (Show)
    deriving anyclass (Exception)
  
  {-------------------------------------------------------------------------------
    Support for half-closing
  -------------------------------------------------------------------------------}
  
  type InboundResult sess =
         Either (NoMessages (Inbound sess))
                (Trailers   (Inbound sess))
  
  -- | Should we allow for a half-clsoed connection state?
  --
  -- In HTTP2, streams are bidirectional and can be half-closed in either
  -- direction. This is however not true for all applications /of/ HTTP2. For
  -- example, in gRPC the stream can be half-closed from the client to the server
  -- (indicating that the client will not send any more messages), but not from
  -- the server to the client: when the server half-closes their connection, it
  -- sends the gRPC trailers and this terminates the call.
  data AllowHalfClosed sess =
      ContinueWhenInboundClosed
    | TerminateWhenInboundClosed (InboundResult sess -> SomeException)
  
  -- | Link outbound thread to the inbound thread
  --
  -- This should be wrapped around the body of the inbound thread. It ensures that
  -- when the inbound thread throws an exception, the outbound thread dies also.
  -- This improves predictability of exceptions: the inbound thread spends most of
  -- its time blocked on messages from the peer, and will therefore notice when
  -- the connection is lost. This is not true for the outbound thread, which
  -- spends most of its time blocked waiting for messages to send to the peer.
  linkOutboundToInbound :: forall sess.
       IsSession sess
    => AllowHalfClosed sess
    -> Channel sess
    -> IO (InboundResult sess)
    -> IO ()
  linkOutboundToInbound allowHalfClosed channel inbound = do
      mResult <- try inbound
  
      -- Implementation note: After cancelThread returns, 'channelOutbound' has
      -- been updated, and considered dead, even if perhaps the thread is still
      -- cleaning up.
  
      case (mResult, allowHalfClosed) of
        (Right _result, ContinueWhenInboundClosed) ->
          return ()
        (Right result, TerminateWhenInboundClosed f) ->
          void $ cancelThread (channelOutbound channel) (f result)
        (Left (exception :: SomeException), _) -> do
          void $ cancelThread (channelOutbound channel) exception
          throwIO exception
    where
      _ = addConstraint @(IsSession sess)
  
  {-------------------------------------------------------------------------------
    Constructing channels
  
    Both 'sendMessageLoop' and 'recvMessageLoop' will be run in newly forked
    threads, using the 'Thread' API from "Network.GRPC.Util.Thread". We are
    therefore not particularly worried about these loops being interrupted by
    asynchronous exceptions: this only happens if the threads are explicitly
    terminated (when the corresponding channels are closed), in which case any
    attempt to interact with them after they have been killed will be handled by
    'getThreadInterface' throwing 'ThreadInterfaceUnavailable'.
  -------------------------------------------------------------------------------}
  
  -- | Send all messages to the node's peer
  sendMessageLoop :: forall sess.
       IsSession sess
    => sess
    -> RegularFlowState (Outbound sess)
    -> OutputStream
    -> IO ()
  sendMessageLoop sess st stream = do
      trailers <- loop
      atomically $ putTMVar (flowTerminated st) trailers
    where
      build :: (Message (Outbound sess) -> Builder)
      build = buildMsg sess (flowHeaders st)
  
      loop :: IO (Trailers (Outbound sess))
      loop = do
          msg <- atomically $ takeTMVar (flowMsg st)
          case msg of
            StreamElem x -> do
              writeChunk stream $ build x
              flush stream
              loop
            FinalElem x trailers -> do
              writeChunkFinal stream $ build x
              return trailers
            NoMoreElems trailers -> do
              -- It is crucial to still 'writeChunkFinal' here to guarantee that
              -- cancellation is a no-op. Without it, cancellation may result in a
              -- @RST_STREAM@ frame may being sent to the peer.
              --
              -- This does not necessarily write a DATA frame, since http2 avoids
              -- writing empty data frames unless they are marked @END_OF_STREAM@.
              writeChunkFinal stream $ mempty
              return trailers
  
  -- | Receive all messages sent by the node's peer
  recvMessageLoop :: forall sess.
       IsSession sess
    => sess
    -> RegularFlowState (Inbound sess)
    -> InputStream
    -> IO (Trailers (Inbound sess))
  recvMessageLoop sess st stream =
      go $ parseMsg sess (flowHeaders st)
    where
      go :: Parser String (Message (Inbound sess)) -> IO (Trailers (Inbound sess))
      go parser = do
          mProcessedFinal <- throwParseErrors =<< Parser.processAll
            (getChunk stream)
            processOne
            processFinal
            parser
          case mProcessedFinal of
            Just trailers ->
              return trailers
            Nothing -> do
              trailers <- processTrailers
              atomically $ putTMVar (flowMsg st) $ NoMoreElems trailers
              return trailers
  
      processOne :: Message (Inbound sess) -> IO ()
      processOne msg = do
          atomically $ putTMVar (flowMsg st) $ StreamElem msg
  
      processFinal :: Message (Inbound sess) -> IO (Trailers (Inbound sess))
      processFinal msg = do
          trailers <- processTrailers
          atomically $ putTMVar (flowMsg st) $ FinalElem msg trailers
          return trailers
  
      processTrailers :: IO (Trailers (Inbound sess))
      processTrailers = do
          trailers <- parseInboundTrailers sess =<< getTrailers stream
          atomically $ putTMVar (flowTerminated st) $ trailers
          return trailers
  
      throwParseErrors :: Parser.ProcessResult String b -> IO (Maybe b)
      throwParseErrors (Parser.ProcessError err) =
          throwIO $ PeerSentMalformedMessage err
      throwParseErrors (Parser.ProcessedWithFinal b leftover) = do
          unless (BS.Lazy.null leftover) $ throwIO PeerSentIncompleteMessage
          return $ Just b
      throwParseErrors (Parser.ProcessedWithoutFinal leftover) = do
          unless (BS.Lazy.null leftover) $ throwIO PeerSentIncompleteMessage
          return $ Nothing
  
  outboundTrailersMaker :: forall sess.
       IsSession sess
    => sess
    -> Channel sess
    -> RegularFlowState (Outbound sess)
    -> HTTP2.TrailersMaker
  outboundTrailersMaker sess Channel{channelOutbound} regular = go
    where
      go :: HTTP2.TrailersMaker
      go (Just _) = return $ HTTP2.NextTrailersMaker go
      go Nothing  = do
          mFlowState <- atomically $
            unlessAbnormallyTerminated channelOutbound $
              readTMVar (flowTerminated regular)
          case mFlowState of
              Right trailers ->
                return $ HTTP2.Trailers $ buildOutboundTrailers sess trailers
              Left _exception ->
                return $ HTTP2.Trailers []
  