use 'atomicModifyIORefCAS' where possible

Author: jkachmar

core/package.yaml

@@ -50,6 +50,7 @@ library:
     # hackage dependencies
     - aeson
     - async
+    - atomic-primops
     - lens-family
     - monad-logger
     - network-bsd

core/src/Temporal/Core/Worker.hs

@@ -49,6 +49,7 @@ import Control.Exception
 import Control.Monad
 import Data.Aeson
 import Data.Aeson.TH
+import Data.Atomics (atomicModifyIORefCAS)
 import Data.ByteString (ByteString)
 import qualified Data.ByteString.Lazy as BL
 import Data.IORef
@@ -258,7 +259,7 @@ instance Exception WorkerAlreadyClosed
 -- After calling this, the worker must not be used again.
 closeWorker :: Worker ty -> IO ()
 closeWorker (Worker w _ _ _) = mask_ $ do
-  wp <- atomicModifyIORef' w $ \wp -> (throw WorkerAlreadyClosed, wp)
+  wp <- liftIO $ atomicModifyIORefCAS w $ \wp -> (throw WorkerAlreadyClosed, wp)
   raw_closeWorker wp
 
 

core/temporal-sdk-core.cabal

@@ -88,6 +88,7 @@ library
   build-depends:
       aeson
     , async
+    , atomic-primops
     , base >=4.14 && <5
     , bytestring
     , containers

sdk/package.yaml

@@ -10,6 +10,7 @@ dependencies:
 - proto-lens-protobuf-types
 - temporal-api-protos
 - temporal-sdk-core
+- atomic-primops
 - bytestring
 - discover-instances
 - mtl

sdk/src/Temporal/Workflow/Internal/Instance.hs

@@ -18,6 +18,7 @@ module Temporal.Workflow.Internal.Instance (
 ) where
 
 import Control.Monad.Reader
+import Data.Atomics (atomicModifyIORefCAS)
 import Data.ProtoLens
 import qualified Data.Text as T
 import GHC.Stack
@@ -62,46 +63,46 @@ flushCommands = do
 nextExternalCancelSequence :: InstanceM Sequence
 nextExternalCancelSequence = do
   inst <- ask
-  atomicModifyIORef' inst.workflowSequences $ \seqs ->
+  liftIO $ atomicModifyIORefCAS inst.workflowSequences $ \seqs ->
     let seq' = externalCancel seqs
     in (seqs {externalCancel = succ seq'}, Sequence seq')
 
 
 nextChildWorkflowSequence :: InstanceM Sequence
 nextChildWorkflowSequence = do
   inst <- ask
-  atomicModifyIORef' inst.workflowSequences $ \seqs ->
+  liftIO $ atomicModifyIORefCAS inst.workflowSequences $ \seqs ->
     let seq' = childWorkflow seqs
     in (seqs {childWorkflow = succ seq'}, Sequence seq')
 
 
 nextExternalSignalSequence :: InstanceM Sequence
 nextExternalSignalSequence = do
   inst <- ask
-  atomicModifyIORef' inst.workflowSequences $ \seqs ->
+  liftIO $ atomicModifyIORefCAS inst.workflowSequences $ \seqs ->
     let seq' = externalSignal seqs
     in (seqs {externalSignal = succ seq'}, Sequence seq')
 
 
 nextTimerSequence :: InstanceM Sequence
 nextTimerSequence = do
   inst <- ask
-  atomicModifyIORef' inst.workflowSequences $ \seqs ->
+  liftIO $ atomicModifyIORefCAS inst.workflowSequences $ \seqs ->
     let seq' = timer seqs
     in (seqs {timer = succ seq'}, Sequence seq')
 
 
 nextActivitySequence :: InstanceM Sequence
 nextActivitySequence = do
   inst <- ask
-  atomicModifyIORef' inst.workflowSequences $ \seqs ->
+  liftIO $ atomicModifyIORefCAS inst.workflowSequences $ \seqs ->
     let seq' = activity seqs
     in (seqs {activity = succ seq'}, Sequence seq')
 
 
 nextConditionSequence :: InstanceM Sequence
 nextConditionSequence = do
   inst <- ask
-  atomicModifyIORef' inst.workflowSequences $ \seqs ->
+  liftIO $ atomicModifyIORefCAS inst.workflowSequences $ \seqs ->
     let seq' = condition seqs
     in (seqs {condition = succ seq'}, Sequence seq')

sdk/src/Temporal/Workflow/Internal/Monad.hs

@@ -8,6 +8,7 @@ import Control.Monad
 import qualified Control.Monad.Catch as Catch
 import Control.Monad.Logger
 import Control.Monad.Reader
+import Data.Atomics (atomicModifyIORefCAS, atomicModifyIORefCAS_)
 import Data.HashMap.Strict (HashMap)
 import qualified Data.HashMap.Strict as HashMap
 import Data.Kind
@@ -19,7 +20,6 @@ import Data.Time.Clock.System (SystemTime)
 import Data.Vault.Strict
 import Data.Vector (Vector)
 import Data.Word (Word32)
--- import Debug.Trace
 import GHC.Stack
 import GHC.TypeLits
 import Proto.Temporal.Sdk.Core.WorkflowActivation.WorkflowActivation
@@ -343,7 +343,7 @@ instance FrozenGen StdGen Workflow where
 {-# INLINE addJob #-}
 addJob :: ContinuationEnv -> Workflow b -> IVar b -> IVar a -> InstanceM ()
 addJob env !wf !resultIVar IVar {ivarRef = !ref} =
-  join $ atomicModifyIORef' ref $ \case
+  join $ liftIO $ atomicModifyIORefCAS ref $ \case
     IVarEmpty list -> (IVarEmpty (JobCons env wf resultIVar list), pure ())
     full -> (full, modifyIORef' env.runQueueRef (JobCons env wf resultIVar))
 
@@ -599,7 +599,7 @@ class MonadWriteStateVar m where
 instance MonadReadStateVar Condition where
   readStateVar var = Condition $ do
     touchedVars <- ask
-    modifyIORef' touchedVars (Set.insert var.stateVarId)
+    liftIO $ atomicModifyIORefCAS_ touchedVars (Set.insert var.stateVarId)
     readIORef var.stateVarRef
 
 
@@ -617,7 +617,7 @@ instance MonadWriteStateVar Workflow where
     reevaluateDependentConditions var
     pure $ Done ()
   modifyStateVar var f = Workflow $ \_ -> do
-    res <- modifyIORef' var.stateVarRef f
+    res <- liftIO $ atomicModifyIORefCAS_ var.stateVarRef f
     reevaluateDependentConditions var
     pure $ Done res
 
@@ -921,6 +921,6 @@ instance Monoid WorkflowOutboundInterceptor where
 nextVarIdSequence :: InstanceM Sequence
 nextVarIdSequence = do
   inst <- ask
-  atomicModifyIORef' inst.workflowSequences $ \seqs ->
+  liftIO $ atomicModifyIORefCAS inst.workflowSequences $ \seqs ->
     let seq' = varId seqs
     in (seqs {varId = succ seq'}, Sequence seq')

sdk/temporal-sdk.cabal

@@ -102,6 +102,7 @@ library
       aeson
     , annotated-exception
     , async
+    , atomic-primops
     , base >=4.7 && <5
     , base64
     , bytestring
@@ -187,6 +188,7 @@ test-suite temporal-sdk-tests
       aeson
     , annotated-exception
     , async
+    , atomic-primops
     , base
     , base64
     , bytestring