content: {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE RecordWildCards #-}
  
  module Network.QUIC.Recovery.LossRecovery (
      onPacketSent,
      onPacketReceived,
      onAckReceived,
      onPacketNumberSpaceDiscarded,
  ) where
  
  import Control.Concurrent.STM
  import Data.Sequence (Seq, ViewR (..), (|>))
  import qualified Data.Sequence as Seq
  
  import Network.QUIC.Connector
  import Network.QUIC.Imports
  import Network.QUIC.Qlog
  import Network.QUIC.Recovery.Constants
  import Network.QUIC.Recovery.Detect
  import Network.QUIC.Recovery.Metrics
  import Network.QUIC.Recovery.Misc
  import Network.QUIC.Recovery.PeerPacketNumbers
  import Network.QUIC.Recovery.Release
  import Network.QUIC.Recovery.Timer
  import Network.QUIC.Recovery.Types
  import Network.QUIC.Recovery.Utils
  import Network.QUIC.Types
  
  ----------------------------------------------------------------
  
  onPacketSent :: LDCC -> SentPacket -> IO ()
  onPacketSent ldcc@LDCC{..} sentPacket = do
      let lvl0 = spEncryptionLevel sentPacket
      let lvl
              | lvl0 == RTT0Level = RTT1Level
              | otherwise = lvl0
      discarded <- getPacketNumberSpaceDiscarded ldcc lvl
      unless discarded $ do
          onPacketSentCC ldcc sentPacket
          when (spAckEliciting sentPacket) $
              atomicModifyIORef'' (lossDetection ! lvl) $ \ld ->
                  ld
                      { timeOfLastAckElicitingPacket = spTimeSent sentPacket
                      }
          atomicModifyIORef'' (sentPackets ! lvl) $
              \(SentPackets db) -> SentPackets (db |> sentPacket)
          setLossDetectionTimer ldcc lvl
  
  onPacketSentCC :: LDCC -> SentPacket -> IO ()
  onPacketSentCC ldcc@LDCC{..} sentPacket = metricsUpdated ldcc $
      atomically $
          modifyTVar' recoveryCC $ \cc ->
              cc
                  { bytesInFlight = bytesInFlight cc + sentBytes
                  , numOfAckEliciting = numOfAckEliciting cc + countAckEli sentPacket
                  }
    where
      sentBytes = spSentBytes sentPacket
  
  ----------------------------------------------------------------
  
  onPacketReceived :: LDCC -> EncryptionLevel -> PacketNumber -> IO ()
  onPacketReceived ldcc lvl pn = do
      discarded <- getPacketNumberSpaceDiscarded ldcc lvl
      unless discarded $ addPeerPacketNumbers ldcc lvl pn
  
  ----------------------------------------------------------------
  
  onAckReceived :: LDCC -> EncryptionLevel -> AckInfo -> Microseconds -> IO ()
  onAckReceived ldcc@LDCC{..} lvl ackInfo@(AckInfo largestAcked _ _) ackDelay = do
      changed <- atomicModifyIORef' (lossDetection ! lvl) update
      when changed $ do
          let predicate = fromAckInfoToPred ackInfo . spPacketNumber
          releaseLostCandidates ldcc lvl predicate >>= updateCConAck
          releaseByPredicate ldcc lvl predicate >>= detectLossUpdateCC
    where
      update ld@LossDetection{..} = (ld', changed)
        where
          ld' =
              ld
                  { largestAckedPacket = max largestAckedPacket largestAcked
                  , previousAckInfo = ackInfo
                  }
          changed = previousAckInfo /= ackInfo
      detectLossUpdateCC newlyAckedPackets = case Seq.viewr newlyAckedPackets of
          EmptyR -> return ()
          _ :> lastPkt -> do
              -- If the largest acknowledged is newly acked and
              -- at least one ack-eliciting was newly acked, update the RTT.
              when
                  ( spPacketNumber lastPkt == largestAcked
                      && any spAckEliciting newlyAckedPackets
                  )
                  $ do
                      rtt <- getElapsedTimeMicrosecond $ spTimeSent lastPkt
                      let latestRtt = max rtt kGranularity
                      updateRTT ldcc lvl latestRtt ackDelay
  
              {- fimxe
              -- Process ECN information if present.
              if (ACK frame contains ECN information):
                 ProcessECN(ack, lvl)
              -}
  
              lostPackets <- detectAndRemoveLostPackets ldcc lvl
              unless (null lostPackets) $ do
                  mode <- ccMode <$> readTVarIO recoveryCC
                  if lvl == RTT1Level && mode /= SlowStart
                      then mergeLostCandidates ldcc lostPackets
                      else do
                          -- just in case
                          lostPackets' <- mergeLostCandidatesAndClear ldcc lostPackets
                          onPacketsLost ldcc lostPackets'
                          retransmit ldcc lostPackets'
              -- setLossDetectionTimer in updateCConAck
              updateCConAck newlyAckedPackets
  
      updateCConAck newlyAckedPackets
          | newlyAckedPackets == Seq.empty = return ()
          | otherwise = do
              onPacketsAcked ldcc newlyAckedPackets
  
              -- Sec 6.2.1. Computing PTO
              -- "The PTO backoff factor is reset when an acknowledgement is
              --  received, except in the following case. A server might
              --  take longer to respond to packets during the handshake
              --  than otherwise. To protect such a server from repeated
              --  client probes, the PTO backoff is not reset at a client
              --  that is not yet certain that the server has finished
              --  validating the client's address."
              validated <- peerCompletedAddressValidation ldcc
              when validated $
                  metricsUpdated ldcc $
                      atomicModifyIORef'' recoveryRTT $
                          \rtt -> rtt{ptoCount = 0}
  
              setLossDetectionTimer ldcc lvl
  
  releaseLostCandidates
      :: LDCC -> EncryptionLevel -> (SentPacket -> Bool) -> IO (Seq SentPacket)
  releaseLostCandidates ldcc@LDCC{..} lvl predicate = do
      packets <- atomically $ do
          SentPackets db <- readTVar lostCandidates
          let (pkts, db') = Seq.partition predicate db
          writeTVar lostCandidates $ SentPackets db'
          return pkts
      removePacketNumbers ldcc lvl packets
      return packets
  
  onPacketsAcked :: LDCC -> Seq SentPacket -> IO ()
  onPacketsAcked ldcc@LDCC{..} ackedPackets = metricsUpdated ldcc $ do
      maxPktSiz <- getMaxPacketSize ldcc
      oldcc <- readTVarIO recoveryCC
      atomically $ modifyTVar' recoveryCC $ modify maxPktSiz
      newcc <- readTVarIO recoveryCC
      when (ccMode oldcc /= ccMode newcc) $
          qlogContestionStateUpdated ldcc $
              ccMode newcc
    where
      modify maxPktSiz cc@CC{..} =
          cc
              { bytesInFlight = bytesInFlight'
              , congestionWindow = congestionWindow'
              , bytesAcked = bytesAcked'
              , ccMode = ccMode'
              , numOfAckEliciting = numOfAckEliciting'
              }
        where
          (bytesInFlight', congestionWindow', bytesAcked', ccMode', numOfAckEliciting') =
              foldl'
                  (.+)
                  (bytesInFlight, congestionWindow, bytesAcked, ccMode, numOfAckEliciting)
                  ackedPackets
          (bytes, cwin, acked, _, cnt) .+ sp@SentPacket{..} = (bytes', cwin', acked', mode', cnt')
            where
              isRecovery = inCongestionRecovery spTimeSent congestionRecoveryStartTime
              bytes' = bytes - spSentBytes
              ackedA = acked + spSentBytes
              cnt' = cnt - countAckEli sp
              (cwin', acked', mode')
                  -- Do not increase congestion window in recovery period.
                  | isRecovery = (cwin, acked, Recovery)
                  -- fixme: Do not increase congestion_window if application
                  -- limited or flow control limited.
                  --
                  -- Slow start.
                  | cwin < ssthresh = (cwin + spSentBytes, acked, SlowStart)
                  -- Congestion avoidance.
                  -- In this implementation, maxPktSiz == spSentBytes.
                  -- spSentBytes is large enough, so we don't care
                  -- the roundup issue of `div`.
                  | ackedA >= cwin = (cwin + maxPktSiz, ackedA - cwin, Avoidance)
                  | otherwise = (cwin, ackedA, Avoidance)
  
  ----------------------------------------------------------------
  
  onPacketNumberSpaceDiscarded :: LDCC -> EncryptionLevel -> IO ()
  onPacketNumberSpaceDiscarded ldcc lvl = do
      let (lvl', label) = case lvl of
              InitialLevel -> (HandshakeLevel, "initial")
              _ -> (RTT1Level, "handshake")
      qlogDebug ldcc $ Debug (label <> " discarded")
      void $ discard ldcc lvl
      setLossDetectionTimer ldcc lvl'
  