fix: target exact sabr buffered range

Author: Priveetee

src/main/kotlin/dev/typetype/server/services/SabrTargetRequestShape.kt

@@ -1,5 +1,6 @@
 package dev.typetype.server.services
 
+import org.schabi.newpipe.extractor.services.youtube.sabr.SabrBufferedRange
 import org.schabi.newpipe.extractor.services.youtube.sabr.SabrSegmentRequest
 import org.schabi.newpipe.extractor.services.youtube.sabr.YoutubeSabrFormat
 
@@ -12,13 +13,11 @@ internal inline fun <T> withTargetedRequestShape(
     val state = holder.session.streamState
     val targetIsAudio = request.format.isAudio
     state.setSelectVideoFormatBeforeAudio(targetIsAudio)
-    state.setLastOnlyRange(request.format, true)
-    state.setFullyBuffered(companion, true)
+    state.setBufferedRangesOverride(listOf(companion.fullRange(), request.targetRange(holder)))
     return try {
         block()
     } finally {
-        state.setFullyBuffered(companion, false)
-        state.setLastOnlyRange(request.format, false)
+        state.setBufferedRangesOverride(null)
         state.setActiveTrackTypes(holder.isVideoActive(), holder.isAudioActive())
         state.setSelectVideoFormatBeforeAudio(holder.playerTimeMs() > SEEK_FORMAT_ORDER_MS)
     }
@@ -27,4 +26,36 @@ internal inline fun <T> withTargetedRequestShape(
 private fun SabrSessionHolder.companionFormat(format: YoutubeSabrFormat): YoutubeSabrFormat =
     if (format.isAudio) videoFormat else audioFormat
 
+private fun YoutubeSabrFormat.fullRange(): SabrBufferedRange =
+    SabrBufferedRange(
+        itag,
+        lastModified,
+        xtags,
+        0L,
+        MAX_RANGE_DURATION_MS,
+        MAX_RANGE_INDEX,
+        MAX_RANGE_INDEX,
+        TIMESCALE,
+    )
+
+private fun SabrSegmentRequest.targetRange(holder: SabrSessionHolder): SabrBufferedRange {
+    val state = holder.session.streamState
+    val startMs = state.getSegmentStartMs(format, sequenceNumber).coerceAtLeast(0L)
+    val endMs = state.getSegmentEndMs(format, sequenceNumber)
+    val durationMs = (endMs - startMs).takeIf { it > 0L } ?: 1L
+    return SabrBufferedRange(
+        format.itag,
+        format.lastModified,
+        format.xtags,
+        0L,
+        durationMs,
+        sequenceNumber,
+        sequenceNumber,
+        TIMESCALE,
+    )
+}
+
 private const val SEEK_FORMAT_ORDER_MS = 1_000L
+private const val MAX_RANGE_DURATION_MS = Int.MAX_VALUE.toLong()
+private const val MAX_RANGE_INDEX = Int.MAX_VALUE
+private const val TIMESCALE = 1_000

src/test/kotlin/dev/typetype/server/services/SabrSessionPumpLoopTest.kt

@@ -32,11 +32,10 @@ class SabrSessionPumpLoopTest {
             every { session.getCachedSegment(request) } returns null
             every { streamState.getMinBufferedEndMs() } returns 329_492L
             every { streamState.getSegmentStartMs(audio, 35) } returns 340_000L
+            every { streamState.getSegmentEndMs(audio, 35) } returns 349_476L
             every { streamState.setPlayerTimeMs(any()) } returns Unit
-            every { streamState.setLastOnlyRange(audio, true) } returns Unit
-            every { streamState.setLastOnlyRange(audio, false) } returns Unit
-            every { streamState.setFullyBuffered(video, true) } returns Unit
-            every { streamState.setFullyBuffered(video, false) } returns Unit
+            every { streamState.setBufferedRangesOverride(any()) } returns Unit
+            every { streamState.setBufferedRangesOverride(null) } returns Unit
             every { streamState.setSelectVideoFormatBeforeAudio(any()) } returns Unit
             every { session.prepareForForwardJump(request) } returns Unit
             every { session.pumpOnceStreamingUntilCached(any(), request) } returns 0
@@ -51,11 +50,13 @@ class SabrSessionPumpLoopTest {
             verify(exactly = 1) { session.pumpOnceStreamingUntilCached(any(), request) }
             verifyOrder {
                 streamState.setSelectVideoFormatBeforeAudio(true)
-                streamState.setLastOnlyRange(audio, true)
-                streamState.setFullyBuffered(video, true)
+                streamState.setBufferedRangesOverride(match { ranges ->
+                    ranges.size == 2 &&
+                        ranges[0].summarize() == "itag=247:seq=2147483647-2147483647:time=0+2147483647:timescale=1000" &&
+                        ranges[1].summarize() == "itag=140:seq=35-35:time=0+9476:timescale=1000"
+                })
                 session.pumpOnceStreamingUntilCached(any(), request)
-                streamState.setFullyBuffered(video, false)
-                streamState.setLastOnlyRange(audio, false)
+                streamState.setBufferedRangesOverride(null)
                 streamState.setActiveTrackTypes(true, true)
             }
         } finally {
@@ -82,6 +83,13 @@ class SabrSessionPumpLoopTest {
         every { format.itag } returns itag
         every { format.isAudio } returns isAudio
         every { format.bitrate } returns if (isAudio) 128_000 else 2_000_000
+        every { format.lastModified } returns if (isAudio) AUDIO_LAST_MODIFIED else VIDEO_LAST_MODIFIED
+        every { format.xtags } returns null
         return format
     }
+
+    private companion object {
+        const val AUDIO_LAST_MODIFIED = 1_765_814_035_331_078L
+        const val VIDEO_LAST_MODIFIED = 1_726_365_891_623_401L
+    }
 }