From 37482d69d79172ac3770d9bffb66f893235eed5e Mon Sep 17 00:00:00 2001 From: "Julian K. Arni" Date: Tue, 27 Feb 2018 15:31:41 +0100 Subject: [PATCH] Test that Stream combinator doesn't blow up memory. --- servant-client/servant-client.cabal | 2 +- .../src/Servant/Client/Internal/HttpClient.hs | 12 ++- servant-client/test/Servant/ClientSpec.hs | 3 + servant-client/test/Servant/StreamSpec.hs | 89 ++++++++++++------- 4 files changed, 70 insertions(+), 36 deletions(-) diff --git a/servant-client/servant-client.cabal b/servant-client/servant-client.cabal index 80da8678..71fe22c4 100644 --- a/servant-client/servant-client.cabal +++ b/servant-client/servant-client.cabal @@ -82,7 +82,7 @@ library test-suite spec type: exitcode-stdio-1.0 - ghc-options: -Wall + ghc-options: -Wall -rtsopts -with-rtsopts=-T default-language: Haskell2010 hs-source-dirs: test main-is: Spec.hs diff --git a/servant-client/src/Servant/Client/Internal/HttpClient.hs b/servant-client/src/Servant/Client/Internal/HttpClient.hs index f976deed..eddb0afc 100644 --- a/servant-client/src/Servant/Client/Internal/HttpClient.hs +++ b/servant-client/src/Servant/Client/Internal/HttpClient.hs @@ -183,10 +183,18 @@ requestToClientRequest burl r = Client.defaultRequest where hs = toList $ requestAccept r + convertBody bd = case bd of + RequestBodyLBS body' -> Client.RequestBodyLBS body' + RequestBodyBS body' -> Client.RequestBodyBS body' + RequestBodyBuilder size body' -> Client.RequestBodyBuilder size body' + RequestBodyStream size body' -> Client.RequestBodyStream size body' + RequestBodyStreamChunked body' -> Client.RequestBodyStreamChunked body' + RequestBodyIO body' -> Client.RequestBodyIO (convertBody <$> body') + (body, contentTypeHdr) = case requestBody r of Nothing -> (Client.RequestBodyLBS "", Nothing) - Just (RequestBodyLBS body', typ) - -> (Client.RequestBodyLBS body', Just (hContentType, renderHeader typ)) + Just (body', typ) + -> (convertBody body', Just (hContentType, renderHeader typ)) isSecure = case baseUrlScheme burl of Http -> False diff --git a/servant-client/test/Servant/ClientSpec.hs b/servant-client/test/Servant/ClientSpec.hs index c9a96cab..6d33cd27 100644 --- a/servant-client/test/Servant/ClientSpec.hs +++ b/servant-client/test/Servant/ClientSpec.hs @@ -103,6 +103,9 @@ instance FromJSON Person instance ToForm Person instance FromForm Person +instance Arbitrary Person where + arbitrary = Person <$> arbitrary <*> arbitrary + alice :: Person alice = Person "Alice" 42 diff --git a/servant-client/test/Servant/StreamSpec.hs b/servant-client/test/Servant/StreamSpec.hs index 2df336da..ad4a2664 100644 --- a/servant-client/test/Servant/StreamSpec.hs +++ b/servant-client/test/Servant/StreamSpec.hs @@ -26,25 +26,26 @@ #include "overlapping-compat.h" module Servant.StreamSpec (spec) where -import Prelude () -import Prelude.Compat +import Control.Monad (replicateM_, void) +import qualified Data.ByteString as BS import Data.Proxy -import qualified Network.HTTP.Client as C -import System.IO.Unsafe (unsafePerformIO) +import GHC.Stats (currentBytesUsed, getGCStats) +import qualified Network.HTTP.Client as C +import Prelude () +import Prelude.Compat +import System.IO (IOMode (ReadMode), withFile) +import System.IO.Unsafe (unsafePerformIO) import Test.Hspec +import Test.QuickCheck -import Servant.API ((:<|>) ((:<|>)), - (:>), - EmptyAPI, JSON, - StreamGet, - NewlineFraming, - NetstringFraming, - ResultStream(..), - StreamGenerator(..)) +import Servant.API ((:<|>) ((:<|>)), (:>), JSON, + NetstringFraming, NewlineFraming, + OctetStream, ResultStream (..), + StreamGenerator (..), StreamGet) import Servant.Client +import Servant.ClientSpec (Person (..)) +import qualified Servant.ClientSpec as CS import Servant.Server -import qualified Servant.ClientSpec as CS -import Servant.ClientSpec (Person(..)) spec :: Spec @@ -54,7 +55,7 @@ spec = describe "Servant.Stream" $ do type StreamApi f = "streamGetNewline" :> StreamGet NewlineFraming JSON (f Person) :<|> "streamGetNetstring" :> StreamGet NetstringFraming JSON (f Person) - :<|> EmptyAPI + :<|> "streamALot" :> StreamGet NewlineFraming OctetStream (f BS.ByteString) capi :: Proxy (StreamApi ResultStream) @@ -63,12 +64,9 @@ capi = Proxy sapi :: Proxy (StreamApi StreamGenerator) sapi = Proxy - -getGetNL :<|> getGetNS :<|> EmptyClient = client capi - - -getGetNL :: ClientM (ResultStream Person) -getGetNS :: ClientM (ResultStream Person) +getGetNL, getGetNS :: ClientM (ResultStream Person) +getGetALot :: ClientM (ResultStream BS.ByteString) +getGetNL :<|> getGetNS :<|> getGetALot = client capi alice :: Person alice = Person "Alice" 42 @@ -77,14 +75,24 @@ bob :: Person bob = Person "Bob" 25 server :: Application -server = serve sapi ( - (return (StreamGenerator (\f r -> f alice >> r bob >> r alice)) - :: Handler (StreamGenerator Person)) - :<|> - (return (StreamGenerator (\f r -> f alice >> r bob >> r alice)) - :: Handler (StreamGenerator Person)) - :<|> - emptyServer) +server = serve sapi + $ return (StreamGenerator (\f r -> f alice >> r bob >> r alice)) + :<|> return (StreamGenerator (\f r -> f alice >> r bob >> r alice)) + :<|> return (StreamGenerator lotsGenerator) + where + lotsGenerator f r = do + f "" + withFile "/dev/urandom" ReadMode $ + \handle -> streamFiveMBNTimes handle 1000 r + return () + + streamFiveMBNTimes handle left sink + | left <= 0 = return "" + | otherwise = do + msg <- BS.hGet handle (megabytes 5) + sink msg + streamFiveMBNTimes handle (left - 1) sink + {-# NOINLINE manager' #-} @@ -94,20 +102,35 @@ manager' = unsafePerformIO $ C.newManager C.defaultManagerSettings runClient :: ClientM a -> BaseUrl -> IO (Either ServantError a) runClient x baseUrl' = runClientM x (mkClientEnv manager' baseUrl') -runResultStream :: ResultStream a -> IO (Maybe (Either String a), Maybe (Either String a), Maybe (Either String a), Maybe (Either String a)) -runResultStream (ResultStream k) = k $ \act -> (,,,) <$> act <*> act <*> act <*> act +runResultStream :: ResultStream a + -> IO ( Maybe (Either String a) + , Maybe (Either String a) + , Maybe (Either String a) + , Maybe (Either String a)) +runResultStream (ResultStream k) + = k $ \act -> (,,,) <$> act <*> act <*> act <*> act streamSpec :: Spec streamSpec = beforeAll (CS.startWaiApp server) $ afterAll CS.endWaiApp $ do - it "Servant.API.StreamGet.Newline" $ \(_, baseUrl) -> do + it "works with Servant.API.StreamGet.Newline" $ \(_, baseUrl) -> do Right res <- runClient getGetNL baseUrl let jra = Just (Right alice) jrb = Just (Right bob) runResultStream res `shouldReturn` (jra, jrb, jra, Nothing) - it "Servant.API.StreamGet.Netstring" $ \(_, baseUrl) -> do + it "works with Servant.API.StreamGet.Netstring" $ \(_, baseUrl) -> do Right res <- runClient getGetNS baseUrl let jra = Just (Right alice) jrb = Just (Right bob) runResultStream res `shouldReturn` (jra, jrb, jra, Nothing) + + it "streams in constant memory" $ \(_, baseUrl) -> do + Right (ResultStream res) <- runClient getGetALot baseUrl + let consumeNChunks n = replicateM_ n (res void) + consumeNChunks 900 + memUsed <- currentBytesUsed <$> getGCStats + memUsed `shouldSatisfy` (< (megabytes 20)) + +megabytes :: Num a => a -> a +megabytes n = n * (1000 ^ 2)