Parallelize checksum validation

javasteam-depotdownloader/src/main/kotlin/in/dragonbra/javasteam/depotdownloader/Util.kt

@@ -7,6 +7,12 @@ import `in`.dragonbra.javasteam.types.DepotManifest
 import `in`.dragonbra.javasteam.util.Adler32
 import `in`.dragonbra.javasteam.util.log.LogManager
 import `in`.dragonbra.javasteam.util.log.Logger
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.async
+import kotlinx.coroutines.awaitAll
+import kotlinx.coroutines.coroutineScope
+import kotlinx.coroutines.sync.Semaphore
+import kotlinx.coroutines.sync.withPermit
 import okio.FileHandle
 import okio.FileSystem
 import okio.Path
@@ -25,6 +31,11 @@ object Util {
 
     private val logger: Logger = LogManager.getLogger<Util>()
 
+    private const val MAX_VALIDATE_CONCURRENCY = 16
+
+    private fun defaultValidateConcurrency(): Int =
+        Runtime.getRuntime().availableProcessors().coerceIn(1, MAX_VALIDATE_CONCURRENCY)
+
     @JvmOverloads
     @JvmStatic
     fun getSteamOS(androidEmulation: Boolean = false): String {
@@ -147,27 +158,32 @@ object Util {
      * @throws IOException If there's an error reading the file
      */
     @JvmStatic
+    @JvmOverloads
     @Throws(IOException::class)
-    fun validateSteam3FileChecksums(handle: FileHandle, chunkData: List<ChunkData>): List<ChunkData> {
-        val neededChunks = mutableListOf<ChunkData>()
-
-        for (data in chunkData) {
-            val chunk = ByteArray(data.uncompressedLength)
-            val read = handle.read(data.offset, chunk, 0, data.uncompressedLength)
-
-            val tempChunk = if (read > 0 && read < data.uncompressedLength) {
-                chunk.copyOf(read)
-            } else {
-                chunk
-            }
-
-            val adler = Adler32.calculate(tempChunk)
-            if (adler != data.checksum) {
-                neededChunks.add(data)
+    suspend fun validateSteam3FileChecksums(
+        handle: FileHandle,
+        chunkData: List<ChunkData>,
+        concurrency: Int = defaultValidateConcurrency(),
+    ): List<ChunkData> = coroutineScope {
+        val semaphore = Semaphore(concurrency.coerceAtLeast(1))
+
+        chunkData.map { data ->
+            async(Dispatchers.IO) {
+                semaphore.withPermit {
+                    val chunk = ByteArray(data.uncompressedLength)
+                    val read = handle.read(data.offset, chunk, 0, data.uncompressedLength)
+
+                    val tempChunk = if (read > 0 && read < data.uncompressedLength) {
+                        chunk.copyOf(read)
+                    } else {
+                        chunk
+                    }
+
+                    val adler = Adler32.calculate(tempChunk)
+                    if (adler != data.checksum) data else null
+                }
             }
-        }
-
-        return neededChunks
+        }.awaitAll().filterNotNull()
     }
 
     @JvmStatic

javasteam-depotdownloader/src/test/kotlin/in/dragonbra/javasteam/depotdownloader/UtilValidateSteam3FileChecksumsTest.kt

@@ -0,0 +1,119 @@
+package `in`.dragonbra.javasteam.depotdownloader
+
+import `in`.dragonbra.javasteam.types.ChunkData
+import `in`.dragonbra.javasteam.util.Adler32
+import kotlinx.coroutines.runBlocking
+import okio.FileSystem
+import okio.Path.Companion.toOkioPath
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Assertions.assertTrue
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.io.TempDir
+import java.io.File
+import kotlin.random.Random
+
+/**
+ * Verifies that [Util.validateSteam3FileChecksums] correctly identifies mismatched chunks,
+ * and that parallelizing the per-chunk checksum work does not change the result.
+ */
+class UtilValidateSteam3FileChecksumsTest {
+
+    companion object {
+        const val CHUNK_SIZE = 64
+    }
+
+    private fun writeChunkedFile(tempDir: File, chunkCount: Int, chunkSize: Int = CHUNK_SIZE): Pair<File, List<ChunkData>> {
+        val file = tempDir.resolve("data.bin")
+        val random = Random(seed = 42)
+        val chunks = mutableListOf<ChunkData>()
+
+        file.outputStream().use { out ->
+            for (i in 0 until chunkCount) {
+                val bytes = ByteArray(chunkSize) { random.nextInt().toByte() }
+                out.write(bytes)
+                chunks.add(
+                    ChunkData(
+                        offset = (i * chunkSize).toLong(),
+                        uncompressedLength = chunkSize,
+                        checksum = Adler32.calculate(bytes),
+                    )
+                )
+            }
+        }
+
+        return file to chunks
+    }
+
+    @Test
+    fun `all chunks match returns empty list`(@TempDir tempDir: File) = runBlocking {
+        val (file, chunks) = writeChunkedFile(tempDir, chunkCount = 20)
+
+        FileSystem.SYSTEM.openReadOnly(file.toOkioPath()).use { handle ->
+            val needed = Util.validateSteam3FileChecksums(handle, chunks)
+            assertTrue(needed.isEmpty(), "Expected no chunks to need re-download, got ${needed.size}")
+        }
+    }
+
+    @Test
+    fun `mismatched chunks are returned`(@TempDir tempDir: File) = runBlocking {
+        val (file, chunks) = writeChunkedFile(tempDir, chunkCount = 20)
+
+        // Corrupt the checksum of a known subset so those chunks appear stale.
+        val corruptedIndices = setOf(0, 5, 19)
+        val chunksWithCorruption = chunks.mapIndexed { index, chunk ->
+            if (index in corruptedIndices) chunk.copy(checksum = chunk.checksum xor 0x1) else chunk
+        }
+
+        FileSystem.SYSTEM.openReadOnly(file.toOkioPath()).use { handle ->
+            val needed = Util.validateSteam3FileChecksums(handle, chunksWithCorruption)
+            val neededOffsets = needed.map { it.offset }.toSet()
+            val expectedOffsets = corruptedIndices.map { chunksWithCorruption[it].offset }.toSet()
+
+            assertEquals(expectedOffsets, neededOffsets)
+        }
+    }
+
+    @Test
+    fun `result order matches input order`(@TempDir tempDir: File) = runBlocking {
+        val (file, chunks) = writeChunkedFile(tempDir, chunkCount = 50)
+        val chunksWithCorruption = chunks.map { it.copy(checksum = it.checksum xor 0x1) }
+
+        FileSystem.SYSTEM.openReadOnly(file.toOkioPath()).use { handle ->
+            val needed = Util.validateSteam3FileChecksums(handle, chunksWithCorruption)
+            assertEquals(chunksWithCorruption.map { it.offset }, needed.map { it.offset })
+        }
+    }
+
+    @Test
+    fun `concurrency parameter does not change correctness`(@TempDir tempDir: File) = runBlocking {
+        val (file, chunks) = writeChunkedFile(tempDir, chunkCount = 40)
+        val corruptedIndices = setOf(1, 2, 3, 30)
+        val chunksWithCorruption = chunks.mapIndexed { index, chunk ->
+            if (index in corruptedIndices) chunk.copy(checksum = chunk.checksum xor 0x1) else chunk
+        }
+
+        FileSystem.SYSTEM.openReadOnly(file.toOkioPath()).use { handle ->
+            val serial = Util.validateSteam3FileChecksums(handle, chunksWithCorruption, concurrency = 1)
+            val parallel = Util.validateSteam3FileChecksums(handle, chunksWithCorruption, concurrency = 8)
+
+            assertEquals(serial.map { it.offset }.toSet(), parallel.map { it.offset }.toSet())
+        }
+    }
+
+    @Test
+    fun `large chunk count completes and is correct`(@TempDir tempDir: File) = runBlocking {
+        val (file, chunks) = writeChunkedFile(tempDir, chunkCount = 1000, chunkSize = 16)
+        val corruptedIndices = (0 until 1000 step 37).toSet()
+        val chunksWithCorruption = chunks.mapIndexed { index, chunk ->
+            if (index in corruptedIndices) chunk.copy(checksum = chunk.checksum xor 0x1) else chunk
+        }
+
+        FileSystem.SYSTEM.openReadOnly(file.toOkioPath()).use { handle ->
+            val needed = Util.validateSteam3FileChecksums(handle, chunksWithCorruption)
+            val neededOffsets = needed.map { it.offset }.toSet()
+            val expectedOffsets = corruptedIndices.map { chunksWithCorruption[it].offset }.toSet()
+
+            assertEquals(expectedOffsets, neededOffsets)
+        }
+    }
+}