312 lines
12 KiB
Rust
312 lines
12 KiB
Rust
use collab::core::awareness;
|
|
use std::sync::Arc;
|
|
|
|
use crate::collaborate::sync_protocol::ServerSyncProtocol;
|
|
use collab::core::awareness::{Awareness, AwarenessUpdate};
|
|
use collab::core::collab::MutexCollab;
|
|
use collab::core::origin::CollabOrigin;
|
|
use futures_util::{SinkExt, StreamExt};
|
|
use realtime_protocol::handle_msg;
|
|
use realtime_protocol::{Message, MessageReader, MSG_SYNC, MSG_SYNC_UPDATE};
|
|
use tokio::select;
|
|
use tokio::sync::broadcast::error::SendError;
|
|
use tokio::sync::broadcast::{channel, Sender};
|
|
use tokio::sync::Mutex;
|
|
use yrs::updates::decoder::DecoderV1;
|
|
use yrs::updates::encoder::{Encode, Encoder, EncoderV1};
|
|
use yrs::UpdateSubscription;
|
|
|
|
use crate::collaborate::retry::SinkCollabMessageAction;
|
|
use crate::error::RealtimeError;
|
|
use realtime_entity::collab_msg::{CollabAck, CollabAwareness, CollabBroadcastData, CollabMessage};
|
|
use tracing::{error, trace, warn};
|
|
use yrs::encoding::write::Write;
|
|
|
|
/// A broadcast can be used to propagate updates produced by yrs [yrs::Doc] and [Awareness]
|
|
/// to subscribes. One broadcast can be used to propagate updates for a single document with
|
|
/// object_id.
|
|
pub struct CollabBroadcast {
|
|
object_id: String,
|
|
collab: MutexCollab,
|
|
sender: Sender<CollabMessage>,
|
|
awareness_sub: Mutex<Option<awareness::UpdateSubscription>>,
|
|
doc_sub: Mutex<Option<UpdateSubscription>>,
|
|
}
|
|
|
|
impl CollabBroadcast {
|
|
/// Creates a new [CollabBroadcast] over a provided `collab` instance. All changes triggered
|
|
/// by this collab will be propagated to all subscribers which have been registered via
|
|
/// [CollabBroadcast::subscribe] method.
|
|
///
|
|
/// The overflow of the incoming events that needs to be propagates will be buffered up to a
|
|
/// provided `buffer_capacity` size.
|
|
pub fn new(object_id: &str, collab: MutexCollab, buffer_capacity: usize) -> Self {
|
|
let object_id = object_id.to_owned();
|
|
// broadcast channel
|
|
let (sender, _) = channel(buffer_capacity);
|
|
CollabBroadcast {
|
|
object_id,
|
|
collab,
|
|
sender,
|
|
awareness_sub: Default::default(),
|
|
doc_sub: Default::default(),
|
|
}
|
|
}
|
|
|
|
pub async fn observe_collab_changes(&self) {
|
|
let (doc_sub, awareness_sub) = {
|
|
let mut mutex_collab = self.collab.lock();
|
|
|
|
// Observer the document's update and broadcast it to all subscribers.
|
|
let cloned_oid = self.object_id.clone();
|
|
let broadcast_sink = self.sender.clone();
|
|
let doc_sub = mutex_collab
|
|
.get_mut_awareness()
|
|
.doc_mut()
|
|
.observe_update_v1(move |txn, event| {
|
|
let update_len = event.update.len();
|
|
let origin = CollabOrigin::from(txn);
|
|
let payload = gen_update_message(&event.update);
|
|
let msg = CollabBroadcastData::new(origin, cloned_oid.clone(), payload);
|
|
|
|
match broadcast_sink.send(msg.into()) {
|
|
Ok(_) => trace!("observe doc update with len:{}", update_len),
|
|
Err(e) => error!(
|
|
"observe doc update with len:{} - broadcast sink fail: {}",
|
|
update_len, e
|
|
),
|
|
}
|
|
})
|
|
.unwrap();
|
|
|
|
let broadcast_sink = self.sender.clone();
|
|
let cloned_oid = self.object_id.clone();
|
|
|
|
// Observer the awareness's update and broadcast it to all subscribers.
|
|
let awareness_sub = mutex_collab
|
|
.get_mut_awareness()
|
|
.on_update(move |awareness, event| {
|
|
if let Ok(awareness_update) = gen_awareness_update_message(awareness, event) {
|
|
let payload = Message::Awareness(awareness_update).encode_v1();
|
|
let msg = CollabAwareness::new(cloned_oid.clone(), payload);
|
|
if let Err(_e) = broadcast_sink.send(msg.into()) {
|
|
trace!("Broadcast group is closed");
|
|
}
|
|
}
|
|
});
|
|
(doc_sub, awareness_sub)
|
|
};
|
|
|
|
*self.doc_sub.lock().await = Some(doc_sub);
|
|
*self.awareness_sub.lock().await = Some(awareness_sub);
|
|
}
|
|
|
|
/// Returns a reference to an underlying [MutexCollab] instance.
|
|
pub fn collab(&self) -> &MutexCollab {
|
|
&self.collab
|
|
}
|
|
|
|
/// Broadcasts user message to all active subscribers. Returns error if message could not have
|
|
/// been broadcast.
|
|
#[allow(clippy::result_large_err)]
|
|
pub fn broadcast_awareness(&self, msg: CollabAwareness) -> Result<(), SendError<CollabMessage>> {
|
|
self.sender.send(msg.into())?;
|
|
Ok(())
|
|
}
|
|
|
|
/// Subscribes a new connection - represented by `sink`/`stream` pair implementing a futures
|
|
/// Sink and Stream protocols - to a current broadcast group.
|
|
///
|
|
/// Returns a subscription structure, which can be dropped in order to unsubscribe or awaited
|
|
/// via [Subscription::stop] method in order to complete of its own volition (due to
|
|
/// an internal connection error or closed connection).
|
|
pub fn subscribe<Sink, Stream, E>(
|
|
&self,
|
|
subscriber_origin: CollabOrigin,
|
|
sink: Sink,
|
|
mut stream: Stream,
|
|
) -> Subscription
|
|
where
|
|
Sink: SinkExt<CollabMessage> + Send + Sync + Unpin + 'static,
|
|
Stream: StreamExt<Item = Result<CollabMessage, E>> + Send + Sync + Unpin + 'static,
|
|
<Sink as futures_util::Sink<CollabMessage>>::Error: std::error::Error + Send + Sync,
|
|
E: Into<anyhow::Error> + Send + Sync + 'static,
|
|
{
|
|
let cloned_origin = subscriber_origin.clone();
|
|
trace!("[realtime]: new subscriber: {}", subscriber_origin);
|
|
let sink = Arc::new(Mutex::new(sink));
|
|
// Receive a update from the document observer and forward the update to all
|
|
// connected subscribers using its Sink.
|
|
let sink_stop_tx = {
|
|
let sink = sink.clone();
|
|
let (stop_tx, mut stop_rx) = tokio::sync::mpsc::channel::<()>(1);
|
|
let mut receiver = self.sender.subscribe();
|
|
tokio::spawn(async move {
|
|
loop {
|
|
select! {
|
|
_ = stop_rx.recv() => break,
|
|
Ok(message) = receiver.recv() => {
|
|
if let Some(msg_origin) = message.origin() {
|
|
if msg_origin == &subscriber_origin {
|
|
continue;
|
|
}
|
|
}
|
|
|
|
trace!("[realtime]: broadcast collab message: {}", message);
|
|
let action = SinkCollabMessageAction {
|
|
sink: &sink,
|
|
message,
|
|
};
|
|
if let Err(err) = action.run().await {
|
|
error!("fail to broadcast message:{}", err);
|
|
}
|
|
},
|
|
}
|
|
}
|
|
});
|
|
stop_tx
|
|
};
|
|
|
|
// Receive messages from clients and reply with the response. The message may alter the
|
|
// document that the current broadcast group is associated with. If the message alter
|
|
// the document state then the document observer will be triggered and the update will be
|
|
// broadcast to all connected subscribers. Check out the [observe_update_v1] and [sink_task]
|
|
// above.
|
|
let stream_stop_tx = {
|
|
let collab = self.collab().clone();
|
|
let object_id = self.object_id.clone();
|
|
let (stream_stop_tx, mut stop_rx) = tokio::sync::mpsc::channel::<()>(1);
|
|
tokio::spawn(async move {
|
|
loop {
|
|
select! {
|
|
_ = stop_rx.recv() => break,
|
|
Some(Ok(collab_msg)) = stream.next() => {
|
|
// Continue if the message is empty
|
|
let payload = collab_msg.payload();
|
|
if payload.is_none() {
|
|
warn!("Unexpected empty payload of collab message:{}", collab_msg);
|
|
continue;
|
|
}
|
|
|
|
let collab_msg_origin = collab_msg.origin();
|
|
if object_id != collab_msg.object_id() {
|
|
error!("[🔴Server]: Incoming message's object id does not match the broadcast group's object id");
|
|
continue;
|
|
}
|
|
|
|
// safety: payload is not none
|
|
let mut decoder = DecoderV1::from(payload.unwrap().as_ref());
|
|
match sink.try_lock() {
|
|
Ok(mut sink) => {
|
|
let reader = MessageReader::new(&mut decoder);
|
|
for msg in reader {
|
|
match msg {
|
|
Ok(msg) => {
|
|
if let Ok(payload) =
|
|
handle_msg(&collab_msg_origin, &ServerSyncProtocol, &collab, msg) {
|
|
// Send the response to the corresponding client
|
|
match collab_msg_origin {
|
|
None => warn!("Client message does not have a origin"),
|
|
Some(collab_msg_origin) => {
|
|
if let Some(msg_id) = collab_msg.msg_id() {
|
|
let resp = CollabAck::new(
|
|
collab_msg_origin.clone(),
|
|
object_id.clone(),
|
|
payload.unwrap_or_default(),
|
|
msg_id,
|
|
collab_msg.type_str()
|
|
);
|
|
|
|
trace!("Send response to client: {}", resp);
|
|
match sink.send(resp.into()).await {
|
|
Ok(_) => {},
|
|
Err(err) => {
|
|
trace!("fail to send response to client: {}", err);
|
|
},
|
|
}
|
|
}
|
|
},
|
|
}
|
|
}
|
|
},
|
|
Err(e) => {
|
|
error!("Parser sync message failed: {:?}", e);
|
|
},
|
|
}
|
|
}
|
|
},
|
|
Err(err) => error!("Requires sink lock failed: {:?}", err),
|
|
}
|
|
}
|
|
}
|
|
}
|
|
});
|
|
stream_stop_tx
|
|
};
|
|
|
|
Subscription {
|
|
origin: cloned_origin,
|
|
sink_stop_tx: Some(sink_stop_tx),
|
|
stream_stop_tx: Some(stream_stop_tx),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// A subscription structure returned from [CollabBroadcast::subscribe], which represents a
|
|
/// subscribed connection. It can be dropped in order to unsubscribe or awaited via
|
|
/// [Subscription::stop] method in order to complete of its own volition (due to an internal
|
|
/// connection error or closed connection).
|
|
#[derive(Debug)]
|
|
pub struct Subscription {
|
|
pub origin: CollabOrigin,
|
|
sink_stop_tx: Option<tokio::sync::mpsc::Sender<()>>,
|
|
stream_stop_tx: Option<tokio::sync::mpsc::Sender<()>>,
|
|
}
|
|
|
|
impl Subscription {
|
|
pub async fn stop(mut self) {
|
|
if let Some(sink_stop_tx) = self.sink_stop_tx.take() {
|
|
let _ = sink_stop_tx.send(()).await;
|
|
}
|
|
if let Some(stream_stop_tx) = self.stream_stop_tx.take() {
|
|
let _ = stream_stop_tx.send(()).await;
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Drop for Subscription {
|
|
fn drop(&mut self) {
|
|
if self.stream_stop_tx.is_some() || self.stream_stop_tx.is_some() {
|
|
error!("Subscription is not stopped before dropping");
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Generates a message: Message::Sync::(SyncMessage::Update(update))
|
|
#[inline]
|
|
fn gen_update_message(update: &[u8]) -> Vec<u8> {
|
|
let mut encoder = EncoderV1::new();
|
|
// write the tag for Message::Sync
|
|
encoder.write_var(MSG_SYNC);
|
|
// write the tag for SyncMessage::Update
|
|
encoder.write_var(MSG_SYNC_UPDATE);
|
|
encoder.write_buf(update);
|
|
encoder.to_vec()
|
|
}
|
|
|
|
#[inline]
|
|
fn gen_awareness_update_message(
|
|
awareness: &Awareness,
|
|
event: &awareness::Event,
|
|
) -> Result<AwarenessUpdate, RealtimeError> {
|
|
let added = event.added();
|
|
let updated = event.updated();
|
|
let removed = event.removed();
|
|
let mut changed = Vec::with_capacity(added.len() + updated.len() + removed.len());
|
|
changed.extend_from_slice(added);
|
|
changed.extend_from_slice(updated);
|
|
changed.extend_from_slice(removed);
|
|
let update = awareness.update_with_clients(changed)?;
|
|
Ok(update)
|
|
}
|