Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
#! /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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment