-- | Perf protocol implementation (specs/perf).
--
-- Protocol ID: /perf/1.0.0
--
-- Wire format: the client sends a single 8-byte big-endian uint64
-- naming the number of bytes it wants the server to send back, then
-- streams its upload payload, then half-closes its write side. The
-- server reads the 8-byte header, drains the upload until EOF, and
-- only then (perf.md: the response "MUST NOT be run concurrently"
-- with the upload) writes the requested number of bytes back and
-- closes the stream.
--
-- Each measurement runs on its own stream; there is no framing and no
-- protobuf. Payload bytes carry no meaning, so both sides send zeros.
module LibP2P.Protocol.Perf
  ( -- * Protocol ID
    perfProtocolId
    -- * Types
  , PerfError (..)
  , PerfResult (..)
    -- * Responder
  , handlePerf
    -- * Initiator
  , perfOnStream
  , runPerf
    -- * Registration
  , registerPerfHandler
  ) where

import Control.Concurrent.STM (atomically, readTVar, writeTVar)
import Control.Exception (SomeException, catch, finally, try)
import qualified Data.ByteString as BS
import qualified Data.Map.Strict as Map
import Data.Text (Text)
import Data.Time.Clock (NominalDiffTime, diffUTCTime, getCurrentTime)
import Data.Word (Word64)
import LibP2P.Core.Binary (readWord64BE, word64BE)
import LibP2P.Crypto.PeerId (PeerId)
import LibP2P.MultistreamSelect.Negotiation
  ( NegotiationResult (..)
  , StreamIO (..)
  , closeQuietly
  , negotiateInitiator
  , readExactBounded
  )
import LibP2P.Switch.Connection (newStream)
import LibP2P.Switch.Types
  ( Connection (..)
  , Switch (..)
  )

-- | Perf protocol ID.
perfProtocolId :: Text
perfProtocolId :: Text
perfProtocolId = Text
"/perf/1.0.0"

-- | Chunk size for bulk sends, matching the 64 KiB block the reference
-- implementations use.
perfBlockSize :: Int
perfBlockSize :: Int
perfBlockSize = Int
65536

-- | Perf error types.
data PerfError
  = PerfNegotiationError !String  -- ^ Stream open or protocol negotiation failed
  | PerfStreamError !String       -- ^ I/O error during the exchange
  deriving (Int -> PerfError -> ShowS
[PerfError] -> ShowS
PerfError -> [Char]
(Int -> PerfError -> ShowS)
-> (PerfError -> [Char])
-> ([PerfError] -> ShowS)
-> Show PerfError
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> PerfError -> ShowS
showsPrec :: Int -> PerfError -> ShowS
$cshow :: PerfError -> [Char]
show :: PerfError -> [Char]
$cshowList :: [PerfError] -> ShowS
showList :: [PerfError] -> ShowS
Show, PerfError -> PerfError -> Bool
(PerfError -> PerfError -> Bool)
-> (PerfError -> PerfError -> Bool) -> Eq PerfError
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: PerfError -> PerfError -> Bool
== :: PerfError -> PerfError -> Bool
$c/= :: PerfError -> PerfError -> Bool
/= :: PerfError -> PerfError -> Bool
Eq)

-- | Successful perf exchange result.
data PerfResult = PerfResult
  { PerfResult -> NominalDiffTime
perfElapsed :: !NominalDiffTime  -- ^ Header write to last byte received
  } deriving (Int -> PerfResult -> ShowS
[PerfResult] -> ShowS
PerfResult -> [Char]
(Int -> PerfResult -> ShowS)
-> (PerfResult -> [Char])
-> ([PerfResult] -> ShowS)
-> Show PerfResult
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> PerfResult -> ShowS
showsPrec :: Int -> PerfResult -> ShowS
$cshow :: PerfResult -> [Char]
show :: PerfResult -> [Char]
$cshowList :: [PerfResult] -> ShowS
showList :: [PerfResult] -> ShowS
Show, PerfResult -> PerfResult -> Bool
(PerfResult -> PerfResult -> Bool)
-> (PerfResult -> PerfResult -> Bool) -> Eq PerfResult
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: PerfResult -> PerfResult -> Bool
== :: PerfResult -> PerfResult -> Bool
$c/= :: PerfResult -> PerfResult -> Bool
/= :: PerfResult -> PerfResult -> Bool
Eq)

-- | A shared zero block for bulk sends.
zeroBlock :: BS.ByteString
zeroBlock :: ByteString
zeroBlock = Int -> Word8 -> ByteString
BS.replicate Int
perfBlockSize Word8
0

-- | Write @n@ zero bytes in 'perfBlockSize' chunks.
writeZeros :: StreamIO -> Word64 -> IO ()
writeZeros :: StreamIO -> Word64 -> IO ()
writeZeros StreamIO
stream = Word64 -> IO ()
forall {t}. Integral t => t -> IO ()
go
  where
    go :: t -> IO ()
go t
0 = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    go t
n = do
      let chunk :: t
chunk = t -> t -> t
forall a. Ord a => a -> a -> a
min t
n (Int -> t
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
perfBlockSize)
      StreamIO -> ByteString -> IO ()
streamWrite StreamIO
stream (Int -> ByteString -> ByteString
BS.take (t -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral t
chunk) ByteString
zeroBlock)
      t -> IO ()
go (t
n t -> t -> t
forall a. Num a => a -> a -> a
- t
chunk)

-- | Read and discard bytes until EOF (the initiator's half-close).
-- Chunk-level reads (#276) keep the drain off the byte-at-a-time path
-- that bounded download throughput.
drainUntilEof :: StreamIO -> IO ()
drainUntilEof :: StreamIO -> IO ()
drainUntilEof StreamIO
stream = IO ()
forall {b}. IO b
loop IO () -> (SomeException -> IO ()) -> IO ()
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` \(SomeException
_ :: SomeException) -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  where
    loop :: IO b
loop = StreamIO -> Int -> IO ByteString
streamReadChunk StreamIO
stream Int
perfBlockSize IO ByteString -> IO b -> IO b
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> IO b
loop

-- | Read and discard exactly @n@ bytes at chunk granularity (#276). The
-- payload carries no meaning, so the chunks are dropped; premature EOF
-- throws. A chunk request never exceeds the bytes still owed, so no
-- byte beyond @n@ is consumed from the stream.
discardExactly :: StreamIO -> Word64 -> IO ()
discardExactly :: StreamIO -> Word64 -> IO ()
discardExactly StreamIO
stream = Word64 -> IO ()
go
  where
    go :: Word64 -> IO ()
    go :: Word64 -> IO ()
go Word64
0 = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    go !Word64
n = do
      chunk <- StreamIO -> Int -> IO ByteString
streamReadChunk StreamIO
stream (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64 -> Word64 -> Word64
forall a. Ord a => a -> a -> a
min Word64
n (Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
perfBlockSize)))
      go (n - fromIntegral (BS.length chunk))

-- | Handle an inbound perf request (responder).
--
-- Reads the 8-byte download size, drains the client's upload until it
-- half-closes, then sends the requested bytes back and closes. A client
-- that closes before sending a full header is dropped silently.
handlePerf :: StreamIO -> PeerId -> IO ()
handlePerf :: StreamIO -> PeerId -> IO ()
handlePerf StreamIO
stream PeerId
_remotePeerId = IO ()
serve IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO a
`finally` StreamIO -> IO ()
closeQuietly StreamIO
stream
  where
    serve :: IO ()
serve = do
      header <- StreamIO -> Int -> Int -> IO (Either [Char] ByteString)
readExactBounded StreamIO
stream Int
8 Int
8 IO (Either [Char] ByteString)
-> (SomeException -> IO (Either [Char] ByteString))
-> IO (Either [Char] ByteString)
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch`
                (\(SomeException
_ :: SomeException) -> Either [Char] ByteString -> IO (Either [Char] ByteString)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Char] -> Either [Char] ByteString
forall a b. a -> Either a b
Left [Char]
"stream closed"))
      case header of
        Left [Char]
_ -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        Right ByteString
sizeBytes -> do
          let downloadSize :: Word64
downloadSize = ByteString -> Word64
readWord64BE ByteString
sizeBytes
          StreamIO -> IO ()
drainUntilEof StreamIO
stream
          StreamIO -> Word64 -> IO ()
writeZeros StreamIO
stream Word64
downloadSize
            IO () -> (SomeException -> IO ()) -> IO ()
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` (\(SomeException
_ :: SomeException) -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())

-- | One perf exchange on an already-negotiated stream (initiator).
--
-- Sends the header and @uploadBytes@ zeros, half-closes the write side,
-- then reads exactly @downloadBytes@ back. The elapsed time covers the
-- full exchange, header write to last byte read.
perfOnStream :: StreamIO -> Word64 -> Word64 -> IO (Either PerfError PerfResult)
perfOnStream :: StreamIO -> Word64 -> Word64 -> IO (Either PerfError PerfResult)
perfOnStream StreamIO
stream Word64
uploadBytes Word64
downloadBytes = do
  t0 <- IO UTCTime
getCurrentTime
  outcome <- try $ do
    streamWrite stream (word64BE downloadBytes)
    writeZeros stream uploadBytes
    streamClose stream
    discardExactly stream downloadBytes
  case outcome of
    Left (SomeException
e :: SomeException) ->
      Either PerfError PerfResult -> IO (Either PerfError PerfResult)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PerfError -> Either PerfError PerfResult
forall a b. a -> Either a b
Left ([Char] -> PerfError
PerfStreamError ([Char]
"perf I/O failed: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ SomeException -> [Char]
forall a. Show a => a -> [Char]
show SomeException
e)))
    Right () -> do
      t1 <- IO UTCTime
getCurrentTime
      pure (Right (PerfResult (diffUTCTime t1 t0)))

-- | Run one perf measurement against a connected peer: open a stream,
-- negotiate /perf/1.0.0, run the exchange, and release the stream.
runPerf :: Switch -> Connection -> Word64 -> Word64 -> IO (Either PerfError PerfResult)
runPerf :: Switch
-> Connection
-> Word64
-> Word64
-> IO (Either PerfError PerfResult)
runPerf Switch
sw Connection
conn Word64
uploadBytes Word64
downloadBytes = do
  streamOrErr <- Switch -> Connection -> IO (Either ResourceError StreamIO)
newStream Switch
sw Connection
conn
  case streamOrErr of
    Left ResourceError
err ->
      Either PerfError PerfResult -> IO (Either PerfError PerfResult)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PerfError -> Either PerfError PerfResult
forall a b. a -> Either a b
Left ([Char] -> PerfError
PerfNegotiationError ([Char]
"stream reservation failed: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ ResourceError -> [Char]
forall a. Show a => a -> [Char]
show ResourceError
err)))
    Right StreamIO
stream -> do
      negotiated <- IO NegotiationResult -> IO (Either SomeException NegotiationResult)
forall e a. Exception e => IO a -> IO (Either e a)
try (StreamIO -> [Text] -> IO NegotiationResult
negotiateInitiator StreamIO
stream [Text
perfProtocolId])
      case negotiated of
        Right (Accepted Text
_) ->
          StreamIO -> Word64 -> Word64 -> IO (Either PerfError PerfResult)
perfOnStream StreamIO
stream Word64
uploadBytes Word64
downloadBytes
            IO (Either PerfError PerfResult)
-> IO () -> IO (Either PerfError PerfResult)
forall a b. IO a -> IO b -> IO a
`finally` StreamIO -> IO ()
closeQuietly StreamIO
stream
        Right NegotiationResult
NoProtocol -> do
          StreamIO -> IO ()
closeQuietly StreamIO
stream
          Either PerfError PerfResult -> IO (Either PerfError PerfResult)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PerfError -> Either PerfError PerfResult
forall a b. a -> Either a b
Left ([Char] -> PerfError
PerfNegotiationError [Char]
"remote does not support perf"))
        Left (SomeException
e :: SomeException) -> do
          StreamIO -> IO ()
closeQuietly StreamIO
stream
          Either PerfError PerfResult -> IO (Either PerfError PerfResult)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PerfError -> Either PerfError PerfResult
forall a b. a -> Either a b
Left ([Char] -> PerfError
PerfNegotiationError ([Char]
"perf negotiation failed: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ SomeException -> [Char]
forall a. Show a => a -> [Char]
show SomeException
e)))

-- | Register the perf handler on the Switch.
registerPerfHandler :: Switch -> IO ()
registerPerfHandler :: Switch -> IO ()
registerPerfHandler Switch
sw = STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
  protos <- TVar (Map Text StreamHandler) -> STM (Map Text StreamHandler)
forall a. TVar a -> STM a
readTVar (Switch -> TVar (Map Text StreamHandler)
swProtocols Switch
sw)
  let handler Connection
conn StreamIO
stream = StreamIO -> PeerId -> IO ()
handlePerf StreamIO
stream (Connection -> PeerId
connPeerId Connection
conn)
  writeTVar (swProtocols sw) (Map.insert perfProtocolId handler protos)