fix: prime sabr seek requests

Author: Priveetee

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

@@ -0,0 +1,52 @@
+package dev.typetype.server.services
+
+import org.schabi.newpipe.extractor.localization.Localization
+import org.schabi.newpipe.extractor.services.youtube.sabr.SabrSegmentRequest
+import org.schabi.newpipe.extractor.services.youtube.sabr.YoutubeSabrFormat
+import org.schabi.newpipe.extractor.services.youtube.sabr.YoutubeSabrSession
+
+internal fun YoutubeSabrSession.prepareForRequestedSegment(
+    holder: SabrSessionHolder,
+    request: SabrSegmentRequest,
+    localization: Localization,
+) {
+    if (request.isInitializationSegment) return
+    if (request.sequenceNumber > 1) primeInitialMedia(holder, localization)
+    prepareForRewind(request)
+    prepareForForwardJump(request)
+    streamState.setBufferedRangesOverride(emptyList())
+    val startMs = streamState.getSegmentStartMs(request.format, request.sequenceNumber).coerceAtLeast(0L)
+    streamState.setPlayerTimeMs(if (startMs == 0L) 0L else startMs + 1L)
+    if (request.format.isAudio) {
+        selectTargetFormat(holder.audioFormat, holder.videoFormat)
+    } else {
+        selectTargetFormat(holder.videoFormat, holder.audioFormat)
+    }
+}
+
+private fun YoutubeSabrSession.primeInitialMedia(holder: SabrSessionHolder, localization: Localization) {
+    val requests = listOf(
+        SabrSegmentRequest.initialization(holder.videoFormat),
+        SabrSegmentRequest.initialization(holder.audioFormat),
+        SabrSegmentRequest.media(holder.videoFormat, 1),
+        SabrSegmentRequest.media(holder.audioFormat, 1),
+    )
+    for (request in requests) {
+        if (getCachedSegment(request) == null && !isBeyondEnd(request)) {
+            runCatching { fetchSegment(request, localization) }
+        }
+    }
+}
+
+private fun YoutubeSabrSession.selectTargetFormat(target: YoutubeSabrFormat, companion: YoutubeSabrFormat) {
+    if (target.isAudio) {
+        streamState.setRequestTrackMode(AUDIO_ONLY_TRACKS, true, false)
+    } else {
+        streamState.setRequestTrackMode(VIDEO_ONLY_TRACKS, false, true)
+    }
+    streamState.setFullyBuffered(companion, true)
+    streamState.setFullyBuffered(target, false)
+}
+
+private const val AUDIO_ONLY_TRACKS = 1
+private const val VIDEO_ONLY_TRACKS = 2

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

@@ -5,7 +5,6 @@ import kotlinx.coroutines.sync.withLock
 import org.schabi.newpipe.extractor.localization.Localization
 import org.schabi.newpipe.extractor.services.youtube.sabr.SabrMediaSegment
 import org.schabi.newpipe.extractor.services.youtube.sabr.SabrSegmentRequest
-import org.schabi.newpipe.extractor.services.youtube.sabr.YoutubeSabrSession
 import java.time.Instant
 
 internal class SabrSessionPump {
@@ -37,7 +36,7 @@ internal class SabrSessionPump {
                 return@withLock segment
             }
             if (holder.session.isBeyondEnd(request)) return@withLock null
-            holder.session.prepareForRequestedSegment(holder, request)
+            holder.session.prepareForRequestedSegment(holder, request, localization)
             runCatching { holder.session.fetchSegment(request, localization) }
                 .getOrNull()
                 ?.also { holder.markServed(it) }
@@ -55,7 +54,7 @@ internal class SabrSessionPump {
                         return@withLock
                     }
                     signal?.let {
-                        holder.session.prepareForRequestedSegment(holder, signal)
+                        holder.session.prepareForRequestedSegment(holder, signal, localization)
                         holder.pendingSignals.remove(signal)
                     }
                     holder.session.evictPlayed()
@@ -90,27 +89,4 @@ internal class SabrSessionPump {
             }
         }
     }
-
-    private fun YoutubeSabrSession.prepareForRequestedSegment(holder: SabrSessionHolder, request: SabrSegmentRequest) {
-        if (request.isInitializationSegment) return
-        prepareForRewind(request)
-        prepareForForwardJump(request)
-        streamState.setBufferedRangesOverride(emptyList())
-        val startMs = streamState.getSegmentStartMs(request.format, request.sequenceNumber).coerceAtLeast(0L)
-        streamState.setPlayerTimeMs(if (startMs == 0L) 0L else startMs + 1L)
-        if (request.format.isAudio) {
-            streamState.setRequestTrackMode(AUDIO_ONLY_TRACKS, true, false)
-            streamState.setFullyBuffered(holder.videoFormat, true)
-            streamState.setFullyBuffered(holder.audioFormat, false)
-        } else {
-            streamState.setRequestTrackMode(VIDEO_ONLY_TRACKS, false, true)
-            streamState.setFullyBuffered(holder.audioFormat, true)
-            streamState.setFullyBuffered(holder.videoFormat, false)
-        }
-    }
-
-    private companion object {
-        const val AUDIO_ONLY_TRACKS = 1
-        const val VIDEO_ONLY_TRACKS = 2
-    }
 }

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

@@ -32,8 +32,8 @@ class SabrSessionPumpTest {
         val session = mockk<YoutubeSabrSession>()
         val streamState = mockk<YoutubeSabrStreamState>()
         val segment = mediaSegment(startMs, durationMs)
-        every { session.getCachedSegment(request) } returns null
-        every { session.isBeyondEnd(request) } returns false
+        every { session.getCachedSegment(any()) } returns null
+        every { session.isBeyondEnd(any()) } returns false
         every { session.prepareForRewind(request) } returns Unit
         every { session.prepareForForwardJump(request) } returns Unit
         every { session.streamState } returns streamState
@@ -43,7 +43,7 @@ class SabrSessionPumpTest {
         every { streamState.setRequestTrackMode(any(), any(), any()) } returns Unit
         every { streamState.setFullyBuffered(companion, true) } returns Unit
         every { streamState.setFullyBuffered(format, false) } returns Unit
-        every { session.fetchSegment(request, any<Localization>()) } returns segment
+        every { session.fetchSegment(any(), any<Localization>()) } returns segment
         every { session.setPlayHeadMs(startMs + durationMs) } returns Unit
         val holder = SabrSessionHolder(
             session = session,