<commit_before>{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE GADTSyntax #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Node.Internal (
    NodeId(..),
    Node(..),
    startNode,
    stopNode,
    sendMsg,
    ChannelIn(..),
    ChannelOut(..),
    connectOutChannel,
    connectInOutChannel,
    closeChannel,
    withOutChannel,
    withInOutChannel,
    writeChannel,
    readChannel,
  ) where

import Data.Binary     as Bin
import qualified Data.ByteString.Lazy as LBS
import qualified Data.ByteString      as BS
import qualified Data.ByteString.Builder as BS
import qualified Data.ByteString.Builder.Extra as BS
import Data.Map (Map)
import qualified Data.Map.Strict as Map
import Data.Monoid
import Data.Typeable
import Control.Exception hiding (bracket, throw)
import qualified Network.Transport.Abstract as NT
import qualified Network.Transport as NT (EventErrorCode(EventConnectionLost, EventEndPointFailed, EventTransportFailed))
import System.Random (StdGen, random)
import Mockable.Class
import Mockable.Concurrent
import Mockable.Exception
import qualified Mockable.Channel as Channel
import Mockable.SharedAtomic


-- A node id wraps a network-transport endpoint address
newtype NodeId = NodeId NT.EndPointAddress
  deriving (Eq, Ord, Show)

data Node (m :: * -> *) = Node {
       nodeEndPoint         :: NT.EndPoint m,
       nodeDispatcherThread :: ThreadId m,
       nodeExpectedIncoming :: SharedAtomicT m (StdGen, Map Nonce (ThreadId m, ChannelIn m))
     }

type Nonce = Word64

data NodeException =
       ProtocolError String
     | InternalError String
  deriving (Show, Typeable)

instance Exception NodeException

newtype ChannelIn m = ChannelIn (Channel.ChannelT m (Maybe BS.ByteString))

newtype ChannelOut m = ChannelOut (NT.Connection m)

startNode :: ( Mockable SharedAtomic m, Mockable Fork m
             , Mockable Channel.Channel m, Mockable Throw m
             , Mockable RunInUnboundThread m )
          => NT.Transport m
          -> StdGen
          -> (NodeId -> ChannelIn m -> m ())
          -> (NodeId -> ChannelIn m -> ChannelOut m -> m ())
          -> m (Node m)
startNode transport prng handlerIn handlerOut = do
    Right endpoint <- NT.newEndPoint transport --TODO: error handling
    incomingVar    <- newSharedAtomic (prng, Map.empty)
    tid  <- fork (nodeDispatcher endpoint incomingVar handlerIn handlerOut)
    --TODO: exceptions in the forkIO
    return Node {
      nodeEndPoint         = endpoint,
      nodeDispatcherThread = tid,
      nodeExpectedIncoming = incomingVar
    }


stopNode :: Node m -> m ()
stopNode Node {..} =
    NT.closeEndPoint nodeEndPoint
    -- This eventually will shut down the dispatcher thread, which in turn
    -- ought to stop the connection handling threads.
    -- It'll also close all TCP connections.


type DispatcherState m = Map NT.ConnectionId (ConnectionState m)
data ConnectionState m =

       -- We got a new connection and are waiting on the first chunk of data
       ConnectionNew !NodeId

       -- We got the first chunk of data, we're now waiting either for more
       -- data or for the connection to be closed. This supports the small
       -- message optimisation.
     | ConnectionNewChunks !NodeId ![BS.ByteString]

       -- We've forked a thread to handle the message. The connection is still
       -- open and data is still arriving. We have a channel to pass the
       -- incoming chunks off to the other thread.
     | ConnectionReceiving !(ThreadId m) !(Channel.ChannelT m (Maybe BS.ByteString))

       -- We've forked a thread to handle the message. The connection is now
       -- closed and we have all the data already, but the thread we forked
       -- to handle it is still active.
     | ConnectionClosed !(ThreadId m)

--TODO: extend this to keep track of the number of active threads and total
-- amount of in flight incoming data. This will be needed to inform the
-- back-pressure policy.

--TODO: need to fill in how threads are cleaned up. The 'finally' handler for
-- the forked thread will post a cleanup event to the dispatcher thread which
-- will eventually get it and remove the ConnectionId from the DispatcherState
-- and update any stats (ie decrement counts of resources in use)

-- | The one thread that handles /all/ incoming messages and dispatches them
-- to various handlers.
--
nodeDispatcher :: forall m g .
                  ( Mockable SharedAtomic m, Mockable Fork m
                  , Mockable Channel.Channel m, Mockable RunInUnboundThread m
                  , Mockable Throw m )
               => NT.EndPoint m
               -> SharedAtomicT m (g, Map Nonce (ThreadId m, ChannelIn m))
               -> (NodeId -> ChannelIn m -> m ())
               -> (NodeId -> ChannelIn m -> ChannelOut m -> m ())
               -> m ()
nodeDispatcher endpoint incomingVar handlerIn handlerInOut =
    runInUnboundThread $ loop Map.empty
  where
    loop :: DispatcherState m -> m ()
--    loop state
--      | overloaded = do ... TODO

    loop !state = do
      event <- NT.receive endpoint
      case event of

        NT.ConnectionOpened connid NT.ReliableOrdered peer ->
          -- Just keep track of the new connection, nothing else
          loop (Map.insert connid (ConnectionNew (NodeId peer)) state)

        -- receiving data for an existing connection (ie a multi-chunk message)
        NT.Received connid chunks ->
          case Map.lookup connid state of
            Nothing ->
              throw (InternalError "received data on unknown connection")

            -- TODO: need policy here on queue size<commit_msg>Special g -> StdGen in dispatcher<commit_after>{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE GADTSyntax #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Node.Internal (
    NodeId(..),
    Node(..),
    startNode,
    stopNode,
    sendMsg,
    ChannelIn(..),
    ChannelOut(..),
    connectOutChannel,
    connectInOutChannel,
    closeChannel,
    withOutChannel,
    withInOutChannel,
    writeChannel,
    readChannel,
  ) where

import Data.Binary     as Bin
import qualified Data.ByteString.Lazy as LBS
import qualified Data.ByteString      as BS
import qualified Data.ByteString.Builder as BS
import qualified Data.ByteString.Builder.Extra as BS
import Data.Map (Map)
import qualified Data.Map.Strict as Map
import Data.Monoid
import Data.Typeable
import Control.Exception hiding (bracket, throw)
import qualified Network.Transport.Abstract as NT
import qualified Network.Transport as NT (EventErrorCode(EventConnectionLost, EventEndPointFailed, EventTransportFailed))
import System.Random (StdGen, random)
import Mockable.Class
import Mockable.Concurrent
import Mockable.Exception
import qualified Mockable.Channel as Channel
import Mockable.SharedAtomic


-- A node id wraps a network-transport endpoint address
newtype NodeId = NodeId NT.EndPointAddress
  deriving (Eq, Ord, Show)

data Node (m :: * -> *) = Node {
       nodeEndPoint         :: NT.EndPoint m,
       nodeDispatcherThread :: ThreadId m,
       nodeExpectedIncoming :: SharedAtomicT m (StdGen, Map Nonce (ThreadId m, ChannelIn m))
     }

type Nonce = Word64

data NodeException =
       ProtocolError String
     | InternalError String
  deriving (Show, Typeable)

instance Exception NodeException

newtype ChannelIn m = ChannelIn (Channel.ChannelT m (Maybe BS.ByteString))

newtype ChannelOut m = ChannelOut (NT.Connection m)

startNode :: ( Mockable SharedAtomic m, Mockable Fork m
             , Mockable Channel.Channel m, Mockable Throw m
             , Mockable RunInUnboundThread m )
          => NT.Transport m
          -> StdGen
          -> (NodeId -> ChannelIn m -> m ())
          -> (NodeId -> ChannelIn m -> ChannelOut m -> m ())
          -> m (Node m)
startNode transport prng handlerIn handlerOut = do
    Right endpoint <- NT.newEndPoint transport --TODO: error handling
    incomingVar    <- newSharedAtomic (prng, Map.empty)
    tid  <- fork (nodeDispatcher endpoint incomingVar handlerIn handlerOut)
    --TODO: exceptions in the forkIO
    return Node {
      nodeEndPoint         = endpoint,
      nodeDispatcherThread = tid,
      nodeExpectedIncoming = incomingVar
    }


stopNode :: Node m -> m ()
stopNode Node {..} =
    NT.closeEndPoint nodeEndPoint
    -- This eventually will shut down the dispatcher thread, which in turn
    -- ought to stop the connection handling threads.
    -- It'll also close all TCP connections.


type DispatcherState m = Map NT.ConnectionId (ConnectionState m)
data ConnectionState m =

       -- We got a new connection and are waiting on the first chunk of data
       ConnectionNew !NodeId

       -- We got the first chunk of data, we're now waiting either for more
       -- data or for the connection to be closed. This supports the small
       -- message optimisation.
     | ConnectionNewChunks !NodeId ![BS.ByteString]

       -- We've forked a thread to handle the message. The connection is still
       -- open and data is still arriving. We have a channel to pass the
       -- incoming chunks off to the other thread.
     | ConnectionReceiving !(ThreadId m) !(Channel.ChannelT m (Maybe BS.ByteString))

       -- We've forked a thread to handle the message. The connection is now
       -- closed and we have all the data already, but the thread we forked
       -- to handle it is still active.
     | ConnectionClosed !(ThreadId m)

--TODO: extend this to keep track of the number of active threads and total
-- amount of in flight incoming data. This will be needed to inform the
-- back-pressure policy.

--TODO: need to fill in how threads are cleaned up. The 'finally' handler for
-- the forked thread will post a cleanup event to the dispatcher thread which
-- will eventually get it and remove the ConnectionId from the DispatcherState
-- and update any stats (ie decrement counts of resources in use)

-- | The one thread that handles /all/ incoming messages and dispatches them
-- to various handlers.
--
nodeDispatcher :: forall m .
                  ( Mockable SharedAtomic m, Mockable Fork m
                  , Mockable Channel.Channel m, Mockable RunInUnboundThread m
                  , Mockable Throw m )
               => NT.EndPoint m
               -> SharedAtomicT m (StdGen, Map Nonce (ThreadId m, ChannelIn m))
               -> (NodeId -> ChannelIn m -> m ())
               -> (NodeId -> ChannelIn m -> ChannelOut m -> m ())
               -> m ()
nodeDispatcher endpoint incomingVar handlerIn handlerInOut =
    runInUnboundThread $ loop Map.empty
  where
    loop :: DispatcherState m -> m ()
--    loop state
--      | overloaded = do ... TODO

    loop !state = do
      event <- NT.receive endpoint
      case event of

        NT.ConnectionOpened connid NT.ReliableOrdered peer ->
          -- Just keep track of the new connection, nothing else
          loop (Map.insert connid (ConnectionNew (NodeId peer)) state)

        -- receiving data for an existing connection (ie a multi-chunk message)
        NT.Received connid chunks ->
          case Map.lookup connid state of
            Nothing ->
              throw (InternalError "received data on unknown connection")

            -- TODO: need policy here on queue size
{'fraction_non_alphanumeric': 0.08008266597778352, 'fraction_numerical': 0.0003444415740979936, 'mean_word_length': 2.6940203562340965, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 56, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5246466', 'n_tokens_mistral': 3183, 'n_tokens_neox': 3098, 'n_words': 1232}