perf: reuse SABR bootstrap preparation

Author: Priveetee

src/main/kotlin/dev/typetype/server/ExtractionServiceRegistry.kt

@@ -134,10 +134,8 @@ internal class ExtractionServiceRegistry(
         ),
         YouTubeSubtitleCache(cache),
     )
-    val youtubeSabrBootstrapStreamService = CachedStreamService(
-        YoutubeScopedStreamService(SabrBootstrapStreamService(sabrSessionStore, tokenYoutubeSessionClient)),
-        cache,
-        "stream-youtube-sabr-bootstrap:v1",
+    val youtubeSabrBootstrapStreamService = YoutubeScopedStreamService(
+        SabrBootstrapStreamService(sabrSessionStore, tokenYoutubeSessionClient),
     )
     val nicoNicoStreamService = CachedStreamService(directPipePipeStreamService, cache, "stream-niconico:v1")
     val bilibiliStreamService = CachedStreamService(directPipePipeStreamService, cache, "stream-bilibili:v1")

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

@@ -2,24 +2,21 @@ package dev.typetype.server.services
 
 import dev.typetype.server.models.ExtractionResult
 import dev.typetype.server.models.StreamResponse
-import kotlinx.coroutines.async
-import kotlinx.coroutines.coroutineScope
 
 internal class SabrBootstrapStreamService(
     private val sessionStore: SabrSessionStore,
     private val tokenSessionClient: TypetypeTokenYoutubeSessionClient,
 ) : StreamService {
-    override suspend fun getStreamInfo(url: String): ExtractionResult<StreamResponse> = coroutineScope {
+    override suspend fun getStreamInfo(url: String): ExtractionResult<StreamResponse> {
         val videoId = youtubeVideoId(url)
-            ?: return@coroutineScope ExtractionResult.BadRequest("Invalid YouTube URL")
-        val prepared = async { sessionStore.fetchInfo(videoId, cachedFirst = true) }
-        val session = async { tokenSessionClient.fetchPlaybackSession(videoId) }
-        val metadata = session.await()
-            ?: return@coroutineScope ExtractionResult.Failure("SABR bootstrap metadata unavailable")
-        if (prepared.await() == null) {
-            return@coroutineScope ExtractionResult.Failure("SABR playback formats unavailable")
-        }
-        ExtractionResult.Success(metadata.toFallbackStreamResponse(videoId))
+            ?: return ExtractionResult.BadRequest("Invalid YouTube URL")
+        val metadata = tokenSessionClient.fetchPlaybackSession(videoId)
+            ?: return ExtractionResult.Failure("SABR bootstrap metadata unavailable")
+        val prepared = metadata.preparedSabrInfo()
+            ?: sessionStore.fetchInfo(videoId, cachedFirst = true)
+            ?: return ExtractionResult.Failure("SABR playback formats unavailable")
+        sessionStore.rememberPreparedInfo(videoId, prepared)
+        return ExtractionResult.Success(metadata.toFallbackStreamResponse(videoId))
     }
 }
 

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

@@ -59,9 +59,12 @@ internal class SabrInfoFetcher(
 
     suspend fun rememberExtractedInfo(videoId: String, info: YoutubeSabrInfo): Unit {
         repository.rememberInitialization(videoId, info)
-        fetchInfoOnce(videoId, startTimeMs = 0L)
-            ?.takeIf { it.hasAudioAndVideoFormats() }
-            ?.let { repository.putPrepared(videoId, startTimeMs = 0L, it) }
+        fetchInfo(videoId, startTimeMs = 0L, cachedFirst = true)
+    }
+
+    suspend fun rememberPreparedInfo(videoId: String, prepared: SabrPreparedInfo): Unit {
+        repository.rememberInitialization(videoId, prepared.info)
+        repository.putPrepared(videoId, startTimeMs = 0L, prepared)
     }
 
     suspend fun invalidatePlayback(videoId: String): Unit = repository.invalidatePlayback(videoId)
@@ -87,13 +90,7 @@ internal class SabrInfoFetcher(
     ): SabrPreparedInfo? =
         withTimeoutOrNull(SabrSessionStoreDefaults.INFO_TIMEOUT_MS) {
             val tokenSession = sessionClient?.fetchPlaybackSession(videoId, isolatedPlayback)
-            tokenSession?.token
-                ?.takeIf { it.visitorData == tokenSession.info.visitorData }
-                ?.let { sessionToken ->
-                    SabrPreparedInfo(tokenSession.info, sessionToken, tokenSession.isLive, tokenSession.isLiveContent)
-                        .takeIf { it.hasAudioAndVideoFormats() }
-                        ?.let { return@withTimeoutOrNull it }
-                }
+            tokenSession?.preparedSabrInfo()?.let { return@withTimeoutOrNull it }
             val token = tokenClient.fetch(videoId)
                 ?: return@withTimeoutOrNull null.also {
                     logger.warn(

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

@@ -151,7 +151,7 @@ internal class SabrSessionStore(
 
     internal suspend fun rememberExtractedInfo(videoId: String, info: YoutubeSabrInfo): Unit =
         infoFetcher.rememberExtractedInfo(videoId, info)
-
+    internal suspend fun rememberPreparedInfo(videoId: String, prepared: SabrPreparedInfo): Unit = infoFetcher.rememberPreparedInfo(videoId, prepared)
     internal suspend fun invalidatePlaybackInfo(videoId: String): Unit = infoFetcher.invalidatePlayback(videoId)
 
     internal suspend fun recoverProtectedPlaybackInfo(holder: SabrSessionHolder): Unit =

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

@@ -18,3 +18,9 @@ internal data class TokenYoutubeSession(
     val isLiveContent: Boolean,
     val hlsUrl: String = "",
 )
+
+internal fun TokenYoutubeSession.preparedSabrInfo(): SabrPreparedInfo? {
+    val boundToken = token?.takeIf { it.visitorData == info.visitorData } ?: return null
+    return SabrPreparedInfo(info, boundToken, isLive, isLiveContent)
+        .takeIf(SabrPreparedInfo::hasAudioAndVideoFormats)
+}

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

@@ -17,8 +17,9 @@ class SabrBootstrapStreamServiceTest {
         val sessionStore = mockk<SabrSessionStore>()
         val tokenClient = mockk<TypetypeTokenYoutubeSessionClient>()
         val prepared = preparedInfo()
-        coEvery { sessionStore.fetchInfo(VIDEO_ID, cachedFirst = true) } returns prepared
-        coEvery { tokenClient.fetchPlaybackSession(VIDEO_ID) } returns tokenSession(prepared.info)
+        val session = tokenSession(prepared.info, tokenBundle())
+        coEvery { sessionStore.rememberPreparedInfo(VIDEO_ID, any()) } returns Unit
+        coEvery { tokenClient.fetchPlaybackSession(VIDEO_ID) } returns session
         val service = SabrBootstrapStreamService(sessionStore, tokenClient)
 
         val result = service.getStreamInfo(YOUTUBE_URL)
@@ -27,7 +28,8 @@ class SabrBootstrapStreamServiceTest {
         assertEquals("Bootstrap title", response.title)
         assertEquals(listOf(137), response.videoOnlyStreams.map { it.itag })
         assertEquals(listOf(140), response.audioStreams.map { it.itag })
-        coVerify(exactly = 1) { sessionStore.fetchInfo(VIDEO_ID, cachedFirst = true) }
+        coVerify(exactly = 0) { sessionStore.fetchInfo(any(), any()) }
+        coVerify(exactly = 1) { sessionStore.rememberPreparedInfo(VIDEO_ID, any()) }
         coVerify(exactly = 1) { tokenClient.fetchPlaybackSession(VIDEO_ID) }
     }
 
@@ -58,12 +60,22 @@ class SabrBootstrapStreamServiceTest {
         every { audio.audioTrackId } returns "en.4"
         val info = mockk<YoutubeSabrInfo>()
         every { info.formats } returns listOf(video, audio)
+        every { info.visitorData } returns VISITOR_DATA
         return SabrPreparedInfo(info, null)
     }
 
-    private fun tokenSession(info: YoutubeSabrInfo): TokenYoutubeSession = TokenYoutubeSession(
+    private fun tokenBundle() = SabrTokenBundle(
+        videoId = VIDEO_ID,
+        visitorBoundPoToken = "player-token",
+        visitorBoundPoTokenBytes = byteArrayOf(1),
+        visitorData = VISITOR_DATA,
+        videoBoundPoToken = "media-token",
+        videoBoundPoTokenBytes = byteArrayOf(2),
+    )
+
+    private fun tokenSession(info: YoutubeSabrInfo, token: SabrTokenBundle? = null): TokenYoutubeSession = TokenYoutubeSession(
         info = info,
-        token = null,
+        token = token,
         title = "Bootstrap title",
         author = "Bootstrap channel",
         channelId = "channel-id",
@@ -80,5 +92,6 @@ class SabrBootstrapStreamServiceTest {
     private companion object {
         const val VIDEO_ID = "f6f3PhauXyg"
         const val YOUTUBE_URL = "https://www.youtube.com/watch?v=$VIDEO_ID"
+        const val VISITOR_DATA = "visitor-data"
     }
 }