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 <- 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 = forall a. Integral a => a -> Integer
toInteger Int
len forall a. Num a => a -> a -> a
+ Integer
bytesSent
      case Next
signal of
        Next
Done -> forall (m :: * -> *) a. Monad m => a -> m a
return Integer
totalBytesSent
        More Int
minSize BufferWriter
next
          | Int
size forall a. Ord a => a -> a -> Bool
< Int
minSize -> do
              forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
minSize forall a. Ord a => a -> a -> Bool
> Int
maxRspBufSize) forall a b. (a -> b) -> a -> b
$
                forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Sending a Builder response required a buffer of size "
                          forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Int
minSize forall a. [a] -> [a] -> [a]
++ [Char]
" which is bigger than the specified maximum of "
                          forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Int
maxRspBufSize 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 <- forall a. IO a -> IO a
mask_ forall a b. (a -> b) -> a -> b
$ do
                WriteBuffer -> IO ()
bufFree WriteBuffer
writeBuffer
                WriteBuffer
biggerWriteBuffer <- Int -> IO WriteBuffer
createWriteBuffer Int
minSize
                forall a. IORef a -> a -> IO ()
writeIORef IORef WriteBuffer
writeBufferRef WriteBuffer
biggerWriteBuffer
                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