Pool decompression by memory budget

2026-09-19 04:59:59 +00:00 · 2026-07-25 09:39:01 +02:00 · 2026-07-25 09:39:01 +02:00 · a491a65b06
commit a491a65b06
parent 6e4bb8daca
4 changed files with 242 additions and 34 deletions
--- a/app/src/main/java/com/bitchat/android/protocol/BinaryProtocol.kt
+++ b/app/src/main/java/com/bitchat/android/protocol/BinaryProtocol.kt
@ -183,10 +183,6 @@ object BinaryProtocol {
    private const val SENDER_ID_SIZE = 8
    private const val RECIPIENT_ID_SIZE = 8
    private const val SIGNATURE_SIZE = 64
-    // Acquire before copying compressed input so queued decodes cannot each retain another
-    // near-limit payload while waiting for the inflater's single-flight lock.
-    private val compressedDecodeLock = Any()
-
    object Flags {
        const val HAS_RECIPIENT: UByte = 0x01u
        const val HAS_SIGNATURE: UByte = 0x02u
@ -337,7 +333,7 @@ object BinaryProtocol {
    }
    
    fun decode(data: ByteArray): BitchatPacket? =
-        decode(data, CompressionUtil::decompress)
+        decode(data, CompressionUtil::decompressWithResourcesReserved)

    /** Test seam used to prove rejected expansion sizes never reach inflation. */
    internal fun decodeForTesting(
@ -485,9 +481,10 @@ object BinaryProtocol {
                    return null
                }
                
-                // Copy and inflate single-flight. Acquiring before the copy bounds both the
-                // compressed input copy and expanded output allocation across concurrent decodes.
-                val expandedPayload = synchronized(compressedDecodeLock) {
+                // Reserve the compressed copy plus expanded output before either allocation.
+                // Small packets share the memory pool; packets wait only while its budget is full.
+                val resourceBytes = compressedSize.toLong() + originalSize.toLong()
+                val expandedPayload = CompressionUtil.withDecompressionResources(resourceBytes) {
                    val compressedPayload = ByteArray(compressedSize)
                    buffer.get(compressedPayload)
                    decompress(compressedPayload, originalSize)
--- a/app/src/main/java/com/bitchat/android/protocol/CompressionUtil.kt
+++ b/app/src/main/java/com/bitchat/android/protocol/CompressionUtil.kt
@ -13,9 +13,7 @@ import java.util.zip.Inflater
 object CompressionUtil {
    private const val COMPRESSION_THRESHOLD = com.bitchat.android.util.AppConstants.Protocol.COMPRESSION_THRESHOLD_BYTES  // bytes - same as iOS

-    // Inflation allocates the full declared output buffer. Keep that allocation single-flight so
-    // concurrent packets cannot multiply the bounded per-packet memory cost.
-    private val decompressionLock = Any()
+    private val decompressionPool = DecompressionResourcePool.forRuntime()
    
    /**
     * Helper to check if compression is worth it - exact same logic as iOS
@ -78,47 +76,70 @@ object CompressionUtil {
     * iOS COMPRESSION_ZLIB produces raw deflate data (no headers)
     */
    fun decompress(compressedData: ByteArray, originalSize: Int): ByteArray? {
+        if (!isValidRequest(compressedData, originalSize)) return null
+        return withDecompressionResources(originalSize.toLong()) {
+            decompressWithResourcesReserved(compressedData, originalSize)
+        }
+    }
+
+    internal fun <T> withDecompressionResources(bytes: Long, block: () -> T): T? =
+        decompressionPool.withReservation(bytes, block)
+
+    /**
+     * Inflate after the caller has reserved all packet-specific allocations.
+     * This avoids nested acquisition when BinaryProtocol reserves both its input copy and output.
+     */
+    internal fun decompressWithResourcesReserved(
+        compressedData: ByteArray,
+        originalSize: Int
+    ): ByteArray? {
+        if (!isValidRequest(compressedData, originalSize)) return null
+        return decompressExact(compressedData, originalSize)
+    }
+
+    private fun isValidRequest(compressedData: ByteArray, originalSize: Int): Boolean {
        val maxExpandedSize = com.bitchat.android.util.AppConstants.Protocol.MAX_PAYLOAD_LENGTH
        if (compressedData.isEmpty()) {
            Log.w("CompressionUtil", "Refusing an empty compressed payload")
-            return null
+            return false
        }
        if (originalSize <= 0 || originalSize > maxExpandedSize) {
            Log.w(
                "CompressionUtil",
                "Refusing expanded payload size $originalSize outside 1..$maxExpandedSize"
            )
-            return null
+            return false
        }
+        return true
+    }

-        return synchronized(decompressionLock) {
-            if (looksLikeZlib(compressedData)) {
-                // A raw stream can coincidentally begin with a valid-looking zlib header. The
-                // header therefore only determines which format to try first; any non-exact zlib
-                // result must still fall back to raw under the same size/completion bounds.
-                val zlibResult = try {
-                    inflateExact(compressedData, originalSize, nowrap = false)
-                } catch (zlibException: DataFormatException) {
-                    null
-                }
-                if (zlibResult != null) {
-                    zlibResult
-                } else {
-                    try {
-                        inflateExact(compressedData, originalSize, nowrap = true)
-                    } catch (rawException: DataFormatException) {
-                        Log.d("CompressionUtil", "Invalid zlib/raw deflate stream")
-                        null
-                    }
-                }
+    private fun decompressExact(compressedData: ByteArray, originalSize: Int): ByteArray? {
+        return if (looksLikeZlib(compressedData)) {
+            // A raw stream can coincidentally begin with a valid-looking zlib header. The
+            // header therefore only determines which format to try first; any non-exact zlib
+            // result must still fall back to raw under the same size/completion bounds.
+            val zlibResult = try {
+                inflateExact(compressedData, originalSize, nowrap = false)
+            } catch (zlibException: DataFormatException) {
+                null
+            }
+            if (zlibResult != null) {
+                zlibResult
            } else {
                try {
                    inflateExact(compressedData, originalSize, nowrap = true)
                } catch (rawException: DataFormatException) {
-                    Log.d("CompressionUtil", "Invalid raw deflate stream")
+                    Log.d("CompressionUtil", "Invalid zlib/raw deflate stream")
                    null
                }
            }
+        } else {
+            try {
+                inflateExact(compressedData, originalSize, nowrap = true)
+            } catch (rawException: DataFormatException) {
+                Log.d("CompressionUtil", "Invalid raw deflate stream")
+                null
+            }
        }
    }

--- a/app/src/main/java/com/bitchat/android/protocol/DecompressionResourcePool.kt
+++ b/app/src/main/java/com/bitchat/android/protocol/DecompressionResourcePool.kt
@ -0,0 +1,73 @@
+package com.bitchat.android.protocol
+
+import java.util.concurrent.Semaphore
+import java.util.concurrent.TimeUnit
+import kotlin.math.ceil
+
+/**
+ * Fair, weighted admission control for decompression allocations.
+ *
+ * Permits represent memory rather than workers: small packets can proceed concurrently while
+ * near-limit packets consume most of the budget. Callers must reserve before allocating any
+ * packet-specific compressed copy or expanded output.
+ */
+internal class DecompressionResourcePool(
+    budgetBytes: Long,
+    private val unitBytes: Int,
+    private val waitTimeoutMs: Long
+) {
+    private val totalPermits = (budgetBytes / unitBytes).toInt().coerceAtLeast(1)
+    private val permits = Semaphore(totalPermits, true)
+
+    fun <T> withReservation(bytes: Long, block: () -> T): T? {
+        val requiredPermits = permitsFor(bytes)
+        val acquired = try {
+            permits.tryAcquire(requiredPermits, waitTimeoutMs, TimeUnit.MILLISECONDS)
+        } catch (_: InterruptedException) {
+            Thread.currentThread().interrupt()
+            false
+        }
+        if (!acquired) return null
+
+        return try {
+            block()
+        } finally {
+            permits.release(requiredPermits)
+        }
+    }
+
+    internal fun permitsFor(bytes: Long): Int =
+        ceil(bytes.coerceAtLeast(1).toDouble() / unitBytes.toDouble())
+            .toInt()
+            .coerceAtMost(totalPermits)
+
+    internal val availablePermits: Int
+        get() = permits.availablePermits()
+
+    companion object {
+        private const val DEFAULT_UNIT_BYTES = 256 * 1024
+        private const val DEFAULT_WAIT_TIMEOUT_MS = 1_000L
+        private const val HEAP_BUDGET_DIVISOR = 8L
+        private const val MAX_BUDGET_BYTES = 64L * 1024 * 1024
+
+        fun forRuntime(
+            maxHeapBytes: Long = Runtime.getRuntime().maxMemory(),
+            maxPacketResourceBytes: Long =
+                2L * com.bitchat.android.util.AppConstants.Protocol.MAX_PAYLOAD_LENGTH
+        ): DecompressionResourcePool {
+            val budget = recommendedBudgetBytes(maxHeapBytes, maxPacketResourceBytes)
+            return DecompressionResourcePool(
+                budgetBytes = budget,
+                unitBytes = DEFAULT_UNIT_BYTES,
+                waitTimeoutMs = DEFAULT_WAIT_TIMEOUT_MS
+            )
+        }
+
+        internal fun recommendedBudgetBytes(
+            maxHeapBytes: Long,
+            maxPacketResourceBytes: Long
+        ): Long = (maxHeapBytes / HEAP_BUDGET_DIVISOR)
+            .coerceAtLeast(maxPacketResourceBytes)
+            .coerceAtMost(MAX_BUDGET_BYTES.coerceAtLeast(maxPacketResourceBytes))
+    }
+}
--- a/app/src/test/java/com/bitchat/android/protocol/DecompressionResourcePoolTest.kt
+++ b/app/src/test/java/com/bitchat/android/protocol/DecompressionResourcePoolTest.kt
@ -0,0 +1,117 @@
+package com.bitchat.android.protocol
+
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertFalse
+import org.junit.Assert.assertNull
+import org.junit.Assert.assertTrue
+import org.junit.Test
+import java.util.concurrent.CountDownLatch
+import java.util.concurrent.Executors
+import java.util.concurrent.TimeUnit
+
+class DecompressionResourcePoolTest {
+    @Test
+    fun `small reservations run concurrently and next waits only when budget is full`() {
+        val pool = DecompressionResourcePool(
+            budgetBytes = 2_000,
+            unitBytes = 1_000,
+            waitTimeoutMs = 2_000
+        )
+        val executor = Executors.newFixedThreadPool(3)
+        val entered = CountDownLatch(2)
+        val release = CountDownLatch(1)
+        val thirdEntered = CountDownLatch(1)
+
+        try {
+            repeat(2) {
+                executor.submit {
+                    pool.withReservation(1_000) {
+                        entered.countDown()
+                        release.await()
+                    }
+                }
+            }
+            assertTrue(entered.await(1, TimeUnit.SECONDS))
+
+            executor.submit {
+                pool.withReservation(1_000) {
+                    thirdEntered.countDown()
+                }
+            }
+            assertFalse("third reservation must wait while budget is full", thirdEntered.await(100, TimeUnit.MILLISECONDS))
+
+            release.countDown()
+            assertTrue("third reservation must proceed after release", thirdEntered.await(1, TimeUnit.SECONDS))
+        } finally {
+            release.countDown()
+            executor.shutdownNow()
+        }
+    }
+
+    @Test
+    fun `timed admission drops work instead of waiting indefinitely`() {
+        val pool = DecompressionResourcePool(
+            budgetBytes = 1_000,
+            unitBytes = 1_000,
+            waitTimeoutMs = 50
+        )
+        val entered = CountDownLatch(1)
+        val release = CountDownLatch(1)
+        val executor = Executors.newSingleThreadExecutor()
+
+        try {
+            executor.submit {
+                pool.withReservation(1_000) {
+                    entered.countDown()
+                    release.await()
+                }
+            }
+            assertTrue(entered.await(1, TimeUnit.SECONDS))
+            assertNull(pool.withReservation(1_000) { "unexpected" })
+        } finally {
+            release.countDown()
+            executor.shutdownNow()
+        }
+    }
+
+    @Test
+    fun `permits are released when decode throws`() {
+        val pool = DecompressionResourcePool(2_000, 1_000, 50)
+
+        try {
+            pool.withReservation(2_000) { error("boom") }
+        } catch (_: IllegalStateException) {
+            // Expected.
+        }
+
+        assertEquals(2, pool.availablePermits)
+        assertEquals("ok", pool.withReservation(2_000) { "ok" })
+    }
+
+    @Test
+    fun `runtime budget is based on heap memory and always admits one maximum packet`() {
+        val maxPacketResources = 20L * 1024 * 1024
+
+        assertEquals(
+            maxPacketResources,
+            DecompressionResourcePool.recommendedBudgetBytes(
+                maxHeapBytes = 64L * 1024 * 1024,
+                maxPacketResourceBytes = maxPacketResources
+            )
+        )
+        assertEquals(
+            32L * 1024 * 1024,
+            DecompressionResourcePool.recommendedBudgetBytes(
+                maxHeapBytes = 256L * 1024 * 1024,
+                maxPacketResourceBytes = maxPacketResources
+            )
+        )
+        assertEquals(
+            64L * 1024 * 1024,
+            DecompressionResourcePool.recommendedBudgetBytes(
+                maxHeapBytes = 2L * 1024 * 1024 * 1024,
+                maxPacketResourceBytes = maxPacketResources
+            )
+        )
+    }
+}