fix: preload adaptive SABR initialization ranges

Author: Priveetee

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

@@ -0,0 +1,43 @@
+package dev.typetype.server.services
+
+import dev.typetype.server.cache.CacheService
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.runInterruptible
+import org.schabi.newpipe.extractor.localization.Localization
+import org.schabi.newpipe.extractor.services.youtube.sabr.YoutubeSabrFormat
+
+internal object SabrAdaptiveInitialization {
+    private val localization = Localization("en", "US")
+
+    suspend fun fetch(
+        holder: SabrSessionHolder,
+        format: YoutubeSabrFormat,
+        cache: CacheService?,
+        timeoutMs: Long = 2_000L,
+    ): ByteArray? {
+        SabrInitializationData.fetch(holder.key.videoId, format, cache)?.let {
+            holder.session.streamState.ingestInitializationData(format, it)
+            return it
+        }
+        val data = fetchRange(holder, format, timeoutMs) ?: return null
+        SabrInitializationData.remember(holder.key.videoId, format, data, cache)
+        return data
+    }
+
+    suspend fun fetchRange(
+        holder: SabrSessionHolder,
+        format: YoutubeSabrFormat,
+        timeoutMs: Long,
+    ): ByteArray? {
+        val poToken = holder.session.streamState.poToken?.takeIf { it.isNotEmpty() }
+            ?: holder.playerContextToken?.streamingPoTokenBytesFor(holder.info)
+            ?: return null
+        return runCatchingNonCancellation {
+            runInterruptible(Dispatchers.IO) {
+                holder.withPlayerContext {
+                    fetchInitializationData(format, localization, timeoutMs, poToken)
+                }
+            }
+        }.getOrNull()
+    }
+}

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

@@ -1,33 +1,15 @@
 package dev.typetype.server.services
 
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.runInterruptible
-import org.schabi.newpipe.extractor.localization.Localization
 import org.schabi.newpipe.extractor.services.youtube.sabr.YoutubeSabrFormat
 
 internal object SabrDownloadInitialization {
-    private val localization = Localization("en", "US")
-
     suspend fun fetch(
         store: SabrSessionStore,
         holder: SabrSessionHolder,
         format: YoutubeSabrFormat,
     ): ByteArray? {
-        SabrInitializationData.fetch(holder.key.videoId, format, store.initCache)?.let {
-            holder.session.streamState.ingestInitializationData(format, it)
-            return it
-        }
-        val poToken = holder.session.streamState.poToken?.takeIf { it.isNotEmpty() }
-            ?: holder.playerContextToken?.streamingPoTokenBytesFor(holder.info)
+        val direct = SabrAdaptiveInitialization.fetch(holder, format, store.initCache, DIRECT_TIMEOUT_MS)
             ?: return store.fetchInitializationData(holder, format)
-        val direct = runCatchingNonCancellation {
-            runInterruptible(Dispatchers.IO) {
-                holder.withPlayerContext {
-                    fetchInitializationData(format, localization, DIRECT_TIMEOUT_MS, poToken)
-                }
-            }
-        }.getOrNull() ?: return store.fetchInitializationData(holder, format)
-        SabrInitializationData.remember(holder.key.videoId, format, direct, store.initCache)
         return direct
     }
 

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

@@ -184,10 +184,7 @@ internal class SabrSessionStore(
         holder.liveInitialization(format)?.let { return it }
         val request = SabrSegmentRequest.initialization(format)
         holder.session.getCachedSegment(request)?.let { segmentCache.put(holder, it); return it.data }
-        SabrInitializationData.fetch(holder.key.videoId, format, initCache)?.let {
-            holder.session.streamState.ingestInitializationData(format, it)
-            return it
-        }
+        SabrAdaptiveInitialization.fetch(holder, format, initCache)?.let { return it }
         SabrInitializationData.bootstrap(holder, format, initCache)?.let { return it }
         val segment = pump.fetchSegment(holder, request) ?: return null
         segmentCache.put(holder, segment)

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

@@ -0,0 +1,79 @@
+package dev.typetype.server.services
+
+import io.mockk.every
+import io.mockk.mockk
+import io.mockk.verify
+import kotlinx.coroutines.runBlocking
+import org.junit.jupiter.api.Assertions.assertArrayEquals
+import org.junit.jupiter.api.Assertions.assertNull
+import org.junit.jupiter.api.Test
+import org.schabi.newpipe.extractor.localization.Localization
+import org.schabi.newpipe.extractor.services.youtube.sabr.YoutubeSabrFormat
+import org.schabi.newpipe.extractor.services.youtube.sabr.YoutubeSabrInfo
+import org.schabi.newpipe.extractor.services.youtube.sabr.YoutubeSabrSession
+import org.schabi.newpipe.extractor.services.youtube.sabr.YoutubeSabrStreamState
+
+class SabrAdaptiveInitializationTest {
+    @Test
+    fun `fetches initialization range with the active streaming token`() = runBlocking {
+        val fixture = fixture(byteArrayOf(1, 2, 3))
+        val expected = byteArrayOf(4, 5, 6)
+        every {
+            fixture.session.fetchInitializationData(
+                fixture.format,
+                any<Localization>(),
+                2_000L,
+                match { it.contentEquals(fixture.poToken) },
+            )
+        } returns expected
+
+        val actual = SabrAdaptiveInitialization.fetchRange(fixture.holder, fixture.format, 2_000L)
+
+        assertArrayEquals(expected, actual)
+    }
+
+    @Test
+    fun `returns null when the adaptive range request fails`() = runBlocking {
+        val fixture = fixture(byteArrayOf(1))
+        every {
+            fixture.session.fetchInitializationData(
+                fixture.format,
+                any<Localization>(),
+                2_000L,
+                any(),
+            )
+        } throws java.io.IOException("range unavailable")
+
+        assertNull(SabrAdaptiveInitialization.fetchRange(fixture.holder, fixture.format, 2_000L))
+    }
+
+    @Test
+    fun `does not request a range without a streaming token`() = runBlocking {
+        val fixture = fixture(null)
+
+        assertNull(SabrAdaptiveInitialization.fetchRange(fixture.holder, fixture.format, 2_000L))
+        verify(exactly = 0) {
+            fixture.session.fetchInitializationData(any(), any(), any(), any())
+        }
+    }
+
+    private fun fixture(poToken: ByteArray?): Fixture {
+        val holder = mockk<SabrSessionHolder>()
+        val session = mockk<YoutubeSabrSession>()
+        val state = mockk<YoutubeSabrStreamState>()
+        val format = mockk<YoutubeSabrFormat>()
+        every { holder.session } returns session
+        every { holder.playerContextToken } returns null
+        every { holder.info } returns mockk<YoutubeSabrInfo>()
+        every { session.streamState } returns state
+        every { state.poToken } returns poToken
+        return Fixture(holder, session, format, poToken ?: byteArrayOf())
+    }
+
+    private data class Fixture(
+        val holder: SabrSessionHolder,
+        val session: YoutubeSabrSession,
+        val format: YoutubeSabrFormat,
+        val poToken: ByteArray,
+    )
+}