Pipe devToolsProxy into stores (PolymerLabs#6037)

* Pipe devToolsProxy into stores

* Fix tests.

* Merge conflict resolution

Co-authored-by: Sarah Heimlich <[email protected]>

java/arcs/android/devtools/DevToolsService.kt

@@ -63,7 +63,7 @@ open class DevToolsService : Service() {
             val service = initialize()
             val proxy = service.devToolsProxy
 
-            forwardProxyMessageToken = proxy.registerBindingContextProxyMessageCallback(
+            forwardProxyMessageToken = proxy.registerRefModeStoreProxyMessageCallback(
                 object : IStorageServiceCallback.Stub() {
                     override fun onProxyMessage(proxyMessage: ByteArray) {
                         scope.launch {
@@ -123,7 +123,7 @@ open class DevToolsService : Service() {
     override fun onDestroy() {
         super.onDestroy()
         devToolsServer.close()
-        devToolsProxy?.deRegisterBindingContextProxyMessageCallback(forwardProxyMessageToken)
+        devToolsProxy?.deRegisterRefModeStoreProxyMessageCallback(forwardProxyMessageToken)
         scope.cancel()
     }
 

java/arcs/android/storage/service/BindingContext.kt

@@ -45,10 +45,11 @@ import kotlinx.coroutines.withTimeout
  */
 @ExperimentalCoroutinesApi
 class DeferredStore<Data : CrdtData, Op : CrdtOperation, T>(
-    options: StoreOptions
+    options: StoreOptions,
+    private val devToolsProxy: DevToolsProxyImpl?
 ) {
     private val store = SuspendableLazy<ActiveStore<Data, Op, T>> {
-        DefaultActivationFactory(options)
+        DefaultActivationFactory(options, devToolsProxy)
     }
     suspend operator fun invoke() = store()
 }
@@ -147,8 +148,6 @@ class BindingContext(
                     resultCallback.takeIf { it.asBinder().isBinderAlive }?.onResult(null)
 
                     val actualMessage = proxyMessage.decodeProxyMessage()
-                    // TODO: (sarahheimlich) remove once we dive into stores (b/162955831)
-                    devToolsProxy?.onBindingContextProxyMessage(proxyMessage)
 
                     (store() as ActiveStore<CrdtData, CrdtOperation, Any?>).let { store ->
                         if (store.onProxyMessage(actualMessage)) {

java/arcs/android/storage/service/DevToolsProxyImpl.kt

@@ -1,33 +1,37 @@
 package arcs.android.storage.service
 
+import arcs.android.storage.toProto
+import arcs.core.storage.DevToolsProxy
+import arcs.core.storage.ProxyMessage
+
 /**
  * Implementation of [IDevToolsProxy] to allow communication between the [StorageService] and
  * [DevToolsService]
  */
-class DevToolsProxyImpl : IDevToolsProxy.Stub() {
-    private val onBindingContextProxyMessageCallbacks = mutableMapOf<Int, IStorageServiceCallback>()
-    private var bindingContextCallbackCounter = 0
+class DevToolsProxyImpl : IDevToolsProxy.Stub(), DevToolsProxy {
+    private val onRefModeStoreProxyMessageCallbacks = mutableMapOf<Int, IStorageServiceCallback>()
+    private var refModeStoreCallbackCounter = 0
 
     /**
      * TODO: (sarahheimlich) remove once we dive into stores (b/162955831)
      *
      * Execute the callbacks to be called with the [BindingContext] receives a [ProxyMessage]
      */
-    fun onBindingContextProxyMessage(proxyMessage: ByteArray) {
-        onBindingContextProxyMessageCallbacks.forEach { (_, callback) ->
-            callback.onProxyMessage(proxyMessage)
+    override fun onRefModeStoreProxyMessage(proxyMessage: ProxyMessage<*, *, *>) {
+        onRefModeStoreProxyMessageCallbacks.forEach { (_, callback) ->
+            callback.onProxyMessage(proxyMessage.toProto().toByteArray())
         }
     }
 
-    override fun registerBindingContextProxyMessageCallback(
+    override fun registerRefModeStoreProxyMessageCallback(
         callback: IStorageServiceCallback
     ): Int {
-        bindingContextCallbackCounter++
-        onBindingContextProxyMessageCallbacks.put(bindingContextCallbackCounter, callback)
-        return bindingContextCallbackCounter
+        refModeStoreCallbackCounter++
+        onRefModeStoreProxyMessageCallbacks.put(refModeStoreCallbackCounter, callback)
+        return refModeStoreCallbackCounter
     }
 
-    override fun deRegisterBindingContextProxyMessageCallback(callbackToken: Int) {
-        onBindingContextProxyMessageCallbacks.remove(callbackToken)
+    override fun deRegisterRefModeStoreProxyMessageCallback(callbackToken: Int) {
+        onRefModeStoreProxyMessageCallbacks.remove(callbackToken)
     }
 }

java/arcs/android/storage/service/IDevToolsProxy.aidl

@@ -8,16 +8,12 @@ import arcs.android.storage.service.IStorageServiceCallback;
 interface IDevToolsProxy {
 
     /**
-     * TODO: (sarahheimlich) remove once we dive into stores (b/162955831)
-     *
-     * Register a callback to be called with the [BindingContext] receives a [ProxyMessage]
+     * Register a callback to be called when the [ReferenceModeStore] receives a [ProxyMessage]
      */
-    int registerBindingContextProxyMessageCallback(in IStorageServiceCallback callback);
+    int registerRefModeStoreProxyMessageCallback(in IStorageServiceCallback callback);
 
     /**
-     * TODO: (sarahheimlich) remove once we dive into stores (b/162955831)
-     *
-     * Remove a callback that is called with the [BindingContext] receives a [ProxyMessage]
+     * Remove a callback that is called when the [ReferenceModeStore] receives a [ProxyMessage]
      */
-    oneway void deRegisterBindingContextProxyMessageCallback(in int callbackToken);
+    oneway void deRegisterRefModeStoreProxyMessageCallback(in int callbackToken);
 }

java/arcs/core/storage/DevToolsProxy.kt

@@ -0,0 +1,12 @@
+package arcs.core.storage
+
+/**
+ * Exposed API to communicate between an [ActiveStore] and [DevToolsService].
+ */
+interface DevToolsProxy {
+
+    /**
+     * Function to call when a [ReferenceModeStore] receives a [ProxyMessage].
+     */
+    fun onRefModeStoreProxyMessage(proxyMessage: ProxyMessage<*, *, *>)
+}

java/arcs/core/storage/ReferenceModeStore.kt

@@ -92,7 +92,8 @@ class ReferenceModeStore private constructor(
     val containerStore: DirectStore<CrdtData, CrdtOperation, Any?>,
     /* internal */
     val backingKey: StorageKey,
-    backingType: Type
+    backingType: Type,
+    private val devToolsProxy: DevToolsProxy?
 ) : ActiveStore<RefModeStoreData, RefModeStoreOp, RefModeStoreOutput>(options) {
     // TODO(#5551): Consider including a hash of the storage key in log prefix.
     private val log = TaggedLog { "ReferenceModeStore" }
@@ -202,6 +203,7 @@ class ReferenceModeStore private constructor(
     ): Boolean {
         log.verbose { "onProxyMessage: $message" }
         val refModeMessage = message.sanitizeForRefModeStore(type)
+        devToolsProxy?.onRefModeStoreProxyMessage(message)
         return receiveQueue.enqueueAndWait {
             handleProxyMessage(refModeMessage)
         }
@@ -689,7 +691,8 @@ class ReferenceModeStore private constructor(
 
         @Suppress("UNCHECKED_CAST")
         suspend fun create(
-            options: StoreOptions
+            options: StoreOptions,
+            devToolsProxy: DevToolsProxy?
         ): ReferenceModeStore {
             val refableOptions =
                 requireNotNull(
@@ -727,7 +730,8 @@ class ReferenceModeStore private constructor(
                 refableOptions,
                 containerStore,
                 storageKey.backingKey,
-                type.containedType
+                type.containedType,
+                devToolsProxy
             ).also { refModeStore ->
                 // Since `on` is a suspending method, we need to setup the container store callback
                 // here in this create method, which is inside of a coroutine.

java/arcs/core/storage/Store.kt

@@ -24,10 +24,11 @@ import kotlinx.coroutines.ExperimentalCoroutinesApi
 object DefaultActivationFactory : ActivationFactory {
     @ExperimentalCoroutinesApi
     override suspend fun <Data : CrdtData, Op : CrdtOperation, T> invoke(
-        options: StoreOptions
+        options: StoreOptions,
+        devToolsProxy: DevToolsProxy?
     ): ActiveStore<Data, Op, T> = when (options.storageKey) {
         is ReferenceModeStorageKey ->
-            ReferenceModeStore.create(options) as ActiveStore<Data, Op, T>
+            ReferenceModeStore.create(options, devToolsProxy) as ActiveStore<Data, Op, T>
         else -> DirectStore.create(options)
     }
 }
@@ -38,6 +39,7 @@ object DefaultActivationFactory : ActivationFactory {
  */
 interface ActivationFactory {
     suspend operator fun <Data : CrdtData, Op : CrdtOperation, T> invoke(
-        options: StoreOptions
+        options: StoreOptions,
+        devToolsProxy: DevToolsProxy?
     ): ActiveStore<Data, Op, T>
 }

java/arcs/core/storage/StoreManager.kt

@@ -37,7 +37,7 @@ class StoreManager(
         storeOptions: StoreOptions
     ) = storesMutex.withLock {
         stores.getOrPut(storeOptions.storageKey) {
-            activationFactory<Data, Op, T>(storeOptions)
+            activationFactory<Data, Op, T>(storeOptions, null)
         } as ActiveStore<Data, Op, T>
     }
 

java/arcs/sdk/android/storage/ServiceStore.kt

@@ -28,6 +28,7 @@ import arcs.core.data.EntityType
 import arcs.core.data.SingletonType
 import arcs.core.storage.ActivationFactory
 import arcs.core.storage.ActiveStore
+import arcs.core.storage.DevToolsProxy
 import arcs.core.storage.ProxyCallback
 import arcs.core.storage.ProxyMessage
 import arcs.core.storage.StoreOptions
@@ -61,7 +62,8 @@ class ServiceStoreFactory(
     private val connectionFactory: ConnectionFactory? = null
 ) : ActivationFactory {
     override suspend operator fun <Data : CrdtData, Op : CrdtOperation, ConsumerData> invoke(
-        options: StoreOptions
+        options: StoreOptions,
+        devToolsProxy: DevToolsProxy?
     ): ServiceStore<Data, Op, ConsumerData> {
         val storeContext = coroutineContext + CoroutineName("ServiceStore(${options.storageKey})")
         val parcelableType = when (options.type) {

java/arcs/sdk/android/storage/service/StorageService.kt

@@ -168,7 +168,7 @@ open class StorageService : ResurrectorService() {
         return BindingContext(
             stores.computeIfAbsent(options.storageKey) {
                 @Suppress("UNCHECKED_CAST")
-                DeferredStore<CrdtData, CrdtOperation, Any>(options)
+                DeferredStore<CrdtData, CrdtOperation, Any>(options, devToolsProxy)
             },
             coroutineContext,
             stats,

javatests/arcs/android/storage/ReferenceModeStoreDatabaseImplIntegrationTest.kt

@@ -701,7 +701,8 @@ class ReferenceModeStoreDatabaseImplIntegrationTest {
             StoreOptions(
                 testKey,
                 CollectionType(EntityType(schema))
-            )
+            ),
+            null
         )
     }
 
@@ -710,7 +711,8 @@ class ReferenceModeStoreDatabaseImplIntegrationTest {
             StoreOptions(
                 testKey,
                 SingletonType(EntityType(schema))
-            )
+            ),
+            null
         )
     }
 

javatests/arcs/android/storage/service/BindingContextTest.kt

@@ -65,7 +65,8 @@ class BindingContextTest {
             StoreOptions(
                 storageKey,
                 CountType()
-            )
+            ),
+            null
         )
     }
 

javatests/arcs/core/storage/ReferenceModeStoreDatabaseIntegrationTest.kt

@@ -663,7 +663,8 @@ class ReferenceModeStoreDatabaseIntegrationTest {
             StoreOptions(
                 testKey,
                 CollectionType(EntityType(schema))
-            )
+            ),
+            null
         )
     }
 
@@ -672,7 +673,8 @@ class ReferenceModeStoreDatabaseIntegrationTest {
             StoreOptions(
                 testKey,
                 SingletonType(EntityType(schema))
-            )
+            ),
+            null
         )
     }
 

javatests/arcs/core/storage/ReferenceModeStoreStabilityTest.kt

@@ -86,7 +86,8 @@ class ReferenceModeStoreStabilityTest {
             StoreOptions(
                 storageKey,
                 SingletonType(EntityType(schema))
-            )
+            ),
+            null
         )
 
         val modelValue = CompletableDeferred<RefModeStoreData.Singleton>()
@@ -128,7 +129,8 @@ class ReferenceModeStoreStabilityTest {
             StoreOptions(
                 storageKey,
                 CollectionType(EntityType(schema))
-            )
+            ),
+            null
         )
 
         val modelValue = CompletableDeferred<RefModeStoreData.Set>()
@@ -192,7 +194,8 @@ class ReferenceModeStoreStabilityTest {
             StoreOptions(
                 storageKey,
                 CollectionType(EntityType(schema))
-            )
+            ),
+            null
         )
 
         val modelValue = CompletableDeferred<RefModeStoreData.Set>()
@@ -254,7 +257,8 @@ class ReferenceModeStoreStabilityTest {
             StoreOptions(
                 storageKey,
                 SingletonType(EntityType(schema))
-            )
+            ),
+            null
         )
 
         val modelValue = CompletableDeferred<RefModeStoreData.Singleton>()
@@ -316,7 +320,8 @@ class ReferenceModeStoreStabilityTest {
             StoreOptions(
                 storageKey,
                 CollectionType(EntityType(schema))
-            )
+            ),
+            null
         )
 
         val modelValue = CompletableDeferred<RefModeStoreData.Set>()

javatests/arcs/core/storage/ReferenceModeStoreTest.kt

@@ -90,7 +90,8 @@ class ReferenceModeStoreTest {
                 StoreOptions(
                     testKey,
                     SingletonType(EntityType(schema))
-                )
+                ),
+                null
             )
         }
     }
@@ -103,7 +104,8 @@ class ReferenceModeStoreTest {
             StoreOptions(
                 testKey,
                 CollectionType(EntityType(schema))
-            )
+            ),
+            null
         )
 
         assertThat(store).isInstanceOf(ReferenceModeStore::class.java)
@@ -779,7 +781,8 @@ class ReferenceModeStoreTest {
             StoreOptions(
                 testKey,
                 CollectionType(EntityType(schema))
-            )
+            ),
+            null
         )
     }
 
@@ -788,7 +791,8 @@ class ReferenceModeStoreTest {
             StoreOptions(
                 testKey,
                 SingletonType(EntityType(schema))
-            )
+            ),
+            null
         )
     }
 

javatests/arcs/core/storage/ReferenceTest.kt

@@ -60,7 +60,7 @@ class ReferenceTest {
                 storageKey = refModeKey,
                 type = CollectionType(EntityType(Person.SCHEMA))
             )
-        val store: CollectionStore<RawEntity> = DefaultActivationFactory(options)
+        val store: CollectionStore<RawEntity> = DefaultActivationFactory(options, null)
 
         val addPeople = listOf(
             CrdtSet.Operation.Add(
@@ -90,7 +90,7 @@ class ReferenceTest {
 
         @Suppress("UNCHECKED_CAST")
         val directCollection: CollectionStore<Reference> =
-            DefaultActivationFactory(collectionOptions)
+            DefaultActivationFactory(collectionOptions, null)
 
         val job = Job()
         val me = directCollection.on(ProxyCallback {