fix: Serialize entry writes

Author: harrytmthy

safebox/src/androidTest/java/com/harrytmthy/safebox/SafeBoxTest.kt

@@ -28,18 +28,20 @@ import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.UnconfinedTestDispatcher
+import kotlinx.coroutines.test.runTest
+import kotlinx.coroutines.withTimeout
 import org.junit.After
 import org.junit.runner.RunWith
 import java.io.File
 import java.security.KeyStore
-import java.util.concurrent.CopyOnWriteArrayList
 import java.util.concurrent.atomic.AtomicBoolean
 import kotlin.test.Test
 import kotlin.test.assertContentEquals
 import kotlin.test.assertEquals
 import kotlin.test.assertFalse
 import kotlin.test.assertNull
 import kotlin.test.assertTrue
+import kotlin.time.Duration.Companion.seconds
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @RunWith(AndroidJUnit4::class)
@@ -179,47 +181,6 @@ class SafeBoxTest {
         assertContentEquals(expectedValueChanges, changedValues)
     }
 
-    @Test
-    fun closeWhenIdle_shouldWaitUntilWritesAreDoneBeforeClosing() {
-        val observedStates = CopyOnWriteArrayList<SafeBoxState>()
-        val closed = AtomicBoolean(false)
-        safeBox = createSafeBox(
-            ioDispatcher = Dispatchers.IO,
-            stateListener = SafeBoxStateListener { state ->
-                observedStates.add(state)
-                closed.set(state == SafeBoxState.CLOSED)
-            },
-        )
-        repeat(5) {
-            safeBox.edit()
-                .putString("key", "value")
-                .apply()
-        }
-        safeBox.closeWhenIdle()
-        repeat(5) {
-            safeBox.edit()
-                .putString("key", "value")
-                .apply()
-        }
-        while (!closed.get()) {
-            Thread.sleep(3)
-        }
-        val expectedOnSlowInit = listOf(
-            SafeBoxState.STARTING,
-            SafeBoxState.WRITING,
-            SafeBoxState.IDLE,
-            SafeBoxState.CLOSED,
-        )
-        val expectedOnFastInit = listOf(
-            SafeBoxState.STARTING,
-            SafeBoxState.IDLE, // finished STARTING before launching any write operation
-            SafeBoxState.WRITING,
-            SafeBoxState.IDLE,
-            SafeBoxState.CLOSED,
-        )
-        assertTrue(observedStates == expectedOnSlowInit || observedStates == expectedOnFastInit)
-    }
-
     @Test
     fun putString_shouldDoNothingAfterClosing() {
         val hasEmissionAfterClose = AtomicBoolean(false)
@@ -241,6 +202,27 @@ class SafeBoxTest {
         assertFalse(hasEmissionAfterClose.get())
     }
 
+    @Test
+    fun apply_then_commit_shouldHaveCorrectOrder() = runTest {
+        safeBox = createSafeBox(ioDispatcher = Dispatchers.IO)
+
+        withTimeout(10.seconds) {
+            safeBox.edit().putInt("0", 0).apply()
+            safeBox.edit().putInt("1", 1).apply()
+            assertTrue(safeBox.edit().clear().commit())
+            safeBox.edit().putInt("2", 2).apply()
+            safeBox.edit().putInt("3", 3).apply()
+            assertTrue(safeBox.edit().clear().commit())
+            safeBox.edit().putInt("4", 4).apply()
+        }
+
+        assertEquals(4, safeBox.getInt("4", -1))
+        assertEquals(-1, safeBox.getInt("3", -1))
+        assertEquals(-1, safeBox.getInt("2", -1))
+        assertEquals(-1, safeBox.getInt("1", -1))
+        assertEquals(-1, safeBox.getInt("0", -1))
+    }
+
     private fun createSafeBox(
         ioDispatcher: CoroutineDispatcher = UnconfinedTestDispatcher(),
         stateListener: SafeBoxStateListener? = null,

safebox/src/main/java/com/harrytmthy/safebox/SafeBox.kt

@@ -41,7 +41,6 @@ import com.harrytmthy.safebox.strategy.ValueFallbackStrategy
 import com.harrytmthy.safebox.strategy.ValueFallbackStrategy.WARN
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.CoroutineDispatcher
-import kotlinx.coroutines.CoroutineExceptionHandler
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.sync.Mutex
 import kotlinx.coroutines.sync.withLock
@@ -86,38 +85,34 @@ public class SafeBox private constructor(
 
     private val listeners = CopyOnWriteArrayList<OnSharedPreferenceChangeListener>()
 
-    private val commitMutex = Mutex()
+    private val writeMutex = Mutex()
 
-    private val applyMutex = Mutex()
-
-    @Volatile
-    private var applyCompleted = CompletableDeferred<Unit>().apply { complete(Unit) }
+    private val writeBarrier = AtomicReference(CompletableDeferred<Unit>().apply { complete(Unit) })
 
     private val delegate = object : Delegate {
 
         private val updateLock = Any()
 
-        private val exceptionHandler = CoroutineExceptionHandler { _, throwable ->
-            Log.e("SafeBox", "Failed to apply changes.", throwable)
-            applyCompleted.complete(Unit)
-        }
-
         override fun commit(entries: LinkedHashMap<String, Action>, cleared: Boolean): Boolean {
             val entriesToWrite = LinkedHashMap(entries)
             synchronized(updateLock) {
                 entries.clear() // Prevents stale mutations on reused editor instance
                 updateEntries(entriesToWrite, cleared)
             }
+            val currentWriteBarrier = CompletableDeferred<Unit>()
+            val previousWriteBarrier = writeBarrier.getAndSet(currentWriteBarrier)
             return stateManager.launchCommitWithWritingState {
                 try {
-                    applyCompleted.await()
-                    commitMutex.withLock {
+                    previousWriteBarrier.await()
+                    writeMutex.withLock {
                         applyChanges(entriesToWrite, cleared)
                     }
                     true
                 } catch (e: Exception) {
                     Log.e("SafeBox", "Failed to commit changes.", e)
                     false
+                } finally {
+                    currentWriteBarrier.complete(Unit)
                 }
             }
         }
@@ -128,12 +123,19 @@ public class SafeBox private constructor(
                 entries.clear() // Prevents stale mutations on reused editor instance
                 updateEntries(entriesToWrite, cleared)
             }
-            stateManager.launchApplyWithWritingState(exceptionHandler) {
-                applyCompleted = CompletableDeferred()
-                applyMutex.withLock {
-                    applyChanges(entriesToWrite, cleared)
+            val currentWriteBarrier = CompletableDeferred<Unit>()
+            val previousWriteBarrier = writeBarrier.getAndSet(currentWriteBarrier)
+            stateManager.launchApplyWithWritingState {
+                try {
+                    previousWriteBarrier.await()
+                    writeMutex.withLock {
+                        applyChanges(entriesToWrite, cleared)
+                    }
+                } catch (e: Exception) {
+                    Log.e("SafeBox", "Failed to commit changes.", e)
+                } finally {
+                    currentWriteBarrier.complete(Unit)
                 }
-                applyCompleted.complete(Unit)
             }
         }
 

safebox/src/main/java/com/harrytmthy/safebox/state/SafeBoxStateManager.kt

@@ -24,7 +24,6 @@ import com.harrytmthy.safebox.state.SafeBoxState.STARTING
 import com.harrytmthy.safebox.state.SafeBoxState.WRITING
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.CoroutineDispatcher
-import kotlinx.coroutines.CoroutineExceptionHandler
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
 import java.util.concurrent.atomic.AtomicBoolean
@@ -64,52 +63,46 @@ internal class SafeBoxStateManager(
     inline fun launchWithStartingState(crossinline block: suspend () -> Unit) {
         updateState(STARTING)
         safeBoxScope.launch(ioDispatcher) {
-            block()
-            if (concurrentWriteCount.get() == 0) {
-                updateState(IDLE)
-            } else {
-                updateState(WRITING)
+            try {
+                block()
+            } finally {
+                initialReadCompleted.complete(Unit)
+                if (concurrentWriteCount.get() == 0) {
+                    updateState(IDLE)
+                }
             }
-            initialReadCompleted.complete(Unit)
         }
     }
 
     inline fun launchCommitWithWritingState(crossinline block: suspend () -> Boolean): Boolean {
         if (closed.get()) {
             return false
         }
-        if (concurrentWriteCount.incrementAndGet() == 1) {
-            writeCompleted.set(CompletableDeferred())
-            if (initialReadCompleted.isCompleted) {
-                updateState(WRITING)
-            }
-        }
         return runBlocking {
-            initialReadCompleted.await()
-            val result = block()
-            finalizeWriting()
-            result
+            withStateTransition(block)
         }
     }
 
-    inline fun launchApplyWithWritingState(
-        exceptionHandler: CoroutineExceptionHandler,
-        crossinline block: suspend () -> Unit,
-    ) {
+    inline fun launchApplyWithWritingState(crossinline block: suspend () -> Unit) {
         if (closed.get()) {
             return
         }
+        safeBoxScope.launch(ioDispatcher) {
+            withStateTransition(block)
+        }
+    }
+
+    private suspend inline fun <T> withStateTransition(crossinline block: suspend () -> T): T {
+        initialReadCompleted.await()
         if (concurrentWriteCount.incrementAndGet() == 1) {
-            writeCompleted.set(CompletableDeferred())
-            if (initialReadCompleted.isCompleted) {
-                updateState(WRITING)
-            }
+            updateState(WRITING)
         }
-        safeBoxScope.launch(ioDispatcher + exceptionHandler) {
-            initialReadCompleted.await()
+        return try {
             block()
-        }.invokeOnCompletion {
-            finalizeWriting()
+        } finally {
+            if (concurrentWriteCount.decrementAndGet() == 0) {
+                updateState(IDLE)
+            }
         }
     }
 
@@ -133,11 +126,4 @@ internal class SafeBoxStateManager(
         stateListener?.onStateChanged(newState)
         SafeBoxGlobalStateObserver.updateState(fileName, newState)
     }
-
-    private fun finalizeWriting() {
-        if (concurrentWriteCount.decrementAndGet() == 0) {
-            updateState(IDLE)
-            writeCompleted.get()?.complete(Unit)
-        }
-    }
 }