playing around with bytestring lazy streaming in servant

Will be needed for git-annex http p2p protocol.

diff --git a/servant.hs b/servant.hs
new file mode 100644 (file)
index 0000000..9cf0dd0
--- /dev/null
+++ b/servant.hs
@@ -0,0 +1,102 @@
+{-# 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)
+
+api :: Proxy API
+api = Proxy
+
+server :: Server API
+server = readme where
+    readme = liftIO $ do
+        putStrLn "/proxy"
+        return $ S.SourceT $ \k -> do
+               k =<< readfilelazy "README.md"
+               k =<< readfilelazy "another"
+
+app :: Application
+app = serve api server
+
+cli :: ClientM (S.SourceT IO BS.ByteString)
+cli = 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 cli (mkClientEnv mgr burl) $ \me -> case me of
+                Left err  -> print err
+                Right src -> do
+                       b <- S.unSourceT src gatherbytestring
+                       liftIO $ print "got it all, writing"
+                       BL.writeFile "got" (BL.init b)
+        _ -> 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"
+
+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)
+