mempool: adapt for generalized validation

Author: crocodile-dentist

ouroboros-network/src/Ouroboros/Network/TxSubmission/Mempool/Simple.hs

@@ -1,8 +1,11 @@
-{-# LANGUAGE DerivingStrategies  #-}
-{-# LANGUAGE GADTs               #-}
-{-# LANGUAGE NamedFieldPuns      #-}
-{-# LANGUAGE ScopedTypeVariables #-}
-{-# LANGUAGE StandaloneDeriving  #-}
+{-# LANGUAGE BangPatterns             #-}
+{-# LANGUAGE DerivingStrategies       #-}
+{-# LANGUAGE DisambiguateRecordFields #-}
+{-# LANGUAGE FlexibleContexts         #-}
+{-# LANGUAGE GADTs                    #-}
+{-# LANGUAGE NamedFieldPuns           #-}
+{-# LANGUAGE ScopedTypeVariables      #-}
+{-# LANGUAGE StandaloneDeriving       #-}
 
 -- | The module should be imported qualified.
 --
@@ -19,24 +22,22 @@ module Ouroboros.Network.TxSubmission.Mempool.Simple
 import Prelude hiding (read, seq)
 
 import Control.Concurrent.Class.MonadSTM.Strict
-import Control.Monad (when)
 import Control.Monad.Class.MonadThrow
-
+import Control.Monad.Trans.Except
 import Data.Bifunctor (bimap)
-import Data.Either (partitionEithers)
+import Data.Either
 import Data.Foldable (toList)
 import Data.Foldable qualified as Foldable
-import Data.Function (on)
-import Data.List (find, nubBy)
-import Data.Maybe (isJust)
+import Data.List (find)
+import Data.Maybe (isJust, isNothing, fromJust)
 import Data.Sequence (Seq)
 import Data.Sequence qualified as Seq
 import Data.Set (Set)
 import Data.Set qualified as Set
 import Data.Typeable (Typeable)
 
+import Ouroboros.Network.Protocol.LocalTxSubmission.Type (SubmitResult (..))
 import Ouroboros.Network.SizeInBytes
-import Ouroboros.Network.TxSubmission.Inbound.V2.Types
 import Ouroboros.Network.TxSubmission.Mempool.Reader
 
 
@@ -120,11 +121,35 @@ deriving instance Show InvalidTxsError
 instance Exception InvalidTxsError
 
 
--- | A simple mempool writer.
+-- | A mempool writer which generalizes the tx submission mempool writer
+-- TODO: We could replace TxSubmissionMempoolWriter with this at some point
+--
+data MempoolWriter txid tx failure idx m =
+     MempoolWriter {
+
+       -- | Compute the transaction id from a transaction.
+       --
+       -- This is used in the protocol handler to verify a full transaction
+       -- matches a previously given transaction id.
+       --
+       txId          :: tx -> txid,
+
+       -- | Supply a batch of transactions to the mempool. They are either
+       -- accepted or rejected individually, but in the order supplied.
+       --
+       -- The 'txid's of all transactions that were added successfully are
+       -- returned.
+       mempoolAddTxs :: [tx] -> m [SubmitResult failure]
+    }
+
+
+-- | A mempool writer with validation harness
+-- PRECONDITION: no duplicates given to mempoolAddTxs
 --
 getWriter :: forall tx txid ctx failure m.
              ( MonadSTM m
-             , MonadThrow m
+             , Exception failure
+             , MonadThrow (STM m)
              , Ord txid
              , Typeable txid
              , Typeable failure
@@ -134,40 +159,38 @@ getWriter :: forall tx txid ctx failure m.
           => (tx -> txid)
           -- ^ get txid of a tx
           -> m ctx
-          -- ^ monadic validation ctx
-          -> (ctx -> tx -> Either failure ())
-          -- ^ validate a tx, any failing `tx` throws an exception.
-          -> (failure -> Bool)
-          -- ^ return `True` when a failure should throw an exception
+          -- ^ acquire validation context
+          -> ([tx] -> ctx -> [Except failure (Either failure ())])
+          -- ^ validation function
+          -> Maybe failure
+          -- ^ replace duplicates if Just
           -> Mempool m txid tx
-          -> TxSubmissionMempoolWriter txid tx Int m
-getWriter getTxId getValidationCtx validateTx failureFilterFn (Mempool mempool) =
-    TxSubmissionMempoolWriter {
-        txId = getTxId,
-
-        mempoolAddTxs = \txs -> do
-          ctx <- getValidationCtx
-          (invalidTxIds, validTxs) <- atomically $ do
-            MempoolSeq { mempoolSet, mempoolSeq } <- readTVar mempool
-            let (invalidTxIds, validTxs) =
-                    bimap (filter (failureFilterFn . snd))
-                          (nubBy (on (==) getTxId))
-                  . partitionEithers
-                  . map (\tx -> case validateTx ctx tx of
-                                 Left  e -> Left (getTxId tx, e)
-                                 Right _ -> Right tx
-                        )
-                  . filter (\tx -> getTxId tx `Set.notMember` mempoolSet)
-                  $ txs
-                mempoolTxs' = MempoolSeq {
-                    mempoolSet = Foldable.foldl' (\s tx -> getTxId tx `Set.insert` s)
-                                                 mempoolSet
-                                                 validTxs,
-                    mempoolSeq = Foldable.foldl' (Seq.|>) mempoolSeq validTxs
-                  }
-            writeTVar mempool mempoolTxs'
-            return (invalidTxIds, map getTxId validTxs)
-          when (not (null invalidTxIds)) $
-            throwIO (InvalidTxsError invalidTxIds)
-          return validTxs
-      }
+          -> MempoolWriter txid tx failure Int m
+getWriter getTxId acquireCtx validateTxs mDuplicate (Mempool mempool) =
+  MempoolWriter {
+    txId = getTxId,
+
+    mempoolAddTxs = \txs -> do
+      ctx <- acquireCtx
+      let !vTxs = validateTxs txs ctx
+      atomically $ do
+        MempoolSeq { mempoolSet, mempoolSeq } <- readTVar mempool
+        result <- sequence
+          [if duplicate then
+             pure . Left . SubmitFail . fromJust $ mDuplicate
+           else case runExcept vtx of
+             Left e       -> throwSTM e
+             Right result -> pure $ bimap SubmitFail (const (txid, tx)) result
+          | (tx, vtx) <- zip txs vTxs
+          , let txid = getTxId tx
+                duplicate = txid `Set.member` mempoolSet
+          , not (duplicate && isNothing mDuplicate)
+          ]
+        let (validIds, validTxs) = unzip . rights $ result
+            mempoolTxs' = MempoolSeq {
+               mempoolSet = Set.union mempoolSet (Set.fromList validIds),
+               mempoolSeq = Foldable.foldl' (Seq.|>) mempoolSeq validTxs
+            }
+        writeTVar mempool mempoolTxs'
+        return $ fromLeft SubmitSuccess <$> result
+  }