Test that Stream combinator doesn't blow up memory.

Author: jkarni

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

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

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
 

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)