haskellconcurrencystmtvar

Haskell: Updating two or more TVars atomically. Possible?


Can one transaction update two different TVars in an atomic way? i.e. can I compose data structures out of lots of TVars to reduce contention? If so, could you provide an example?


Solution

  • Can one transaction update two different TVars in an atomic way?

    Yes, you can update multiple TVars atomically in one transaction. That's sort of the whole point of STM. It wouldn't be very useful if you couldn't.

    Can I compose data structures out of lots of TVars to reduce contention? If so, could you provide an example?

    Here is a (somewhat silly) example of storing TVars in a data structure. It simulates a bunch of random concurrent transactions between accounts in a bank, where each account is just a TVar Integer. The account TVars are kept in a map from account IDs, which is itself kept in a TVar so that new accounts can be created on the fly.

    import Control.Concurrent
    import Control.Concurrent.MVar
    import Control.Concurrent.STM
    import Control.Monad
    import System.Random
    
    import qualified Data.Map as Map
    
    type AccountId = Int
    type Account = TVar Dollars
    type Dollars = Integer
    type Bank = TVar (Map.Map AccountId Account)
    
    numberOfAccounts = 20
    threads = 100
    transactionsPerThread = 100
    maxAmount = 1000
    
    -- Get account by ID, create new empty account if it didn't exist
    getAccount :: Bank -> AccountId -> STM Account
    getAccount bank accountId = do
      accounts <- readTVar bank
      case Map.lookup accountId accounts of
        Just account -> return account
        Nothing -> do
          account <- newTVar 0
          writeTVar bank $ Map.insert accountId account accounts
          return account
    
    -- Transfer amount between two accounts (accounts can go negative)
    transfer :: Dollars -> Account -> Account -> STM ()
    transfer amount from to = when (from /= to) $ do
      balanceFrom <- readTVar from
      balanceTo <- readTVar to
      writeTVar from $! balanceFrom - amount
      writeTVar to $! balanceTo + amount
    
    randomTransaction :: Bank -> IO ()
    randomTransaction bank = do
      -- Make a random transaction
      fromId <- randomRIO (1, numberOfAccounts)
      toId   <- randomRIO (1, numberOfAccounts)
      amount <- randomRIO (1, maxAmount)
    
      -- Perform it atomically
      atomically $ do
        from <- getAccount bank fromId
        to   <- getAccount bank toId
        transfer amount from to
    
    main = do
      bank <- newTVarIO Map.empty
    
      -- Start some worker threads to each do a number of random transactions
      workers <- replicateM threads $ do
        done <- newEmptyMVar
        forkIO $ do
          replicateM_ transactionsPerThread $ randomTransaction bank
          putMVar done ()
        return done
    
      -- Wait for worker threads to finish
      mapM_ takeMVar workers
    
      -- Print list of accounts and total bank balance (which should be zero)
      summary <- atomically $ do
        accounts <- readTVar bank
        forM (Map.assocs accounts) $ \(accountId, account) -> do
          balance <- readTVar account
          return (accountId, balance)
    
      mapM_ print summary
      putStrLn "----------------"
      putStrLn $ "TOTAL BALANCE: " ++ show (sum $ map snd summary)
    

    This should print a total balance of zero at the end if there were no race conditions during the transfers.