Expose detailed sync stats through Database::get_sync_usage_stats

Signed-off-by: Piotr Jastrzebski <[email protected]>

Author: haaawk

libsql/src/local/database.rs

@@ -10,6 +10,7 @@ cfg_replication!(
     use crate::replication::remote_client::RemoteClient;
     use crate::replication::EmbeddedReplicator;
     pub use crate::replication::Frames;
+    pub use crate::replication::SyncUsageStats;
 
     pub struct ReplicationContext {
         pub(crate) replicator: EmbeddedReplicator,
@@ -277,6 +278,20 @@ impl Database {
         Ok(self.sync_oneshot().await?)
     }
 
+    #[cfg(feature = "replication")]
+    /// Return detailed logs about bytes synced with primary
+    pub async fn get_sync_usage_stats(&self) -> Result<SyncUsageStats> {
+        if let Some(ctx) = &self.replication_ctx {
+            let sync_stats = ctx.replicator.get_sync_usage_stats().await?;
+            Ok(sync_stats)
+        } else {
+            Err(crate::errors::Error::Misuse(
+                "No replicator available. Use Database::with_replicator() to enable replication"
+                    .to_string(),
+            ))
+        }
+    }
+
     #[cfg(feature = "replication")]
     /// Sync with primary at least to a given replication index
     pub async fn sync_until(&self, replication_index: FrameNo) -> Result<crate::replication::Replicated> {

libsql/src/replication/mod.rs

@@ -65,6 +65,60 @@ pub enum Frames {
     Snapshot(SnapshotFile),
 }
 
+/// Detailed logs about bytes synced with primary
+pub struct SyncUsageStats {
+    prefetched_bytes: u64,
+    prefetched_bytes_discarded_due_to_new_session: u64,
+    prefetched_bytes_discarded_due_to_consecutive_handshake: u64,
+    prefetched_bytes_discarded_due_to_invalid_frame_header: u64,
+    synced_bytes_discarded_due_to_invalid_frame_header: u64,
+    prefetched_bytes_used: u64,
+    synced_bytes_used: u64,
+    snapshot_bytes: u64,
+}
+
+impl SyncUsageStats {
+    /// Number of bytes prefetched while doing handshake
+    pub fn prefetched_bytes(&self) -> u64 {
+        self.prefetched_bytes
+    }
+
+    /// Number of bytes prefetched and discarded due to the change of the client session
+    pub fn prefetched_bytes_discarded_due_to_new_session(&self) -> u64 {
+        self.prefetched_bytes_discarded_due_to_new_session
+    }
+
+    /// Number of bytes prefetched and discarded due to consecutive handshake with new prefetch
+    pub fn prefetched_bytes_discarded_due_to_consecutive_handshake(&self) -> u64 {
+        self.prefetched_bytes_discarded_due_to_consecutive_handshake
+    }
+
+    /// Number of bytes prefetched and discarded due to invalid frame header in received frames
+    pub fn prefetched_bytes_discarded_due_to_invalid_frame_header(&self) -> u64 {
+        self.prefetched_bytes_discarded_due_to_invalid_frame_header
+    }
+
+    /// Number of bytes synced and discarded due to invalid frame header in received frames
+    pub fn synced_bytes_discarded_due_to_invalid_frame_header(&self) -> u64 {
+        self.synced_bytes_discarded_due_to_invalid_frame_header
+    }
+
+    /// Number of bytes prefetched and used
+    pub fn prefetched_bytes_used(&self) -> u64 {
+        self.prefetched_bytes_used
+    }
+
+    /// Number of bytes synced and used
+    pub fn synced_bytes_used(&self) -> u64 {
+        self.synced_bytes_used
+    }
+
+    /// Number of bytes downloaded as snapshots
+    pub fn snapshot_bytes(&self) -> u64 {
+        self.snapshot_bytes
+    }
+}
+
 #[derive(Clone)]
 pub(crate) struct Writer {
     pub(crate) client: client::Client,
@@ -210,6 +264,35 @@ impl EmbeddedReplicator {
         })
     }
 
+    pub async fn get_sync_usage_stats(&self) -> Result<SyncUsageStats> {
+        let mut replicator = self.replicator.lock().await;
+        match replicator.client_mut() {
+            Either::Right(_) => {
+                Err(crate::errors::Error::Misuse(
+                    "Trying to get sync usage stats, but this is a local replicator".into(),
+                ))
+            }
+            Either::Left(c) => {
+                let stats = c.sync_stats();
+                Ok(SyncUsageStats {
+                    prefetched_bytes: stats.prefetched_bytes.load(std::sync::atomic::Ordering::SeqCst),
+                    prefetched_bytes_discarded_due_to_new_session: stats
+                        .prefetched_bytes_discarded_due_to_new_session.load(std::sync::atomic::Ordering::SeqCst),
+                    prefetched_bytes_discarded_due_to_consecutive_handshake: stats
+                        .prefetched_bytes_discarded_due_to_consecutive_handshake.load(std::sync::atomic::Ordering::SeqCst),
+                    prefetched_bytes_discarded_due_to_invalid_frame_header: stats
+                        .prefetched_bytes_discarded_due_to_invalid_frame_header.load(std::sync::atomic::Ordering::SeqCst),
+                    synced_bytes_discarded_due_to_invalid_frame_header: stats
+                        .synced_bytes_discarded_due_to_invalid_frame_header.load(std::sync::atomic::Ordering::SeqCst),
+                    prefetched_bytes_used: stats.prefetched_bytes_used.load(std::sync::atomic::Ordering::SeqCst),
+                    synced_bytes_used: stats.synced_bytes_used.load(std::sync::atomic::Ordering::SeqCst),
+                    snapshot_bytes: stats.snapshot_bytes.load(std::sync::atomic::Ordering::SeqCst),
+                })
+            }
+        }
+
+    }
+
     pub async fn sync_oneshot(&self) -> Result<Replicated> {
         use libsql_replication::replicator::ReplicatorClient;
 

libsql/src/replication/remote_client.rs

@@ -24,7 +24,7 @@ async fn time<O>(fut: impl Future<Output = O>) -> (O, Duration) {
     (out, before.elapsed())
 }
 
-struct SyncStats {
+pub(crate) struct SyncStats {
     pub prefetched_bytes: AtomicU64,
     pub prefetched_bytes_discarded_due_to_new_session: AtomicU64,
     pub prefetched_bytes_discarded_due_to_consecutive_handshake: AtomicU64,
@@ -121,6 +121,10 @@ impl RemoteClient {
         })
     }
 
+    pub(crate) fn sync_stats(&self) -> Arc<SyncStats> {
+        self.sync_stats.clone()
+    }
+
     fn next_offset(&self) -> FrameNo {
         match self.last_received {
             Some(fno) => fno + 1,