feat(state): Add SafeBoxStateManager for better lifecycle support

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

@@ -22,12 +22,16 @@ import androidx.test.core.app.ApplicationProvider
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.harrytmthy.safebox.SafeBox.Companion.DEFAULT_KEY_ALIAS
 import com.harrytmthy.safebox.SafeBox.Companion.DEFAULT_VALUE_KEYSTORE_ALIAS
+import com.harrytmthy.safebox.state.SafeBoxState
+import com.harrytmthy.safebox.state.SafeBoxStateListener
+import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.UnconfinedTestDispatcher
 import org.junit.After
 import org.junit.runner.RunWith
 import java.io.File
 import java.security.KeyStore
+import java.util.concurrent.CopyOnWriteArrayList
 import kotlin.test.Test
 import kotlin.test.assertContentEquals
 import kotlin.test.assertEquals
@@ -158,4 +162,40 @@ class SafeBoxTest {
         assertContentEquals(expectedKeyChanges, changedKeys)
         assertContentEquals(expectedValueChanges, changedValues)
     }
+
+    @Test
+    fun closeWhenIdle_shouldWaitUntilWritesAreDoneBeforeClosing() {
+        val observedStates = CopyOnWriteArrayList<SafeBoxState>()
+        var shouldLoop = true
+        val safeBox = SafeBox.create(
+            context = context,
+            fileName = "${fileName}2",
+            ioDispatcher = Dispatchers.IO,
+            stateListener = SafeBoxStateListener {
+                observedStates.add(it)
+                shouldLoop = it != SafeBoxState.CLOSED
+            },
+        )
+        repeat(5) {
+            safeBox.edit()
+                .putString("key", "value")
+                .apply()
+        }
+        safeBox.closeWhenIdle()
+        repeat(5) {
+            safeBox.edit()
+                .putString("key", "value")
+                .apply()
+        }
+        while (shouldLoop) {
+            Thread.sleep(3)
+        }
+        val expected = listOf(
+            SafeBoxState.STARTING,
+            SafeBoxState.WRITING,
+            SafeBoxState.IDLE,
+            SafeBoxState.CLOSED,
+        )
+        assertEquals(expected, observedStates)
+    }
 }

safebox/src/androidTest/java/com/harrytmthy/safebox/storage/SafeBoxBlobStoreTest.kt

@@ -22,10 +22,8 @@ import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.harrytmthy.safebox.extensions.toBytes
 import com.harrytmthy.safebox.state.SafeBoxState
 import com.harrytmthy.safebox.state.SafeBoxStateListener
-import kotlinx.coroutines.Dispatchers
+import com.harrytmthy.safebox.state.SafeBoxStateManager
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.joinAll
-import kotlinx.coroutines.launch
 import kotlinx.coroutines.test.UnconfinedTestDispatcher
 import kotlinx.coroutines.test.runTest
 import org.junit.After
@@ -53,8 +51,7 @@ class SafeBoxBlobStoreTest {
     private val blobStore = SafeBoxBlobStore.create(
         context,
         fileName,
-        UnconfinedTestDispatcher(),
-        stateListener,
+        SafeBoxStateManager(fileName, stateListener, UnconfinedTestDispatcher()),
     )
 
     @After
@@ -221,64 +218,10 @@ class SafeBoxBlobStoreTest {
     }
 
     @Test
-    fun write_shouldEmitWritingAndIdleStates() = runTest {
-        buildList {
-            repeat(10) {
-                add(
-                    launch(Dispatchers.Default) {
-                        blobStore.write("alpha".toByteArray().toBytes(), "123".toByteArray())
-                    },
-                )
-            }
-        }.joinAll()
-
-        val expected = listOf(
-            SafeBoxState.STARTING,
-            SafeBoxState.IDLE,
-            SafeBoxState.WRITING,
-            SafeBoxState.IDLE,
-        )
-        assertEquals(expected, observedStates)
-    }
-
-    @Test
-    fun close_shouldEmitClosedState() {
-        blobStore.close()
-
+    fun init_shouldEmitStartingState() {
         val expected = listOf(
             SafeBoxState.STARTING,
             SafeBoxState.IDLE,
-            SafeBoxState.CLOSED,
-        )
-        assertEquals(expected, observedStates)
-    }
-
-    @Test
-    fun closeWhenIdle_shouldWaitUntilWritesAreDone() = runTest {
-        buildList {
-            repeat(5) {
-                add(
-                    launch(Dispatchers.Default) {
-                        blobStore.write("alpha".toByteArray().toBytes(), "123".toByteArray())
-                    },
-                )
-            }
-            add(launch(Dispatchers.Default) { blobStore.closeWhenIdle() })
-            repeat(5) {
-                add(
-                    launch(Dispatchers.Default) {
-                        blobStore.write("alpha".toByteArray().toBytes(), "123".toByteArray())
-                    },
-                )
-            }
-        }.joinAll()
-
-        val expected = listOf(
-            SafeBoxState.STARTING,
-            SafeBoxState.IDLE,
-            SafeBoxState.WRITING,
-            SafeBoxState.IDLE,
-            SafeBoxState.CLOSED,
         )
         assertEquals(expected, observedStates)
     }

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

@@ -33,8 +33,8 @@ import com.harrytmthy.safebox.extensions.toBytes
 import com.harrytmthy.safebox.extensions.toEncodedByteArray
 import com.harrytmthy.safebox.keystore.SecureRandomKeyProvider
 import com.harrytmthy.safebox.registry.SafeBoxBlobFileRegistry
-import com.harrytmthy.safebox.state.SafeBoxState
 import com.harrytmthy.safebox.state.SafeBoxStateListener
+import com.harrytmthy.safebox.state.SafeBoxStateManager
 import com.harrytmthy.safebox.storage.Bytes
 import com.harrytmthy.safebox.storage.SafeBoxBlobStore
 import com.harrytmthy.safebox.strategy.ValueFallbackStrategy
@@ -43,8 +43,6 @@ import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineExceptionHandler
 import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.launch
-import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.sync.Mutex
 import kotlinx.coroutines.sync.withLock
 import java.util.concurrent.CopyOnWriteArrayList
@@ -70,13 +68,13 @@ import java.util.concurrent.atomic.AtomicReference
  * @param blobStore Internal storage engine managing encrypted key-value pairs.
  * @param keyCipherProvider Cipher used for encrypting and decrypting keys (deterministic).
  * @param valueCipherProvider Cipher used for encrypting and decrypting values (randomized).
- * @param ioDispatcher Coroutine dispatcher for IO operations.
+ * @param stateManager Responsible for managing SafeBox lifecycle states and its concurrency.
  */
 public class SafeBox private constructor(
     private val blobStore: SafeBoxBlobStore,
     private val keyCipherProvider: CipherProvider,
     private val valueCipherProvider: CipherProvider,
-    private val ioDispatcher: CoroutineDispatcher,
+    private val stateManager: SafeBoxStateManager,
 ) : SharedPreferences {
 
     private val castFailureStrategy = AtomicReference<ValueFallbackStrategy>(WARN)
@@ -100,7 +98,7 @@ public class SafeBox private constructor(
         }
 
         override fun commit(entries: LinkedHashMap<String, Action>, cleared: Boolean): Boolean =
-            runBlocking {
+            stateManager.launchCommitWithWritingState {
                 try {
                     applyCompleted.await()
                     commitMutex.withLock {
@@ -114,7 +112,7 @@ public class SafeBox private constructor(
             }
 
         override fun apply(entries: LinkedHashMap<String, Action>, cleared: Boolean) {
-            safeBoxScope.launch(ioDispatcher + exceptionHandler) {
+            stateManager.launchApplyWithWritingState(exceptionHandler) {
                 applyCompleted = CompletableDeferred()
                 applyMutex.withLock {
                     applyChanges(entries, cleared)
@@ -181,9 +179,7 @@ public class SafeBox private constructor(
      * becomes idle before releasing resources.
      */
     public fun closeWhenIdle() {
-        safeBoxScope.launch(ioDispatcher) {
-            blobStore.closeWhenIdle()
-        }
+        stateManager.closeWhenIdle(::close)
     }
 
     override fun getAll(): Map<String, Any?> {
@@ -372,7 +368,6 @@ public class SafeBox private constructor(
             stateListener: SafeBoxStateListener? = null,
         ): SafeBox {
             SafeBoxBlobFileRegistry.register(fileName)
-            stateListener?.onStateChanged(SafeBoxState.IDLE)
             val aesGcmCipherProvider = AesGcmCipherProvider.create(
                 alias = valueKeyStoreAlias,
                 aad = additionalAuthenticatedData,
@@ -386,8 +381,9 @@ public class SafeBox private constructor(
             )
             val keyCipherProvider = ChaCha20CipherProvider(keyProvider, deterministic = true)
             val valueCipherProvider = ChaCha20CipherProvider(keyProvider, deterministic = false)
-            val blobStore = SafeBoxBlobStore.create(context, fileName, ioDispatcher, stateListener)
-            return SafeBox(blobStore, keyCipherProvider, valueCipherProvider, ioDispatcher)
+            val stateManager = SafeBoxStateManager(fileName, stateListener, ioDispatcher)
+            val blobStore = SafeBoxBlobStore.create(context, fileName, stateManager)
+            return SafeBox(blobStore, keyCipherProvider, valueCipherProvider, stateManager)
         }
 
         /**
@@ -422,8 +418,9 @@ public class SafeBox private constructor(
             stateListener: SafeBoxStateListener? = null,
         ): SafeBox {
             SafeBoxBlobFileRegistry.register(fileName)
-            val blobStore = SafeBoxBlobStore.create(context, fileName, ioDispatcher, stateListener)
-            return SafeBox(blobStore, keyCipherProvider, valueCipherProvider, ioDispatcher)
+            val stateManager = SafeBoxStateManager(fileName, stateListener, ioDispatcher)
+            val blobStore = SafeBoxBlobStore.create(context, fileName, stateManager)
+            return SafeBox(blobStore, keyCipherProvider, valueCipherProvider, stateManager)
         }
     }
 }

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

@@ -0,0 +1,132 @@
+/*
+ * Copyright 2025 Harry Timothy Tumalewa
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.harrytmthy.safebox.state
+
+import com.harrytmthy.safebox.SafeBox
+import com.harrytmthy.safebox.extensions.safeBoxScope
+import com.harrytmthy.safebox.state.SafeBoxState.CLOSED
+import com.harrytmthy.safebox.state.SafeBoxState.IDLE
+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.AtomicInteger
+import java.util.concurrent.atomic.AtomicReference
+
+/**
+ * Manages the lifecycle state of a [SafeBox] instance and coordinates concurrent read/write access.
+ *
+ * This class emits state changes to both the instance-bound [SafeBoxStateListener] and the global
+ * [SafeBoxGlobalStateObserver].
+ *
+ * Key behaviors:
+ * - Tracks concurrent writes using an atomic counter.
+ * - Waits for the blob store's initial read before permitting writes or close.
+ * - Guarantees transition to [IDLE] after all writes complete.
+ * - Supports safe, deferred closing via [closeWhenIdle], ensuring no writes are in progress.
+ *
+ * @param fileName The unique file identifier associated with this SafeBox instance.
+ * @param stateListener Optional listener for observing state transitions on this instance.
+ * @param ioDispatcher Dispatcher used for coroutine-based I/O tasks.
+ */
+internal class SafeBoxStateManager(
+    private val fileName: String,
+    private val stateListener: SafeBoxStateListener?,
+    private val ioDispatcher: CoroutineDispatcher,
+) {
+
+    private val concurrentWriteCount = AtomicInteger(0)
+
+    private val initialReadCompleted = CompletableDeferred<Unit>()
+
+    private val writeCompleted = AtomicReference<CompletableDeferred<Unit>>()
+
+    inline fun launchWithStartingState(crossinline block: suspend () -> Unit) {
+        updateState(STARTING)
+        safeBoxScope.launch(ioDispatcher) {
+            block()
+            if (concurrentWriteCount.get() == 0) {
+                updateState(IDLE)
+            } else {
+                updateState(WRITING)
+            }
+            initialReadCompleted.complete(Unit)
+        }
+    }
+
+    inline fun launchCommitWithWritingState(crossinline block: suspend () -> Boolean): Boolean {
+        if (concurrentWriteCount.incrementAndGet() == 1) {
+            writeCompleted.set(CompletableDeferred())
+            if (initialReadCompleted.isCompleted) {
+                updateState(WRITING)
+            }
+        }
+        return runBlocking {
+            initialReadCompleted.await()
+            val result = block()
+            finalizeWriting()
+            result
+        }
+    }
+
+    inline fun launchApplyWithWritingState(
+        exceptionHandler: CoroutineExceptionHandler,
+        crossinline block: suspend () -> Unit,
+    ) {
+        if (concurrentWriteCount.incrementAndGet() == 1) {
+            writeCompleted.set(CompletableDeferred())
+            if (initialReadCompleted.isCompleted) {
+                updateState(WRITING)
+            }
+        }
+        safeBoxScope.launch(ioDispatcher + exceptionHandler) {
+            initialReadCompleted.await()
+            block()
+        }.invokeOnCompletion {
+            finalizeWriting()
+        }
+    }
+
+    inline fun closeWhenIdle(crossinline block: () -> Unit) {
+        if (concurrentWriteCount.get() == 0 && initialReadCompleted.isCompleted) {
+            block()
+            updateState(CLOSED)
+            return
+        }
+        safeBoxScope.launch(ioDispatcher) {
+            initialReadCompleted.await()
+            writeCompleted.get()?.await()
+            block()
+            updateState(CLOSED)
+        }
+    }
+
+    private fun updateState(newState: SafeBoxState) {
+        stateListener?.onStateChanged(newState)
+        SafeBoxGlobalStateObserver.updateState(fileName, newState)
+    }
+
+    private fun finalizeWriting() {
+        if (concurrentWriteCount.decrementAndGet() == 0) {
+            updateState(IDLE)
+            writeCompleted.get()?.complete(Unit)
+        }
+    }
+}