performance of stm

stm-test.hs

module Main where

import Control.Monad
import Control.Concurrent
import Control.Concurrent.STM

-- STM (TVar [STM (TVar Int)])

type Data = TVar Int
type DataList = TVar [STM Data]

scale :: Int
scale = 10000

(...) :: Int -> Int -> [Int]
a ... b = [a..b-1]

newDataList :: IO DataList
newDataList = newTVarIO . map newTVar $ 0...scale

-- DataListに100000個のSTM (TVar Int)を追加
push :: DataList -> TVar Int -> IO ()
push tv c = forM_ (scale...(scale*2)) proc >> atomically (modifyTVar c succ)
  where
  ¦ proc i = do
  ¦ ¦ atomically $ do
  ¦ ¦ ¦ modifyTVar tv (newTVar i:)

-- 0...100000番目の各要素をsuccで更新
update :: DataList -> TVar Int -> IO ()
update tv c = forM_ (0...scale) proc >> atomically (modifyTVar c succ)
  where
  ¦ proc i = do
  ¦ ¦ ¦ atomically $ do
  ¦ ¦ ¦ ¦ modifyTVar tv change -- TVar [STM (TVar Int)]をmodify
  ¦ ¦ where
  ¦ ¦ ¦ -- N番目の要素のTVar Intをsucc
  ¦ ¦ ¦ change xs = modifyNth i (\x -> do t <- x; modifyTVar t succ; return t) xs

modifyNth n f (x:xs)
  | n == 0 = f x : xs
  | otherwise = x : modifyNth (n - 1) f xs

main :: IO ()
main = do
  dataList <- newDataList
  p <- newTVarIO 0
  u <- newTVarIO 0
  replicateM_ 10 $ forkIO $ push dataList p -- pushを10スレッド
  replicateM_ 40 $ forkIO $ update dataList u -- updateを40スレッド
  atomically $ do
  ¦ p' <- readTVar p
  ¦ u' <- readTVar u
  ¦ unless (p' == 10 && u' == 40) retry -- pushとupdateがすべて終わるまで待つ