<commit_before>{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE BangPatterns, CPP #-}
{-# LANGUAGE NamedFieldPuns #-}

module Network.Wai.Handler.Warp.HTTP2.Sender (frameSender) where

#if __GLASGOW_HASKELL__ < 709
import Control.Applicative
#endif
import Control.Concurrent.STM
import qualified Control.Exception as E
import Control.Monad (unless, void, when)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Builder.Extra as B
import Data.Maybe (isNothing)
import Foreign.Ptr
import Network.HPACK (setLimitForEncoding)
import Network.HTTP2
import Network.HTTP2.Priority (dequeueSTM)
import Network.Wai
import Network.Wai.Handler.Warp.Buffer
import Network.Wai.Handler.Warp.HTTP2.EncodeFrame
import Network.Wai.Handler.Warp.HTTP2.HPACK
import Network.Wai.Handler.Warp.HTTP2.Types
import Network.Wai.Handler.Warp.IORef
import qualified Network.Wai.Handler.Warp.Settings as S
#ifndef WINDOWS
import qualified Network.Wai.Handler.Warp.Timeout as T
#endif
import Network.Wai.Handler.Warp.Types
import Network.Wai.Internal (Response(..))

#ifdef WINDOWS
import qualified System.IO as IO
#else
import Network.Wai.Handler.Warp.FdCache
import Network.Wai.Handler.Warp.SendFile (positionRead)
import System.Posix.IO (openFd, OpenFileFlags(..), defaultFileFlags, OpenMode(ReadOnly), closeFd)
import System.Posix.Types
#endif

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

data Leftover = LZero
              | LOne B.BufferWriter
              | LTwo BS.ByteString B.BufferWriter

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

getStreamWindowSize :: Stream -> IO WindowSize
getStreamWindowSize Stream{streamWindow} = atomically $ readTVar streamWindow

waitStreamWindowSize :: Stream -> STM ()
waitStreamWindowSize Stream{streamWindow} = do
    w <- readTVar streamWindow
    check (w > 0)

waitStreaming :: TBQueue a -> STM ()
waitStreaming tbq = do
    isEmpty <- isEmptyTBQueue tbq
    check (isEmpty == False)

frameSender :: Context -> Connection -> InternalInfo -> S.Settings -> IO ()
frameSender ctx@Context{outputQ,controlQ,connectionWindow,encodeDynamicTable}
            conn@Connection{connWriteBuffer,connBufferSize,connSendAll}
            ii settings = loop `E.catch` ignore
  where
    dequeueControl = Left <$> readTQueue controlQ
    dequeueData    = Right <$> do
        w <- readTVar connectionWindow
        check (w > 0)
        dequeueSTM outputQ

    loop = do
        ex <- atomically (dequeueControl `orElse` dequeueData)
        case ex of
            Left  ctl         -> control ctl
            Right (_,pre,out) -> do
                writeIORef (streamPrecedence (outputStream out)) pre
                switch out

    control CFinish         = return ()
    control (CGoaway frame) = connSendAll frame
    control (CFrame frame)  = do
        connSendAll frame
        loop
    control (CSettings frame alist) = do
        connSendAll frame
        case lookup SettingsHeaderTableSize alist of
            Nothing  -> return ()
            Just siz -> do
                dyntbl <- readIORef encodeDynamicTable
                setLimitForEncoding siz dyntbl
        loop

    switch out@(ONext strm curr binfo) = do
        whenReadyOrEnqueueAgain strm binfo out $ \sws -> do
            cws <- atomically $ readTVar connectionWindow
            let !lim = min cws sws
            -- Data frame payload
            Next datPayloadLen mnext <- curr lim
            fillDataHeaderSend strm 0 datPayloadLen mnext
            maybeEnqueueNext strm mnext binfo
        loop
    switch out@(OResponse strm rsp binfo) = do
        whenReadyOrEnqueueAgain strm binfo out $ \sws -> do
            -- Header frame and Continuation frame
            let sid = streamNumber strm
                endOfStream = case binfo of
                    Persist{}          -> False
                    OneshotWithBody    -> False
                    OneshotWithoutBody -> True
            len <- headerContinue sid rsp endOfStream
            let total = len + frameHeaderLength
            case binfo of
                OneshotWithBody -> do
                    -- Data frame payload
                    (off, _) <- sendHeadersIfNecessary total
                    let payloadOff = off + frameHeaderLength
                    cws <- atomically $ readTVar connectionWindow
                    let !lim = min cws sws
                    Next datPayloadLen mnext <-
                        fillResponseBodyGetNext conn ii payloadOff lim rsp
                    fillDataHeaderSend strm total datPayloadLen mnext
                    maybeEnqueueNext strm mnext binfo
                OneshotWithoutBody -> do
                    -- "closed" must be before "connSendAll". If not,
                    -- the context would be switched to the receiver,
                    -- resulting the inconsistency of concurrency.
                    closed ctx strm Finished
                    flushN total
                Persist sq -> do
                    (off, needSend) <- sendHeadersIfNecessary total
                    let payloadOff = off + frameHeaderLength
                    cws <- atomically $ readTVar connectionWindow
                    let !lim = min cws sws
                    Next datPayloadLen mnext <-
                        fillStreamBodyGetNext conn payloadOff lim sq strm
                    -- If no data was immediately available, avoid sending an
                    -- empty data frame.
                    if datPayloadLen > 0 then
                        fillDataHeaderSend strm total datPayloadLen mnext
                      else
                        when needSend $ flushN off
                    maybeEnqueueNext strm mnext binfo
        loop

    whenReadyOrEnqueueAgain strm binfo out body = E.handle resetStream $ do
        state <- readIORef $ streamState strm
        unless (isClosed state) $ case binfo of
            Persist tbq -> checkStreaming tbq
            _           -> checkStreamWindowSize
      where
        checkStreaming tbq = do
            isEmpty <- atomically $ isEmptyTBQueue tbq
            if isEmpty then
                enqueueWhenReady (waitStreaming tbq) outputQ out
              else
                checkStreamWindowSize
        checkStreamWindowSize = do
            sws <- getStreamWindowSize strm
            if sws == 0 then do
                enqueueWhenReady (waitStreamWindowSize strm) outputQ out
              else
                body sws
        resetStream e = do
            closed ctx strm (ResetByMe e)
            let !rst = resetFrame InternalError $ streamNumber strm
            enqueueControl controlQ $ CFrame rst

    -- Flush the connection buffer to the socket, where the first 'n' bytes of
    -- the buffer are filled.
    flushN :: Int -> IO ()
    flushN n = bufferIO connWriteBuffer n connSendAll

    headerContinue sid rsp endOfStream = do
        builder <- hpackEncodeHeader ctx ii settings rsp
        (len, signal) <- B.runBuilder builder bufHeaderPayload headerPayloadLim
        let flag0 = case signal of
                B.Done -> setEndHeader defaultFlags
                _      -> defaultFlags
            flag = if endOfStream then setEndStream flag0 else flag0
        fillFrameHeader FrameHeaders len sid flag connWriteBuffer
        continue sid len signal

    continue _   len B.Done = return len<commit_msg>fixing a bug of waiter.<commit_after>{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE BangPatterns, CPP #-}
{-# LANGUAGE NamedFieldPuns #-}

module Network.Wai.Handler.Warp.HTTP2.Sender (frameSender) where

#if __GLASGOW_HASKELL__ < 709
import Control.Applicative
#endif
import Control.Concurrent (forkIO)
import Control.Concurrent.STM
import qualified Control.Exception as E
import Control.Monad (unless, void, when)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Builder.Extra as B
import Data.Maybe (isNothing)
import Foreign.Ptr
import Network.HPACK (setLimitForEncoding)
import Network.HTTP2
import Network.HTTP2.Priority (dequeueSTM)
import Network.Wai
import Network.Wai.Handler.Warp.Buffer
import Network.Wai.Handler.Warp.HTTP2.EncodeFrame
import Network.Wai.Handler.Warp.HTTP2.HPACK
import Network.Wai.Handler.Warp.HTTP2.Types
import Network.Wai.Handler.Warp.IORef
import qualified Network.Wai.Handler.Warp.Settings as S
#ifndef WINDOWS
import qualified Network.Wai.Handler.Warp.Timeout as T
#endif
import Network.Wai.Handler.Warp.Types
import Network.Wai.Internal (Response(..))

#ifdef WINDOWS
import qualified System.IO as IO
#else
import Network.Wai.Handler.Warp.FdCache
import Network.Wai.Handler.Warp.SendFile (positionRead)
import System.Posix.IO (openFd, OpenFileFlags(..), defaultFileFlags, OpenMode(ReadOnly), closeFd)
import System.Posix.Types
#endif

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

data Leftover = LZero
              | LOne B.BufferWriter
              | LTwo BS.ByteString B.BufferWriter

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

getStreamWindowSize :: Stream -> IO WindowSize
getStreamWindowSize Stream{streamWindow} = atomically $ readTVar streamWindow

waitStreamWindowSize :: Stream -> STM ()
waitStreamWindowSize Stream{streamWindow} = do
    w <- readTVar streamWindow
    check (w > 0)

waitStreaming :: TBQueue a -> STM ()
waitStreaming tbq = do
    isEmpty <- isEmptyTBQueue tbq
    check (isEmpty == False)

frameSender :: Context -> Connection -> InternalInfo -> S.Settings -> IO ()
frameSender ctx@Context{outputQ,controlQ,connectionWindow,encodeDynamicTable}
            conn@Connection{connWriteBuffer,connBufferSize,connSendAll}
            ii settings = loop `E.catch` ignore
  where
    dequeueControl = Left <$> readTQueue controlQ
    dequeueData    = Right <$> do
        w <- readTVar connectionWindow
        check (w > 0)
        dequeueSTM outputQ

    loop = do
        ex <- atomically (dequeueControl `orElse` dequeueData)
        case ex of
            Left  ctl         -> control ctl
            Right (_,pre,out) -> do
                writeIORef (streamPrecedence (outputStream out)) pre
                switch out

    control CFinish         = return ()
    control (CGoaway frame) = connSendAll frame
    control (CFrame frame)  = do
        connSendAll frame
        loop
    control (CSettings frame alist) = do
        connSendAll frame
        case lookup SettingsHeaderTableSize alist of
            Nothing  -> return ()
            Just siz -> do
                dyntbl <- readIORef encodeDynamicTable
                setLimitForEncoding siz dyntbl
        loop

    switch out@(ONext strm curr binfo) = do
        whenReadyOrEnqueueAgain strm binfo out $ \sws -> do
            cws <- atomically $ readTVar connectionWindow
            let !lim = min cws sws
            -- Data frame payload
            Next datPayloadLen mnext <- curr lim
            fillDataHeaderSend strm 0 datPayloadLen mnext
            maybeEnqueueNext strm mnext binfo
        loop
    switch out@(OResponse strm rsp binfo) = do
        whenReadyOrEnqueueAgain strm binfo out $ \sws -> do
            -- Header frame and Continuation frame
            let sid = streamNumber strm
                endOfStream = case binfo of
                    Persist{}          -> False
                    OneshotWithBody    -> False
                    OneshotWithoutBody -> True
            len <- headerContinue sid rsp endOfStream
            let total = len + frameHeaderLength
            case binfo of
                OneshotWithBody -> do
                    -- Data frame payload
                    (off, _) <- sendHeadersIfNecessary total
                    let payloadOff = off + frameHeaderLength
                    cws <- atomically $ readTVar connectionWindow
                    let !lim = min cws sws
                    Next datPayloadLen mnext <-
                        fillResponseBodyGetNext conn ii payloadOff lim rsp
                    fillDataHeaderSend strm total datPayloadLen mnext
                    maybeEnqueueNext strm mnext binfo
                OneshotWithoutBody -> do
                    -- "closed" must be before "connSendAll". If not,
                    -- the context would be switched to the receiver,
                    -- resulting the inconsistency of concurrency.
                    closed ctx strm Finished
                    flushN total
                Persist sq -> do
                    (off, needSend) <- sendHeadersIfNecessary total
                    let payloadOff = off + frameHeaderLength
                    cws <- atomically $ readTVar connectionWindow
                    let !lim = min cws sws
                    Next datPayloadLen mnext <-
                        fillStreamBodyGetNext conn payloadOff lim sq strm
                    -- If no data was immediately available, avoid sending an
                    -- empty data frame.
                    if datPayloadLen > 0 then
                        fillDataHeaderSend strm total datPayloadLen mnext
                      else
                        when needSend $ flushN off
                    maybeEnqueueNext strm mnext binfo
        loop

    whenReadyOrEnqueueAgain strm binfo out body = E.handle resetStream $ do
        state <- readIORef $ streamState strm
        unless (isClosed state) $ case binfo of
            Persist tbq -> checkStreaming tbq
            _           -> checkStreamWindowSize
      where
        checkStreaming tbq = do
            isEmpty <- atomically $ isEmptyTBQueue tbq
            if isEmpty then
                void . forkIO $ enqueueWhenReady (waitStreaming tbq) outputQ out
              else
                checkStreamWindowSize
        checkStreamWindowSize = do
            sws <- getStreamWindowSize strm
            if sws == 0 then do
                void . forkIO $ enqueueWhenReady (waitStreamWindowSize strm) outputQ out
              else
                body sws
        resetStream e = do
            closed ctx strm (ResetByMe e)
            let !rst = resetFrame InternalError $ streamNumber strm
            enqueueControl controlQ $ CFrame rst

    -- Flush the connection buffer to the socket, where the first 'n' bytes of
    -- the buffer are filled.
    flushN :: Int -> IO ()
    flushN n = bufferIO connWriteBuffer n connSendAll

    headerContinue sid rsp endOfStream = do
        builder <- hpackEncodeHeader ctx ii settings rsp
        (len, signal) <- B.runBuilder builder bufHeaderPayload headerPayloadLim
        let flag0 = case signal of
                B.Done -> setEndHeader defaultFlags
                _      -> defaultFlags
            flag = if endOfStream then setEndStream flag0 else flag0
        fillFrameHeader FrameHeaders len sid flag connWriteBuffer
        continue sid len signal

    continue _   len B.Done = return len
{'fraction_non_alphanumeric': 0.07434392498144775, 'fraction_numerical': 0.0022937327126762463, 'mean_word_length': 1.8496732026143792, 'pattern_counts': {'":': 0, '<': 47, '<?xml version=': 0, '>': 61, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5003240', 'n_tokens_mistral': 3904, 'n_tokens_neox': 3695, 'n_words': 1214}