refactor(engine): Improve recovery scheduling implementation (#144)

Author: harrytmthy

safebox/src/main/java/com/harrytmthy/safebox/engine/SafeBoxEngine.kt

@@ -38,7 +38,6 @@ import com.harrytmthy.safebox.storage.SafeBoxRecoveryBlobStore
 import com.harrytmthy.safebox.strategy.ValueFallbackStrategy
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.CoroutineDispatcher
-import kotlinx.coroutines.Job
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
@@ -78,11 +77,9 @@ internal class SafeBoxEngine private constructor(
 
     private val scanScheduled = AtomicBoolean(false)
 
-    private var callback: Callback? = null
-
-    private var recoveryJob: Job? = null
+    private val recoveryScheduled = AtomicBoolean(false)
 
-    private var recoveryBackoffMs = 1_000L
+    private var callback: Callback? = null
 
     init {
         launchWithStartingState {
@@ -93,10 +90,6 @@ internal class SafeBoxEngine private constructor(
                     entries += it
                     recoveryEntries += it
                 }
-        }.invokeOnCompletion {
-            if (recoveryEntries.isNotEmpty()) {
-                scheduleRecoveryEntriesWrite()
-            }
         }
     }
 
@@ -215,7 +208,6 @@ internal class SafeBoxEngine private constructor(
                             encryptedValue = encryptedValue,
                         )
                         recoveryEntries[encryptedKey] = encryptedValue
-                        scheduleRecoveryEntriesWrite()
                     }
                 }
                 is Remove -> {
@@ -228,46 +220,32 @@ internal class SafeBoxEngine private constructor(
         }
     }
 
-    private fun scheduleRecoveryEntriesWrite() {
-        if (recoveryJob?.isActive == true) {
-            return
-        }
-        recoveryJob = launchWriteAsync {
-            val fileId = blobStore.getFileName().toBytes()
-            val snapshot = ArrayList(recoveryEntries.entries)
-            var failed = false
-            for ((encryptedKey, encryptedValue) in snapshot) {
-                try {
-                    blobStore.write(encryptedKey, encryptedValue)
-                    recoveryEntries.remove(encryptedKey)
-                } catch (_: Exception) {
-                    failed = true
-                    break
-                }
-            }
-            if (!failed && recoveryEntries.isEmpty()) {
-                recoveryBlobStore.delete(fileId)
-                recoveryBackoffMs = 1_000L
-            }
-        }.apply {
-            invokeOnCompletion { cause ->
-                if (recoveryEntries.isEmpty()) {
-                    return@invokeOnCompletion
-                }
-                if (cause != null) {
-                    recoveryBackoffMs = (recoveryBackoffMs * 2).coerceAtMost(30_000L)
-                }
-                safeBoxScope.launch(ioDispatcher) {
-                    delay(recoveryBackoffMs)
-                    scheduleRecoveryEntriesWrite()
+    private fun scheduleRecoveryEntriesWrite(recoveryBackoffMs: Long) {
+        safeBoxScope.launch(ioDispatcher) {
+            delay(recoveryBackoffMs)
+        }.invokeOnCompletion {
+            val nextBackoffMs = (recoveryBackoffMs * 2).coerceAtMost(MAX_BACKOFF_MS)
+            launchWriteAsync(nextBackoffMs) {
+                val fileName = blobStore.getFileName().toBytes()
+                val snapshot = ArrayList(recoveryEntries.entries)
+                val removedKeys = ArrayList<Bytes>()
+                for ((encryptedKey, encryptedValue) in snapshot) {
+                    try {
+                        blobStore.write(encryptedKey, encryptedValue)
+                        recoveryEntries.remove(encryptedKey)
+                        removedKeys += encryptedKey
+                    } catch (_: Exception) {
+                        continue
+                    }
                 }
+                recoveryBlobStore.delete(fileName, *removedKeys.toTypedArray())
             }
         }
     }
 
-    private inline fun launchWithStartingState(crossinline block: suspend () -> Unit): Job {
+    private inline fun launchWithStartingState(crossinline block: suspend () -> Unit) {
         updateState(STARTING)
-        return safeBoxScope.launch(ioDispatcher) {
+        safeBoxScope.launch(ioDispatcher) {
             try {
                 block()
             } finally {
@@ -276,12 +254,15 @@ internal class SafeBoxEngine private constructor(
                     updateState(IDLE)
                 }
             }
+        }.invokeOnCompletion {
+            if (recoveryEntries.isNotEmpty()) {
+                recoveryScheduled.set(true)
+                scheduleRecoveryEntriesWrite(DEFAULT_BACKOFF_MS)
+            }
         }
     }
 
-    private inline fun launchWriteBlocking(
-        crossinline block: suspend () -> Unit,
-    ): Boolean {
+    private inline fun launchWriteBlocking(crossinline block: suspend () -> Unit): Boolean {
         val currentWriteBarrier = CompletableDeferred<Unit>()
         val previousWriteBarrier = writeBarrier.getAndSet(currentWriteBarrier)
         return runBlocking {
@@ -293,21 +274,37 @@ internal class SafeBoxEngine private constructor(
                 false
             } finally {
                 currentWriteBarrier.complete(Unit)
+                if (recoveryEntries.isNotEmpty() && recoveryScheduled.compareAndSet(false, true)) {
+                    scheduleRecoveryEntriesWrite(DEFAULT_BACKOFF_MS)
+                }
             }
         }
     }
 
-    private inline fun launchWriteAsync(crossinline block: suspend () -> Unit): Job {
+    private inline fun launchWriteAsync(
+        recoveryBackoffMs: Long = DEFAULT_BACKOFF_MS,
+        crossinline block: suspend () -> Unit,
+    ) {
         val currentWriteBarrier = CompletableDeferred<Unit>()
         val previousWriteBarrier = writeBarrier.getAndSet(currentWriteBarrier)
-        return safeBoxScope.launch(ioDispatcher) {
+        safeBoxScope.launch(ioDispatcher) {
             try {
                 withStateTransition(previousWriteBarrier, block)
             } catch (e: Exception) {
                 Log.e("SafeBox", "Failed to commit changes.", e)
             } finally {
                 currentWriteBarrier.complete(Unit)
             }
+        }.invokeOnCompletion {
+            if (recoveryBackoffMs == DEFAULT_BACKOFF_MS && recoveryScheduled.get()) {
+                return@invokeOnCompletion
+            }
+            if (recoveryEntries.isNotEmpty()) {
+                recoveryScheduled.set(true)
+                scheduleRecoveryEntriesWrite(recoveryBackoffMs)
+            } else {
+                recoveryScheduled.set(false)
+            }
         }
     }
 
@@ -390,6 +387,10 @@ internal class SafeBoxEngine private constructor(
 
     internal companion object {
 
+        private const val DEFAULT_BACKOFF_MS = 1000L
+
+        private const val MAX_BACKOFF_MS = 30000L
+
         fun create(
             context: Context,
             fileName: String,

safebox/src/main/java/com/harrytmthy/safebox/storage/SafeBoxBlobStore.kt

@@ -180,12 +180,12 @@ internal class SafeBoxBlobStore private constructor(private val file: File) {
      * @param encryptedKeys Vararg array of keys to delete.
      */
     internal suspend fun delete(vararg encryptedKeys: Bytes) {
+        if (encryptedKeys.isEmpty() || entryMetas.isEmpty()) {
+            return
+        }
         writeMutex.withLock {
             for (encryptedKey in encryptedKeys) {
-                if (!entryMetas.containsKey(encryptedKey)) {
-                    continue
-                }
-                val entry = entryMetas.getValue(encryptedKey)
+                val entry = entryMetas[encryptedKey] ?: continue
                 nextWritePositions[entry.page] = buffers[entry.page].shiftLeft(
                     currentTail = nextWritePositions[entry.page],
                     fromOffset = entry.offset + entry.size,

safebox/src/main/java/com/harrytmthy/safebox/storage/SafeBoxRecoveryBlobStore.kt

@@ -43,7 +43,7 @@ import java.nio.channels.FileChannel.MapMode.READ_WRITE
  * - `keyBytes: ByteArray`
  * - `valueBytes: ByteArray`
  */
-internal class SafeBoxRecoveryBlobStore private constructor(private val file: File) {
+internal class SafeBoxRecoveryBlobStore private constructor(file: File) {
 
     private val channel = RandomAccessFile(file, "rw").channel
 
@@ -131,13 +131,18 @@ internal class SafeBoxRecoveryBlobStore private constructor(private val file: Fi
     }
 
     internal suspend fun delete(fileName: Bytes) {
+        val encryptedKeys = encryptedKeysByFileName[fileName] ?: return
+        delete(fileName, *encryptedKeys.toTypedArray())
+    }
+
+    internal suspend fun delete(fileName: Bytes, vararg encryptedKeys: Bytes) {
+        if (encryptedKeys.isEmpty() || entryMetasByFileName[fileName].isNullOrEmpty()) {
+            return
+        }
         writeMutex.withLock {
-            encryptedKeysByFileName[fileName]?.forEach { encryptedKey ->
-                val entryMeta = entryMetasByFileName.getOrPut(fileName) { HashMap() }
-                if (!entryMeta.containsKey(encryptedKey)) {
-                    return@forEach
-                }
-                val entry = entryMeta.getValue(encryptedKey)
+            val entryMeta = entryMetasByFileName[fileName] ?: return@withLock
+            for (encryptedKey in encryptedKeys) {
+                val entry = entryMeta[encryptedKey] ?: continue
                 nextWritePosition = buffer.shiftLeft(
                     currentTail = nextWritePosition,
                     fromOffset = entry.offset + entry.size,
@@ -151,8 +156,14 @@ internal class SafeBoxRecoveryBlobStore private constructor(private val file: Fi
                     )
                 }
                 entryMeta -= encryptedKey
+                encryptedKeysByFileName.getValue(fileName).remove(encryptedKey)
+            }
+            if (entryMeta.isEmpty()) {
+                entryMetasByFileName.remove(fileName)
+            }
+            if (encryptedKeysByFileName.getValue(fileName).isEmpty()) {
+                encryptedKeysByFileName.remove(fileName)
             }
-            encryptedKeysByFileName.remove(fileName)
         }
     }