Fix "identity not yet recieved" intermittent failure in tests (#445)

This error turns out to have been a legitimate, albeit rare, race condition,
which could have (but apparently didn't) manifested in user code.

The cause, ultimately, was holding locks for too-short scopes
while initializing the connection.
It was possible for a new connection to receive and completely process
its `IdentityToken` message before calling `maybe_set_credentials`,
which would then clobber the received identity.

This is now fixed (hopefully) by holding the `CredentialStore` lock
for the duration of `BackgroundDbConnection::connect`.
Two other locks, on `ClientCache` and `ReducerCallbacks`,
are also newly held for the duration to avoid similar issues.

In addition, this commit adds a call to `disconnect` to the beginning of `connect`.
This will prevent pathological user code
which connects twice in a row without an intervening disconnect
from encountering bizarre errors
caused by the first connection's background workers staying active.

This commit also bumps the `RUST_LOG` env variable for the SDK test clients to `trace`,
and adds several trace-level logs to the SDK,
which were instrumental in pinning down the failure.
This does mean that log output of failed SDK tests will now be much noisier.

crates/sdk/src/background_connection.rs

@@ -313,10 +313,37 @@ impl BackgroundDbConnection {
         IntoUri: TryInto<http::Uri>,
         <IntoUri as TryInto<http::Uri>>::Error: std::error::Error + Send + Sync + 'static,
     {
-        let client_address = {
-            let mut lock = self.credentials.lock().expect("CredentialStore Mutex is poisoned");
-            lock.get_or_init_address()
-        };
+        // Disconnect to prevent any outstanding messages from contending for unique resources,
+        // i.e. the reducer callbacks, credential store and client cache.
+        self.disconnect();
+
+        // Hold all internal locks for the duration of this method,
+        // to prevent races if the connection starts receiving messages
+        // before this method returns.
+        let mut reducer_callbacks_lock = self
+            .reducer_callbacks
+            .lock()
+            .expect("ReducerCallbacks Mutex is poisoned");
+        let mut credentials_lock = self.credentials.lock().expect("CredentialStore Mutex is poisoned");
+        let mut client_cache_lock = self.client_cache.lock().expect("ClientCache Mutex is poisoned");
+
+        // Specialize the reducer callbacks for this module.
+        // Registering callbacks doesn't require the module, but firing them does,
+        // in order to parse arguments with appropriate types.
+        reducer_callbacks_lock.set_module(module.clone());
+
+        // If credentials were passed to connect, store them in the credential store.
+        // If not, unset any stale credentials in the credential store.
+        credentials_lock.maybe_set_credentials(credentials.clone());
+
+        let client_address = credentials_lock.get_or_init_address();
+
+        // Construct a new client cache specialized for this module.
+        // The client cache needs to know the module
+        // in order to parse rows and issue callbacks with appropriate types.
+        let client_cache = Arc::new(ClientCache::new(module.clone()));
+        *client_cache_lock = Some(client_cache);
+
         // `block_in_place` is required here, as tokio won't allow us to call
         // `block_on` if it would block the current thread of an outer runtime
         let connection = tokio::task::block_in_place(|| {
@@ -328,41 +355,26 @@ impl BackgroundDbConnection {
             ))
         })?;
 
-        let client_cache = Arc::new(ClientCache::new(module.clone()));
-
-        {
-            // Set the global `CLIENT_CACHE` to our newly-constructed cache.
-            // Do this inside a short scope to avoid holding the lock unnecessarily.
-            let mut client_cache_lock = self.client_cache.lock().expect("ClientCache mutex is poisoned");
-            *client_cache_lock = Some(client_cache);
-        }
-
         let (websocket_loop_handle, recv_chan, send_chan) = connection.spawn_message_loop(&self.handle);
         let recv_handle = self.spawn_receiver(recv_chan, self.client_cache.clone());
 
         self.send_chan = Some(send_chan);
         self.websocket_loop_handle = Some(websocket_loop_handle);
         self.recv_handle = Some(recv_handle);
 
-        self.reducer_callbacks
-            .lock()
-            .expect("ReducerCallbacks Mutex is poisoned")
-            .set_module(module);
-        self.credentials
-            .lock()
-            .expect("CredentialStore Mutex is poisoned")
-            .maybe_set_credentials(credentials);
-
         Ok(())
     }
 
     pub fn disconnect(&mut self) {
         self.send_chan = None;
+
+        // `block_in_place` over `block_on` allows us to wait for a future to complete
+        // regardless of whether we're currently running in an `async` task or not.
         if let Some(h) = self.websocket_loop_handle.take() {
-            let _ = self.handle.block_on(h);
+            let _ = tokio::task::block_in_place(|| self.handle.block_on(h));
         }
         if let Some(h) = self.recv_handle.take() {
-            let _ = self.handle.block_on(h);
+            let _ = tokio::task::block_in_place(|| self.handle.block_on(h));
         }
     }
 

crates/sdk/src/callbacks.rs

@@ -287,6 +287,10 @@ impl<Args: OwnedArgs> CallbackMap<Args> {
                     // Enter a dynamic context where `CURRENT_STATE`
                     // is the state which caused this callback invocation.
                     let _guard = CurrentStateGuard::with_current_state(db_state);
+                    log::trace!(
+                        "CallbackMap worker: received invoke message for args {} and got state guard",
+                        std::any::type_name::<Args>(),
+                    );
 
                     // We're invoking several callbacks, so none of them may consume `args`.
                     let borrowed = args.borrow();
@@ -408,6 +412,7 @@ impl<Args: OwnedArgs> CallbackMap<Args> {
 
     /// Invoke each currently-registered callback in `self` with `args`.
     fn invoke(&self, args: Args, db_state: ClientCacheView) {
+        log::trace!("CallbackMap: invoking for args type {}", std::any::type_name::<Args>());
         self.send
             .unbounded_send(CallbackMessage::Invoke { args, db_state })
             // TODO: properly handle this error somehow
@@ -857,6 +862,15 @@ impl CredentialStore {
     ///
     /// Do not invoke on-connect callbacks.
     pub(crate) fn maybe_set_credentials(&mut self, credentials: Option<Credentials>) {
+        if let Some(creds) = &credentials {
+            log::trace!(
+                "maybe_set_credentials: Setting stored credentials to Some({:?})",
+                creds.identity,
+            );
+        } else {
+            log::trace!("maybe_set_credentials: Setting stored credentials to None");
+        }
+
         self.credentials = credentials;
     }
 
@@ -915,12 +929,16 @@ impl CredentialStore {
             return;
         }
 
+        let identity = Identity::from_bytes(identity);
+        log::trace!("handle_identity_token: received identity {:?}", identity);
+
         let creds = Credentials {
-            identity: Identity::from_bytes(identity),
+            identity,
             token: Token { string: token },
         };
 
         let address = Address::from_slice(&address);
+        log::trace!("handle_identity_token: received address {:?}", address);
 
         if Some(address) != self.address {
             log::error!(
@@ -930,8 +948,6 @@ impl CredentialStore {
             );
         }
 
-        self.callbacks.invoke((creds.clone(), address), state);
-
         if let Some(existing_creds) = &self.credentials {
             // If we already have credentials, make sure that they match. Log an error if
             // they don't.
@@ -944,10 +960,15 @@ impl CredentialStore {
                     creds,
                     existing_creds,
                 );
+            } else {
+                log::trace!("handle_identity_token: received creds match local record.");
             }
         } else {
-            self.credentials = Some(creds);
+            self.credentials = Some(creds.clone());
+            log::trace!("handle_identity_token: storing new credentials");
         }
+
+        self.callbacks.invoke((creds.clone(), address), state);
     }
 
     /// Return the current connection's `Identity`, if one is stored.

crates/testing/src/sdk.rs

@@ -211,7 +211,7 @@ fn run_client(run_command: &str, client_project: &str, db_name: &str, test_name:
         .dir(client_project)
         .env(TEST_CLIENT_PROJECT_ENV_VAR, client_project)
         .env(TEST_DB_NAME_ENV_VAR, db_name)
-        .env("RUST_LOG", "info")
+        .env("RUST_LOG", "trace")
         .stderr_to_stdout()
         .stdout_capture()
         .unchecked()