@@ -26,25 +26,26 @@
2626#include "overlapping-compat.h"
2727module Servant.StreamSpec (spec) where
2828
29−import Prelude ()
30−import Prelude.Compat
29+import Control.Monad (replicateM_, void)
30+import qualified Data.ByteString as BS
3131import Data.Proxy
32−import qualified Network.HTTP.Client as C
33−import System.IO.Unsafe (unsafePerformIO)
32+import GHC.Stats (currentBytesUsed, getGCStats)
33+import qualified Network.HTTP.Client as C
34+import Prelude ()
35+import Prelude.Compat
36+import System.IO (IOMode (ReadMode), withFile)
37+import System.IO.Unsafe (unsafePerformIO)
3438import Test.Hspec
39+import Test.QuickCheck
3540
36−import Servant.API ((:<|>) ((:<|>)),
37− (:>),
38− EmptyAPI, JSON,
39− StreamGet,
40− NewlineFraming,
41− NetstringFraming,
42− ResultStream(..),
43− StreamGenerator(..))
41+import Servant.API ((:<|>) ((:<|>)), (:>), JSON,
42+ NetstringFraming, NewlineFraming,
43+ OctetStream, ResultStream (..),
44+ StreamGenerator (..), StreamGet)
4445import Servant.Client
46+import Servant.ClientSpec (Person (..))
47+import qualified Servant.ClientSpec as CS
4548import Servant.Server
46−import qualified Servant.ClientSpec as CS
47−import Servant.ClientSpec (Person(..))
4849
4950
5051spec :: Spec
@@ -54,7 +55,7 @@ spec = describe "Servant.Stream" $ do
5455type StreamApi f =
5556 "streamGetNewline" :> StreamGet NewlineFraming JSON (f Person)
5657 :<|> "streamGetNetstring" :> StreamGet NetstringFraming JSON (f Person)
57− :<|> EmptyAPI
58+ :<|> "streamALot" :> StreamGet NewlineFraming OctetStream (f BS.ByteString)
5859
5960
6061capi :: Proxy (StreamApi ResultStream)
@@ -63,12 +64,9 @@ capi = Proxy
6364sapi :: Proxy (StreamApi StreamGenerator)
6465sapi = Proxy
6566
66−
67−getGetNL :<|> getGetNS :<|> EmptyClient = client capi
68−
69−
70−getGetNL :: ClientM (ResultStream Person)
71−getGetNS :: ClientM (ResultStream Person)
67+getGetNL, getGetNS :: ClientM (ResultStream Person)
68+getGetALot :: ClientM (ResultStream BS.ByteString)
69+getGetNL :<|> getGetNS :<|> getGetALot = client capi
7270
7371alice :: Person
7472alice = Person "Alice" 42
@@ -77,14 +75,24 @@ bob :: Person
7775bob = Person "Bob" 25
7876
7977server :: Application
80−server = serve sapi (
81− (return (StreamGenerator (\f r -> f alice >> r bob >> r alice))
82− :: Handler (StreamGenerator Person))
83− :<|>
84− (return (StreamGenerator (\f r -> f alice >> r bob >> r alice))
85− :: Handler (StreamGenerator Person))
86− :<|>
87− emptyServer)
78+server = serve sapi
79+ $ return (StreamGenerator (\f r -> f alice >> r bob >> r alice))
80+ :<|> return (StreamGenerator (\f r -> f alice >> r bob >> r alice))
81+ :<|> return (StreamGenerator lotsGenerator)
82+ where
83+ lotsGenerator f r = do
84+ f ""
85+ withFile "/dev/urandom" ReadMode $
86+ \handle -> streamFiveMBNTimes handle 1000 r
87+ return ()
88+
89+ streamFiveMBNTimes handle left sink
90+ | left <= 0 = return ""
91+ | otherwise = do
92+ msg <- BS.hGet handle (megabytes 5)
93+ sink msg
94+ streamFiveMBNTimes handle (left - 1) sink
95+
8896
8997
9098{-# NOINLINE manager' #-}
@@ -94,20 +102,35 @@ manager' = unsafePerformIO $ C.newManager C.defaultManagerSettings
94102runClient :: ClientM a -> BaseUrl -> IO (Either ServantError a)
95103runClient x baseUrl' = runClientM x (mkClientEnv manager' baseUrl')
96104
97−runResultStream :: ResultStream a -> IO (Maybe (Either String a), Maybe (Either String a), Maybe (Either String a), Maybe (Either String a))
98−runResultStream (ResultStream k) = k $ \act -> (,,,) <$> act <*> act <*> act <*> act
105+runResultStream :: ResultStream a
106+ -> IO ( Maybe (Either String a)
107+ , Maybe (Either String a)
108+ , Maybe (Either String a)
109+ , Maybe (Either String a))
110+runResultStream (ResultStream k)
111+ = k $ \act -> (,,,) <$> act <*> act <*> act <*> act
99112
100113streamSpec :: Spec
101114streamSpec = beforeAll (CS.startWaiApp server) $ afterAll CS.endWaiApp $ do
102115
103− it "Servant.API.StreamGet.Newline" $ \(_, baseUrl) -> do
116+ it "works with Servant.API.StreamGet.Newline" $ \(_, baseUrl) -> do
104117 Right res <- runClient getGetNL baseUrl
105118 let jra = Just (Right alice)
106119 jrb = Just (Right bob)
107120 runResultStream res `shouldReturn` (jra, jrb, jra, Nothing)
108121
109− it "Servant.API.StreamGet.Netstring" $ \(_, baseUrl) -> do
122+ it "works with Servant.API.StreamGet.Netstring" $ \(_, baseUrl) -> do
110123 Right res <- runClient getGetNS baseUrl
111124 let jra = Just (Right alice)
112125 jrb = Just (Right bob)
113126 runResultStream res `shouldReturn` (jra, jrb, jra, Nothing)
127+
128+ it "streams in constant memory" $ \(_, baseUrl) -> do
129+ Right (ResultStream res) <- runClient getGetALot baseUrl
130+ let consumeNChunks n = replicateM_ n (res void)
131+ consumeNChunks 900
132+ memUsed <- currentBytesUsed <$> getGCStats
133+ memUsed `shouldSatisfy` (< (megabytes 20))
134+
135+megabytes :: Num a => a -> a
136+megabytes n = n * (1000 ^ 2)