{-# LANGUAGE BlockArguments           #-}
{-# LANGUAGE CPP                      #-}
{-# LANGUAGE DisambiguateRecordFields #-}
{-# LANGUAGE FlexibleContexts         #-}
{-# LANGUAGE NamedFieldPuns           #-}
{-# LANGUAGE RankNTypes               #-}
{-# LANGUAGE ScopedTypeVariables      #-}

module Ouroboros.Network.PeerSelection.PeerSelectionActions
  ( PeerSelectionActions (..)
  , withPeerSelectionActions
  , requestPeerSharingResult
  , requestPublicRootPeers
  ) where


import Control.Applicative (Alternative)
import Control.Concurrent.Class.MonadMVar (MonadMVar (..))
import Control.Concurrent.Class.MonadSTM.Strict
import Control.Monad.Class.MonadAsync
import Control.Monad.Class.MonadFork
import Control.Monad.Class.MonadThrow
import Control.Monad.Class.MonadTime.SI
import Control.Monad.Class.MonadTimer.SI
import Control.Tracer (Tracer)

import Data.Bifunctor (first)
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Set (Set)
import Data.Void (Void)

import Network.DNS qualified as DNS

import Ouroboros.Network.PeerSelection.Governor.Types
           (PeerSelectionActions (PeerSelectionActions, readLocalRootPeersFromFile))
import Ouroboros.Network.PeerSelection.LedgerPeers hiding (getLedgerPeers)
import Ouroboros.Network.PeerSelection.PeerAdvertise (PeerAdvertise)
import Ouroboros.Network.PeerSelection.PublicRootPeers (PublicRootPeers)
import Ouroboros.Network.PeerSelection.PublicRootPeers qualified as PublicRootPeers
import Ouroboros.Network.PeerSelection.RootPeersDNS
import Ouroboros.Network.PeerSelection.State.LocalRootPeers
import Ouroboros.Network.PeerSharing (PeerSharingController,
           PeerSharingResult (..), requestPeers)
import Ouroboros.Network.Protocol.PeerSharing.Type (PeerSharingAmount (..))

import System.Random

withPeerSelectionActions
  :: forall extraState extraFlags extraPeers extraAPI extraCounters peeraddr peerconn resolver exception m a.
     ( Alternative (STM m)
     , MonadAsync m
     , MonadDelay m
     , MonadThrow m
     , Ord peeraddr
     , Exception exception
     , Eq extraFlags
     )
  => Tracer m (TraceLocalRootPeers extraFlags peeraddr exception)
  -> StrictTVar m (Config extraFlags peeraddr)
  -> PeerActionsDNS peeraddr resolver exception m
  -> ( (NumberOfPeers -> LedgerPeersKind -> m (Maybe (Set peeraddr, DiffTime)))
     -> PeerSelectionActions extraState extraFlags extraPeers extraAPI extraCounters peeraddr peerconn m)
  -> WithLedgerPeersArgs extraAPI m
  -> StdGen
  -> (   (Async m Void, Async m Void)
      -> PeerSelectionActions extraState extraFlags extraPeers extraAPI extraCounters peeraddr peerconn m
      -> m a)
  -- ^ continuation, receives a handle to the local roots peer provider thread
  -- (only if local root peers were non-empty).
  -> m a
withPeerSelectionActions :: forall extraState extraFlags extraPeers extraAPI extraCounters
       peeraddr peerconn resolver exception (m :: * -> *) a.
(Alternative (STM m), MonadAsync m, MonadDelay m, MonadThrow m,
 Ord peeraddr, Exception exception, Eq extraFlags) =>
Tracer m (TraceLocalRootPeers extraFlags peeraddr exception)
-> StrictTVar m (Config extraFlags peeraddr)
-> PeerActionsDNS peeraddr resolver exception m
-> ((NumberOfPeers
     -> LedgerPeersKind -> m (Maybe (Set peeraddr, DiffTime)))
    -> PeerSelectionActions
         extraState
         extraFlags
         extraPeers
         extraAPI
         extraCounters
         peeraddr
         peerconn
         m)
-> WithLedgerPeersArgs extraAPI m
-> StdGen
-> ((Async m Void, Async m Void)
    -> PeerSelectionActions
         extraState
         extraFlags
         extraPeers
         extraAPI
         extraCounters
         peeraddr
         peerconn
         m
    -> m a)
-> m a
withPeerSelectionActions
  Tracer m (TraceLocalRootPeers extraFlags peeraddr exception)
localTracer
  StrictTVar m (Config extraFlags peeraddr)
localRootsVar
  PeerActionsDNS peeraddr resolver exception m
peerActionsDNS
  (NumberOfPeers
 -> LedgerPeersKind -> m (Maybe (Set peeraddr, DiffTime)))
-> PeerSelectionActions
     extraState
     extraFlags
     extraPeers
     extraAPI
     extraCounters
     peeraddr
     peerconn
     m
getPeerSelectionActions
  WithLedgerPeersArgs extraAPI m
ledgerPeersArgs
  StdGen
rng0
  (Async m Void, Async m Void)
-> PeerSelectionActions
     extraState
     extraFlags
     extraPeers
     extraAPI
     extraCounters
     peeraddr
     peerconn
     m
-> m a
k = do
    PeerActionsDNS peeraddr resolver exception m
-> WithLedgerPeersArgs extraAPI m
-> ((NumberOfPeers
     -> LedgerPeersKind -> m (Maybe (Set peeraddr, DiffTime)))
    -> Async m Void -> m a)
-> m a
forall peerAddr resolver exception extraAPI (m :: * -> *) a.
(MonadAsync m, MonadThrow m, MonadMonotonicTime m,
 Exception exception, Ord peerAddr) =>
PeerActionsDNS peerAddr resolver exception m
-> WithLedgerPeersArgs extraAPI m
-> ((NumberOfPeers
     -> LedgerPeersKind -> m (Maybe (Set peerAddr, DiffTime)))
    -> Async m Void -> m a)
-> m a
withLedgerPeers
      PeerActionsDNS peeraddr resolver exception m
peerActionsDNS
      WithLedgerPeersArgs extraAPI m
ledgerPeersArgs
      (\NumberOfPeers
-> LedgerPeersKind -> m (Maybe (Set peeraddr, DiffTime))
getLedgerPeers Async m Void
lpThread -> do
          let peerSelectionActions :: PeerSelectionActions
  extraState
  extraFlags
  extraPeers
  extraAPI
  extraCounters
  peeraddr
  peerconn
  m
peerSelectionActions@PeerSelectionActions
                { STM m (Config extraFlags RelayAccessPoint)
readLocalRootPeersFromFile :: forall extraState extraFlags extraPeers extraAPI extraCounters
       peeraddr peerconn (m :: * -> *).
PeerSelectionActions
  extraState
  extraFlags
  extraPeers
  extraAPI
  extraCounters
  peeraddr
  peerconn
  m
-> STM m (Config extraFlags RelayAccessPoint)
readLocalRootPeersFromFile :: STM m (Config extraFlags RelayAccessPoint)
readLocalRootPeersFromFile
                } = (NumberOfPeers
 -> LedgerPeersKind -> m (Maybe (Set peeraddr, DiffTime)))
-> PeerSelectionActions
     extraState
     extraFlags
     extraPeers
     extraAPI
     extraCounters
     peeraddr
     peerconn
     m
getPeerSelectionActions NumberOfPeers
-> LedgerPeersKind -> m (Maybe (Set peeraddr, DiffTime))
getLedgerPeers
          m Void -> (Async m Void -> m a) -> m a
forall a b. m a -> (Async m a -> m b) -> m b
forall (m :: * -> *) a b.
MonadAsync m =>
m a -> (Async m a -> m b) -> m b
withAsync do
              String -> m ()
forall (m :: * -> *). MonadThread m => String -> m ()
labelThisThread String
"local-roots-peers"
              Tracer m (TraceLocalRootPeers extraFlags peeraddr exception)
-> PeerActionsDNS peeraddr resolver exception m
-> ResolvConf
-> StdGen
-> STM m (Config extraFlags RelayAccessPoint)
-> StrictTVar m (Config extraFlags peeraddr)
-> m Void
forall (m :: * -> *) extraFlags peerAddr resolver exception.
(Alternative (STM m), MonadAsync m, MonadDelay m, MonadThrow m,
 Ord peerAddr, Eq extraFlags) =>
Tracer m (TraceLocalRootPeers extraFlags peerAddr exception)
-> PeerActionsDNS peerAddr resolver exception m
-> ResolvConf
-> StdGen
-> STM
     m
     [(HotValency, WarmValency,
       Map RelayAccessPoint (LocalRootConfig extraFlags))]
-> StrictTVar
     m
     [(HotValency, WarmValency,
       Map peerAddr (LocalRootConfig extraFlags))]
-> m Void
localRootPeersProvider
                Tracer m (TraceLocalRootPeers extraFlags peeraddr exception)
localTracer
                PeerActionsDNS peeraddr resolver exception m
peerActionsDNS
                -- NOTE: we don't set `resolvConcurrent` because
                -- of https://github.com/kazu-yamamoto/dns/issues/174
                ResolvConf
DNS.defaultResolvConf
                StdGen
rng0
                STM m (Config extraFlags RelayAccessPoint)
readLocalRootPeersFromFile
                StrictTVar m (Config extraFlags peeraddr)
localRootsVar
            (\Async m Void
lrppThread -> (Async m Void, Async m Void)
-> PeerSelectionActions
     extraState
     extraFlags
     extraPeers
     extraAPI
     extraCounters
     peeraddr
     peerconn
     m
-> m a
k (Async m Void
lpThread, Async m Void
lrppThread) PeerSelectionActions
  extraState
  extraFlags
  extraPeers
  extraAPI
  extraCounters
  peeraddr
  peerconn
  m
peerSelectionActions))

requestPeerSharingResult :: ( MonadSTM m
                            , MonadMVar m
                            , Ord peeraddr
                            )
                         => STM m (Map peeraddr (PeerSharingController peeraddr m))
                         -> PeerSharingAmount
                         -> peeraddr
                         -> m (PeerSharingResult peeraddr)
requestPeerSharingResult :: forall (m :: * -> *) peeraddr.
(MonadSTM m, MonadMVar m, Ord peeraddr) =>
STM m (Map peeraddr (PeerSharingController peeraddr m))
-> PeerSharingAmount -> peeraddr -> m (PeerSharingResult peeraddr)
requestPeerSharingResult STM m (Map peeraddr (PeerSharingController peeraddr m))
sharingController PeerSharingAmount
amount peeraddr
peer = do
  controllerMap <- STM m (Map peeraddr (PeerSharingController peeraddr m))
-> m (Map peeraddr (PeerSharingController peeraddr m))
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically STM m (Map peeraddr (PeerSharingController peeraddr m))
sharingController
  case Map.lookup peer controllerMap of
    -- Peer Registering happens asynchronously with respect to
    -- requestPeerShare. This means that there's a possible race where the
    -- Peer Selection Governor can decide to peer share request to a peer
    -- for the peer is registered. When this happens this map lookup is
    -- going to fail, so instead of erroring we report this to the governor
    -- so it can deal with this particular case accordingly.
    Maybe (PeerSharingController peeraddr m)
Nothing -> PeerSharingResult peeraddr -> m (PeerSharingResult peeraddr)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return PeerSharingResult peeraddr
forall peerAddress. PeerSharingResult peerAddress
PeerSharingNotRegisteredYet
    Just PeerSharingController peeraddr m
psController ->
      [peeraddr] -> PeerSharingResult peeraddr
forall peerAddress. [peerAddress] -> PeerSharingResult peerAddress
PeerSharingResult ([peeraddr] -> PeerSharingResult peeraddr)
-> m [peeraddr] -> m (PeerSharingResult peeraddr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PeerSharingController peeraddr m
-> PeerSharingAmount -> m [peeraddr]
forall (m :: * -> *) peer.
(MonadMVar m, MonadSTM m) =>
PeerSharingController peer m -> PeerSharingAmount -> m [peer]
requestPeers PeerSharingController peeraddr m
psController PeerSharingAmount
amount


-- | Retrieves public root peers
--
-- This function attempts to fetch ledger peers of a specified kind. If no ledger peers
-- are found, it retrieves extra peers instead. The result includes the
-- public root peers and the time taken for the operation.
--
requestPublicRootPeers
  :: forall m peeraddr extraPeers resolver exception .
    ( MonadThrow m
    , MonadAsync m
    , Exception exception
    , Monoid extraPeers
    , Ord peeraddr
    )
  => Tracer m TracePublicRootPeers
  -> STM m (Map RelayAccessPoint PeerAdvertise)
  -> PeerActionsDNS peeraddr resolver exception m
  -> DNSSemaphore m
  -> (Map peeraddr PeerAdvertise -> extraPeers)
  -- ^ Function to convert DNS result into extra peers
  -> (NumberOfPeers -> LedgerPeersKind -> m (Maybe (Set peeraddr, DiffTime)))
  -> LedgerPeersKind
  -> StdGen
  -> Int
  -> m (PublicRootPeers extraPeers peeraddr, DiffTime)
requestPublicRootPeers :: forall (m :: * -> *) peeraddr extraPeers resolver exception.
(MonadThrow m, MonadAsync m, Exception exception,
 Monoid extraPeers, Ord peeraddr) =>
Tracer m TracePublicRootPeers
-> STM m (Map RelayAccessPoint PeerAdvertise)
-> PeerActionsDNS peeraddr resolver exception m
-> DNSSemaphore m
-> (Map peeraddr PeerAdvertise -> extraPeers)
-> (NumberOfPeers
    -> LedgerPeersKind -> m (Maybe (Set peeraddr, DiffTime)))
-> LedgerPeersKind
-> StdGen
-> Int
-> m (PublicRootPeers extraPeers peeraddr, DiffTime)
requestPublicRootPeers
  Tracer m TracePublicRootPeers
publicTracer STM m (Map RelayAccessPoint PeerAdvertise)
readPublicRootPeers
  PeerActionsDNS { paToPeerAddr :: forall peeraddr resolver exception (m :: * -> *).
PeerActionsDNS peeraddr resolver exception m
-> IP -> PortNumber -> peeraddr
paToPeerAddr = IP -> PortNumber -> peeraddr
toPeerAddr
                 , paDnsActions :: forall peeraddr resolver exception (m :: * -> *).
PeerActionsDNS peeraddr resolver exception m
-> DNSActions peeraddr resolver exception m
paDnsActions = DNSActions peeraddr resolver exception m
dnsActions
                 }
  DNSSemaphore m
dnsSemaphore Map peeraddr PeerAdvertise -> extraPeers
toExtraPeers NumberOfPeers
-> LedgerPeersKind -> m (Maybe (Set peeraddr, DiffTime))
getLedgerPeers LedgerPeersKind
ledgerPeersKind StdGen
rng Int
n = do
  mbLedgerPeers <- NumberOfPeers
-> LedgerPeersKind -> m (Maybe (Set peeraddr, DiffTime))
getLedgerPeers (Word16 -> NumberOfPeers
NumberOfPeers (Word16 -> NumberOfPeers) -> Word16 -> NumberOfPeers
forall a b. (a -> b) -> a -> b
$ Int -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n)
                                   LedgerPeersKind
ledgerPeersKind
  case mbLedgerPeers of
    -- no peers from the ledger
    Maybe (Set peeraddr, DiffTime)
Nothing -> do
      (extraPeers, dt) <- (Map peeraddr PeerAdvertise -> extraPeers)
-> (Map peeraddr PeerAdvertise, DiffTime) -> (extraPeers, DiffTime)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first Map peeraddr PeerAdvertise -> extraPeers
toExtraPeers
                      ((Map peeraddr PeerAdvertise, DiffTime) -> (extraPeers, DiffTime))
-> m (Map peeraddr PeerAdvertise, DiffTime)
-> m (extraPeers, DiffTime)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> m (Map peeraddr PeerAdvertise, DiffTime)
getExtraPeers Int
n
      pure (PublicRootPeers.empty extraPeers, dt)
    Just (Set peeraddr
ledgerPeers, DiffTime
dt) ->
      case LedgerPeersKind
ledgerPeersKind of
        LedgerPeersKind
AllLedgerPeers ->
          (PublicRootPeers extraPeers peeraddr, DiffTime)
-> m (PublicRootPeers extraPeers peeraddr, DiffTime)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Set peeraddr -> PublicRootPeers extraPeers peeraddr
forall extraPeers peeraddr.
Monoid extraPeers =>
Set peeraddr -> PublicRootPeers extraPeers peeraddr
PublicRootPeers.fromLedgerPeers Set peeraddr
ledgerPeers, DiffTime
dt)
        LedgerPeersKind
BigLedgerPeers ->
          (PublicRootPeers extraPeers peeraddr, DiffTime)
-> m (PublicRootPeers extraPeers peeraddr, DiffTime)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Set peeraddr -> PublicRootPeers extraPeers peeraddr
forall extraPeers peeraddr.
Monoid extraPeers =>
Set peeraddr -> PublicRootPeers extraPeers peeraddr
PublicRootPeers.fromBigLedgerPeers Set peeraddr
ledgerPeers, DiffTime
dt)
  where
    getExtraPeers :: Int -> m (Map peeraddr PeerAdvertise, DiffTime)
    getExtraPeers :: Int -> m (Map peeraddr PeerAdvertise, DiffTime)
getExtraPeers Int
x =
      -- NOTE: we don't set `resolvConcurrent` because of
      -- https://github.com/kazu-yamamoto/dns/issues/174
      Tracer m TracePublicRootPeers
-> (IP -> PortNumber -> peeraddr)
-> DNSSemaphore m
-> ResolvConf
-> STM m (Map RelayAccessPoint PeerAdvertise)
-> DNSActions peeraddr resolver exception m
-> StdGen
-> ((Int -> m (Map peeraddr PeerAdvertise, DiffTime))
    -> m (Map peeraddr PeerAdvertise, DiffTime))
-> m (Map peeraddr PeerAdvertise, DiffTime)
forall peerAddr resolver exception a (m :: * -> *).
(MonadThrow m, MonadAsync m, Exception exception, Ord peerAddr) =>
Tracer m TracePublicRootPeers
-> (IP -> PortNumber -> peerAddr)
-> DNSSemaphore m
-> ResolvConf
-> STM m (Map RelayAccessPoint PeerAdvertise)
-> DNSActions peerAddr resolver exception m
-> StdGen
-> ((Int -> m (Map peerAddr PeerAdvertise, DiffTime)) -> m a)
-> m a
publicRootPeersProvider Tracer m TracePublicRootPeers
publicTracer
                              IP -> PortNumber -> peeraddr
toPeerAddr
                              DNSSemaphore m
dnsSemaphore
                              -- NOTE: we don't set `resolveConcurrent` because
                              -- of https://github.com/kazu-yamamoto/dns/issues/174
                              ResolvConf
DNS.defaultResolvConf
                              STM m (Map RelayAccessPoint PeerAdvertise)
readPublicRootPeers
                              DNSActions peeraddr resolver exception m
dnsActions
                              StdGen
rng
                              ((Int -> m (Map peeraddr PeerAdvertise, DiffTime))
-> Int -> m (Map peeraddr PeerAdvertise, DiffTime)
forall a b. (a -> b) -> a -> b
$ Int
x)