perf: tune video streaming path

Author: Priveetee

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

@@ -41,6 +41,8 @@ import dev.typetype.server.services.YoutubeSessionCrypto
 import dev.typetype.server.services.YoutubeSessionHlsManifestService
 import dev.typetype.server.services.YoutubeSessionService
 import dev.typetype.server.services.YoutubeSessionStreamService
+import okhttp3.ConnectionPool
+import okhttp3.Dispatcher
 import okhttp3.OkHttpClient
 import java.util.concurrent.TimeUnit
 
@@ -53,6 +55,8 @@ internal class ExtractionServiceRegistry(
         ?.takeIf { it.length >= MIN_YOUTUBE_SESSION_SECRET_LENGTH }
     val httpClient = OkHttpClient()
     val proxyHttpClient: OkHttpClient = httpClient.newBuilder()
+        .dispatcher(proxyDispatcher())
+        .connectionPool(ConnectionPool(64, 5, TimeUnit.MINUTES))
         .connectTimeout(10, TimeUnit.SECONDS)
         .readTimeout(30, TimeUnit.SECONDS)
         .followRedirects(true)
@@ -85,7 +89,7 @@ internal class ExtractionServiceRegistry(
     val nicoVideoProxyService = NicoVideoProxyService()
     val manifestService = CachedManifestService(ManifestService(streamService), cache)
     val nativeManifestService = CachedNativeManifestService(NativeManifestService(), cache)
-    val hlsManifestService = HlsManifestService(streamService, proxyHttpClient)
+    val hlsManifestService = HlsManifestService(streamService, proxyHttpClient, cache)
     val youtubeSessionHlsManifestService = hlsTokenService?.let { tokenService ->
         youtubeSessionStreamService?.let {
             YoutubeSessionHlsManifestService(youtubeSessionService, it, hlsManifestService, tokenService)
@@ -95,5 +99,10 @@ internal class ExtractionServiceRegistry(
 
     private companion object {
         const val MIN_YOUTUBE_SESSION_SECRET_LENGTH = 32
+
+        fun proxyDispatcher(): Dispatcher = Dispatcher().apply {
+            maxRequests = 256
+            maxRequestsPerHost = 64
+        }
     }
 }

src/main/kotlin/dev/typetype/server/models/ProxyResponse.kt

@@ -10,4 +10,5 @@ class ProxyResponse(
     val acceptRanges: String?,
     val stream: InputStream,
     val close: () -> Unit,
+    val cacheControl: String? = null,
 )

src/main/kotlin/dev/typetype/server/routes/ProxyRoutes.kt

@@ -28,8 +28,9 @@ fun Route.proxyRoutes(proxyService: ProxyService) {
                     val contentType = ContentType.parse(proxy.contentType)
                     proxy.contentRange?.let { call.response.headers.append("Content-Range", it) }
                     proxy.acceptRanges?.let { call.response.headers.append("Accept-Ranges", it) }
+                    proxy.cacheControl?.let { call.response.headers.append("Cache-Control", it, safeOnly = false) }
                     call.respondOutputStream(contentType, status, proxy.contentLength) {
-                        withContext(Dispatchers.IO) { proxy.stream.copyTo(this@respondOutputStream) }
+                        withContext(Dispatchers.IO) { proxy.stream.copyTo(this@respondOutputStream, PROXY_COPY_BUFFER_SIZE) }
                     }
                 } finally {
                     proxy.close()
@@ -42,3 +43,5 @@ fun Route.proxyRoutes(proxyService: ProxyService) {
         }
     }
 }
+
+private const val PROXY_COPY_BUFFER_SIZE = 64 * 1024

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

@@ -0,0 +1,19 @@
+package dev.typetype.server.services
+
+import dev.typetype.server.cache.CacheService
+
+internal class HlsManifestCache(private val cache: CacheService) {
+    suspend fun get(manifestUrl: String): String? = runCatching {
+        cache.get(key(manifestUrl))
+    }.getOrNull()
+
+    suspend fun set(manifestUrl: String, manifest: String): Unit {
+        runCatching { cache.set(key(manifestUrl), manifest, TTL_SECONDS) }
+    }
+
+    private fun key(manifestUrl: String): String = PublicCacheKey.of("hls-manifest", manifestUrl)
+
+    private companion object {
+        const val TTL_SECONDS = 30L
+    }
+}

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

@@ -1,13 +1,16 @@
 package dev.typetype.server.services
 
+import dev.typetype.server.cache.CacheService
 import dev.typetype.server.models.ExtractionResult
 import dev.typetype.server.models.StreamResponse
+import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.withContext
 import okhttp3.OkHttpClient
 import okhttp3.Request
 import java.net.URLEncoder
 import java.nio.charset.StandardCharsets
+import java.util.concurrent.ConcurrentHashMap
 
 internal fun isManifestUrl(url: String): Boolean {
     if (!url.startsWith("http")) return false
@@ -38,7 +41,10 @@ private fun toHlsProxyUrl(url: String): String {
 class HlsManifestService(
     private val streamService: StreamService,
     private val httpClient: OkHttpClient,
+    cache: CacheService? = null,
 ) {
+    private val manifestCache = cache?.let(::HlsManifestCache)
+    private val inFlight = ConcurrentHashMap<String, CompletableDeferred<ExtractionResult<String>>>()
 
     suspend fun hlsManifest(url: String): ExtractionResult<String> {
         val manifestUrl = if (isManifestUrl(url)) {
@@ -50,7 +56,7 @@ class HlsManifestService(
                 is ExtractionResult.Failure -> return resolved
             }
         }
-        return fetchAndRewrite(manifestUrl)
+        return cachedOrFetch(manifestUrl)
     }
 
     suspend fun hlsManifestFromStreamInfo(result: ExtractionResult<StreamResponse>): ExtractionResult<String> {
@@ -59,7 +65,30 @@ class HlsManifestService(
             is ExtractionResult.BadRequest -> return resolved
             is ExtractionResult.Failure -> return resolved
         }
-        return fetchAndRewrite(manifestUrl)
+        return cachedOrFetch(manifestUrl)
+    }
+
+    private suspend fun cachedOrFetch(manifestUrl: String): ExtractionResult<String> {
+        manifestCache?.get(manifestUrl)?.let { return ExtractionResult.Success(it) }
+        val pending = CompletableDeferred<ExtractionResult<String>>()
+        val existing = inFlight.putIfAbsent(manifestUrl, pending)
+        if (existing != null) return existing.await()
+        return try {
+            manifestCache?.get(manifestUrl)?.let {
+                val result = ExtractionResult.Success(it)
+                pending.complete(result)
+                return result
+            }
+            val result = fetchAndRewrite(manifestUrl)
+            if (result is ExtractionResult.Success) manifestCache?.set(manifestUrl, result.data)
+            pending.complete(result)
+            result
+        } catch (error: Throwable) {
+            pending.completeExceptionally(error)
+            throw error
+        } finally {
+            inFlight.remove(manifestUrl, pending)
+        }
     }
 
     private suspend fun resolveHlsUrl(videoUrl: String): ExtractionResult<String> {

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

@@ -57,6 +57,7 @@ class OkHttpProxyService(private val client: OkHttpClient) : ProxyService {
                         val contentType = response.header("Content-Type") ?: "application/octet-stream"
                         val contentRange = response.header("Content-Range")
                         val acceptRanges = response.header("Accept-Ranges")
+                        val cacheControl = response.header("Cache-Control")
                         val contentLength = response.header("Content-Length")?.toLongOrNull()
                         if (isHls(contentType)) {
                             val rewritten = if (isNicoNico(stripTrackingParams(fetchUrl))) {
@@ -71,6 +72,7 @@ class OkHttpProxyService(private val client: OkHttpClient) : ProxyService {
                                 contentLength = null,
                                 contentRange = null,
                                 acceptRanges = null,
+                                cacheControl = cacheControl,
                                 stream = ByteArrayInputStream(rewritten.toByteArray(StandardCharsets.UTF_8)),
                                 close = {},
                             ))
@@ -81,6 +83,7 @@ class OkHttpProxyService(private val client: OkHttpClient) : ProxyService {
                                 contentLength = contentLength,
                                 contentRange = contentRange,
                                 acceptRanges = acceptRanges,
+                                cacheControl = cacheControl,
                                 stream = body.byteStream(),
                                 close = response::close,
                             ))

src/test/kotlin/dev/typetype/server/HlsManifestServiceCacheTest.kt

@@ -0,0 +1,52 @@
+package dev.typetype.server
+
+import dev.typetype.server.cache.CacheService
+import dev.typetype.server.models.ExtractionResult
+import dev.typetype.server.models.StreamResponse
+import dev.typetype.server.services.HlsManifestService
+import dev.typetype.server.services.StreamService
+import kotlinx.coroutines.test.runTest
+import okhttp3.Interceptor
+import okhttp3.MediaType.Companion.toMediaType
+import okhttp3.OkHttpClient
+import okhttp3.Protocol
+import okhttp3.Response
+import okhttp3.ResponseBody.Companion.toResponseBody
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Test
+
+class HlsManifestServiceCacheTest {
+    @Test
+    fun `hls manifests are cached briefly by manifest url`() = runTest {
+        var calls = 0
+        val client = OkHttpClient.Builder().addInterceptor(Interceptor { chain ->
+            calls += 1
+            Response.Builder()
+                .request(chain.request())
+                .protocol(Protocol.HTTP_1_1)
+                .code(200)
+                .message("OK")
+                .body("#EXTM3U\nsegment.ts".toResponseBody("application/vnd.apple.mpegurl".toMediaType()))
+                .build()
+        }).build()
+        val service = HlsManifestService(NoopStreamService, client, InMemoryCache())
+        val url = "https://example.com/master.m3u8"
+
+        service.hlsManifest(url)
+        service.hlsManifest(url)
+
+        assertEquals(1, calls)
+    }
+}
+
+private object NoopStreamService : StreamService {
+    override suspend fun getStreamInfo(url: String): ExtractionResult<StreamResponse> =
+        ExtractionResult.Failure("unused")
+}
+
+private class InMemoryCache : CacheService {
+    private val values = mutableMapOf<String, String>()
+    override suspend fun get(key: String): String? = values[key]
+    override suspend fun set(key: String, value: String, ttlSeconds: Long) { values[key] = value }
+    override suspend fun delete(key: String) { values.remove(key) }
+}