Run at most one async job per user

locks.hs

#!/usr/bin/env stack
{- stack script
--resolver lts-18.14
--package async,containers,say,stm
-}

import Control.Concurrent
import Control.Concurrent.Async
import Control.Concurrent.STM (TVar, atomically, modifyTVar, newTVarIO, readTVar)
import qualified Data.IntSet as Set
import Say (sayString)


main :: IO ()
main = do
    userLocks <- newTVarIO Set.empty
    runInBackgroundAtMostOncePerUser userLocks 1 (longRunningAction 1) -- Run
    runInBackgroundAtMostOncePerUser userLocks 2 (longRunningAction 2) -- Run, because user 2 has no tasks running
    runInBackgroundAtMostOncePerUser userLocks 1 (longRunningAction 1) -- Skipped because user 1 has task running
    sleep 3
    runInBackgroundAtMostOncePerUser userLocks 1 (longRunningAction 1) -- Run because user 1 has no tasks running
    sleep 3


type UserId = Int
type UserLocks = TVar Set.IntSet


runInBackgroundAtMostOncePerUser :: UserLocks -> UserId -> IO () -> IO ()
runInBackgroundAtMostOncePerUser locks userId action = do
    lockSuceeded <- tryLock locks userId
    if lockSuceeded
        then void . async
            . withAsync action -- TODO is there a better way to run "finalization logic" (unlocking) than nesting 2 asyncs?
            $ \a -> do
                res <- waitCatch a -- TODO potentially do something with exception etc.
                unlock locks userId
        else sayString $ "Not running action for user " <> show userId <> " because it's already running"


-- Return True if lock for given user acquired suceeded
tryLock :: UserLocks -> UserId -> IO Bool
tryLock locks userId =
    atomically $ do
        userLocked <- Set.member userId <$> readTVar locks
        if userLocked
            then pure False
            else modifyTVar locks (Set.insert userId) >> pure True


unlock :: UserLocks -> UserId -> IO ()
unlock locks userId =
    atomically $ modifyTVar locks (Set.delete userId)


longRunningAction :: UserId -> IO ()
longRunningAction userId = do
    sayString $ "Action started for user " <> show userId
    sleep 2
    sayString $ "Action finished for user " <> show userId


sleep :: Int -> IO ()
sleep seconds = threadDelay (seconds * 1000000)