stm.hs

module CdnPurge where

type Host = T.Text
type Addr = T.Text
type Key = T.Text
type Response = T.Text

data WorkStats = WorkStats
  { startTime :: UTCTime
  , threadsWorking :: TVar Word
  , threadsFinished :: TVar Word
  , responseStats :: TVar (HM.HashMap Response Word)
  , errorStats :: TVar (HM.HashMap Host SomeException)
  }

purgeHost :: WorkStats -> PurgeConfig -> [Key] -> Host -> Addr -> IO ()
purgeHost WorkStats{..} PurgeConfig{..} keys host addr =
  -- Update thread counters no matter what.
  bracket_ enterStats exitStats $ do
    -- results :: Either SomeException [Response]
    results <- tryAny $ doSomeCrap

    case results of
      Right responses ->
        -- Summarize responses and merge them into stats.
        -- Marginally faster when there are many requests.
        let respMap = HM.fromListWith (+) [(resp, 1) | resp <- responses]
         in atomically $ modifyTVar' responseStats (HM.unionWith (+) respMap)
      Left err -> atomically $ modifyTVar' errorStats (HM.insert host err)
 where
  enterStats = atomically $ modifyTVar' threadsWorking (+1)
  exitStats = atomically $ do
    modifyTVar' threadsWorking ((-)1)
    modifyTVar' threadsFinished (+1)
    
main :: IO ()
main = do
    stats <- do
        hosts <- getTierHosts

        startTime <- getCurrentTime
        stats <- atomically $ do
          threadsWorking <- newTVar 0
          threadsFinished <- newTVar 0
          responseStats <- newTVar HM.empty
          errorStats <- newTVar HM.empty
          return WorkStats{..}

        let printer = forever $ do
              threadDelay 3000000
              printStats stats

        -- Run thread to print stats every 3 seconds.
        withAsync printer $ \_ ->
        -- Run all requests in parallel.
          forConcurrently_ hosts $
            uncurry (purgeHost stats purgeConfig keys)
        return stats

    -- Do something with run results.
    errs <- readTVarIO (errorStats stats)
    print errs
    putStrLn " == Final stats =="
    printStats stats
where
  printVar msg var = do
    val <- readTVarIO var
    putStrLn $ msg ++ show val
  printStats WorkStats{..} = do
    currentTime <- getCurrentTime
    let runtime = diffUTCTime currentTime startTime
    putStrLn $ "Fanout time: " ++ show runtime
    printVar "Threads running: " threadsWorking
    printVar "Threads finished: " threadsFinished
    printVar "Responses: " responseStats
    errs <- readTVarIO errorStats
    putStrLn $ "Errors: " ++ show (HM.size errs)