module Network.Wai.Handler.Warp.IO where

import Control.Exception (mask_)
import Data.ByteString.Builder (Builder)
import Data.ByteString.Builder.Extra (Next (Chunk, Done, More), runBuilder)
import Data.IORef (IORef, readIORef, writeIORef)
import Network.Wai.Handler.Warp.Buffer
import Network.Wai.Handler.Warp.Imports
import Network.Wai.Handler.Warp.Types

toBufIOWith
    :: Int -> IORef WriteBuffer -> (ByteString -> IO ()) -> Builder -> IO Integer
toBufIOWith :: Int
-> IORef WriteBuffer
-> (ByteString -> IO ())
-> Builder
-> IO Integer
toBufIOWith Int
maxRspBufSize IORef WriteBuffer
writeBufferRef ByteString -> IO ()
io Builder
builder = do
    WriteBuffer
writeBuffer <- IORef WriteBuffer -> IO WriteBuffer
forall a. IORef a -> IO a
readIORef IORef WriteBuffer
writeBufferRef
    WriteBuffer -> BufferWriter -> Integer -> IO Integer
loop WriteBuffer
writeBuffer BufferWriter
firstWriter Integer
0
  where
    firstWriter :: BufferWriter
firstWriter = Builder -> BufferWriter
runBuilder Builder
builder
    loop :: WriteBuffer -> BufferWriter -> Integer -> IO Integer
loop WriteBuffer
writeBuffer BufferWriter
writer Integer
bytesSent = do
        let buf :: Buffer
buf = WriteBuffer -> Buffer
bufBuffer WriteBuffer
writeBuffer
            size :: Int
size = WriteBuffer -> Int
bufSize WriteBuffer
writeBuffer
        (Int
len, Next
signal) <- BufferWriter
writer Buffer
buf Int
size
        Buffer -> Int -> (ByteString -> IO ()) -> IO ()
bufferIO Buffer
buf Int
len ByteString -> IO ()
io
        let totalBytesSent :: Integer
totalBytesSent = Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
len Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
bytesSent
        case Next
signal of
            Next
Done -> Integer -> IO Integer
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Integer
totalBytesSent
            More Int
minSize BufferWriter
next
                | Int
size Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
minSize -> do
                    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
minSize Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
maxRspBufSize) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
                        [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error ([Char] -> IO ()) -> [Char] -> IO ()
forall a b. (a -> b) -> a -> b
$
                            [Char]
"Sending a Builder response required a buffer of size "
                                [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
minSize
                                [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" which is bigger than the specified maximum of "
                                [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
maxRspBufSize
                                [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"!"
                    -- The current WriteBuffer is too small to fit the next
                    -- batch of bytes from the Builder so we free it and
                    -- create a new bigger one. Freeing the current buffer,
                    -- creating a new one and writing it to the IORef need
                    -- to be performed atomically to prevent both double
                    -- frees and missed frees. So we mask async exceptions:
                    WriteBuffer
biggerWriteBuffer <- IO WriteBuffer -> IO WriteBuffer
forall a. IO a -> IO a
mask_ (IO WriteBuffer -> IO WriteBuffer)
-> IO WriteBuffer -> IO WriteBuffer
forall a b. (a -> b) -> a -> b
$ do
                        WriteBuffer -> IO ()
bufFree WriteBuffer
writeBuffer
                        WriteBuffer
biggerWriteBuffer <- Int -> IO WriteBuffer
createWriteBuffer Int
minSize
                        IORef WriteBuffer -> WriteBuffer -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef WriteBuffer
writeBufferRef WriteBuffer
biggerWriteBuffer
                        WriteBuffer -> IO WriteBuffer
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return WriteBuffer
biggerWriteBuffer
                    WriteBuffer -> BufferWriter -> Integer -> IO Integer
loop WriteBuffer
biggerWriteBuffer BufferWriter
next Integer
totalBytesSent
                | Bool
otherwise -> WriteBuffer -> BufferWriter -> Integer -> IO Integer
loop WriteBuffer
writeBuffer BufferWriter
next Integer
totalBytesSent
            Chunk ByteString
bs BufferWriter
next -> do
                ByteString -> IO ()
io ByteString
bs
                WriteBuffer -> BufferWriter -> Integer -> IO Integer
loop WriteBuffer
writeBuffer BufferWriter
next Integer
totalBytesSent