pass RpcFrame to client methods

necessary to pass different underlying frames

Author: MarinPostma

bottomless/src/replicator.rs

@@ -18,6 +18,7 @@ use aws_sdk_s3::{Client, Config};
 use bytes::{Buf, Bytes};
 use chrono::{DateTime, NaiveDateTime, TimeZone, Utc};
 use libsql_replication::injector::Injector as _;
+use libsql_replication::rpc::replication::Frame as RpcFrame;
 use libsql_sys::{Cipher, EncryptionConfig};
 use std::ops::Deref;
 use std::path::{Path, PathBuf};
@@ -1554,7 +1555,11 @@ impl Replicator {
                         },
                         page_buf.as_slice(),
                     );
-                    injector.inject_frame(frame_to_inject).await?;
+                    let frame = RpcFrame {
+                        data: frame_to_inject.bytes(),
+                        timestamp: None,
+                    };
+                    injector.inject_frame(frame).await?;
                     applied_wal_frame = true;
                 }
             }

libsql-replication/src/injector/libsql_injector.rs

@@ -5,7 +5,8 @@ use libsql_wal::replication::injector::Injector;
 use libsql_wal::segment::Frame as WalFrame;
 use zerocopy::{AsBytes, FromZeroes};
 
-use crate::frame::{Frame, FrameNo};
+use crate::frame::FrameNo;
+use crate::rpc::replication::Frame as RpcFrame;
 
 use super::error::{Error, Result};
 
@@ -14,15 +15,15 @@ pub struct LibsqlInjector {
 }
 
 impl super::Injector for LibsqlInjector {
-    async fn inject_frame(&mut self, frame: Frame) -> Result<Option<FrameNo>> {
+    async fn inject_frame(&mut self, frame: RpcFrame) -> Result<Option<FrameNo>> {
         // this is a bit annoying be we want to read the frame, and it has to be aligned, so we
         // must copy it...
         // FIXME: optimize this.
         let mut wal_frame = WalFrame::new_box_zeroed();
-        if frame.bytes().len() != size_of::<WalFrame>() {
+        if frame.data.len() != size_of::<WalFrame>() {
             todo!("invalid frame");
         }
-        wal_frame.as_bytes_mut().copy_from_slice(&frame.bytes()[..]);
+        wal_frame.as_bytes_mut().copy_from_slice(&frame.data[..]);
         Ok(self
             .injector
             .insert_frame(wal_frame)

libsql-replication/src/injector/mod.rs

@@ -1,9 +1,10 @@
 use std::future::Future;
 
+use super::rpc::replication::Frame as RpcFrame;
 pub use libsql_injector::LibsqlInjector;
 pub use sqlite_injector::SqliteInjector;
 
-use crate::frame::{Frame, FrameNo};
+use crate::frame::FrameNo;
 
 pub use error::Error;
 use error::Result;
@@ -16,7 +17,7 @@ pub trait Injector {
     /// Inject a singular frame.
     fn inject_frame(
         &mut self,
-        frame: Frame,
+        frame: RpcFrame,
     ) -> impl Future<Output = Result<Option<FrameNo>>> + Send;
 
     /// Discard any uncommintted frames.

libsql-replication/src/injector/sqlite_injector/mod.rs

@@ -7,6 +7,7 @@ use rusqlite::OpenFlags;
 use tokio::task::spawn_blocking;
 
 use crate::frame::{Frame, FrameNo};
+use crate::rpc::replication::Frame as RpcFrame;
 
 use self::injector_wal::{
     InjectorWal, InjectorWalManager, LIBSQL_INJECT_FATAL, LIBSQL_INJECT_OK, LIBSQL_INJECT_OK_TXN,
@@ -25,8 +26,10 @@ pub struct SqliteInjector {
 }
 
 impl Injector for SqliteInjector {
-    async fn inject_frame(&mut self, frame: Frame) -> Result<Option<FrameNo>> {
+    async fn inject_frame(&mut self, frame: RpcFrame) -> Result<Option<FrameNo>> {
         let inner = self.inner.clone();
+        let frame =
+            Frame::try_from(&frame.data[..]).map_err(|e| Error::FatalInjectError(e.into()))?;
         spawn_blocking(move || inner.lock().inject_frame(frame))
             .await
             .unwrap()

libsql-replication/src/replicator.rs

@@ -63,7 +63,7 @@ impl From<tokio::task::JoinError> for Error {
 
 #[async_trait::async_trait]
 pub trait ReplicatorClient {
-    type FrameStream: Stream<Item = Result<Frame, Error>> + Unpin + Send;
+    type FrameStream: Stream<Item = Result<RpcFrame, Error>> + Unpin + Send;
 
     /// Perform handshake with remote
     async fn handshake(&mut self) -> Result<(), Error>;
@@ -318,7 +318,7 @@ where
         }
     }
 
-    async fn inject_frame(&mut self, frame: Frame) -> Result<(), Error> {
+    async fn inject_frame(&mut self, frame: RpcFrame) -> Result<(), Error> {
         self.frames_synced += 1;
 
         match self.injector.inject_frame(frame).await? {
@@ -360,6 +360,7 @@ mod test {
     use async_stream::stream;
 
     use crate::frame::{FrameBorrowed, FrameMut};
+    use crate::rpc::replication::Frame as RpcFrame;
 
     use super::*;
 
@@ -370,7 +371,8 @@ mod test {
 
         #[async_trait::async_trait]
         impl ReplicatorClient for Client {
-            type FrameStream = Pin<Box<dyn Stream<Item = Result<Frame, Error>> + Send + 'static>>;
+            type FrameStream =
+                Pin<Box<dyn Stream<Item = Result<RpcFrame, Error>> + Send + 'static>>;
 
             /// Perform handshake with remote
             async fn handshake(&mut self) -> Result<(), Error> {
@@ -414,7 +416,8 @@ mod test {
 
         #[async_trait::async_trait]
         impl ReplicatorClient for Client {
-            type FrameStream = Pin<Box<dyn Stream<Item = Result<Frame, Error>> + Send + 'static>>;
+            type FrameStream =
+                Pin<Box<dyn Stream<Item = Result<RpcFrame, Error>> + Send + 'static>>;
 
             /// Perform handshake with remote
             async fn handshake(&mut self) -> Result<(), Error> {
@@ -456,7 +459,8 @@ mod test {
 
         #[async_trait::async_trait]
         impl ReplicatorClient for Client {
-            type FrameStream = Pin<Box<dyn Stream<Item = Result<Frame, Error>> + Send + 'static>>;
+            type FrameStream =
+                Pin<Box<dyn Stream<Item = Result<RpcFrame, Error>> + Send + 'static>>;
 
             /// Perform handshake with remote
             async fn handshake(&mut self) -> Result<(), Error> {
@@ -500,7 +504,8 @@ mod test {
 
         #[async_trait::async_trait]
         impl ReplicatorClient for Client {
-            type FrameStream = Pin<Box<dyn Stream<Item = Result<Frame, Error>> + Send + 'static>>;
+            type FrameStream =
+                Pin<Box<dyn Stream<Item = Result<RpcFrame, Error>> + Send + 'static>>;
 
             /// Perform handshake with remote
             async fn handshake(&mut self) -> Result<(), Error> {
@@ -544,7 +549,8 @@ mod test {
 
         #[async_trait::async_trait]
         impl ReplicatorClient for Client {
-            type FrameStream = Pin<Box<dyn Stream<Item = Result<Frame, Error>> + Send + 'static>>;
+            type FrameStream =
+                Pin<Box<dyn Stream<Item = Result<RpcFrame, Error>> + Send + 'static>>;
 
             /// Perform handshake with remote
             async fn handshake(&mut self) -> Result<(), Error> {
@@ -586,7 +592,8 @@ mod test {
 
         #[async_trait::async_trait]
         impl ReplicatorClient for Client {
-            type FrameStream = Pin<Box<dyn Stream<Item = Result<Frame, Error>> + Send + 'static>>;
+            type FrameStream =
+                Pin<Box<dyn Stream<Item = Result<RpcFrame, Error>> + Send + 'static>>;
 
             /// Perform handshake with remote
             async fn handshake(&mut self) -> Result<(), Error> {
@@ -627,7 +634,8 @@ mod test {
 
         #[async_trait::async_trait]
         impl ReplicatorClient for Client {
-            type FrameStream = Pin<Box<dyn Stream<Item = Result<Frame, Error>> + Send + 'static>>;
+            type FrameStream =
+                Pin<Box<dyn Stream<Item = Result<RpcFrame, Error>> + Send + 'static>>;
 
             /// Perform handshake with remote
             async fn handshake(&mut self) -> Result<(), Error> {
@@ -672,7 +680,8 @@ mod test {
 
         #[async_trait::async_trait]
         impl ReplicatorClient for Client {
-            type FrameStream = Pin<Box<dyn Stream<Item = Result<Frame, Error>> + Send + 'static>>;
+            type FrameStream =
+                Pin<Box<dyn Stream<Item = Result<RpcFrame, Error>> + Send + 'static>>;
 
             /// Perform handshake with remote
             async fn handshake(&mut self) -> Result<(), Error> {
@@ -740,7 +749,8 @@ mod test {
 
         #[async_trait::async_trait]
         impl ReplicatorClient for Client {
-            type FrameStream = Pin<Box<dyn Stream<Item = Result<Frame, Error>> + Send + 'static>>;
+            type FrameStream =
+                Pin<Box<dyn Stream<Item = Result<RpcFrame, Error>> + Send + 'static>>;
 
             /// Perform handshake with remote
             async fn handshake(&mut self) -> Result<(), Error> {
@@ -752,15 +762,26 @@ mod test {
                     let frames = self
                         .frames
                         .iter()
+                        .map(|f| RpcFrame {
+                            data: f.bytes(),
+                            timestamp: None,
+                        })
                         .take(2)
-                        .cloned()
                         .map(Ok)
                         .chain(Some(Err(Error::Client("some client error".into()))))
                         .collect::<Vec<_>>();
                     Ok(Box::pin(tokio_stream::iter(frames)))
                 } else {
-                    let stream = tokio_stream::iter(self.frames.clone().into_iter().map(Ok));
-                    Ok(Box::pin(stream))
+                    let iter = self
+                        .frames
+                        .iter()
+                        .map(|f| RpcFrame {
+                            data: f.bytes(),
+                            timestamp: None,
+                        })
+                        .map(Ok)
+                        .collect::<Vec<_>>();
+                    Ok(Box::pin(tokio_stream::iter(iter)))
                 }
             }
             /// Return a snapshot for the current replication index. Called after next_frame has returned a

libsql-server/src/replication/replicator_client.rs

@@ -4,16 +4,17 @@ use std::pin::Pin;
 use bytes::Bytes;
 use chrono::{DateTime, Utc};
 use futures::TryStreamExt;
-use libsql_replication::frame::Frame;
 use libsql_replication::meta::WalIndexMeta;
-use libsql_replication::replicator::{map_frame_err, Error, ReplicatorClient};
+use libsql_replication::replicator::{Error, ReplicatorClient};
 use libsql_replication::rpc::replication::hello_request::WalFlavor;
 use libsql_replication::rpc::replication::replication_log_client::ReplicationLogClient;
 use libsql_replication::rpc::replication::{
-    verify_session_token, HelloRequest, LogOffset, NAMESPACE_METADATA_KEY, SESSION_TOKEN_KEY,
+    verify_session_token, Frame as RpcFrame, HelloRequest, LogOffset, NAMESPACE_METADATA_KEY,
+    SESSION_TOKEN_KEY,
 };
 use tokio::sync::watch;
-use tokio_stream::{Stream, StreamExt};
+use tokio_stream::Stream;
+
 use tonic::metadata::{AsciiMetadataValue, BinaryMetadataValue};
 use tonic::transport::Channel;
 use tonic::{Code, Request, Status};
@@ -95,7 +96,7 @@ impl Client {
 
 #[async_trait::async_trait]
 impl ReplicatorClient for Client {
-    type FrameStream = Pin<Box<dyn Stream<Item = Result<Frame, Error>> + Send + 'static>>;
+    type FrameStream = Pin<Box<dyn Stream<Item = Result<RpcFrame, Error>> + Send + 'static>>;
 
     #[tracing::instrument(skip(self))]
     async fn handshake(&mut self) -> Result<(), Error> {
@@ -169,7 +170,7 @@ impl ReplicatorClient for Client {
                     None => REPLICATION_LATENCY_CACHE_MISS.increment(1),
                 }
             })
-            .map(map_frame_err);
+            .map_err(Into::into);
 
         Ok(Box::pin(stream))
     }
@@ -181,7 +182,7 @@ impl ReplicatorClient for Client {
         let req = self.make_request(offset);
         match self.client.snapshot(req).await {
             Ok(resp) => {
-                let stream = resp.into_inner().map(map_frame_err);
+                let stream = resp.into_inner().map_err(Into::into);
                 Ok(Box::pin(stream))
             }
             Err(e) if e.code() == Code::Unavailable => Err(Error::SnapshotPending),

libsql/src/replication/local_client.rs

@@ -3,6 +3,7 @@ use std::pin::Pin;
 
 use futures::{StreamExt, TryStreamExt};
 use libsql_replication::{
+    rpc::replication::Frame as RpcFrame,
     frame::{Frame, FrameNo},
     meta::WalIndexMeta,
     replicator::{Error, ReplicatorClient},
@@ -35,7 +36,7 @@ impl LocalClient {
 
 #[async_trait::async_trait]
 impl ReplicatorClient for LocalClient {
-    type FrameStream = Pin<Box<dyn Stream<Item = Result<Frame, Error>> + Send + 'static>>;
+    type FrameStream = Pin<Box<dyn Stream<Item = Result<RpcFrame, Error>> + Send + 'static>>;
 
     /// Perform handshake with remote
     async fn handshake(&mut self) -> Result<(), Error> {
@@ -46,7 +47,7 @@ impl ReplicatorClient for LocalClient {
     async fn next_frames(&mut self) -> Result<Self::FrameStream, Error> {
         match self.frames.take() {
             Some(Frames::Vec(f)) => {
-                let iter = f.into_iter().map(Ok);
+                let iter = f.into_iter().map(|f| RpcFrame { data: f.bytes(), timestamp: None }).map(Ok);
                 Ok(Box::pin(tokio_stream::iter(iter)))
             }
             Some(f @ Frames::Snapshot(_)) => {
@@ -70,7 +71,8 @@ impl ReplicatorClient for LocalClient {
                         if s.as_mut().peek().await.is_none() {
                             next.header_mut().size_after = size_after.into();
                         }
-                        yield Frame::from(next);
+                        let frame = Frame::from(next);
+                        yield RpcFrame { data: frame.bytes(), timestamp: None };
                     }
                 };
 
@@ -95,8 +97,9 @@ impl ReplicatorClient for LocalClient {
 
 #[cfg(test)]
 mod test {
-    use libsql_replication::snapshot::SnapshotFile;
+    use libsql_replication::{frame::FrameHeader, snapshot::SnapshotFile};
     use tempfile::tempdir;
+    use zerocopy::FromBytes;
 
     use super::*;
 
@@ -111,7 +114,8 @@ mod test {
         let mut s = client.snapshot().await.unwrap();
         assert!(matches!(s.next().await, Some(Ok(_))));
         let last = s.next().await.unwrap().unwrap();
-        assert_eq!(last.header().size_after.get(), 2);
+        let header: FrameHeader = FrameHeader::read_from_prefix(&last.data[..]).unwrap();
+        assert_eq!(header.size_after.get(), 2);
         assert!(s.next().await.is_none());
     }
 }

libsql/src/replication/remote_client.rs

@@ -4,13 +4,13 @@ use std::pin::Pin;
 use std::time::{Duration, Instant};
 
 use bytes::Bytes;
-use futures::StreamExt as _;
-use libsql_replication::frame::{Frame, FrameHeader, FrameNo};
+use futures::{StreamExt as _, TryStreamExt};
+use libsql_replication::frame::{FrameHeader, FrameNo};
 use libsql_replication::meta::WalIndexMeta;
-use libsql_replication::replicator::{map_frame_err, Error, ReplicatorClient};
+use libsql_replication::replicator::{Error, ReplicatorClient};
 use libsql_replication::rpc::replication::hello_request::WalFlavor;
 use libsql_replication::rpc::replication::{
-    verify_session_token, Frames, HelloRequest, HelloResponse, LogOffset, SESSION_TOKEN_KEY,
+    Frame as RpcFrame, verify_session_token, Frames, HelloRequest, HelloResponse, LogOffset, SESSION_TOKEN_KEY,
 };
 use tokio_stream::Stream;
 use tonic::metadata::AsciiMetadataValue;
@@ -161,7 +161,7 @@ impl RemoteClient {
 
         let frames_iter = frames
             .into_iter()
-            .map(|f| Frame::try_from(&*f.data).map_err(|e| Error::Client(e.into())));
+            .map(Ok);
 
         let stream = tokio_stream::iter(frames_iter);
 
@@ -197,15 +197,16 @@ impl RemoteClient {
             .snapshot(req)
             .await?
             .into_inner()
-            .map(map_frame_err)
+            .map_err(|e| e.into())
             .peekable();
 
         {
             let frames = Pin::new(&mut frames);
 
             // the first frame is the one with the highest frame_no in the snapshot
             if let Some(Ok(f)) = frames.peek().await {
-                self.last_received = Some(f.header().frame_no.get());
+                let header: FrameHeader = FrameHeader::read_from_prefix(&f.data[..]).unwrap();
+                self.last_received = Some(header.frame_no.get());
             }
         }
 
@@ -240,7 +241,7 @@ fn maybe_log<T>(
 
 #[async_trait::async_trait]
 impl ReplicatorClient for RemoteClient {
-    type FrameStream = Pin<Box<dyn Stream<Item = Result<Frame, Error>> + Send + 'static>>;
+    type FrameStream = Pin<Box<dyn Stream<Item = Result<RpcFrame, Error>> + Send + 'static>>;
 
     /// Perform handshake with remote
     async fn handshake(&mut self) -> Result<(), Error> {