Docstrings and unit tests

Author: maciejdudko

client/src/replaceable.rs

@@ -4,10 +4,13 @@ use std::sync::atomic::{AtomicU32, Ordering};
 use std::sync::{Arc, RwLock};
 
 /// A client wrapper that allows replacing the underlying client at a later point in time.
-/// Clones of [SharedReplaceableClient] have a shared reference to the underlying client, and also
-/// a local cached clone of the underlying client. Before every service call, the cached clone is
-/// updated if the shared client was replaced.
-#[derive(Debug, Clone)]
+/// Clones of this struct have a shared reference to the underlying client, and each clone also
+/// has its own cached clone of the underlying client. Before every service call, a check is made
+/// whether the shared client was replaced, and the cached clone is updated accordingly.
+///
+/// This struct is fully thread-safe, and it works in a lock-free manner except when the client is
+/// being replaced. A read-write lock is used then, with minimal locking time.
+#[derive(Debug)]
 pub struct SharedReplaceableClient<C>
 where
     C: Clone + Send + Sync,
@@ -30,14 +33,17 @@ impl<C> SharedClientData<C>
 where
     C: Clone + Send + Sync,
 {
+    fn fetch(&self) -> (C, u32) {
+        let lock = self.client.read().unwrap();
+        let client = lock.clone();
+        // Loading generation under lock to ensure the client won't be updated in the meantime.
+        let generation = self.generation.load(Ordering::Acquire);
+        (client, generation)
+    }
+
     fn fetch_newer_than(&self, current_generation: u32) -> Option<(C, u32)> {
-        (current_generation != self.generation.load(Ordering::Acquire)).then(|| {
-            let lock = self.client.read().unwrap();
-            let client = lock.clone();
-            // Loading generation under lock to ensure the client won't be updated in the meantime.
-            let generation = self.generation.load(Ordering::Acquire);
-            (client, generation)
-        })
+        // fetch() will do a second atomic load, but it's necessary to avoid a race condition.
+        (current_generation != self.generation.load(Ordering::Acquire)).then(|| self.fetch())
     }
 
     fn replace_client(&self, client: C) {
@@ -54,7 +60,7 @@ where
     C: Clone + Send + Sync,
 {
     /// Creates the initial instance of replaceable client with the provided underlying client.
-    /// Use `clone` method to create more instances that share the same underlying client.
+    /// Use [`clone()`](Self::clone) method to create more instances that share the same underlying client.
     pub fn new(client: C) -> Self {
         let cloned_client = client.clone();
         Self {
@@ -77,7 +83,10 @@ where
         self.inner_cow().into_owned()
     }
 
-    /// Returns a reference to the underlying client if possible, or its clone otherwise.
+    /// Returns a reference to this instance's cached clone of the underlying client if it's up to
+    /// date, or a fresh clone of the shared client otherwise. Because it's an immutable method,
+    /// it will not update this instance's cached clone. For this reason, prefer to use
+    /// [`refresh_inner()`](Self::refresh_inner) when possible.
     pub fn inner_cow(&self) -> Cow<'_, C> {
         self.shared_data
             .fetch_newer_than(self.cloned_generation)
@@ -86,7 +95,11 @@ where
     }
 
     /// Refreshes this instance's cached clone of the underlying client. Returns a mutable reference
-    /// to it. Called automatically by other mutable methods, in particular by all client calls.
+    /// to it. Called automatically by other mutable methods, in particular by all RPC calls.
+    ///
+    /// While this method allows mutable access to the underlying client, any configuration changes
+    /// will not be shared with other instances, and will be lost if the client gets replaced from
+    /// anywhere. To make configuration changes, use [`replace_client()`](Self::refresh_client) instead.
     pub fn refresh_inner(&mut self) -> &mut C {
         if let Some((client, generation)) =
             self.shared_data.fetch_newer_than(self.cloned_generation)
@@ -98,6 +111,25 @@ where
     }
 }
 
+impl<C> Clone for SharedReplaceableClient<C>
+where
+    C: Clone + Send + Sync,
+{
+    /// Creates a new instance of replaceable client that shares the underlying client with this
+    /// instance. Replacing a client in either instance will replace it for both instances, and all
+    /// other clones too.
+    fn clone(&self) -> Self {
+        // self's cloned_client could've been modified through a mutable reference,
+        // so for consistent behavior, we need to fetch it from shared_data.
+        let (client, generation) = self.shared_data.fetch();
+        Self {
+            shared_data: self.shared_data.clone(),
+            cloned_client: client,
+            cloned_generation: generation,
+        }
+    }
+}
+
 impl<C> NamespacedClient for SharedReplaceableClient<C>
 where
     C: NamespacedClient + Clone + Send + Sync,
@@ -110,3 +142,76 @@ where
         self.inner_cow().identity()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::{NamespacedClient, SharedReplaceableClient};
+    use std::borrow::Cow;
+
+    #[derive(Debug, Clone)]
+    struct StubClient {
+        identity: String,
+    }
+
+    impl StubClient {
+        fn new(identity: &str) -> Self {
+            Self {
+                identity: identity.to_owned(),
+            }
+        }
+    }
+
+    impl NamespacedClient for StubClient {
+        fn namespace(&self) -> String {
+            "default".into()
+        }
+
+        fn identity(&self) -> String {
+            self.identity.clone()
+        }
+    }
+
+    #[test]
+    fn cow_returns_reference_before_and_clone_after_refresh() {
+        let mut client = SharedReplaceableClient::new(StubClient::new("1"));
+        let Cow::Borrowed(inner) = client.inner_cow() else {
+            panic!("expected borrowed inner");
+        };
+        assert_eq!(inner.identity, "1");
+
+        client.replace_client(StubClient::new("2"));
+        let Cow::Owned(inner) = client.inner_cow() else {
+            panic!("expected owned inner");
+        };
+        assert_eq!(inner.identity, "2");
+
+        assert_eq!(client.refresh_inner().identity, "2");
+        let Cow::Borrowed(inner) = client.inner_cow() else {
+            panic!("expected borrowed inner");
+        };
+        assert_eq!(inner.identity, "2");
+    }
+
+    #[test]
+    fn client_replaced_in_clones() {
+        let original1 = SharedReplaceableClient::new(StubClient::new("1"));
+        let clone1 = original1.clone();
+        assert_eq!(original1.identity(), "1");
+        assert_eq!(clone1.identity(), "1");
+
+        original1.replace_client(StubClient::new("2"));
+        assert_eq!(original1.identity(), "2");
+        assert_eq!(clone1.identity(), "2");
+
+        let original2 = SharedReplaceableClient::new(StubClient::new("3"));
+        let clone2 = original2.clone();
+        assert_eq!(original2.identity(), "3");
+        assert_eq!(clone2.identity(), "3");
+
+        clone2.replace_client(StubClient::new("4"));
+        assert_eq!(original2.identity(), "4");
+        assert_eq!(clone2.identity(), "4");
+        assert_eq!(original1.identity(), "2");
+        assert_eq!(clone1.identity(), "2");
+    }
+}