進捗を取りながらファイルをダウンロードする

＿人人人人人人人人人人＿
＞　進捗どうですか！？　＜
￣Y^Y^Y^Y^Y^Y^Y^Y^￣

僕はダメです。

進捗、大事ですよね。毎日進捗に追われる日々。

ダウンロードでも進捗は大事です。小さなファイルを落とすぐらいならすぐに終わって問題ないのですが大きなファイル、例えばLinuxのディストリのisoファイルなんかはギガバイト級だったりするので自分で作ったプログラムで軽々しくダウンロードするといつ終わるのか分からなくてドキドキします。ルータのランプを眺めてダウンロードが行われていることを予想したりします。

そういうわけで進捗をとりながらダウンロードするhaskellのプログラムを書きます。

主に使うパッケージはみんな大好きhttp-conduitさんとstmさん。

まずは進捗を入れるためのデータ型から。並行処理ができるようにSTMを使って実装します。

type URL = String

data Progress = Progress {
  url :: URL
, size :: (TVar Integer)
, total :: (TVar Integer)
, cancel :: (TVar Bool)
}

newProgress :: URL -> IO Progress
newProgress u = Progress u <$> newTVarIO 0 <*> newTVarIO 0 <*> newTVarIO False

packProgress :: Progress -> IO (URL, Integer, Integer, Bool)
packProgress pg = atomically $ do
  s <- readTVar . size $ pg
  t <- readTVar . total $ pg
  c <- readTVar . cancel $ pg
  return (url pg, s, t, c)

packProgressではreadTVarIOで実装した方がパフォーマンスが上がるかもしれませんが、s,t,cの間で齟齬が生じる可能性があるのでatomically内でまとめて処理します。進捗を見たいだけならあまり気にしなくてもいいでしょうが。

今回はpackProgressの結果はタプルにしていますが、新しくデータ型を作ってアプリカティブスタイルにしたいですね。

続いてダウンロードする部分。

import qualified Data.ByteString as B (length)
import qualified Data.ByteString.Char8 as B (unpack)

残りのimportは適当にhackage見てください。

download :: Progress -> IO ()
download pg = do
  req <- parseUrl . url $ pg
  man <- newManager tlsManagerSettings
  runResourceT $ do
    res <- http req man
    let cl = maybe 0 (read . B.unpack) . lookup hContentLength . responseHeaders $ res
    liftIO . atomically . flip writeTVar cl . total $ pg
    responseBody res $=+ updateProgress pg $$+- sinkNull

updateProgress pg = await >>= maybe (return ()) go
  where
    update len = atomically . flip modifyTVar (+ len) . size $ pg
    go chunk = do
      let len = B.length chunk
      liftIO . update . toInteger $ len
      flag <- liftIO . readTVarIO . cancel $ pg
      if flag then return () else do { yield chunk; updateProgress pg }

末尾再帰最適化とかよくわからないです。

コンソールに出力したい時はlet len = B.length chunkの次辺りでliftIO $ packProgress pg >>= printとかすればいいんじゃないでしょうか。

sinkNullに流し込んでいるのでファイルへ保存はしていません。
保存したい場合はsinkFileなりsinkHandleにしてください(conduit-extraにあります)。

依存パッケージは

• base
• bytestring
• transformers
• conduit
• http-conduit
• resourcet
• http-types
• stm

辺り。stackのlts-3.6で書きました。

これを使えばダウンロードの
＿人人人人人人人人人＿
＞　進捗大丈夫です　＜
￣Y^Y^Y^Y^Y^Y^Y^￣