servant/servant-client/test/Servant/StreamSpec.hs

{-# LANGUAGE CPP                    #-}
{-# LANGUAGE ConstraintKinds        #-}
{-# LANGUAGE DataKinds              #-}
{-# LANGUAGE DeriveGeneric          #-}
{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE FlexibleInstances      #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs                  #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE OverloadedStrings      #-}
{-# LANGUAGE PolyKinds              #-}
{-# LANGUAGE RecordWildCards        #-}
{-# LANGUAGE ScopedTypeVariables    #-}
{-# LANGUAGE StandaloneDeriving     #-}
{-# LANGUAGE TypeFamilies           #-}
{-# LANGUAGE TypeOperators          #-}
{-# LANGUAGE UndecidableInstances   #-}
{-# OPTIONS_GHC -freduction-depth=100 #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
{-# OPTIONS_GHC -fno-warn-name-shadowing #-}

module Servant.StreamSpec (spec) where

import           Control.Monad
                 (when)
import           Control.Monad.Codensity
                 (Codensity (..))
import           Control.Monad.IO.Class
                 (MonadIO (..))
import           Control.Monad.Trans.Except
import qualified Data.ByteString               as BS
import           Data.Proxy
import qualified Data.TDigest                  as TD
import qualified Network.HTTP.Client           as C
import           Prelude ()
import           Prelude.Compat
import           Servant.API
                 ((:<|>) ((:<|>)), (:>), JSON, NetstringFraming, StreamBody,
                 NewlineFraming, NoFraming, OctetStream, SourceIO, StreamGet,
                 )
import           Servant.Client.Streaming
import           Servant.ClientSpec
                 (Person (..))
import qualified Servant.ClientSpec            as CS
import           Servant.Server
import           Servant.Test.ComprehensiveAPI
import           Servant.Types.SourceT
import           System.Entropy
                 (getEntropy, getHardwareEntropy)
import           System.IO.Unsafe
                 (unsafePerformIO)
import           System.Mem
                 (performGC)
import           Test.Hspec

#if MIN_VERSION_base(4,10,0)
import           GHC.Stats
                 (gc, gcdetails_live_bytes, getRTSStats)
#else
import           GHC.Stats
                 (currentBytesUsed, getGCStats)
#endif

-- This declaration simply checks that all instances are in place.
-- Note: this is streaming client
_ = client comprehensiveAPI

spec :: Spec
spec = describe "Servant.Client.Streaming" $ do
    streamSpec

type StreamApi =
         "streamGetNewline" :> StreamGet NewlineFraming JSON (SourceIO Person)
    :<|> "streamGetNetstring" :> StreamGet NetstringFraming JSON (SourceIO Person)
    :<|> "streamALot" :> StreamGet NoFraming OctetStream (SourceIO BS.ByteString)
    :<|> "streamBody" :> StreamBody NoFraming OctetStream (SourceIO BS.ByteString) :> StreamGet NoFraming OctetStream (SourceIO BS.ByteString)

api :: Proxy StreamApi
api = Proxy

getGetNL, getGetNS :: ClientM (SourceIO Person)
getGetALot :: ClientM (SourceIO BS.ByteString)
getStreamBody :: SourceT IO BS.ByteString -> ClientM (SourceIO BS.ByteString)
getGetNL :<|> getGetNS :<|> getGetALot :<|> getStreamBody = client api

alice :: Person
alice = Person "Alice" 42

bob :: Person
bob = Person "Bob" 25

server :: Application
server = serve api
    $    return (source [alice, bob, alice])
    :<|> return (source [alice, bob, alice])
    -- 2 ^ (18 + 10) = 256M
    :<|> return (SourceT ($ lots (powerOfTwo 18)))
    :<|> return
  where
    lots n
        | n < 0     = Stop
        | otherwise = Effect $ do
            let size = powerOfTwo 10
            mbs <- getHardwareEntropy size
            bs <- maybe (getEntropy size) pure mbs
            return (Yield bs (lots (n - 1)))

powerOfTwo :: Int -> Int
powerOfTwo = (2 ^)

{-# NOINLINE manager' #-}
manager' :: C.Manager
manager' = unsafePerformIO $ C.newManager C.defaultManagerSettings

withClient :: ClientM a -> BaseUrl -> (Either ServantError a -> IO r) -> IO r
withClient x baseUrl' = withClientM x (mkClientEnv manager' baseUrl')

testRunSourceIO :: SourceIO a
    -> IO (Either String [a])
testRunSourceIO = runExceptT . runSourceT

streamSpec :: Spec
streamSpec = beforeAll (CS.startWaiApp server) $ afterAll CS.endWaiApp $ do
    it "works with Servant.API.StreamGet.Newline" $ \(_, baseUrl) -> do
        withClient getGetNL baseUrl $ \(Right res) ->
            testRunSourceIO res `shouldReturn` Right [alice, bob, alice]

    it "works with Servant.API.StreamGet.Netstring" $ \(_, baseUrl) -> do
        withClient getGetNS baseUrl $ \(Right res) ->
            testRunSourceIO res `shouldReturn` Right [alice, bob, alice]

    it "works with Servant.API.StreamBody" $ \(_, baseUrl) -> do
        withClient (getStreamBody (source input)) baseUrl $ \(Right res) ->
            testRunSourceIO res `shouldReturn` Right output
        where
          input = ["foo", "", "bar"]
          output = ["foo", "bar"]

{-
    it "streams in constant memory" $ \(_, baseUrl) -> do
        Right rs <- runClient getGetALot baseUrl
        performGC
        -- usage0 <- getUsage
        -- putStrLn $ "Start:  " ++ show usage0
        tdigest <- memoryUsage $ joinCodensitySourceT rs

        -- putStrLn $ "Median: " ++ show (TD.median tdigest)
        -- putStrLn $ "Mean:   " ++ show (TD.mean tdigest)
        -- putStrLn $ "Stddev: " ++ show (TD.stddev tdigest)

        -- forM_ [0.01, 0.1, 0.2, 0.5, 0.8, 0.9, 0.99] $ \q ->
        --    putStrLn $ "q" ++ show q ++ ": " ++ show (TD.quantile q tdigest)

        let Just stddev = TD.stddev tdigest

        -- standard deviation of 100k is ok, we generate 256M of data after all.
        -- On my machine deviation is 40k-50k
        stddev `shouldSatisfy` (< 100000)

memoryUsage :: SourceT IO BS.ByteString -> IO (TD.TDigest 25)
memoryUsage src = unSourceT src $ loop mempty (0 :: Int)
  where
    loop !acc !_ Stop          = return acc
    loop !_   !_ (Error err)   = fail err -- !
    loop !acc !n (Skip s)      = loop acc n s
    loop !acc !n (Effect ms)   = ms >>= loop acc n
    loop !acc !n (Yield _bs s) =  do
        usage  <- liftIO getUsage
        -- We perform GC in between as we generate garbage.
        when (n `mod` 1024 == 0) $ liftIO performGC
        loop (TD.insert usage acc) (n + 1) s

getUsage :: IO Double
getUsage = fromIntegral .
#if MIN_VERSION_base(4,10,0)
    gcdetails_live_bytes . gc <$> getRTSStats
#else
    currentBytesUsed <$> getGCStats
#endif
       memUsed `shouldSatisfy` (< megabytes 22)

megabytes :: Num a => a -> a
megabytes n = n * (1000 ^ (2 :: Int))
-}