Used Script of Script for BearerInfo Attenuation

- Modified Snocket's connect to use one script per connection

ouroboros-network-framework/src/Simulation/Network/Snocket.hs

@@ -1,4 +1,3 @@
-{-# LANGUAGE BangPatterns        #-}
 {-# LANGUAGE DeriveTraversable   #-}
 {-# LANGUAGE DerivingStrategies  #-}
 {-# LANGUAGE FlexibleContexts    #-}
@@ -9,6 +8,8 @@
 {-# LANGUAGE NamedFieldPuns      #-}
 {-# LANGUAGE RankNTypes          #-}
 {-# LANGUAGE ScopedTypeVariables #-}
+{-# OPTIONS_GHC -Wno-deferred-out-of-scope-variables #-}
+{-# LANGUAGE TypeApplications #-}
 
 -- | This module provides simulation environment and a snocket implementation
 -- suitable for 'IOSim'.
@@ -25,8 +26,10 @@ module Simulation.Network.Snocket
     withSnocket
   , ObservableNetworkState (..)
   , ResourceException (..)
+  , SDUSize
+  , Script (..)
+  , Size
   , SnocketTrace (..)
-  , TimeoutDetail (..)
   , SockType (..)
   , OpenType (..)
 
@@ -35,10 +38,9 @@ module Simulation.Network.Snocket
   , IOErrType (..)
   , IOErrThrowOrReturn (..)
   , SuccessOrFailure (..)
-  , Size
+  , TimeoutDetail (..)
   , noAttenuation
   , FD
-  , SDUSize
 
   , GlobalAddressScheme (..)
   , AddressType (..)
@@ -78,7 +80,7 @@ import           Ouroboros.Network.ConnectionManager.Types (AddressType (..))
 import           Ouroboros.Network.Snocket
 
 import           Ouroboros.Network.Testing.Data.Script
-                  (Script(..), stepScriptSTM, initScript)
+                  (Script(..), initScript, stepScriptSTM, stepScriptSTMTx, stepScript)
 
 data Connection m addr = Connection
     { -- | Attenuated channels of a connection.
@@ -199,13 +201,23 @@ data NetworkState m addr = NetworkState {
 
       -- | Registry of active connections.
       --
-      nsConnections       :: StrictTVar m (Map (NormalisedId addr) (Connection m addr)),
+      nsConnections       :: StrictTVar
+                              m
+                              (Map (NormalisedId addr) (Connection m addr)),
 
       -- | Get an unused ephemeral address.
       --
       nsNextEphemeralAddr :: AddressType -> STM m addr,
 
-      nsBearerInfo        :: LazySTM.TVar m (Script BearerInfo)
+      nsBearerInfo        :: LazySTM.TVar
+                              m
+                              (Script (LazySTM.STM m (LazySTM.TVar m (Script BearerInfo)))),
+
+      -- | Get the BearerInfo Script for a given connection.
+      --
+      nsAttenuationMap    :: StrictTVar
+                              m (Map (ConnectionId addr)
+                                     (LazySTM.TVar m (Script BearerInfo)))
 
     }
 
@@ -312,17 +324,18 @@ newNetworkState
        ( MonadLabelledSTM m
        , GlobalAddressScheme peerAddr
        )
-    => Script BearerInfo
+    => Script (Script BearerInfo)
     -- ^ the largest ephemeral address
     -> m (NetworkState m (TestAddress peerAddr))
-newNetworkState bearerInfoScript = atomically $ do
+newNetworkState script = atomically $ do
   (v :: StrictTVar m Natural) <- newTVar 0
   let nextEphemeralAddr :: AddressType -> STM m (TestAddress peerAddr)
       nextEphemeralAddr addrType = do
         -- TODO: we should use `(\s -> (succ s, s)` but p2p-master does not
         -- include PR #3172.
          a <- stateTVar v (\s -> let s' = succ s in (s', s'))
          return (ephemeralAddress addrType a)
+
   s <- NetworkState
     -- nsListeningFDs
     <$> newTVar Map.empty
@@ -331,7 +344,10 @@ newNetworkState bearerInfoScript = atomically $ do
     -- nsNextEphemeralAddr
     <*> pure nextEphemeralAddr
     -- nsBearerInfo
-    <*> initScript bearerInfoScript
+    <*> LazySTM.newTVar (initScript <$> script)
+    -- attenuationMap
+    <*> newTVar Map.empty
+
   labelTVar (nsListeningFDs s)   "nsListeningFDs"
   labelTVar (nsConnections s)    "nsConnections"
   return s
@@ -387,7 +403,7 @@ withSnocket
        )
     => Tracer m (WithAddr (TestAddress peerAddr)
                           (SnocketTrace m (TestAddress peerAddr)))
-    -> Script BearerInfo
+    -> Script (Script BearerInfo)
     -> (Snocket m (FD m (TestAddress peerAddr)) (TestAddress peerAddr)
         -> m (ObservableNetworkState (TestAddress peerAddr))
         -> m a)
@@ -721,13 +737,23 @@ mkSnocket state tr = Snocket { getLocalAddr
           -- accepted.
           FDUninitialised mbLocalAddr -> mask $ \unmask -> do
             (connId, bearerInfo, simOpen) <- atomically $ do
-              bearerInfo <- stepScriptSTM (nsBearerInfo state)
               localAddress <-
                 case mbLocalAddr of
                   Just addr -> return addr
                   Nothing   -> nsNextEphemeralAddr state (getAddressType remoteAddress)
               let connId = ConnectionId { localAddress, remoteAddress }
 
+              attenuationMap <- readTVar (nsAttenuationMap state)
+
+              bearerInfo <- case Map.lookup connId attenuationMap of
+                Nothing     -> do
+                  script <- stepScriptSTMTx (nsBearerInfo state)
+                  writeTVar (nsAttenuationMap state)
+                            (Map.insert connId script attenuationMap)
+                  stepScriptSTM script
+
+                Just script -> stepScriptSTM script
+
               connMap <- readTVar (nsConnections state)
               case Map.lookup (normaliseId connId) connMap of
                 Just      Connection { connState = ESTABLISHED } ->
@@ -997,7 +1023,9 @@ mkSnocket state tr = Snocket { getLocalAddr
            -> m (Accept m (FD m (TestAddress addr))
                                 (TestAddress addr))
     accept FD { fdVar } = do time <- getMonotonicTime
-                             deltaAndIOErr <- biAcceptFailures <$> atomically (stepScriptSTM $ nsBearerInfo state)
+                             script <- atomically . stepScriptSTMTx $ nsBearerInfo state
+                             bearerInfo <- stepScript script
+                             let deltaAndIOErr = biAcceptFailures bearerInfo
                              return $ accept_ time deltaAndIOErr
       where
         -- non-blocking; return 'True' if a connection is in 'SYN_SENT' state