perf: bound SABR memory caches

Author: Priveetee

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

@@ -0,0 +1,97 @@
+package dev.typetype.server.services
+
+import java.time.Duration
+import java.util.LinkedHashMap
+
+internal class BoundedExpiringCache<K, V>(
+    private val maxEntries: Int,
+    private val maxWeight: Long = Long.MAX_VALUE,
+    ttl: Duration,
+    private val weigher: (V) -> Long = { 1L },
+    private val clock: () -> Long = System::currentTimeMillis,
+) {
+    private val ttlMs = ttl.toMillis()
+    private val entries = LinkedHashMap<K, Entry<V>>(maxEntries.coerceAtMost(64), 0.75f, true)
+    private var weight = 0L
+
+    init {
+        require(maxEntries > 0) { "maxEntries must be positive" }
+        require(maxWeight > 0L) { "maxWeight must be positive" }
+        require(!ttl.isNegative && !ttl.isZero) { "ttl must be positive" }
+    }
+
+    @Synchronized
+    fun get(key: K): V? {
+        evictExpired(clock())
+        return entries[key]?.value
+    }
+
+    @Synchronized
+    fun put(key: K, value: V) {
+        val now = clock()
+        evictExpired(now)
+        removeEntry(key)
+        val entryWeight = weigher(value).coerceAtLeast(0L)
+        if (entryWeight > maxWeight) return
+        entries[key] = Entry(value, expiresAt(now), entryWeight)
+        weight += entryWeight
+        trim()
+    }
+
+    @Synchronized
+    fun remove(key: K): V? = removeEntry(key)?.value
+
+    @Synchronized
+    fun removeIf(predicate: (K) -> Boolean) {
+        val iterator = entries.iterator()
+        while (iterator.hasNext()) {
+            val entry = iterator.next()
+            if (!predicate(entry.key)) continue
+            weight -= entry.value.weight
+            iterator.remove()
+        }
+    }
+
+    @Synchronized
+    fun evictExpired() {
+        evictExpired(clock())
+    }
+
+    @Synchronized
+    fun clear() {
+        entries.clear()
+        weight = 0L
+    }
+
+    @Synchronized
+    internal fun size(): Int = entries.size
+
+    @Synchronized
+    internal fun weight(): Long = weight
+
+    private fun expiresAt(now: Long): Long =
+        if (Long.MAX_VALUE - now < ttlMs) Long.MAX_VALUE else now + ttlMs
+
+    private fun evictExpired(now: Long) {
+        val iterator = entries.iterator()
+        while (iterator.hasNext()) {
+            val entry = iterator.next().value
+            if (entry.expiresAtMs > now) continue
+            weight -= entry.weight
+            iterator.remove()
+        }
+    }
+
+    private fun trim() {
+        val iterator = entries.iterator()
+        while ((entries.size > maxEntries || weight > maxWeight) && iterator.hasNext()) {
+            weight -= iterator.next().value.weight
+            iterator.remove()
+        }
+    }
+
+    private fun removeEntry(key: K): Entry<V>? =
+        entries.remove(key)?.also { weight -= it.weight }
+
+    private data class Entry<V>(val value: V, val expiresAtMs: Long, val weight: Long)
+}

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

@@ -1,25 +1,39 @@
 package dev.typetype.server.services
 
-import kotlinx.serialization.Serializable
-import kotlinx.serialization.Transient
 import org.schabi.newpipe.extractor.services.youtube.sabr.SabrMediaSegment
 import java.util.Base64
 
-@Serializable
-internal data class CachedSabrSegment(
+internal class CachedSabrSegment(
     val itag: Int,
     val sequence: Int,
     val init: Boolean,
     val startMs: Long,
     val durationMs: Long,
     val mimeType: String,
-    val bytesBase64: String,
-    val byteLength: Int = -1,
+    val bytes: ByteArray,
 ) {
-    @Transient
-    private var decodedBytes: ByteArray? = null
-    val bytes: ByteArray get() = decodedBytes ?: Base64.getDecoder().decode(bytesBase64).also { decodedBytes = it }
-    val length: Int get() = byteLength.takeIf { it >= 0 } ?: bytes.size
+    constructor(
+        itag: Int,
+        sequence: Int,
+        init: Boolean,
+        startMs: Long,
+        durationMs: Long,
+        mimeType: String,
+        bytesBase64: String,
+        byteLength: Int = -1,
+    ) : this(
+        itag,
+        sequence,
+        init,
+        startMs,
+        durationMs,
+        mimeType,
+        Base64.getDecoder().decode(bytesBase64).also {
+            require(byteLength < 0 || byteLength == it.size) { "byteLength does not match decoded bytes" }
+        },
+    )
+
+    val length: Int get() = bytes.size
 }
 
 internal fun SabrMediaSegment.toCachedSabrSegment(
@@ -32,6 +46,5 @@ internal fun SabrMediaSegment.toCachedSabrSegment(
     startMs = header.startMs,
     durationMs = header.durationMs,
     mimeType = mimeType,
-    bytesBase64 = Base64.getEncoder().encodeToString(bytes),
-    byteLength = bytes.size,
+    bytes = bytes,
 )

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

@@ -68,6 +68,8 @@ internal class SabrInfoFetcher(
     fun initializationFormat(videoId: String, target: YoutubeSabrFormat): YoutubeSabrFormat? =
         repository.initializationFormat(videoId, target)
 
+    fun evictExpired(): Unit = repository.evictExpired()
+
     private suspend fun fetchPlayable(videoId: String, startTimeMs: Long): SabrPreparedInfo? =
         fetchInfoOnce(videoId, startTimeMs)?.let { repository.putPrepared(videoId, startTimeMs, it) }
 

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

@@ -3,15 +3,18 @@ package dev.typetype.server.services
 import dev.typetype.server.cache.CacheService
 import org.schabi.newpipe.extractor.services.youtube.sabr.YoutubeSabrFormat
 import org.schabi.newpipe.extractor.services.youtube.sabr.YoutubeSabrInfo
-import java.util.concurrent.ConcurrentHashMap
+import java.time.Duration
 
 internal class SabrInfoRepository(
     infoCache: SabrPreparedInfoCache,
     sharedCache: CacheService?,
 ) {
     private val preparedInfos = infoCache
     private val sharedInfos = SabrInfoSharedCache(sharedCache)
-    private val initializationInfos = ConcurrentHashMap<String, YoutubeSabrInfo>()
+    private val initializationInfos = BoundedExpiringCache<String, YoutubeSabrInfo>(
+        maxEntries = 256,
+        ttl = Duration.ofHours(6),
+    )
 
     fun local(videoId: String, startTimeMs: Long): SabrPreparedInfo? {
         val cachedAtStart = preparedInfos.get(videoId, startTimeMs)
@@ -23,15 +26,15 @@ internal class SabrInfoRepository(
     }
 
     suspend fun shared(videoId: String, token: SabrTokenBundle): SabrPreparedInfo? {
-        sharedInfos.getInitialization(videoId)?.let { initializationInfos[videoId] = it }
+        sharedInfos.getInitialization(videoId)?.let { initializationInfos.put(videoId, it) }
         val info = sharedInfos.getPlayback(videoId) ?: return null
         if (info.visitorData != token.visitorData) return null
         return putPrepared(videoId, startTimeMs = 0L, SabrPreparedInfo(info, token), share = false)
     }
 
     suspend fun rememberInitialization(videoId: String, info: YoutubeSabrInfo): Unit {
         if (!SabrPreparedInfo(info, null).hasAudioAndVideoFormats()) return
-        initializationInfos[videoId] = info
+        initializationInfos.put(videoId, info)
         sharedInfos.putInitialization(videoId, info)
     }
 
@@ -52,7 +55,12 @@ internal class SabrInfoRepository(
     }
 
     fun initializationFormat(videoId: String, target: YoutubeSabrFormat): YoutubeSabrFormat? =
-        initializationInfos[videoId]?.formats?.firstOrNull {
+        initializationInfos.get(videoId)?.formats?.firstOrNull {
             it.itag == target.itag && it.audioTrackId == target.audioTrackId && it.xtags == target.xtags
         }
+
+    fun evictExpired() {
+        preparedInfos.evictExpired()
+        initializationInfos.evictExpired()
+    }
 }

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

@@ -6,14 +6,19 @@ 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 java.security.MessageDigest
+import java.time.Duration
 import java.util.Base64
 import java.util.Collections
 import java.util.WeakHashMap
-import java.util.concurrent.ConcurrentHashMap
 
 internal object SabrInitializationData {
     private const val CACHE_TTL_SECONDS = 21_600L
-    private val memoryCache = ConcurrentHashMap<String, ByteArray>()
+    private val memoryCache = BoundedExpiringCache<String, ByteArray>(
+        maxEntries = 512,
+        maxWeight = 32L * 1024L * 1024L,
+        ttl = Duration.ofSeconds(CACHE_TTL_SECONDS),
+        weigher = { it.size.toLong() },
+    )
     private val formatCache = Collections.synchronizedMap(WeakHashMap<YoutubeSabrFormat, ByteArray>())
 
     suspend fun ingest(
@@ -28,12 +33,12 @@ internal object SabrInitializationData {
     suspend fun fetch(videoId: String, format: YoutubeSabrFormat, cache: CacheService? = null): ByteArray? {
         val key = cacheKey(videoId, format)
         formatCache[format]?.let { return it }
-        memoryCache[key]?.let {
+        memoryCache.get(key)?.let {
             formatCache[format] = it
             return it
         }
         cache?.getBytes(key)?.let { bytes ->
-            memoryCache[key] = bytes
+            memoryCache.put(key, bytes)
             formatCache[format] = bytes
             return bytes
         }
@@ -47,7 +52,7 @@ internal object SabrInitializationData {
         cache: CacheService? = null,
     ): Unit {
         val key = cacheKey(videoId, format)
-        memoryCache[key] = bytes
+        memoryCache.put(key, bytes)
         formatCache[format] = bytes
         cache?.setBytes(key, bytes, CACHE_TTL_SECONDS)
     }
@@ -57,6 +62,8 @@ internal object SabrInitializationData {
         return holder.session.streamState.ingestInitializationData(format, bytes)
     }
 
+    fun evictExpired(): Unit = memoryCache.evictExpired()
+
     suspend fun bootstrap(
         holder: SabrSessionHolder,
         format: YoutubeSabrFormat,

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

@@ -1,43 +1,41 @@
 package dev.typetype.server.services
 
 import java.time.Duration
-import java.time.Instant
-import java.util.concurrent.ConcurrentHashMap
 
 internal class SabrPreparedInfoCache(
     private val ttl: Duration = Duration.ofMinutes(10),
+    maxEntries: Int = 256,
+    clock: () -> Long = System::currentTimeMillis,
 ) {
-    private val items = ConcurrentHashMap<Key, Entry>()
+    private val items = BoundedExpiringCache<Key, SabrPreparedInfo>(
+        maxEntries = maxEntries,
+        ttl = ttl,
+        clock = clock,
+    )
 
     fun get(videoId: String, startTimeMs: Long): SabrPreparedInfo? {
-        val key = Key(videoId, startBucket(startTimeMs))
-        val entry = items[key] ?: return null
-        if (entry.createdAt.plus(ttl).isBefore(Instant.now())) {
-            items.remove(key, entry)
-            return null
-        }
-        return entry.value
+        return items.get(Key(videoId, startBucket(startTimeMs)))
     }
 
     fun remove(videoId: String, startTimeMs: Long): Unit {
         items.remove(Key(videoId, startBucket(startTimeMs)))
     }
 
     fun remove(videoId: String): Unit {
-        items.keys.removeIf { it.videoId == videoId }
+        items.removeIf { it.videoId == videoId }
     }
 
     fun put(videoId: String, startTimeMs: Long, value: SabrPreparedInfo): SabrPreparedInfo {
-        items[Key(videoId, startBucket(startTimeMs))] = Entry(value, Instant.now())
+        items.put(Key(videoId, startBucket(startTimeMs)), value)
         return value
     }
 
+    fun evictExpired(): Unit = items.evictExpired()
+
     private fun startBucket(startTimeMs: Long): Long = startTimeMs.coerceAtLeast(0L) / START_BUCKET_MS
 
     private data class Key(val videoId: String, val startBucket: Long)
 
-    private data class Entry(val value: SabrPreparedInfo, val createdAt: Instant)
-
     private companion object {
         const val START_BUCKET_MS = 30_000L
     }

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

@@ -206,6 +206,8 @@ internal class SabrSessionStore(
         while (true) {
             delay(15_000)
             registry.evictIdle(Instant.now().minus(idleEviction))
+            infoFetcher.evictExpired()
+            SabrInitializationData.evictExpired()
         }
     }
 }

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

@@ -0,0 +1,59 @@
+package dev.typetype.server.services
+
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Assertions.assertNull
+import org.junit.jupiter.api.Test
+import java.time.Duration
+
+class BoundedExpiringCacheTest {
+    @Test
+    fun `least recently used entry is removed at capacity`() {
+        val cache = BoundedExpiringCache<String, String>(
+            maxEntries = 2,
+            ttl = Duration.ofMinutes(1),
+        )
+        cache.put("first", "1")
+        cache.put("second", "2")
+        cache.get("first")
+
+        cache.put("third", "3")
+
+        assertEquals("1", cache.get("first"))
+        assertNull(cache.get("second"))
+        assertEquals("3", cache.get("third"))
+    }
+
+    @Test
+    fun `weight limit removes oldest entries`() {
+        val cache = BoundedExpiringCache<String, String>(
+            maxEntries = 10,
+            maxWeight = 5,
+            ttl = Duration.ofMinutes(1),
+            weigher = { it.length.toLong() },
+        )
+        cache.put("first", "123")
+
+        cache.put("second", "456")
+
+        assertNull(cache.get("first"))
+        assertEquals("456", cache.get("second"))
+        assertEquals(3, cache.weight())
+    }
+
+    @Test
+    fun `new writes purge expired entries without reading their keys`() {
+        var now = 0L
+        val cache = BoundedExpiringCache<String, String>(
+            maxEntries = 10,
+            ttl = Duration.ofMillis(10),
+            clock = { now },
+        )
+        cache.put("expired", "old")
+        now = 10L
+
+        cache.put("current", "new")
+
+        assertNull(cache.get("expired"))
+        assertEquals(1, cache.size())
+    }
+}