Plumb WebSocket observation functionality through the FFI

This will allow Android and iOS clients to observe the underlying state of the WebSocket connection to the Convex backend.

Author: dowski

android/convexmobile/src/androidTest/java/dev/convex/android/ConvexClientInstrumentationTest.kt

@@ -212,6 +212,28 @@ class ConvexClientInstrumentationTest {
         assertEquals(expected, result)
         assertEquals(100, result.aPlainInt)
     }
+
+    @Test
+    fun can_observe_websocket_state() = runTest {
+        val client = ConvexClient(DEPLOYMENT_URL)
+
+        val states = mutableListOf<WebSocketState>()
+
+        val receiveJob = launch {
+            client.webSocketStateFlow.take(2).collect { state ->
+                states.add(state)
+            }
+        }
+
+        // It doesn't really matter which Convex function we call - but calling one should trigger
+        // a connection.
+        client.mutation<Unit?>("messages:clearAll")
+
+        receiveJob.join()
+
+        assertEquals(2, states.size)
+        assertEquals(listOf(WebSocketState.CONNECTING, WebSocketState.CONNECTED), states)
+    }
 }
 
 @Serializable

android/convexmobile/src/androidTest/java/dev/convex/android/ConvexClientWithAuthTest.kt

@@ -26,7 +26,7 @@ class ConvexClientWithAuthTest {
         client = ConvexClientWithAuth(
             "foo://bar",
             authProvider
-        ) { _, _ -> ffiClient }
+        ) { _, _, _ -> ffiClient }
     }
 
     @Test

android/convexmobile/src/main/java/dev/convex/android/ConvexClient.kt

@@ -7,6 +7,7 @@ import kotlinx.coroutines.channels.awaitClose
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.callbackFlow
 import kotlinx.serialization.Contextual
 import kotlinx.serialization.Serializable
@@ -46,12 +47,29 @@ typealias ConvexNum = Contextual
  */
 open class ConvexClient(
     deploymentUrl: String,
-    ffiClientFactory: (deploymentUrl: String, clientId: String) -> MobileConvexClientInterface = ::MobileConvexClient
+    ffiClientFactory: (deploymentUrl: String, clientId: String, webSocketStateSubscriber: WebSocketStateSubscriber?) -> MobileConvexClientInterface = ::MobileConvexClient
 ) {
+    private val _webSocketStateFlow = MutableStateFlow(WebSocketState.CONNECTING)
+    private val webSocketStateSubscriber = object : WebSocketStateSubscriber {
+        override fun onStateChange(state: WebSocketState) {
+            _webSocketStateFlow.value = state
+        }
+    }
 
     @PublishedApi
     internal val ffiClient =
-        ffiClientFactory(deploymentUrl, "kotlin-${BuildConfig.LIBRARY_VERSION}")
+        ffiClientFactory(
+            deploymentUrl,
+            "kotlin-${BuildConfig.LIBRARY_VERSION}",
+            webSocketStateSubscriber
+        )
+
+    /**
+     * A [Flow] of [WebSocketState].
+     *
+     * Will change states as the underlying state of the WebSocket connection to Convex changes.
+     */
+    val webSocketStateFlow: StateFlow<WebSocketState> = _webSocketStateFlow.asStateFlow()
 
     /**
      * Subscribes to the query with the given [name] and converts data from the subscription into a
@@ -195,7 +213,7 @@ open class ConvexClient(
 class ConvexClientWithAuth<T>(
     deploymentUrl: String,
     private val authProvider: AuthProvider<T>,
-    ffiClientFactory: (deploymentUrl: String, clientId: String) -> MobileConvexClientInterface = ::MobileConvexClient
+    ffiClientFactory: (deploymentUrl: String, clientId: String, webSocketSocketStateSubscriber: WebSocketStateSubscriber?) -> MobileConvexClientInterface = ::MobileConvexClient
 ) : ConvexClient(deploymentUrl, ffiClientFactory) {
     private val _authState = MutableStateFlow<AuthState<T>>(AuthState.Unauthenticated())
 

android/convexmobile/src/test/java/dev/convex/android/ConvexClientTest.kt

@@ -39,7 +39,7 @@ class ConvexClientTest {
     @Before
     fun setup() {
         ffiClient = FakeFfiClient()
-        client = ConvexClient("foo://bar") { _, _ -> ffiClient }
+        client = ConvexClient("foo://bar") { _, _, _ -> ffiClient }
     }
 
     @Test

rust/Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "convex-mobile"
-version = "0.2.2"
+version = "0.3.0"
 edition = "2021"
 
 [dependencies]

rust/Cargo.toml

@@ -9,8 +9,18 @@ interface ClientError {
     ServerError(string msg);
 };
 
+enum WebSocketState {
+    "Connected",
+    "Connecting",
+};
+
+[Trait, WithForeign]
+interface WebSocketStateSubscriber {        
+    void on_state_change(WebSocketState state);
+};
+
 interface MobileConvexClient {
-    constructor(string deployment_url, string client_id);
+    constructor(string deployment_url, string client_id, WebSocketStateSubscriber? web_socket_state_subscriber);
 
     [Async, Throws=ClientError]
     string query(string name, record<string, string> args);
@@ -37,4 +47,4 @@ interface SubscriptionHandle {
 interface QuerySubscriber {
     void on_update(string value);
     void on_error(string message, string? value);
-};
+};

rust/src/convex-mobile.udl

@@ -4,12 +4,13 @@ use std::{
 };
 
 use async_once_cell::OnceCell;
-use convex::{ConvexClient, ConvexClientBuilder, FunctionResult, Value};
+use convex::{ConvexClient, ConvexClientBuilder, FunctionResult, Value, WebSocketState};
 use futures::{
     channel::oneshot::{self, Sender},
     pin_mut, select_biased, FutureExt, StreamExt,
 };
 use parking_lot::Mutex;
+use tokio::sync::mpsc;
 use tracing::{debug, info};
 
 mod logging;
@@ -43,6 +44,10 @@ pub trait QuerySubscriber: Send + Sync {
     fn on_error(&self, message: String, value: Option<String>) -> ();
 }
 
+pub trait WebSocketStateSubscriber: Send + Sync {
+    fn on_state_change(&self, state: WebSocketState) -> ();
+}
+
 pub struct SubscriptionHandle {
     cancel_sender: Mutex<Option<Sender<()>>>,
 }
@@ -62,9 +67,9 @@ impl SubscriptionHandle {
 }
 
 /// Initializes logging.
-/// 
+///
 /// Call this early in the life of your application to enable logging from
-/// [MobileConvexClient] and its dependencies. 
+/// [MobileConvexClient] and its dependencies.
 pub fn init_convex_logging() {
     use std::sync::Once;
     static INIT: Once = Once::new();
@@ -84,6 +89,7 @@ pub fn init_convex_logging() {
 struct MobileConvexClient {
     deployment_url: String,
     client_id: String,
+    web_socket_state_subscriber: Option<Arc<dyn WebSocketStateSubscriber>>,
     client: OnceCell<ConvexClient>,
     rt: tokio::runtime::Runtime,
 }
@@ -96,14 +102,19 @@ impl MobileConvexClient {
     ///
     /// The `client_id` should be a string representing the name and version of
     /// the foreign client.
-    pub fn new(deployment_url: String, client_id: String) -> MobileConvexClient {
+    pub fn new(
+        deployment_url: String,
+        client_id: String,
+        web_socket_state_subscriber: Option<Arc<dyn WebSocketStateSubscriber>>,
+    ) -> MobileConvexClient {
         let rt = tokio::runtime::Builder::new_multi_thread()
             .enable_all()
             .build()
             .unwrap();
         MobileConvexClient {
             deployment_url,
             client_id,
+            web_socket_state_subscriber,
             client: OnceCell::new(),
             rt,
         }
@@ -122,12 +133,25 @@ impl MobileConvexClient {
         self.client
             .get_or_try_init(async {
                 let client_id = self.client_id.to_owned();
+                let (tx, mut rx) = mpsc::channel(1);
+                let possible_subscriber = self.web_socket_state_subscriber.clone();
+                if let Some(subscriber) = possible_subscriber.clone() {
+                    self.rt.spawn(async move {
+                        while let Some(state) = rx.recv().await {
+                            subscriber.on_state_change(state);
+                        }
+                    });
+                }
+
+                let has_subscriber = possible_subscriber.is_some();
                 self.rt
                     .spawn(async move {
-                        ConvexClientBuilder::new(url.as_str())
-                            .with_client_id(&client_id)
-                            .build()
-                            .await
+                        let mut builder =
+                            ConvexClientBuilder::new(url.as_str()).with_client_id(&client_id);
+                        if has_subscriber {
+                            builder = builder.with_on_state_change(tx);
+                        }
+                        builder.build().await
                     })
                     .await?
             })