Middleware to record request and response sizes

[MaxGabriel]

I am interested in having WAI middleware that records the request body size and the response body size (in my case to prometheus). I don't think this exists today.

I haven't tested it, but I got this compiling which I think is roughly that. Would WAI be interested in a PR with something similar to this?

-- Largely modelled off of requestSizeCheck from wai-extra
recordRequestSize :: (Word64 -> IO ()) -> Wai.Request -> IO Wai.Request
recordRequestSize recordLength req =
  case Wai.requestBodyLength req of
    Wai.KnownLength len -> do
      recordLength len
      pure req
    Wai.ChunkedBody -> do
      currentSizeRef <- newIORef (0, False)
      let yieldNextBodyChunk :: IO ByteString
          yieldNextBodyChunk = do
            bs <- Wai.getRequestBodyChunk req
            if BS.length bs == 0 -- Out of chunks
              then do
                (currentSize, hasReported :: Bool) <- readIORef currentSizeRef
                -- (Just in case someone repeatedly calls for the next chunk when there is none)
                unless hasReported do
                  atomicModifyIORef' currentSizeRef (\(prevSize, _) -> ((prevSize, True), ()))
                  recordLength currentSize
              else do
                atomicModifyIORef' currentSizeRef $ \(prevSize, reported) ->
                  let nextSize = prevSize + fromIntegral (BS.length bs)
                   in ((nextSize, reported), ())
                pure ()

            pure bs
      -- Each time a caller downstream gets the next chunk, it will call yieldNextBodyChunk
      -- Which gives + increments the counter on the request size, until it's out of chunks
      pure $ Wai.setRequestBodyChunks yieldNextBodyChunk req

recordResponseBody :: IORef Int64 -> Wai.Response -> IO Wai.Response
recordResponseBody outgoingRef response = case response of
  Wai.ResponseStream status headers streamingBody -> do
    let wrappedBody :: (Builder.Builder -> IO ()) -> IO () -> IO ()
        wrappedBody originalSend flush =
          let modifiedSend builder = do
                let chunkSize = LBS.length (Builder.toLazyByteString builder)
                atomicModifyIORef' outgoingRef (\x -> (x + chunkSize, ()))
                originalSend builder
           in streamingBody modifiedSend flush
    pure $ Wai.ResponseStream status headers wrappedBody
  Wai.ResponseBuilder _ _ builder -> do
    let total = LBS.length (Builder.toLazyByteString builder)
    atomicWriteIORef outgoingRef total
    pure response
  Wai.ResponseRaw _ _ -> pure response -- Not sure if it's doable to measure raw responses
  Wai.ResponseFile _ _ filePath _ -> do
    size <- getFileSize filePath
    atomicWriteIORef outgoingRef (fromIntegral size)
    pure response

[Vlix]

It seems like a reasonable thing to want (to do), but I'm not sure this is something for wai.
It sounds like it would fit better in wai-extra as actual middleware, or just as a feature of warp. 🤔

Other idea:
Would it be a good idea in the Request case, to add an IORef to ChunkedBody? And that it will contain the amount of bytes read so far? This is a breaking change, but seems like a worthwile idea. Maybe with the Bool in the example that will be True when the entire body has been read?
We might even be able to mitigate breakage a bit with a pattern?

data RequestBodyLength
    = -- The 'Word64' contains the amount of bytes read up until that moment.
      -- The 'Bool' is 'True' after the entire body has been consumed.
      DelayedLength (IORef (Word64, Bool))
    | KnownLength Word64
  deriving (Show)

pattern ChunkedBody <- DelayedLength _

---

As for the Response part, we might be able to add an IORef to all the constructors, because we don't export them from wai. Only issue is that the functions creating the Response constructors are not in IO, so we'd have to create IORefs with unsafePerformIO (though it's not impossible)
One of the things I see is that recordResponseBody, for example, reads the builder into a lazy ByteString, when it might already have been one if the user used responseLBS.

And might it be possible for a file to have a different size on the system than what will actually be sent? Because then, we're not giving the actual size of the response for the ResponseFile case.
And if someone records the file size (or builder size) and after that, the response is turned into a streaming response that stops half-way through?

I can't think of a good alternative at the moment, though :/

[MaxGabriel]

Agreed I would put this in WAI-extra

[MaxGabriel]

Re: the IORef for request length: one benefit I like of the original solution is that the length can be recorded as soon as it is known. That’s nice because if I wait until eg the end of the HTTP request, that might be too late (eg I had a huge request body and I OOMed because of it, so I never got to the recording) (speculating a bit here)

Also seems just unnecessary to break for this?

[MaxGabriel]

Good call on file part handling, should be this I think:

  Wai.ResponseFile _ _ filePath mFilePart -> do
    size <- case mFilePart of
      Just part -> pure $ filePartByteCount part
      Nothing -> getFileSize filePath

[Vlix]

Fair, breakage is probably unnecessary, yeah.

Looking back at the Response example: anything working with a Response recurses over the default Response in the ResponseRaw constructor, so you might want to at least get the length of that if it's used? 🤷

I'm also wondering when you'll be reading the IORef of the Response example?

---

I've looked at the internals of warp and I think a better location might be sendResponse.
If we add your example's way of counting the ResponseStream, we'll have all (except ResponseRaw) content lengths of responses and the settingsLogger will always be called (except for ResponseRaw) so you'll be able to put your size record function in there 🤔

[parsonsmatt]

This is a bit of an aside, but is there any documentation or explanation of what ResponseRaw is? It's not constructed in this repository, yesod, or servant, and one of the few mentions of ResponseRaw here is that it's totally unsupported in the HTTP2 side of things

[Vlix]

@parsonsmatt It's one of the ways websockets get handled