Try different approach for stopping workers

lib/Echidna/Campaign.hs

@@ -7,7 +7,7 @@ import Optics.Core hiding ((|>))
 import Control.Concurrent (writeChan)
 import Control.DeepSeq (force)
 import Control.Monad (replicateM, when, void, forM_)
-import Control.Monad.Catch (MonadCatch(..), MonadThrow(..), catchAll)
+import Control.Monad.Catch (MonadCatch(..), MonadThrow(..))
 import Control.Monad.Random.Strict (MonadRandom, RandT, evalRandT)
 import Control.Monad.Reader (MonadReader, asks, liftIO, ask)
 import Control.Monad.State.Strict
@@ -77,15 +77,15 @@ replayCorpus vm txSeqs =
 -- we can't solve or shrink anything.
 runWorker
   :: (MonadIO m, MonadCatch m, MonadRandom m, MonadReader Env m)
-  => StateT WorkerState m (Maybe WorkerStopReason)
+  => StateT WorkerState m ()
   -- ^ Callback to run after each state update (for instrumentation)
   -> VM      -- ^ Initial VM state
   -> World   -- ^ Initial world state
   -> GenDict -- ^ Generation dictionary
   -> Int     -- ^ Worker id starting from 0
   -> [[Tx]]  -- ^ Initial corpus of transactions
   -> Int     -- ^ Test limit for this worker
-  -> m WorkerState
+  -> m (WorkerStopReason, WorkerState)
 runWorker callback vm world dict workerId initialCorpus testLimit = do
   metaCacheRef <- asks (.metadataCache)
   fetchContractCacheRef <- asks (.fetchContractCache)
@@ -104,15 +104,11 @@ runWorker callback vm world dict workerId initialCorpus testLimit = do
                   , ncalls = 0
                   }
 
-  flip execStateT initialState $ do
+  flip runStateT initialState $ do
     flip evalRandT (mkStdGen effectiveSeed) $ do
-      catchAll
-        (do void $ lift callback
-            void $ replayCorpus vm initialCorpus
-            stopReason <- run
-            pushEvent $ WorkerStopped stopReason
-        )
-        (pushEvent . WorkerStopped . Crashed . show)
+      lift callback
+      void $ replayCorpus vm initialCorpus
+      run
 
   where
   run = do
@@ -131,7 +127,7 @@ runWorker callback vm world dict workerId initialCorpus testLimit = do
                           _       -> False
 
     if | stopOnFail && any final tests -> do
-         void $ lift callback
+         lift callback
          pure FastFailed
 
        | (null tests || any isOpen tests) && ncalls < testLimit ->
@@ -141,15 +137,12 @@ runWorker callback vm world dict workerId initialCorpus testLimit = do
          continue
 
        | otherwise -> do
-         void $ lift callback
+         lift callback
          pure TestLimitReached
 
   fuzz = randseq vm.env.contracts world >>= callseq vm
 
-  continue = do
-    runUpdate (shrinkTest vm)
-    maybeStop <- lift callback
-    maybe run pure maybeStop
+  continue = runUpdate (shrinkTest vm) >> lift callback >> run
 
   mkMemo = makeBytecodeCache . map (forceBuf . (^. bytecode)) . Map.elems
 

lib/Echidna/UI.hs

@@ -15,6 +15,7 @@ import Echidna.UI.Widgets
 #endif
 
 import Control.Concurrent (killThread, threadDelay)
+import Control.Exception (AsyncException)
 import Control.Monad
 import Control.Monad.Catch
 import Control.Monad.Random.Strict (MonadRandom)
@@ -31,7 +32,7 @@ import EVM (VM, Contract)
 import EVM.Types (Addr, W256)
 
 import Echidna.ABI
-import Echidna.Campaign (runWorker, pushEvent)
+import Echidna.Campaign (runWorker)
 import Echidna.Output.JSON qualified
 import Echidna.Types.Campaign
 import Echidna.Types.Config
@@ -49,8 +50,6 @@ data UIEvent =
                       (Map Addr (Map W256 (Maybe W256)))
   | WorkerEvent (Int, LocalTime, CampaignEvent)
 
-type Worker = (IORef WorkerState, MVar ())
-
 -- | Set up and run an Echidna 'Campaign' and display interactive UI or
 -- print non-interactive output in desired format at the end
 ui
@@ -78,7 +77,7 @@ ui vm world dict initialCorpus = do
     perWorkerTestLimit = ceiling
       (fromIntegral conf.campaignConf.testLimit / fromIntegral jobs :: Double)
 
-  workers <- forM [0..(jobs-1)] (spawnWorker perWorkerTestLimit)
+  workers <- forM [0..(jobs-1)] (spawnWorker env.eventQueue perWorkerTestLimit)
 
   -- Run a thread that will order all workers to exit when timeout is reached
   case conf.uiConf.maxTime of
@@ -198,26 +197,35 @@ ui vm world dict initialCorpus = do
 
   where
 
-  spawnWorker testLimit workerId = do
+  spawnWorker eventQueue testLimit workerId = do
     stateRef <- newIORef initialWorkerState
-    stopWorker <- newEmptyMVar
 
-    -- Is ThreadId useful for anything?
-    void . forkIO . void $ do
-      -- TODO: split corpus into chunks and make each worker replay a chunk
-      runWorker (workerCallback stateRef stopWorker)
-                vm world dict workerId initialCorpus testLimit
+    threadId <- forkIO . void $ do
+      stopReason <- catches
+        (do
+          -- TODO: split corpus into chunks and make each worker replay a chunk
+          (stopReason, _finalState) <-
+            runWorker (workerCallback stateRef)
+                      vm world dict workerId initialCorpus testLimit
+          pure stopReason
+        )
+        [ Handler $ \(_ :: AsyncException) -> pure Killed
+        , Handler $ \(e :: SomeException)  -> pure $ Crashed (show e)
+        ]
+
+      time <- liftIO getTimestamp
+      writeChan eventQueue (workerId, time, WorkerStopped stopReason)
 
-    pure (stateRef, stopWorker)
+    pure (threadId, stateRef)
 
   -- | This function is idempotent and can be called many times. This is
   -- important in case there is a race to stop workers, the first reason wins.
-  stopWorkers workers reason =
-    forM_ workers $ \(_, stopWorker) -> tryPutMVar stopWorker reason
+  stopWorkers workers _reason =
+    forM_ workers $ \(threadId, _) -> liftIO $ killThread threadId
 
   -- | Get a snapshot of all worker states
   workerStates workers =
-    forM workers $ \(stateRef, _) -> readIORef stateRef
+    forM workers $ \(_, stateRef) -> readIORef stateRef
 
   spawnListener
     :: Env
@@ -239,9 +247,8 @@ ui vm world dict initialCorpus = do
           (_, _, WorkerStopped _) -> loop (workersAlive - 1)
           _                       -> loop workersAlive
 
-  workerCallback stateRef stopWorker = do
+  workerCallback stateRef = do
     get >>= writeIORef stateRef
-    tryTakeMVar stopWorker
 
 #ifdef INTERACTIVE_UI
 vtyConfig :: IO Config

src/test/Common.hs

@@ -117,8 +117,8 @@ runContract f selectedContract cfg = do
   (vm, world, dict) <- prepareContract env contracts (f :| []) selectedContract seed
 
   let corpus = []
-  finalState <- flip runReaderT env $
-    runWorker (pure Nothing) vm world dict 0 corpus cfg.campaignConf.testLimit
+  (_stopReason, finalState) <- flip runReaderT env $
+    runWorker (pure ()) vm world dict 0 corpus cfg.campaignConf.testLimit
 
   -- TODO: consider snapshotting the state so checking function don't need to
   -- be IO