git-annex/servant.hs

{-# LANGUAGE BangPatterns  #-}
{-# LANGUAGE DataKinds     #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE LambdaCase    #-}
module Main (main) where

import           Control.Concurrent
                 (threadDelay)
import           Control.Monad.IO.Class
                 (MonadIO (..))
import qualified Data.ByteString              as BS
import qualified Data.ByteString.Lazy         as BL
import qualified Data.ByteString.Lazy.Internal as BLI
import           Data.Maybe
                 (fromMaybe)
import           Network.HTTP.Client
                 (defaultManagerSettings, newManager)
import           Network.Wai
                 (Application)
import           System.Environment
                 (getArgs, lookupEnv)
import           Text.Read
                 (readMaybe)

import           Servant
import           Servant.Client.Streaming
import qualified Servant.Types.SourceT as S
import Control.Concurrent.MVar
import System.IO.Unsafe

import qualified Network.Wai.Handler.Warp     as Warp

type API = "readme" :> StreamGet NoFraming OctetStream (SourceIO BS.ByteString)
	:<|> "writeme" :> StreamBody NoFraming OctetStream (SourceIO BS.ByteString) :> Post '[JSON] Bool

api :: Proxy API
api = Proxy

server :: Server API
server = readme :<|> writeme where
    readme = liftIO $ do
        putStrLn "/proxy"
        return $ S.SourceT $ \k -> do
		k =<< readfilelazy "README.md"
		k =<< readfilelazy "another"

writeme :: SourceIO BS.ByteString -> Handler Bool
writeme src = do
	liftIO $ print "gathering lazy bytestring"
	b <- liftIO $ S.unSourceT src gatherbytestring
	liftIO $ print "got lazy bytestring, writing to file"
	liftIO$ BL.writeFile "writem" b
	liftIO$ print "write complete"
	return True

app :: Application
app = serve api server

cli :: ClientM (S.SourceT IO BS.ByteString)
cli :<|> writecli = client api

main :: IO ()
main = do
    args <- getArgs
    case args of
        ("server":_) -> do
            putStrLn "Starting cookbook-basic-streaming at http://localhost:8000"
            port <- fromMaybe 8000 . (>>= readMaybe) <$> lookupEnv "PORT"
            Warp.run port app
        ("client":ns:_) -> do
            mgr <- newManager defaultManagerSettings
            burl <- parseBaseUrl "http://localhost:8000/"
            withClientM (writecli getit) (mkClientEnv mgr burl) $ \me -> case me of
                Left err  -> print err
                Right src -> print src
        _ -> do
            putStrLn "Try:"
            putStrLn "cabal new-run cookbook-basic-streaming server"
            putStrLn "cabal new-run cookbook-basic-streaming client 10"
            putStrLn "time curl -H 'Accept: application/json' localhost:8000/slow/5"
  where
	getit = S.SourceT $ \k -> do
                k =<< readfilelazy "/home/joey/README.md"

readfilelazy :: FilePath -> IO (S.StepT IO BS.ByteString)
readfilelazy file = do
	bl <- BL.readFile file
	v <- newMVar (BL.toChunks bl)
 	return (go v)
  where
    go :: MVar [BS.ByteString] -> S.StepT IO BS.ByteString
    go v = S.fromActionStep BS.null $ do
	print "chunk"
	modifyMVar v $ pure . \case
		[] -> ([], BS.empty)
		(b:bs) -> (bs, b)

gatherbytestring :: S.StepT IO BS.ByteString -> IO BL.ByteString
gatherbytestring x = do
	l <- unsafeInterleaveIO $ go x
	return l
  where
	go S.Stop        = return BLI.Empty
	go (S.Error err) = error $ show ("ERROR", err)
	go (S.Skip s)    = do
		go s
	go (S.Effect ms) = do
		ms >>= go
	go (S.Yield v s) = do
		BLI.Chunk v <$> unsafeInterleaveIO (go s)