mongo-batch-performance.hs

#! /usr/bin/env stack
-- stack --resolver lts-17.0 script
{-# language OverloadedStrings #-}
import qualified Control.Exception as Exception
import qualified Control.Monad as Monad
import qualified Data.Function as Function
import qualified Data.IORef as IORef
import qualified Database.MongoDB as Mongo
import qualified Data.Text as Text
import qualified GHC.Clock as Clock
import qualified System.Environment as Environment
import qualified System.Random as Random
import qualified Text.Printf as Printf
import qualified Text.Read as Read
main = do
  arguments <- Environment.getArgs
  [ documentCount, batchSize ] <- either fail pure $ traverse Read.readEither arguments
  database <- fmap Text.pack . Monad.replicateM 8 $ Random.randomRIO ('a', 'z')
  collection <- fmap Text.pack . Monad.replicateM 8 $ Random.randomRIO ('a', 'z')
  let acquire = Mongo.connect $ Mongo.host "localhost"
  let release pipe = Mongo.access pipe Mongo.master database $ Mongo.dropDatabase database
  Exception.bracket acquire release $ \ pipe -> do
    let mongo accessMode = Mongo.access pipe accessMode database
    mongo Mongo.master . Mongo.insertMany_ collection $ replicate documentCount ["x" Mongo.=: (1 :: Int)]
    cursor <- mongo Mongo.slaveOk . Mongo.find $ Mongo.select [] collection
    sumRef <- IORef.newIORef (0 :: Int)
    before <- Clock.getMonotonicTime
    Function.fix $ \ loop -> do
      batch <- mongo Mongo.slaveOk $ Mongo.nextN batchSize cursor
      Monad.forM_ batch $ \ document -> do
        x <- Mongo.lookup "x" document
        IORef.modifyIORef' sumRef (+ x)
      Monad.unless (null batch) loop
    after <- Clock.getMonotonicTime
    sum <- IORef.readIORef sumRef
    print sum
    Printf.printf "%d total, %d per, %.3f elapsed\n"
      documentCount batchSize (after - before)