Private
Public Access
1
0

daemon: update monitor: implements ping/pong (required server changes)

This commit is contained in:
2025-06-13 16:45:28 -07:00
parent 4f40be205d
commit dece6f1abc
12 changed files with 202 additions and 55 deletions

View File

@@ -2,15 +2,31 @@ use crate::model::event::Event;
use crate::model::update::UpdateItem;
use async_trait::async_trait;
use futures_util::stream::Stream;
use futures_util::Sink;
#[derive(Debug, Eq, PartialEq, Clone)]
pub enum SinkMessage {
Ping,
}
pub enum SocketUpdate {
Update(Vec<UpdateItem>),
Pong,
}
pub enum SocketEvent {
Update(Event),
Pong,
}
#[async_trait]
pub trait EventSocket {
type Error;
type EventStream: Stream<Item = Result<Event, Self::Error>>;
type UpdateStream: Stream<Item = Result<Vec<UpdateItem>, Self::Error>>;
type EventStream: Stream<Item = Result<SocketEvent, Self::Error>>;
type UpdateStream: Stream<Item = Result<SocketUpdate, Self::Error>>;
/// Modern event pipeline
async fn events(self) -> Self::EventStream;
async fn events(self) -> (Self::EventStream, impl Sink<SinkMessage, Error = Self::Error>);
/// Raw update items from the v1 API.
async fn raw_updates(self) -> Self::UpdateStream;

View File

@@ -3,10 +3,9 @@ extern crate serde;
use std::{path::PathBuf, pin::Pin, str, task::Poll};
use crate::api::event_socket::EventSocket;
use crate::api::event_socket::{EventSocket, SinkMessage, SocketEvent, SocketUpdate};
use crate::api::AuthenticationStore;
use bytes::Bytes;
use hyper::body::HttpBody;
use hyper::{Body, Client, Method, Request, Uri};
use async_trait::async_trait;
@@ -16,7 +15,7 @@ use tokio::net::TcpStream;
use futures_util::stream::{BoxStream, Stream};
use futures_util::task::Context;
use futures_util::{StreamExt, TryStreamExt};
use futures_util::{Sink, SinkExt, StreamExt, TryStreamExt};
use tokio_tungstenite::connect_async;
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
@@ -49,6 +48,7 @@ pub enum Error {
HTTPError(hyper::Error),
SerdeError(serde_json::Error),
DecodeError(String),
PongError(tungstenite::Error),
Unauthorized,
}
@@ -124,34 +124,44 @@ impl<B> AuthSetting for hyper::http::Request<B> {
}
}
type WebsocketSink = futures_util::stream::SplitSink<WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>, tungstenite::Message>;
type WebsocketStream = futures_util::stream::SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>;
pub struct WebsocketEventSocket {
socket: WebSocketStream<MaybeTlsStream<TcpStream>>,
sink: Option<WebsocketSink>,
stream: WebsocketStream
}
impl WebsocketEventSocket {
pub fn new(socket: WebSocketStream<MaybeTlsStream<TcpStream>>) -> Self {
Self { socket }
let (sink, stream) = socket.split();
Self { sink: Some(sink), stream }
}
}
impl WebsocketEventSocket {
fn raw_update_stream(self) -> impl Stream<Item = Result<Vec<UpdateItem>, Error>> {
let (_, stream) = self.socket.split();
stream
fn raw_update_stream(self) -> impl Stream<Item = Result<SocketUpdate, Error>> {
self.stream
.map_err(Error::from)
.try_filter_map(|msg| async move {
match msg {
tungstenite::Message::Text(text) => {
serde_json::from_str::<Vec<UpdateItem>>(&text)
.map(Some)
.map_err(Error::from)
match serde_json::from_str::<Vec<UpdateItem>>(&text) {
Ok(updates) => Ok(Some(SocketUpdate::Update(updates))),
Err(e) => {
log::error!("Error parsing update: {:?}", e);
Err(Error::from(e))
}
}
}
tungstenite::Message::Ping(_) => {
// Borrowing issue here with the sink, need to handle pings at the client level (whomever
// is consuming these updateitems, should be a union type of updateitem | ping).
// We don't expect the server to send us pings.
Ok(None)
}
tungstenite::Message::Pong(_) => {
Ok(Some(SocketUpdate::Pong))
}
tungstenite::Message::Close(_) => {
// Connection was closed cleanly
Err(Error::ClientError("WebSocket connection closed".into()))
@@ -165,16 +175,37 @@ impl WebsocketEventSocket {
#[async_trait]
impl EventSocket for WebsocketEventSocket {
type Error = Error;
type EventStream = BoxStream<'static, Result<Event, Error>>;
type UpdateStream = BoxStream<'static, Result<Vec<UpdateItem>, Error>>;
type EventStream = BoxStream<'static, Result<SocketEvent, Error>>;
type UpdateStream = BoxStream<'static, Result<SocketUpdate, Error>>;
async fn events(self) -> Self::EventStream {
async fn events(mut self) -> (Self::EventStream, impl Sink<SinkMessage, Error = Self::Error>) {
use futures_util::stream::iter;
self.raw_update_stream()
.map_ok(|updates| iter(updates.into_iter().map(|update| Ok(Event::from(update)))))
let sink = self.sink.take().unwrap().with(|f| {
match f {
SinkMessage::Ping => futures_util::future::ready(Ok::<tungstenite::Message, Error>(tungstenite::Message::Ping(Bytes::new())))
}
});
let stream = self.raw_update_stream()
.map_ok(|updates| -> BoxStream<'static, Result<SocketEvent, Error>> {
match updates {
SocketUpdate::Update(updates) => {
let iter_stream = iter(
updates.into_iter().map(|u| Ok(SocketEvent::Update(Event::from(u))))
);
iter_stream.boxed()
}
SocketUpdate::Pong => {
iter(std::iter::once(Ok(SocketEvent::Pong))).boxed()
}
}
})
.try_flatten()
.boxed()
.boxed();
(stream, sink)
}
async fn raw_updates(self) -> Self::UpdateStream {

View File

@@ -12,6 +12,9 @@ pub struct UpdateItem {
#[serde(rename = "message")]
pub message: Option<Message>,
#[serde(default)]
pub pong: bool,
}
impl Default for UpdateItem {
@@ -20,6 +23,7 @@ impl Default for UpdateItem {
seq: 0,
conversation: None,
message: None,
pong: false,
}
}
}

View File

@@ -6,7 +6,7 @@ use uuid::Uuid;
pub use crate::APIInterface;
use crate::{
api::event_socket::EventSocket,
api::event_socket::{EventSocket, SinkMessage},
api::http_client::Credentials,
model::{
Conversation, ConversationID, Event, JwtToken, Message, MessageID, OutgoingMessage,
@@ -63,6 +63,10 @@ impl EventSocket for TestEventSocket {
let results: Vec<Result<Vec<UpdateItem>, TestError>> = vec![];
futures_util::stream::iter(results.into_iter()).boxed()
}
fn get_sink(&mut self) -> impl futures_util::Sink<SinkMessage> {
todo!("")
}
}
#[async_trait]