Benchmark queries in yesod

hack_the_reader.hs

instance YesodPersist App where
  type YesodPersistBackend App = SqlBackend
  runDB action = do
    master <- getYesod
    runSqlPool (hackTheReader action) $ appConnPool master

newtype BenchmarkResults = BenchmarkResults [(Text,TimeSpec)]
  deriving Typeable
type TypeMap      = HashMap TypeRep Dynamic

getBenchmarkResults :: HandlerT site IO BenchmarkResults
getBenchmarkResults = do
  gh <- handlerGetGHState >>= lift . readIORef
  return $ case HashMap.lookup (typeOf (BenchmarkResults [])) (ghsCache gh) of
    Nothing -> BenchmarkResults []
    Just d -> fromDyn d
      $ error $ "getBenchmarkResults: mistake made: " <> show d

addBenchmarkResult :: Text -> TimeSpec -> TypeMap -> TypeMap
addBenchmarkResult name duration = HashMap.insertWith concatBr brType (toDyn br)
  where
  concatBr a b = fromMaybe (error "addBenchmarkResult: mistake made") $ do
    BenchmarkResults xs <- fromDynamic a
    BenchmarkResults ys <- fromDynamic b
    return $ toDyn $ BenchmarkResults $ xs ++ ys
  brType = typeOf br
  br = BenchmarkResults [(name,duration)]

handlerGetGHState :: HandlerT site IO (IORef GHState)
handlerGetGHState = HandlerT (return . handlerState)


hackTheReader :: forall a site. ReaderT SqlBackend (HandlerT site IO) a -> ReaderT SqlBackend (HandlerT site IO) a
hackTheReader original = do
  ghRef <- lift handlerGetGHState
  flip local original $ \backend -> backend
    { connPrepare = \t -> do
        now <- getTime Monotonic
        startTime <- newIORef now
        stmt <- connPrepare backend t
        return $ stmt
          { stmtReset    =  do
              now2 <- getTime Monotonic
              start <- readIORef startTime
              let duration = diffTimeSpec now2 start
              modifyIORef' ghRef $ \r -> r
                { ghsCache = addBenchmarkResult t duration (ghsCache r) }
              print duration
              stmtReset stmt
          , stmtQuery    = injectRelease startTime . stmtQuery stmt
          }
    }



injectRelease :: IORef TimeSpec -> Acquire a -> Acquire a
injectRelease startTime (Acquire f) = Acquire $ \b -> do
  getTime Monotonic >>= writeIORef startTime
  f b