Optimize watch sync/cache CPU, memory, and battery usage

Author: celestialtaha

app/src/main/java/com/wapp/wearmessage/presentation/MessagingViewModel.kt

@@ -11,6 +11,7 @@ import com.wapp.wearmessage.sync.SyncInboundEvent
 import com.wapp.wearmessage.sync.WearSyncBus
 import com.wapp.wearmessage.sync.contract.SyncMessageStatus
 import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.Job
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
@@ -29,6 +30,7 @@ class MessagingViewModel(
     val uiState: StateFlow<WearMessagingUiState> = _uiState.asStateFlow()
     private val secureCache = SecureMessageCache(application.applicationContext)
     private var hasReceivedSyncPayload = false
+    private var persistCacheJob: Job? = null
 
     init {
         bootstrapData()
@@ -309,9 +311,14 @@ class MessagingViewModel(
                 .distinctUntilChanged()
                 .collect { snapshot ->
                     if (snapshot != null) {
-                        withContext(Dispatchers.IO) {
-                            secureCache.save(snapshot)
-                        }
+                        persistCacheJob?.cancel()
+                        persistCacheJob =
+                            viewModelScope.launch {
+                                delay(CACHE_PERSIST_DEBOUNCE_MS)
+                                withContext(Dispatchers.IO) {
+                                    secureCache.save(snapshot)
+                                }
+                            }
                     }
                 }
         }
@@ -363,6 +370,8 @@ class MessagingViewModel(
             val conversations =
                 mergedById.values
                     .sortedByDescending { it.lastUpdatedAtEpochMillis }
+                    .take(MAX_CONVERSATIONS_IN_MEMORY)
+            val retainedConversationIds = conversations.map { it.id }.toSet()
 
             state.copy(
                 conversationsState =
@@ -371,6 +380,10 @@ class MessagingViewModel(
                     } else {
                         ConversationsUiState.Success(conversations)
                     },
+                messagesByConversation =
+                    state.messagesByConversation.filterKeys { conversationId ->
+                        retainedConversationIds.contains(conversationId)
+                    },
                 syncStatus = SyncStatus.Idle,
             )
         }
@@ -399,9 +412,18 @@ class MessagingViewModel(
                             )
                         }
                 }
+            val mergedMessagesByConversation =
+                (state.messagesByConversation + mappedByConversation)
+                    .mapValues { (_, messages) ->
+                        if (messages.size <= MAX_MESSAGES_PER_CONVERSATION_IN_MEMORY) {
+                            messages
+                        } else {
+                            messages.takeLast(MAX_MESSAGES_PER_CONVERSATION_IN_MEMORY)
+                        }
+                    }
 
             state.copy(
-                messagesByConversation = state.messagesByConversation + mappedByConversation,
+                messagesByConversation = mergedMessagesByConversation,
                 syncStatus = SyncStatus.Idle,
             )
         }
@@ -423,6 +445,11 @@ class MessagingViewModel(
             )
         }
     }
+
+    override fun onCleared() {
+        persistCacheJob?.cancel()
+        super.onCleared()
+    }
 }
 
 private fun ConversationsUiState.withConversationList(
@@ -490,10 +517,24 @@ private fun WearMessagingUiState.toCachedSyncSnapshotOrNull(): CachedSyncSnapsho
     val messages =
         messagesByConversation
             .values
+            .map { conversationMessages ->
+                if (conversationMessages.size <= MAX_MESSAGES_PER_CONVERSATION_CACHE) {
+                    conversationMessages
+                } else {
+                    conversationMessages.takeLast(MAX_MESSAGES_PER_CONVERSATION_CACHE)
+                }
+            }
             .flatten()
             .sortedWith(
                 compareBy<Message>({ it.conversationId }, { it.id })
             )
+            .let { sortedMessages ->
+                if (sortedMessages.size <= MAX_MESSAGES_TOTAL_CACHE) {
+                    sortedMessages
+                } else {
+                    sortedMessages.takeLast(MAX_MESSAGES_TOTAL_CACHE)
+                }
+            }
             .map { message ->
                 CachedMessage(
                     id = message.id,
@@ -511,7 +552,6 @@ private fun WearMessagingUiState.toCachedSyncSnapshotOrNull(): CachedSyncSnapsho
     return CachedSyncSnapshot(
         conversations = conversations,
         messages = messages,
-        savedAtEpochMillis = System.currentTimeMillis(),
     )
 }
 
@@ -537,3 +577,9 @@ private fun Long.toClockLabel(): String {
     val minute = (minutes % 60).toInt().absoluteValue
     return "%02d:%02d".format(hour, minute)
 }
+
+private const val CACHE_PERSIST_DEBOUNCE_MS = 500L
+private const val MAX_CONVERSATIONS_IN_MEMORY = 300
+private const val MAX_MESSAGES_PER_CONVERSATION_IN_MEMORY = 250
+private const val MAX_MESSAGES_PER_CONVERSATION_CACHE = 120
+private const val MAX_MESSAGES_TOTAL_CACHE = 1500

app/src/main/java/com/wapp/wearmessage/storage/SecureMessageCache.kt

@@ -7,6 +7,7 @@ import android.util.Base64
 import org.json.JSONArray
 import org.json.JSONObject
 import java.nio.ByteBuffer
+import java.security.MessageDigest
 import java.security.KeyStore
 import javax.crypto.Cipher
 import javax.crypto.KeyGenerator
@@ -19,16 +20,25 @@ class SecureMessageCache(
     private val prefs by lazy {
         appContext.getSharedPreferences(PREFS_NAME, Context.MODE_PRIVATE)
     }
+    private val secretKey by lazy { getOrCreateSecretKey() }
+    @Volatile
+    private var lastSavedDigest: String? = prefs.getString(PREF_KEY_CACHE_DIGEST, null)
 
     fun save(snapshot: CachedSyncSnapshot) {
         val plainPayload = snapshot.toJson().toString().toByteArray(Charsets.UTF_8)
+        val digest = plainPayload.sha256B64()
+        if (digest == lastSavedDigest) {
+            return
+        }
         val encryptedPayload = encrypt(plainPayload) ?: return
         prefs.edit()
             .putString(
                 PREF_KEY_CACHE_BLOB,
                 Base64.encodeToString(encryptedPayload, Base64.NO_WRAP),
             )
+            .putString(PREF_KEY_CACHE_DIGEST, digest)
             .apply()
+        lastSavedDigest = digest
     }
 
     fun load(): CachedSyncSnapshot? {
@@ -42,18 +52,23 @@ class SecureMessageCache(
             runCatching { JSONObject(String(plainPayload, Charsets.UTF_8)) }
                 .getOrNull()
                 ?: return clearAndNull()
+        lastSavedDigest = prefs.getString(PREF_KEY_CACHE_DIGEST, null)
         return CachedSyncSnapshot.fromJson(json) ?: clearAndNull()
     }
 
     private fun clearAndNull(): CachedSyncSnapshot? {
-        prefs.edit().remove(PREF_KEY_CACHE_BLOB).apply()
+        prefs.edit()
+            .remove(PREF_KEY_CACHE_BLOB)
+            .remove(PREF_KEY_CACHE_DIGEST)
+            .apply()
+        lastSavedDigest = null
         return null
     }
 
     private fun encrypt(plain: ByteArray): ByteArray? =
         runCatching {
             val cipher = Cipher.getInstance(TRANSFORMATION)
-            cipher.init(Cipher.ENCRYPT_MODE, getOrCreateSecretKey())
+            cipher.init(Cipher.ENCRYPT_MODE, secretKey)
             val iv = cipher.iv ?: return null
             val encrypted = cipher.doFinal(plain)
             ByteBuffer.allocate(1 + iv.size + encrypted.size)
@@ -74,7 +89,7 @@ class SecureMessageCache(
             val encrypted = ByteArray(payloadBuffer.remaining())
             payloadBuffer.get(encrypted)
             val cipher = Cipher.getInstance(TRANSFORMATION)
-            cipher.init(Cipher.DECRYPT_MODE, getOrCreateSecretKey(), javax.crypto.spec.GCMParameterSpec(TAG_SIZE_BITS, iv))
+            cipher.init(Cipher.DECRYPT_MODE, secretKey, javax.crypto.spec.GCMParameterSpec(TAG_SIZE_BITS, iv))
             cipher.doFinal(encrypted)
         }.getOrNull()
 
@@ -105,6 +120,7 @@ class SecureMessageCache(
     companion object {
         private const val PREFS_NAME = "wessage_secure_message_cache"
         private const val PREF_KEY_CACHE_BLOB = "cache_blob_v1"
+        private const val PREF_KEY_CACHE_DIGEST = "cache_blob_digest_v1"
         private const val KEY_ALIAS = "wessage.message.cache.key.v1"
         private const val ANDROID_KEYSTORE = "AndroidKeyStore"
         private const val TRANSFORMATION = "AES/GCM/NoPadding"
@@ -115,12 +131,11 @@ class SecureMessageCache(
 data class CachedSyncSnapshot(
     val conversations: List<CachedConversation>,
     val messages: List<CachedMessage>,
-    val savedAtEpochMillis: Long,
 ) {
     fun toJson(): JSONObject =
         JSONObject()
             .put("schemaVersion", SCHEMA_VERSION)
-            .put("savedAtEpochMillis", savedAtEpochMillis)
+            .put("savedAtEpochMillis", System.currentTimeMillis())
             .put(
                 "conversations",
                 JSONArray().apply {
@@ -207,7 +222,6 @@ data class CachedSyncSnapshot(
             return CachedSyncSnapshot(
                 conversations = conversations,
                 messages = messages,
-                savedAtEpochMillis = json.optLong("savedAtEpochMillis", 0L),
             )
         }
     }
@@ -241,3 +255,8 @@ private fun JSONArray.toStringList(): List<String> =
             add(optString(index))
         }
     }
+
+private fun ByteArray.sha256B64(): String {
+    val digest = MessageDigest.getInstance("SHA-256").digest(this)
+    return Base64.encodeToString(digest, Base64.NO_WRAP)
+}

app/src/main/java/com/wapp/wearmessage/sync/WearSyncTransport.kt

@@ -36,13 +36,15 @@ class WearSyncTransport(
                     )
                 var sent = false
                 nodes.forEach { node ->
-                    Tasks.await(
-                        messageClient.sendMessage(
-                            node.id,
-                            SyncPaths.KEY_EXCHANGE_REQUEST,
-                            secureCodec.createKeyExchangePayload(),
+                    if (!secureCodec.hasPeerKey(node.id)) {
+                        Tasks.await(
+                            messageClient.sendMessage(
+                                node.id,
+                                SyncPaths.KEY_EXCHANGE_REQUEST,
+                                secureCodec.createKeyExchangePayload(),
+                            )
                         )
-                    )
+                    }
                     Tasks.await(
                         messageClient.sendMessage(
                             node.id,
@@ -65,6 +67,16 @@ class WearSyncTransport(
                 val plainPayload = SyncJsonCodec.encodeWatchMutation(mutation)
                 var sent = false
                 nodes.forEach { node ->
+                    if (!secureCodec.hasPeerKey(node.id)) {
+                        Tasks.await(
+                            messageClient.sendMessage(
+                                node.id,
+                                SyncPaths.KEY_EXCHANGE_REQUEST,
+                                secureCodec.createKeyExchangePayload(),
+                            )
+                        )
+                        return@forEach
+                    }
                     val encrypted =
                         secureCodec.encrypt(
                             nodeId = node.id,
@@ -79,6 +91,7 @@ class WearSyncTransport(
                                 secureCodec.createKeyExchangePayload(),
                             )
                         )
+                        return@forEach
                     } else {
                         Tasks.await(messageClient.sendMessage(node.id, SyncPaths.MUTATION, encrypted))
                         sent = true