Multicore WIP

lib/Echidna/Campaign.hs

@@ -5,7 +5,7 @@ module Echidna.Campaign where
 
 import Control.DeepSeq (force)
 import Control.Lens
-import Control.Monad (replicateM, when)
+import Control.Monad (replicateM, when, unless, void)
 import Control.Monad.Catch (MonadCatch(..), MonadThrow(..))
 import Control.Monad.Random.Strict (MonadRandom, RandT, evalRandT)
 import Control.Monad.Reader (MonadReader, asks, liftIO)
@@ -267,7 +267,7 @@ callseq initialCorpus vm world seqLen = do
 -- to generate calls with. Return the 'Campaign' state once we can't solve or shrink anything.
 campaign
   :: (MonadIO m, MonadCatch m, MonadRandom m, MonadReader Env m)
-  => StateT Campaign m a -- ^ Callback to run after each state update (for instrumentation)
+  => StateT Campaign m Bool -- ^ Callback to run after each state update (for instrumentation)
   -> VM                  -- ^ Initial VM state
   -> World               -- ^ Initial world state
   -> [EchidnaTest]       -- ^ Tests to evaluate
@@ -294,13 +294,16 @@ campaign u vm world ts dict initialCorpus = do
       CampaignConf{testLimit, stopOnFail, seqLen, shrinkLimit} <- asks (.cfg.campaignConf)
       Campaign{ncallseqs} <- get
       if | stopOnFail && any (\case Solved -> True; Failed _ -> True; _ -> False) testStates ->
-           lift u
+           void $ lift u
          | any (\case Open n -> n <= testLimit; _ -> False) testStates ->
            callseq initialCorpus vm world seqLen >> step
          | any (\case Large n -> n < shrinkLimit; _ -> False) testStates ->
            step
          | null testStates && (seqLen * ncallseqs) <= testLimit ->
            callseq initialCorpus vm world seqLen >> step
          | otherwise ->
-           lift u
-    step = runUpdate (shrinkTest vm) >> lift u >> runCampaign
+           void $ lift u
+    step = do
+      runUpdate (shrinkTest vm)
+      stop <- lift u -- callback can instruct the campaign to stop running
+      unless stop runCampaign

lib/Echidna/UI.hs

@@ -6,32 +6,30 @@ module Echidna.UI where
 import Brick
 import Brick.BChan
 import Brick.Widgets.Dialog qualified as B
-import Control.Monad.Catch (MonadCatch(..), catchAll)
-import Control.Monad.Reader (MonadReader (ask), runReader, asks)
-import Control.Monad.State (modify')
+import Control.Monad.Reader (ask)
 import Graphics.Vty qualified as V
 import Graphics.Vty (Config, Event(..), Key(..), Modifier(..), defaultConfig, inputMap, mkVty)
 import System.Posix.Terminal (queryTerminal)
 import System.Posix.Types (Fd(..))
 
 import Echidna.UI.Widgets
-#else /* !INTERACTIVE_UI */
-import Control.Monad.Catch (MonadCatch(..))
-import Control.Monad.Reader (MonadReader, runReader, asks)
-import Control.Monad.State.Strict (get)
 #endif
 
-import Control.Monad
 import Control.Concurrent (killThread, threadDelay)
-import Control.Monad.IO.Class (MonadIO(..))
-import Control.Monad.Random.Strict (MonadRandom)
+import Control.Monad
+import Control.Monad.Catch (MonadCatch)
+import Control.Monad.Random.Strict (MonadRandom (getRandom))
+import Control.Monad.Reader (MonadReader, runReader, asks)
+import Control.Monad.State.Strict hiding (state)
 import Data.ByteString.Lazy qualified as BS
-import Data.IORef
+import Data.List (transpose)
+import Data.Set qualified as Set
 import Data.Map (Map)
+import Data.Map qualified as Map
 import Data.Maybe (fromMaybe)
-import UnliftIO (MonadUnliftIO)
+import UnliftIO (MonadUnliftIO, newIORef, readIORef, atomicWriteIORef)
 import UnliftIO.Timeout (timeout)
-import UnliftIO.Concurrent hiding (killThread, threadDelay)
+import UnliftIO.Concurrent (forkIO, getNumCapabilities, newEmptyMVar, putMVar, tryTakeMVar)
 
 import EVM (VM, Contract)
 import EVM.Types (Addr, W256)
@@ -55,6 +53,13 @@ data UIEvent =
   | CampaignCrashed String
   | FetchCacheUpdated (Map Addr (Maybe Contract)) (Map Addr (Map W256 (Maybe W256)))
 
+-- TODO: sync corpus from time to time?
+data WorkerSyncMessage
+  = TestFalsified Int EchidnaTest -- TODO: rethink if whole test should be pushed
+  | TestLimitReached Int
+  | ShrinkLimitReached Int
+  | TimeLimitReached -- timed out, exit
+
 -- | Set up and run an Echidna 'Campaign' and display interactive UI or
 -- print non-interactive output in desired format at the end
 ui :: (MonadCatch m, MonadRandom m, MonadReader Env m, MonadUnliftIO m)
@@ -68,7 +73,7 @@ ui vm world ts dict initialCorpus = do
   conf <- asks (.cfg)
   let uiConf = conf.uiConf
   ref <- liftIO $ newIORef defaultCampaign
-  let updateRef = get >>= liftIO . atomicWriteIORef ref
+  let updateRef = get >>= liftIO . atomicWriteIORef ref >> pure False
       secToUsec = (* 1000000)
       timeoutUsec = secToUsec $ fromMaybe (-1) uiConf.maxTime
       runCampaign = timeout timeoutUsec (campaign updateRef vm world ts dict initialCorpus)
@@ -80,8 +85,8 @@ ui vm world ts dict initialCorpus = do
   let effectiveMode = case uiConf.operationMode of
         Interactive | not terminalPresent -> NonInteractive Text
         other -> other
+          {-
   case effectiveMode of
-#ifdef INTERACTIVE_UI
     Interactive -> do
       bc <- liftIO $ newBChan 100
       let updateUI e = readIORef ref >>= writeBChan bc . e
@@ -94,7 +99,24 @@ ui vm world ts dict initialCorpus = do
           c <- readIORef env.fetchContractCache
           s <- readIORef env.fetchSlotCache
           writeBChan bc (FetchCacheUpdated c s)
-      _ <- forkFinally -- run worker
+        -- _ <- forkFinally -- run worker
+        -}
+
+  numCapabilities <- getNumCapabilities
+  -- TODO: Performance peaks around 12 jobs, even if there are 24 CPU threads
+  -- available (Ryzen 5900X). Is it possible to tweak scheduler/GC?
+  let njobs = numCapabilities -- max 12 numCapabilities
+  -- Communication channels with workers
+  ioRefs <- replicateM njobs $ newIORef defaultCampaign
+  mVars <- replicateM njobs newEmptyMVar
+
+  let worker (ioRef, mVar) = forkIO $ void $ do
+        -- Generate a new seed for each worker
+        seed <- getRandom
+        let dict' = dict { defSeed = seed }
+        campaign (syncWorker ioRef mVar ioRefs) vm world ts dict' initialCorpus
+
+  {-
         (void $ do
           catchAll
             (runCampaign >>= \case
@@ -103,6 +125,48 @@ ui vm world ts dict initialCorpus = do
             (liftIO . writeBChan bc . CampaignCrashed . show)
         )
         (const $ liftIO $ killThread ticker)
+        -}
+
+      maybeTimeout = secToUsec <$> uiConf.maxTime
+
+  {-workerSyncer = forkIO $ void $ do
+        campaigns <- sequence $ readIORef <$> ioRefs
+        pure undefined-}
+
+  case effectiveMode of
+#ifdef INTERACTIVE_UI
+    Interactive -> do
+      -- Channel to push events to update UI
+      bc <- liftIO $ newBChan 100
+
+      env <- ask
+      let updateUI e = do
+            campaigns <- mapM readIORef ioRefs
+            -- TODO: push MVar messages
+            writeBChan bc $ e (mergeCampaigns campaigns)
+
+      ticker <- liftIO $ forkIO $
+        -- run UI update every 100ms
+        forever $ do
+          threadDelay 100000
+          updateUI CampaignUpdated
+          c <- readIORef env.fetchContractCache
+          s <- readIORef env.fetchSlotCache
+          writeBChan bc (FetchCacheUpdated c s)
+
+      -- Timeouter thread, sleep for the timeout then order all workers
+      -- to exit and update the UI
+      case maybeTimeout of
+        Just t -> liftIO $ void $ forkIO $ do
+          threadDelay t
+          killThread ticker
+          mapM_ (`putMVar` TimeLimitReached) mVars
+          updateUI CampaignTimedout
+        Nothing -> pure ()
+
+      _threadIds <- mapM worker (zip ioRefs mVars)
+
+      -- UI initialization
       let buildVty = do
             v <- mkVty =<< vtyConfig
             V.setMode (V.outputIface v) V.Mouse True
@@ -117,7 +181,11 @@ ui vm world ts dict initialCorpus = do
         , fetchedDialog = B.dialog (Just "Fetched contracts/slots") Nothing 80
         , displayFetchedDialog = False
         }
-      final <- liftIO $ readIORef ref
+
+      mapM_ (`putMVar` TimeLimitReached) mVars
+
+      final <- mergeCampaigns <$> mapM readIORef ioRefs
+      -- final <- liftIO $ readIORef ref
       liftIO . putStrLn $ runReader (ppCampaign final) conf
       pure final
 #else
@@ -134,24 +202,108 @@ ui vm world ts dict initialCorpus = do
           putStrLn $ time <> "[status] " <> statusLine conf.campaignConf camp
       result <- runCampaign
       liftIO $ killThread ticker
-      (final, timedout) <- case result of
+      (_final, _timedout) <- case result of
         Nothing -> do
           final <- liftIO $ readIORef ref
           pure (final, True)
         Just final ->
           pure (final, False)
+
+      -- Timeouter thread, sleep for the timeout then order all workers to exit
+      -- TODO: this is similar to the UI one, think about extracting it?
+      didTimeout <- newIORef False
+
+      case maybeTimeout of
+        Just t -> liftIO $ void $ forkIO $ do
+          threadDelay t
+          mapM_ (`putMVar` TimeLimitReached) mVars
+          atomicWriteIORef didTimeout True
+        Nothing -> pure ()
+
+      _threadIds <- mapM worker (zip ioRefs mVars)
+
+      -- TODO wait for threads
+      liftIO $ threadDelay 1000000
+
+      campaigns <- mapM readIORef ioRefs
+      let final = mergeCampaigns campaigns
+
       case outputFormat of
         JSON ->
           liftIO . BS.putStr $ Echidna.Output.JSON.encodeCampaign final
         Text -> do
           liftIO . putStrLn $ runReader (ppCampaign final) conf
+          timedout <- readIORef didTimeout
           when timedout $ liftIO $ putStrLn "TIMEOUT!"
         None ->
           pure ()
       pure final
 
-#ifdef INTERACTIVE_UI
+  where
+  syncWorker ioRef mVar _ioRefs = do
+    c <- get
+    -- push campaign update
+    liftIO $ atomicWriteIORef ioRef c
+    -- read a message if a breakthrough happened in another worker
+    -- TODO use a channel instead of MVar as it could be block the UI thread a bit
+    maybeMessage <- liftIO $ tryTakeMVar mVar
+    case maybeMessage of
+      Nothing -> pure False
+      Just message -> do
+        let (c', stop) = updateCampaign c message
+        put c'
+        pure stop
+    where
+    updateCampaign c = \case
+      TestFalsified _i _t ->
+        -- NOTE: the first worker wins, here we overwrite work that was
+        -- done by the current worker, TODO: rethink this
+        -- (c { tests = c.tests
+        -- (c & tests . ix i .~ t, False)
+        error "implement me"
+      TestLimitReached _i ->
+        -- bump the all? Open trials to max to stop fuzzing
+        --case c ^. tests . at i of
+        --  Nothing -> undefined
+        --  Just t -> undefined
+        error "implement me"
+      ShrinkLimitReached _i ->
+        -- bump the all? Open trials to max to stop fuzzing
+        error "implement me"
+      TimeLimitReached -> (c, True)
+
 
+-- Summarize all campaigns from workers as a single campaign
+-- TODO: this should return a richer data structure, good enough for now
+mergeCampaigns :: [Campaign] -> Campaign
+mergeCampaigns [] = error "won't happen, fix me with NonEmpty"
+mergeCampaigns [c] = c -- don't even try
+mergeCampaigns campaigns =
+  defaultCampaign
+    { tests = mergeTests <$> transpose ((.tests) <$> campaigns)
+    , coverage = Map.unionsWith Set.union ((.coverage) <$> campaigns)
+    , gasInfo = mempty -- TODO
+    , genDict = emptyDict -- TODO
+    , corpus = mempty -- TODO
+    , ncallseqs = sum ((.ncallseqs) <$> campaigns)
+    }
+  where
+  mergeTests :: [EchidnaTest] -> EchidnaTest
+  mergeTests [] = error "won't happen, fix me with NonEmpty"
+  mergeTests (f:ts) =
+    foldl (\t acc ->
+      case (t.state, acc.state) of
+        -- update if better what we have so far
+        (Solved, _) -> t
+        (Large i, Large j) -> t { state = Large (i+j) }
+        (Large _, Open _) -> t
+        (Large _, Passed) -> t -- shoudn't happen but just in case
+        (Open i, Open j) -> t { state = Open (i+j) }
+        -- skip otherwise
+        _ -> acc
+      ) f ts
+
+#ifdef INTERACTIVE_UI
 vtyConfig :: IO Config
 vtyConfig = do
   config <- V.standardIOConfig

package.yaml

@@ -6,7 +6,7 @@ maintainer: Trail of Bits <echidna-dev@trailofbits.com>
 version: 2.1.0
 
 # https://github.com/haskell/cabal/issues/4739
-ghc-options: -Wall -fno-warn-orphans -O2 -threaded +RTS -N -RTS -optP-Wno-nonportable-include-path
+ghc-options: -Wall -fno-warn-orphans -O2 -threaded +RTS -N -RTS -optP-Wno-nonportable-include-path -rtsopts
 
 dependencies:
   - base

src/test/Common.hs

@@ -107,7 +107,7 @@ runContract f selectedContract cfg = do
   (vm, world, echidnaTests, dict) <- prepareContract env contracts (f :| []) selectedContract seed
   let corpus = []
   -- start ui and run tests
-  runReaderT (campaign (pure ()) vm world echidnaTests dict corpus) env
+  runReaderT (campaign (pure False) vm world echidnaTests dict corpus) env
 
 testContract :: FilePath -> Maybe FilePath -> [(String, Campaign -> Bool)] -> TestTree
 testContract fp cfg = testContract' fp Nothing Nothing cfg True