network-mux: provide a way to run bind a mini-protocol to a capability

Co-authored-by: Karl Knutsson (@karknu)
Co-authored-by: Marcin Szamotulski (@coot)

network-mux/demo/mux-demo.hs

@@ -137,9 +137,10 @@ serverWorker bearer = do
   where
     ptcls :: [MiniProtocolInfo ResponderMode]
     ptcls = [ MiniProtocolInfo {
-                miniProtocolNum    = MiniProtocolNum 2,
-                miniProtocolDir    = ResponderDirectionOnly,
-                miniProtocolLimits = defaultProtocolLimits
+                miniProtocolNum        = MiniProtocolNum 2,
+                miniProtocolDir        = ResponderDirectionOnly,
+                miniProtocolLimits     = defaultProtocolLimits,
+                miniProtocolCapability = Nothing
               }
             ]
 
@@ -196,9 +197,10 @@ clientWorker bearer n msg = do
   where
     ptcls :: [MiniProtocolInfo Mx.InitiatorMode]
     ptcls = [ MiniProtocolInfo {
-                miniProtocolNum    = MiniProtocolNum 2,
-                miniProtocolDir    = InitiatorDirectionOnly,
-                miniProtocolLimits = defaultProtocolLimits
+                miniProtocolNum        = MiniProtocolNum 2,
+                miniProtocolDir        = InitiatorDirectionOnly,
+                miniProtocolLimits     = defaultProtocolLimits,
+                miniProtocolCapability = Nothing
               }
             ]
 

network-mux/src/Control/Concurrent/JobPool.hs

@@ -12,6 +12,7 @@ module Control.Concurrent.JobPool
   , Job (..)
   , withJobPool
   , forkJob
+  , forkJobOn
   , readSize
   , readGroupSize
   , waitForJob
@@ -29,6 +30,7 @@ import Control.Monad.Class.MonadAsync
 import Control.Monad.Class.MonadFork (MonadThread (..))
 import Control.Monad.Class.MonadThrow
 
+
 -- | JobPool allows to submit asynchronous jobs, wait for their completion or
 -- cancel.  Jobs are grouped, each group can be cancelled separately.
 --
@@ -69,16 +71,19 @@ withJobPool =
       jobs <- readTVarIO jobsVar
       mapM_ uninterruptibleCancel jobs
 
-forkJob :: forall group m a.
-           ( MonadAsync m, MonadMask m
-           , Ord group
-           )
-        => JobPool group m a
-        -> Job     group m a
-        -> m ()
-forkJob JobPool{jobsVar, completionQueue} (Job action handler group label) =
+
+forkJob' :: forall group m a.
+            ( MonadAsync m, MonadMask m
+            , Ord group
+            )
+         => (m () -> m (Async m ()))
+         -- ^ how to fork a thread, e.g. `async`, `asyncOn`.
+         -> JobPool group m a
+         -> Job     group m a
+         -> m ()
+forkJob' doFork JobPool{jobsVar, completionQueue} (Job action handler group label) =
     mask $ \restore -> do
-      jobAsync <- async $ do
+      jobAsync <- doFork $ do
         tid <- myThreadId
         io tid restore
           `onException`
@@ -104,6 +109,33 @@ forkJob JobPool{jobsVar, completionQueue} (Job action handler group label) =
               restore action
       atomically $ writeTQueue completionQueue res
 
+
+
+-- | Fork a `Job` using `async`.
+--
+forkJob :: forall group m a.
+           ( MonadAsync m, MonadMask m
+           , Ord group
+           )
+        => JobPool group m a
+        -> Job     group m a
+        -> m ()
+forkJob = forkJob' async
+
+
+-- | Fork a `Job` using `asyncOn`.
+--
+forkJobOn :: forall group m a.
+             ( MonadAsync m, MonadMask m
+             , Ord group
+             )
+          => Int
+          -> JobPool group m a
+          -> Job     group m a
+          -> m ()
+forkJobOn cap = forkJob' (asyncOn cap)
+
+
 readSize :: MonadSTM m => JobPool group m a -> STM m Int
 readSize JobPool{jobsVar} = Map.size <$> readTVar jobsVar
 

network-mux/src/Network/Mux.hs

@@ -201,7 +201,7 @@ data Group = MuxJob
 -- * at any given time each @TranslocationServiceRequest@ contains a non-empty
 -- 'Wanton'
 --
-run :: forall m mode.
+run :: forall m (mode :: Mode).
        ( MonadAsync m
        , MonadFork m
        , MonadLabelledSTM m
@@ -214,7 +214,12 @@ run :: forall m mode.
     -> Mux mode m
     -> Bearer m
     -> m ()
-run tracer Mux {muxMiniProtocols, muxControlCmdQueue, muxStatus} bearer@Bearer {name} = do
+run tracer
+    Mux { muxMiniProtocols,
+          muxControlCmdQueue,
+          muxStatus
+        }
+    bearer@Bearer{name} = do
     egressQueue <- atomically $ newTBQueue 100
 
     -- label shared variables
@@ -459,18 +464,28 @@ monitor tracer timeout jobpool egressQueue cmdQueue muxStatus =
                            ptclState@MiniProtocolState {
                              miniProtocolInfo = MiniProtocolInfo {
                                miniProtocolNum,
-                               miniProtocolDir
+                               miniProtocolDir,
+                               miniProtocolCapability
                              }
                            }
                            ptclAction) -> do
           traceWith tracer (TraceStartEagerly miniProtocolNum
-                             (protocolDirEnum miniProtocolDir))
-          JobPool.forkJob jobpool $
-            miniProtocolJob
-              tracer
-              egressQueue
-              ptclState
-              ptclAction
+                                              (protocolDirEnum miniProtocolDir))
+          case miniProtocolCapability of
+            Nothing ->
+              JobPool.forkJob jobpool $
+                miniProtocolJob
+                  tracer
+                  egressQueue
+                  ptclState
+                  ptclAction
+            Just cap ->
+              JobPool.forkJobOn cap jobpool $
+                miniProtocolJob
+                  tracer
+                  egressQueue
+                  ptclState
+                  ptclAction
           go monitorCtx
 
         EventControlCmd (CmdStartProtocolThread
@@ -554,20 +569,30 @@ monitor tracer timeout jobpool egressQueue cmdQueue muxStatus =
     doStartOnDemand ptclState@MiniProtocolState {
                       miniProtocolInfo = MiniProtocolInfo {
                            miniProtocolNum,
-                           miniProtocolDir
+                           miniProtocolDir,
+                           miniProtocolCapability
                       },
                       miniProtocolStatusVar
                     }
                     ptclAction = do
       traceWith tracer (TraceStartedOnDemand miniProtocolNum
-                         (protocolDirEnum miniProtocolDir))
+                                             (protocolDirEnum miniProtocolDir))
       atomically $ modifyTVar miniProtocolStatusVar (\a -> assert (a /= StatusRunning) StatusRunning)
-      JobPool.forkJob jobpool $
-        miniProtocolJob
-          tracer
-          egressQueue
-          ptclState
-          ptclAction
+      case miniProtocolCapability of
+        Nothing ->
+          JobPool.forkJob jobpool $
+            miniProtocolJob
+              tracer
+              egressQueue
+              ptclState
+              ptclAction
+        Just cap ->
+          JobPool.forkJobOn cap jobpool $
+            miniProtocolJob
+              tracer
+              egressQueue
+              ptclState
+              ptclAction
 
     checkNonEmptyQueue :: IngressQueue m -> STM m ()
     checkNonEmptyQueue q = do

network-mux/src/Network/Mux/Types.hs

@@ -127,12 +127,14 @@ type family HasResponder (mode :: Mode) :: Bool where
 --
 data MiniProtocolInfo (mode :: Mode) =
      MiniProtocolInfo {
-       miniProtocolNum    :: !MiniProtocolNum,
+       miniProtocolNum        :: !MiniProtocolNum,
        -- ^ Unique mini-protocol number.
-       miniProtocolDir    :: !(MiniProtocolDirection mode),
+       miniProtocolDir        :: !(MiniProtocolDirection mode),
        -- ^ Mini-protocol direction.
-       miniProtocolLimits :: !MiniProtocolLimits
+       miniProtocolLimits     :: !MiniProtocolLimits,
        -- ^ ingress queue limits for the protocol
+       miniProtocolCapability :: !(Maybe Int)
+       -- ^ capability on which the mini-protocol should run
      }
 
 data MiniProtocolDirection (mode :: Mode) where