srt: limit retransmissions to avoid a NAK-driven retransmit storm

library/src/main/java/com/pedro/library/util/streamclient/SrtStreamClient.kt

@@ -183,4 +183,17 @@ class SrtStreamClient(
    * Packets lost reported by NAK command. Increment each time a NAK is received.
    */
   fun getPacketsLost() = srtClient.packetsLost
+
+  /**
+   * Unique lost sequence numbers reported by NAK. Each sequence is counted once.
+   */
+  fun getPacketsLostUnique() = srtClient.packetsLostUnique
+
+  /**
+   * Max retransmit bandwidth as a percentage of the estimated media rate (libsrt SRTO_OHEADBW).
+   * Default 25. Values <= 0 disable the limit.
+   */
+  fun setRetransmitOverhead(percent: Int) {
+    srtClient.setRetransmitOverhead(percent)
+  }
 }

srt/src/main/java/com/pedro/srt/srt/CommandsManager.kt

@@ -36,12 +36,14 @@ import com.pedro.srt.utils.SrtSocket
 import kotlinx.coroutines.sync.Mutex
 import kotlinx.coroutines.sync.withLock
 import java.io.IOException
+import kotlin.math.max
+import kotlin.math.min
 import kotlin.random.Random
 
 /**
  * Created by pedro on 23/8/23.
  */
-class CommandsManager {
+class CommandsManager(private val timeProvider: () -> Long = { TimeUtils.getCurrentTimeMicro() }) {
 
   private val TAG = "CommandsManager"
   //used for packet lost
@@ -58,12 +60,26 @@ class CommandsManager {
   var videoDisabled = false
   var host = ""
   var latency = 120 //in millis
+  /**
+   * Max retransmit bandwidth as a percentage of the estimated media rate.
+   * Values <= 0 disable the limit (legacy behavior).
+   */
+  var retransmitOverheadPercent: Int = 25
   //Avoid write a packet in middle of other.
   private val writeSync = Mutex(locked = false)
   private var encryptor: EncryptionUtil? = null
   var videoCodec = VideoCodec.H264
   var audioCodec = AudioCodec.AAC
 
+  private var rtt = 0
+  private var rttVariance = 0
+  private var mediaBytesPerSecond = 0.0
+  private var mediaWindowStartUs = 0L
+  private var mediaWindowBytes = 0L
+  private var retransmitTokens = 0.0
+  private var lastTokenRefillUs = 0L
+  private var bucketInitialized = false
+
   fun setPassphrase(passphrase: String, type: EncryptionType) {
     encryptor = if (passphrase.isEmpty() || type == EncryptionType.NONE) null else EncryptionUtil(type, passphrase)
   }
@@ -79,12 +95,19 @@ class CommandsManager {
   fun encryptionEnabled() = encryptor != null
 
   fun loadStartTs() {
-    startTS = TimeUtils.getCurrentTimeMicro()
+    startTS = timeProvider()
     localSocketId = generateSocketId()
   }
 
   fun getTs(): Int {
-    return (TimeUtils.getCurrentTimeMicro() - startTS).toInt()
+    return (timeProvider() - startTS).toInt()
+  }
+
+  suspend fun updateRtt(rtt: Int, rttVariance: Int) {
+    writeSync.withLock {
+      this.rtt = rtt
+      this.rttVariance = rttVariance
+    }
   }
 
   @Throws(IOException::class)
@@ -127,25 +150,53 @@ class CommandsManager {
       sequenceNumber++
       packetHandlingQueue.add(dataPacket)
       dropTooLatePackets(dataPacket.ts)
+      trackMediaBytes(packet.buffer.size, timeProvider())
       dataPacket.write()
       socket?.write(dataPacket)
       return dataPacket.getSize()
     }
   }
 
   @Throws(IOException::class)
-  suspend fun reSendPackets(lostRanges: List<Pair<Int, Int>>, socket: SrtSocket?) {
+  suspend fun reSendPackets(lostRanges: List<Pair<Int, Int>>, socket: SrtSocket?): Int {
     writeSync.withLock {
-      val dataPackets = packetHandlingQueue.filter { packet ->
-        lostRanges.any { (min, max) ->
-          ((packet.sequenceNumber - min) and 0x7FFFFFFF) <= ((max - min) and 0x7FFFFFFF)
-        }
+      val unlimited = retransmitOverheadPercent <= 0
+      val nowTs = getTs()
+      val nowUs = timeProvider()
+      if (!unlimited) {
+        refillRetransmitTokens(nowUs)
       }
-      dataPackets.forEach { packet ->
+      val latencyUs = latency * 1000
+      val minResendInterval = if (unlimited) 0 else {
+        min(max(rtt + 4 * rttVariance, MIN_RESEND_INTERVAL_US), latencyUs / 4)
+      }
+      var newlyReported = 0
+      var budgetExhausted = false
+      for (packet in packetHandlingQueue) {
+        if (!isInLostRange(packet.sequenceNumber, lostRanges)) continue
+        if (!packet.nakReported) {
+          packet.nakReported = true
+          newlyReported++
+        }
+        if (!unlimited) {
+          if ((nowTs - packet.ts + rtt / 2) >= latencyUs) continue
+          // The gate only suppresses repeated reports of a packet that was already retransmitted;
+          // the first NAK is always honored even when it arrives within minResendInterval of the original send.
+          if (packet.retransmitted && (nowTs - packet.lastSentTs) < minResendInterval) continue
+          if (budgetExhausted) continue
+          val packetSize = dataPacketWireSize(packet)
+          if (retransmitTokens < packetSize) {
+            budgetExhausted = true
+            continue
+          }
+          retransmitTokens -= packetSize
+        }
         packet.retransmitted = true
         packet.write()
         socket?.write(packet)
+        packet.lastSentTs = nowTs
       }
+      return newlyReported
     }
   }
 
@@ -159,6 +210,63 @@ class CommandsManager {
     }
   }
 
+  private fun isInLostRange(sequenceNumber: Int, lostRanges: List<Pair<Int, Int>>): Boolean {
+    return lostRanges.any { (min, max) ->
+      ((sequenceNumber - min) and 0x7FFFFFFF) <= ((max - min) and 0x7FFFFFFF)
+    }
+  }
+
+  private fun trackMediaBytes(bytes: Int, nowUs: Long) {
+    if (mediaWindowStartUs == 0L) mediaWindowStartUs = nowUs
+    mediaWindowBytes += bytes
+    val elapsed = nowUs - mediaWindowStartUs
+    if (elapsed >= MEDIA_RATE_WINDOW_US) {
+      val rate = mediaWindowBytes.toDouble() * MEDIA_RATE_WINDOW_US / elapsed
+      mediaBytesPerSecond = if (mediaBytesPerSecond == 0.0) rate else {
+        mediaBytesPerSecond * MEDIA_RATE_EWMA_OLD + rate * MEDIA_RATE_EWMA_NEW
+      }
+      mediaWindowStartUs = nowUs
+      mediaWindowBytes = 0
+    }
+  }
+
+  private fun getRetransmitRate(): Double {
+    val percent = retransmitOverheadPercent
+    val mediaRate = if (mediaBytesPerSecond > 0.0) mediaBytesPerSecond else MIN_RETRANSMIT_BYTES_PER_SECOND.toDouble()
+    return max(mediaRate * percent / 100.0, MIN_RETRANSMIT_BYTES_PER_SECOND.toDouble())
+  }
+
+  private fun getRetransmitCapacity(rate: Double): Int {
+    // Allow an immediate burst up to half a second of media so short loss events
+    // (e.g. a brief link flap) are not retried one packet at a time on healthy links.
+    val burstCapacity = if (mediaBytesPerSecond > 0.0) {
+      (mediaBytesPerSecond * RETRANSMIT_BURST_WINDOW_US / MEDIA_RATE_WINDOW_US).toInt()
+    } else {
+      0
+    }
+    return max(burstCapacity, max((rate * latency / 1000.0).toInt(), MTU))
+  }
+
+  private fun refillRetransmitTokens(nowUs: Long) {
+    if (!bucketInitialized) {
+      val rate = getRetransmitRate()
+      retransmitTokens = getRetransmitCapacity(rate).toDouble()
+      lastTokenRefillUs = nowUs
+      bucketInitialized = true
+      return
+    }
+    val elapsedUs = nowUs - lastTokenRefillUs
+    if (elapsedUs <= 0) return
+    val rate = getRetransmitRate()
+    val capacity = getRetransmitCapacity(rate)
+    retransmitTokens = min(retransmitTokens + rate * elapsedUs / MEDIA_RATE_WINDOW_US, capacity.toDouble())
+    lastTokenRefillUs = nowUs
+  }
+
+  private fun dataPacketWireSize(packet: DataPacket): Int {
+    return packet.payload.size + DATA_HEADER_SIZE
+  }
+
   private fun dropTooLatePackets(nowTs: Int) {
     val thresholdUs = latency * 1000
     val firstKept = packetHandlingQueue.indexOfFirst { (nowTs - it.ts) <= thresholdUs }
@@ -200,6 +308,14 @@ class CommandsManager {
     startTS = 0L
     host = ""
     packetHandlingQueue.clear()
+    rtt = 0
+    rttVariance = 0
+    mediaBytesPerSecond = 0.0
+    mediaWindowStartUs = 0L
+    mediaWindowBytes = 0L
+    retransmitTokens = 0.0
+    lastTokenRefillUs = 0L
+    bucketInitialized = false
   }
 
   private fun generateInitialSequence(): Int {
@@ -209,4 +325,14 @@ class CommandsManager {
   private fun generateSocketId(): Int {
     return Random.nextInt(1, Int.MAX_VALUE)
   }
-}
+
+  companion object {
+    private const val DATA_HEADER_SIZE = 16
+    private const val MIN_RETRANSMIT_BYTES_PER_SECOND = 8_000
+    private const val MIN_RESEND_INTERVAL_US = 20_000
+    private const val MEDIA_RATE_WINDOW_US = 1_000_000L
+    private const val RETRANSMIT_BURST_WINDOW_US = 500_000L
+    private const val MEDIA_RATE_EWMA_OLD = 0.8
+    private const val MEDIA_RATE_EWMA_NEW = 0.2
+  }
+}

srt/src/main/java/com/pedro/srt/srt/SrtClient.kt

@@ -107,6 +107,11 @@ class SrtClient(private val connectChecker: ConnectChecker) {
     private set
   var packetsLost = 0
     private set
+  /**
+   * Unique lost sequence numbers reported by NAK. Each sequence is counted once.
+   */
+  var packetsLostUnique = 0
+    private set
   var socketType = SocketType.JAVA
   var socketTimeout = StreamSocket.DEFAULT_TIMEOUT
 
@@ -132,6 +137,14 @@ class SrtClient(private val connectChecker: ConnectChecker) {
     commandsManager.latency = latency
   }
 
+  /**
+   * Max retransmit bandwidth as a percentage of the estimated media rate (libsrt SRTO_OHEADBW).
+   * Default 25. Values <= 0 disable the limit.
+   */
+  fun setRetransmitOverhead(percent: Int) {
+    commandsManager.retransmitOverheadPercent = percent
+  }
+
   fun setDelay(millis: Long) {
     srtSender.setDelay(millis)
   }
@@ -326,6 +339,7 @@ class SrtClient(private val connectChecker: ConnectChecker) {
     commandsManager.reset()
     rtt = 0
     packetsLost = 0
+    packetsLostUnique = 0
     job?.cancelAndJoin()
     job = null
     scope.cancel()
@@ -405,14 +419,15 @@ class SrtClient(private val connectChecker: ConnectChecker) {
             commandsManager.updateHandlingQueue(lastPacketSequence)
             if (ackSequence != 0) {
               rtt = srtPacket.rtt
+              commandsManager.updateRtt(srtPacket.rtt, srtPacket.rttVariance)
               commandsManager.writeAck2(ackSequence, socket)
             }
           }
           is Nak -> {
             //packet lost reported, we should resend it
             val lostRanges = srtPacket.getNakRanges()
             this.packetsLost += srtPacket.getLostCount()
-            commandsManager.reSendPackets(lostRanges, socket)
+            packetsLostUnique += commandsManager.reSendPackets(lostRanges, socket)
           }
           is Shutdown -> {
             onMainThread {

srt/src/main/java/com/pedro/srt/srt/packets/DataPacket.kt

@@ -41,6 +41,9 @@ class DataPacket(
   var payload: ByteArray = byteArrayOf()
 ): SrtPacket() {
 
+  var lastSentTs: Int = 0
+  var nakReported: Boolean = false
+
   fun write() {
     resetBuffer()
     val headerData = (PacketType.DATA.value shl 31) or (sequenceNumber and 0x7FFFFFFF)