Asset Modules: Introduce session manager with one id per download request

vending-app/src/main/kotlin/com/google/android/finsky/assetmoduleservice/AssetModuleService.kt

@@ -57,6 +57,7 @@ class AssetModuleServiceImpl(
 ) : AbstractAssetModuleServiceImpl(context, lifecycle) {
     private val fileDescriptorMap = mutableMapOf<File, ParcelFileDescriptor>()
     private val lock = Any()
+    private val sessionManager = AssetModuleSessionManager(context)
 
     private fun checkSessionValid(packageName: String, sessionId: Int) {
         Log.d(TAG, "checkSessionValid: $packageName $sessionId ${packageDownloadData[packageName]?.sessionIds}")
@@ -92,6 +93,21 @@ class AssetModuleServiceImpl(
             Log.d(TAG, "startDownload: resetAllModuleStatus ")
             packageDownloadData[packageName]?.resetAllModuleStatus()
         }
+
+        params.moduleNames.forEach { moduleName ->
+            val status = synchronized(lock) {
+                packageDownloadData[packageName]?.getModuleData(moduleName)?.status
+            }
+            if (status == AssetPackStatus.DOWNLOADING) {
+                Log.w(TAG, "startDownload: module $moduleName still DOWNLOADING; overlapping cancel deferred to #2673")
+            }
+        }
+
+        val session = sessionManager.createSession(packageName, params.moduleNames)
+        synchronized(lock) {
+            packageDownloadData[packageName]?.bindModulesToSession(context, session.sessionId, params.moduleNames)
+        }
+
         params.moduleNames.forEach {
             val moduleData = packageDownloadData[packageName]?.getModuleData(it)
             if (moduleData?.status != AssetPackStatus.DOWNLOADING && moduleData?.status != AssetPackStatus.COMPLETED) {
@@ -110,8 +126,8 @@ class AssetModuleServiceImpl(
             }
         }
         val bundleData = buildDownloadBundle(packageDownloadData[packageName]!!, params.moduleNames)
-        Log.d(TAG, "startDownload: $bundleData")
-        callback?.onStartDownload(-1, bundleData)
+        Log.d(TAG, "startDownload: sessionId=${session.sessionId} $bundleData")
+        callback?.onStartDownload(session.sessionId, bundleData)
         params.moduleNames.forEach {
             val packData = packageDownloadData[packageName]?.getModuleData(it)
             if (packData?.status == AssetPackStatus.PENDING) {
@@ -134,15 +150,24 @@ class AssetModuleServiceImpl(
                 return
             }
 
-            packageDownloadData[packageName]?.moduleNames?.forEach { moduleName ->
-                if (moduleName in params.installedAssetModules) return@forEach
-
-                listBundleData.add(sendBroadcastForExistingFile(context, packageDownloadData[packageName]!!, moduleName, null, null))
+            val downloadData = packageDownloadData[packageName]
+            if (downloadData != null) {
+                val modulesBySession = downloadData.moduleNames
+                    .filter { it !in params.installedAssetModules }
+                    .groupBy { downloadData.sessionIds[it] ?: downloadData.sessionId }
 
-                packageDownloadData[packageName]?.getModuleData(moduleName)?.chunks?.forEach { chunkData ->
-                    val destination = chunkData.getChunkFile(context)
-                    if (destination.exists() && destination.length() == chunkData.chunkBytesToDownload) {
-                        sendBroadcastForExistingFile(context, packageDownloadData[packageName]!!, moduleName, chunkData, destination)
+                modulesBySession.forEach { (sessionId, modules) ->
+                    if (sessionId == 0 && modules.all { downloadData.getModuleData(it).status == AssetPackStatus.NOT_INSTALLED }) {
+                        return@forEach
+                    }
+                    listBundleData.add(buildDownloadBundle(downloadData, modules))
+                    modules.forEach { moduleName ->
+                        downloadData.getModuleData(moduleName).chunks.forEach { chunkData ->
+                            val destination = chunkData.getChunkFile(context)
+                            if (destination.exists() && destination.length() == chunkData.chunkBytesToDownload) {
+                                sendBroadcastForExistingFile(context, downloadData, moduleName, chunkData, destination)
+                            }
+                        }
                     }
                 }
             }
@@ -176,6 +201,7 @@ class AssetModuleServiceImpl(
         synchronized(lock) {
             packageDownloadData[packageName]?.updateDownloadStatus(params.moduleName, AssetPackStatus.COMPLETED)
             sendBroadcastForExistingFile(context, packageDownloadData[packageName]!!, params.moduleName, null, null)
+            packageDownloadData[packageName]?.let { sessionManager.markTerminalIfDone(params.sessionId, it) }
         }
 
         val directory = context.getModuleDir(params.sessionId, params.moduleName)
@@ -195,9 +221,19 @@ class AssetModuleServiceImpl(
     override suspend fun notifySessionFailed(params: NotifySessionFailedParameters, packageName: String, callback: IAssetModuleServiceCallback?) {
         checkSessionValid(packageName, params.sessionId)
 
-        // TODO: Implement
+        synchronized(lock) {
+            val downloadData = packageDownloadData[packageName]
+            if (downloadData != null) {
+                val modulesInSession = downloadData.sessionIds.filter { it.value == params.sessionId }.keys
+                modulesInSession.forEach { moduleName ->
+                    downloadData.updateDownloadStatus(moduleName, AssetPackStatus.FAILED)
+                    sendBroadcastForExistingFile(context, downloadData, moduleName, null, null)
+                }
+                sessionManager.markTerminalIfDone(params.sessionId, downloadData)
+            }
+        }
+        // Full cancel of in-flight downloads is #2673
         callback?.onNotifySessionFailed(bundleOf(BundleKeys.SESSION_ID to params.sessionId))
-        //throw UnsupportedOperationException()
     }
 
     override suspend fun keepAlive(params: KeepAliveParameters, packageName: String, callback: IAssetModuleServiceCallback?) {
@@ -251,6 +287,7 @@ class AssetModuleServiceImpl(
                 callback?.onError(Bundle().apply { put(BundleKeys.ERROR_CODE, AssetPackErrorCode.NETWORK_ERROR) })
             }
         }
+        // Do not mint a new session here — only startDownload allocates session ids.
         val bundleData = buildDownloadBundle(packageDownloadData[packageName]!!, params.moduleNames)
         Log.d(TAG, "requestDownloadInfo -> $bundleData")
         callback?.onRequestDownloadInfo(bundleData, bundleData)
@@ -267,4 +304,4 @@ class AssetModuleServiceImpl(
         callback?.onCancelDownloads(bundleOf(BundleKeys.ERROR_CODE to AssetPackErrorCode.NO_ERROR))
         //throw UnsupportedOperationException()
     }
-}
+}

vending-app/src/main/kotlin/com/google/android/finsky/assetmoduleservice/AssetModuleSessionManager.kt

@@ -0,0 +1,94 @@
+/*
+ * SPDX-FileCopyrightText: 2023 microG Project Team
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package com.google.android.finsky.assetmoduleservice
+
+import android.content.Context
+import android.os.SystemClock
+import android.util.Log
+import com.google.android.play.core.assetpacks.model.AssetPackStatus
+
+private const val SESSION_MGR_TAG = "AssetModuleSessionMgr"
+private const val PREFS_NAME = "AssetModuleSessionIdGenerator"
+private const val PREFS_KEY_LATEST = "Latest"
+
+data class AssetModuleSession(
+    val sessionId: Int,
+    val packageName: String,
+    val moduleNames: Set<String>,
+    val createdAtElapsedMs: Long,
+    var terminal: Boolean = false
+)
+
+/**
+ * Play-like session registry for asset module downloads.
+ * One ever-increasing session id covers all modules in a single startDownload request.
+ */
+class AssetModuleSessionManager(private val context: Context) {
+    private val lock = Any()
+    private val sessionsById = mutableMapOf<Int, AssetModuleSession>()
+    private val sessionIdsByPackage = mutableMapOf<String, MutableSet<Int>>()
+
+    fun nextSessionId(): Int = synchronized(lock) {
+        allocateSessionIdLocked()
+    }
+
+    fun createSession(packageName: String, modules: Collection<String>): AssetModuleSession {
+        synchronized(lock) {
+            val sessionId = allocateSessionIdLocked()
+            val session = AssetModuleSession(
+                sessionId = sessionId,
+                packageName = packageName,
+                moduleNames = modules.toSet(),
+                createdAtElapsedMs = SystemClock.elapsedRealtime()
+            )
+            sessionsById[sessionId] = session
+            sessionIdsByPackage.getOrPut(packageName) { mutableSetOf() }.add(sessionId)
+            Log.d(SESSION_MGR_TAG, "createSession: package=$packageName sessionId=$sessionId modules=${session.moduleNames}")
+            return session
+        }
+    }
+
+    fun getSession(sessionId: Int): AssetModuleSession? = synchronized(lock) {
+        sessionsById[sessionId]
+    }
+
+    fun findActiveSession(packageName: String, module: String): AssetModuleSession? = synchronized(lock) {
+        val ids = sessionIdsByPackage[packageName] ?: return null
+        ids.mapNotNull { sessionsById[it] }
+            .filter { !it.terminal && module in it.moduleNames }
+            .maxByOrNull { it.createdAtElapsedMs }
+    }
+
+    fun sessionsForPackage(packageName: String): List<AssetModuleSession> = synchronized(lock) {
+        sessionIdsByPackage[packageName]
+            ?.mapNotNull { sessionsById[it] }
+            ?.sortedBy { it.createdAtElapsedMs }
+            ?: emptyList()
+    }
+
+    fun markTerminalIfDone(sessionId: Int, downloadData: DownloadData) {
+        synchronized(lock) {
+            val session = sessionsById[sessionId] ?: return
+            val allTerminal = session.moduleNames.all { moduleName ->
+                val status = runCatching { downloadData.getModuleData(moduleName).status }.getOrNull()
+                status == AssetPackStatus.COMPLETED ||
+                    status == AssetPackStatus.FAILED ||
+                    status == AssetPackStatus.CANCELED
+            }
+            if (allTerminal) {
+                session.terminal = true
+                Log.d(SESSION_MGR_TAG, "markTerminalIfDone: sessionId=$sessionId terminal=true")
+            }
+        }
+    }
+
+    private fun allocateSessionIdLocked(): Int {
+        val sharedPreferences = context.getSharedPreferences(PREFS_NAME, Context.MODE_PRIVATE)
+        val latest = sharedPreferences.getInt(PREFS_KEY_LATEST, 0) + 1
+        sharedPreferences.edit().putInt(PREFS_KEY_LATEST, latest).commit()
+        return latest
+    }
+}

vending-app/src/main/kotlin/com/google/android/finsky/assetmoduleservice/DownloadData.kt

@@ -6,16 +6,15 @@
 package com.google.android.finsky.assetmoduleservice
 
 import android.content.Context
-import android.util.Log
 import com.google.android.finsky.getChunkFile
 import com.google.android.play.core.assetpacks.model.AssetPackStatus
 import com.google.android.play.core.assetpacks.protocol.CompressionFormat
-import java.io.File
 import java.io.Serializable
 
 data class DownloadData(
     var packageName: String = "",
     var errorCode: Int = 0,
+    var sessionId: Int = 0,
     var sessionIds: Map<String, Int> = emptyMap(),
     var status: Int = 0,
     var moduleNames: Set<String> = emptySet(),
@@ -48,7 +47,11 @@ fun DownloadData?.merge(data: DownloadData?): DownloadData? {
     if (this == null) return data
     if (data == null) return this
     moduleNames += data.moduleNames
-    sessionIds += data.sessionIds.filter { it.key !in sessionIds.keys }
+    // Overwrite session mapping for modules present in incoming data (new startDownload session)
+    sessionIds = sessionIds + data.sessionIds
+    if (data.sessionId != 0) {
+        sessionId = data.sessionId
+    }
     moduleDataMap += data.moduleDataMap.filter { it.key !in moduleDataMap.keys }
     return this
 }
@@ -85,4 +88,4 @@ data class ChunkData(
     val numberOfChunksInSlice: Int
 )
 
-fun ChunkData.getChunkFile(context: Context) = context.getChunkFile(sessionId, moduleName, sliceId, chunkIndex)
+fun ChunkData.getChunkFile(context: Context) = context.getChunkFile(sessionId, moduleName, sliceId, chunkIndex)