Stream.hs

module Network.QUIC.Simple.Stream
  ( -- $intro

    -- * Message queues
    MessageQueues
  , serialise
  , serialiseFrom
  , codec
  , codecFrom
    -- * One-shot messages
  , sendMessage
  , recvMessage
    -- * Raw bytes
  , consume
    -- * Multiple streams
  , open
  , acceptLoop
  ) where

import Codec.Serialise (Serialise, deserialiseIncremental)
import Codec.Serialise qualified as CBOR
import Codec.Serialise qualified as IDecode (IDecode(..))
import Control.Concurrent.Async (Async, async, cancel, link, poll, race_)
import Control.Concurrent.STM
import Control.Exception (finally, throwIO)
import Control.Monad (filterM, forever, unless)
import Control.Monad.ST (stToIO)
import Data.ByteString qualified as BS
import Data.ByteString.Lazy qualified as BSL
import Data.IORef
import Data.Maybe (isNothing)
import Network.QUIC qualified as QUIC

{- $intro
This module is meant to be imported qualified:

> import Network.QUIC.Simple.Stream qualified as Stream

The connection wrappers in "Network.QUIC.Simple" hand out the first stream.
The functions here run protocols over it, or open and accept more streams.
-}

{- | A pair of bounded queues wrapping a stream.
-}
type MessageQueues sendMsg recvMsg = (TBQueue sendMsg, TBQueue recvMsg)

{- | Wrap the stream with the CBOR codec for both incoming and outgoing messages.

The decoder will perform incremental parsing and emit complete messages.

No extra framing is required since CBOR is self-delimiting.
-}
serialise
  :: (Serialise sendMsg, Serialise recvMsg)
  => QUIC.Stream
  -> IO (Async (), MessageQueues sendMsg recvMsg)
serialise = serialiseFrom ""

{- | Same as 'serialise', but starts decoding from the bytes already received.

Use with the leftovers from 'recvMessage' when a stream starts with a header.
-}
serialiseFrom
  :: forall sendMsg recvMsg
  . (Serialise sendMsg, Serialise recvMsg)
  => BS.ByteString
  -> QUIC.Stream
  -> IO (Async (), MessageQueues sendMsg recvMsg)
serialiseFrom leftovers stream = do
  initial <- stToIO $ deserialiseIncremental @recvMsg
  state <- newIORef initial
  let
    decode starting chunk = do
      decoder <- readIORef state
      case decoder of
        IDecode.Fail _leftovers _offset err ->
          throwIO err -- crash writer (thus the stream, and the reader/writer etc)
        IDecode.Done leftovers' _consumed msg -> do
          stToIO deserialiseIncremental >>= writeIORef state -- restart decoder
          pure (leftovers', Just msg)
        IDecode.Partial consume_ -> do
          -- want more data (initial state?)
          stToIO (consume_ $ Just chunk) >>= writeIORef state -- step decoder
          if starting then
            -- re-check if done
            decode False ""
          else
            -- suspend and wait for next chunk
            pure ("", Nothing)
  codecFrom CBOR.serialise (decode True) leftovers stream

{- | Wrap the stream with a codec to provide a TBQueue interface to it.

The decoder loop is stateless.
But it runs in IO so you can use external state and terminate the stream by erroring out.

The worker finishes when the peer closes the stream, or when either the codec or the stream fails.
Cancel it to close the stream from this side.
Messages still queued for sending at that moment are dropped.
-}
codec
  :: (sendMsg -> BSL.ByteString) -- ^ Encoder for outgoing messages
  -> (BS.ByteString -> IO (BS.ByteString, Maybe recvMsg)) -- ^ Decoder for incomming chunks
  -> QUIC.Stream
  -> IO (Async (), MessageQueues sendMsg recvMsg)
codec encode decode = codecFrom encode decode ""

{- | Same as 'codec', but starts decoding from the bytes already received.
-}
codecFrom
  :: (sendMsg -> BSL.ByteString) -- ^ Encoder for outgoing messages
  -> (BS.ByteString -> IO (BS.ByteString, Maybe recvMsg)) -- ^ Decoder for incomming chunks
  -> BS.ByteString -- ^ Bytes received before the codec took over
  -> QUIC.Stream
  -> IO (Async (), MessageQueues sendMsg recvMsg)
codecFrom encode decode leftovers stream = do
  readQ <- newTBQueueIO 1024
  writeQ <- newTBQueueIO 1024
  worker <- async $
    race_ (reader leftovers readQ) (writer writeQ) `finally` QUIC.closeStream stream
  pure (worker, (writeQ, readQ))
  where
    reader buffered readQ = do
      chunk <-
        if BS.null buffered then
          QUIC.recvStream stream 4096
        else
          pure buffered
      unless (BS.null chunk) do
        (buffered', message_) <- decode chunk
        mapM_ (atomically . writeTBQueue readQ) message_
        reader buffered' readQ

    writer writeQ = do
      message <- atomically $ readTBQueue writeQ
      let chunks = BSL.toChunks $ encode message
      QUIC.sendStreamMany stream chunks
      writer writeQ

{- | Send a single CBOR message to a stream.
-}
sendMessage :: Serialise msg => QUIC.Stream -> msg -> IO ()
sendMessage stream = QUIC.sendStreamMany stream . BSL.toChunks . CBOR.serialise

{- | Receive a single CBOR message from a stream.

Bytes received past the end of the message are returned along with it.
They belong to whatever comes next on the stream,
so feed them to 'serialiseFrom' or 'consume'.
In a lockstep exchange there will be none.

Throws a decoding error if the stream ends or gets garbled before the message is complete.
-}
recvMessage :: forall msg. Serialise msg => QUIC.Stream -> IO (msg, BS.ByteString)
recvMessage stream = stToIO (deserialiseIncremental @msg) >>= step
  where
    step = \case
      IDecode.Done leftovers _offset msg ->
        pure (msg, leftovers)
      IDecode.Fail _leftovers _offset err ->
        throwIO err
      IDecode.Partial consume_ -> do
        chunk <- QUIC.recvStream stream 4096
        stToIO (consume_ $ if BS.null chunk then Nothing else Just chunk) >>= step

{- | Feed the incoming bytes to a sink until the peer closes the stream.

The bytes received before, e.g. the leftovers from 'recvMessage', go first.
-}
consume :: BS.ByteString -> QUIC.Stream -> (BS.ByteString -> IO ()) -> IO ()
consume leftovers stream sink = do
  unless (BS.null leftovers) $ sink leftovers
  drain
  where
    drain = do
      chunk <- QUIC.recvStream stream 65536
      unless (BS.null chunk) do
        sink chunk
        drain

{- | Request a new bidirectional stream and announce its purpose with a header message.

The peer is expected to run 'acceptLoop' with a matching header type.
Anything sent right after the header will arrive as leftovers on the other side.
-}
open :: Serialise header => QUIC.Connection -> header -> IO QUIC.Stream
open conn header = do
  new <- QUIC.stream conn
  sendMessage new header
  pure new

{- | Accept streams opened by the peer and dispatch them by their header message.

Each stream gets a thread of its own that closes the stream when the handler returns.
The handler receives the bytes that arrived along with the header,
to be passed along to 'serialiseFrom' or 'consume'.

A crashing handler takes down the accept loop, and thus the connection.
Leaving the loop, by cancelling it or by connection failure, cancels the handlers still running.
-}
acceptLoop
  :: Serialise header
  => QUIC.Connection
  -> (header -> BS.ByteString -> QUIC.Stream -> IO ())
  -> IO ()
acceptLoop conn handler = do
  running <- newTVarIO []
  let
    cancelAll = readTVarIO running >>= mapM_ cancel
    prune = readTVarIO running >>= filterM (fmap isNothing . poll) >>= atomically . writeTVar running
    handleStream new = do
      (header, leftovers) <- recvMessage new
      handler header leftovers new
    accept = forever do
      new <- QUIC.acceptStream conn
      worker <- async $ handleStream new `finally` QUIC.closeStream new
      link worker
      prune
      atomically $ modifyTVar' running (worker :)
  accept `finally` cancelAll