use crate::error::Result;
use crate::protocol::Message;
use bytes::Bytes;
use futures::stream::Stream;
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::sync::mpsc;
pub struct MessageStream {
receiver: mpsc::UnboundedReceiver<Message>,
}
impl MessageStream {
pub fn new(receiver: mpsc::UnboundedReceiver<Message>) -> Self {
Self { receiver }
}
pub async fn next_message(&mut self) -> Option<Message> {
self.receiver.recv().await
}
}
impl Stream for MessageStream {
type Item = Message;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.receiver.poll_recv(cx)
}
}
pub struct TileStream {
receiver: mpsc::UnboundedReceiver<TileData>,
}
impl TileStream {
pub fn new(receiver: mpsc::UnboundedReceiver<TileData>) -> Self {
Self { receiver }
}
pub async fn next_tile(&mut self) -> Option<TileData> {
self.receiver.recv().await
}
}
impl Stream for TileStream {
type Item = TileData;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.receiver.poll_recv(cx)
}
}
#[derive(Debug, Clone)]
pub struct TileData {
pub x: u32,
pub y: u32,
pub zoom: u8,
pub data: Bytes,
pub mime_type: String,
}
impl TileData {
pub fn new(x: u32, y: u32, zoom: u8, data: Vec<u8>, mime_type: String) -> Self {
Self {
x,
y,
zoom,
data: Bytes::from(data),
mime_type,
}
}
pub fn coords(&self) -> (u32, u32, u8) {
(self.x, self.y, self.zoom)
}
pub fn size(&self) -> usize {
self.data.len()
}
}
pub struct FeatureStream {
receiver: mpsc::UnboundedReceiver<FeatureData>,
}
impl FeatureStream {
pub fn new(receiver: mpsc::UnboundedReceiver<FeatureData>) -> Self {
Self { receiver }
}
pub async fn next_feature(&mut self) -> Option<FeatureData> {
self.receiver.recv().await
}
}
impl Stream for FeatureStream {
type Item = FeatureData;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.receiver.poll_recv(cx)
}
}
#[derive(Debug, Clone)]
pub struct FeatureData {
pub geojson: String,
pub change_type: crate::protocol::ChangeType,
pub layer: Option<String>,
}
impl FeatureData {
pub fn new(
geojson: String,
change_type: crate::protocol::ChangeType,
layer: Option<String>,
) -> Self {
Self {
geojson,
change_type,
layer,
}
}
pub fn parse_json(&self) -> Result<serde_json::Value> {
serde_json::from_str(&self.geojson).map_err(Into::into)
}
}
pub struct EventStream {
receiver: mpsc::UnboundedReceiver<EventData>,
}
impl EventStream {
pub fn new(receiver: mpsc::UnboundedReceiver<EventData>) -> Self {
Self { receiver }
}
pub async fn next_event(&mut self) -> Option<EventData> {
self.receiver.recv().await
}
}
impl Stream for EventStream {
type Item = EventData;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.receiver.poll_recv(cx)
}
}
#[derive(Debug, Clone)]
pub struct EventData {
pub event_type: crate::protocol::EventType,
pub payload: serde_json::Value,
pub timestamp: chrono::DateTime<chrono::Utc>,
}
impl EventData {
pub fn new(event_type: crate::protocol::EventType, payload: serde_json::Value) -> Self {
Self {
event_type,
payload,
timestamp: chrono::Utc::now(),
}
}
pub fn with_timestamp(
event_type: crate::protocol::EventType,
payload: serde_json::Value,
timestamp: chrono::DateTime<chrono::Utc>,
) -> Self {
Self {
event_type,
payload,
timestamp,
}
}
}
pub struct BackpressureController {
max_buffer_size: usize,
current_buffer_size: usize,
high_watermark: f64,
low_watermark: f64,
state: BackpressureState,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BackpressureState {
Normal,
High,
Critical,
}
impl BackpressureController {
pub fn new(max_buffer_size: usize) -> Self {
Self {
max_buffer_size,
current_buffer_size: 0,
high_watermark: 0.7,
low_watermark: 0.3,
state: BackpressureState::Normal,
}
}
pub fn update(&mut self, buffer_size: usize) -> BackpressureState {
self.current_buffer_size = buffer_size;
let ratio = buffer_size as f64 / self.max_buffer_size as f64;
self.state = if ratio >= 0.9 {
BackpressureState::Critical
} else if ratio >= self.high_watermark {
BackpressureState::High
} else if ratio <= self.low_watermark {
BackpressureState::Normal
} else {
self.state
};
self.state
}
pub fn state(&self) -> BackpressureState {
self.state
}
pub fn should_throttle(&self) -> bool {
matches!(
self.state,
BackpressureState::High | BackpressureState::Critical
)
}
pub fn should_drop(&self) -> bool {
self.state == BackpressureState::Critical
}
}
pub struct DeltaEncoder {
cache: dashmap::DashMap<(u32, u32, u8), Bytes>,
}
impl DeltaEncoder {
pub fn new() -> Self {
Self {
cache: dashmap::DashMap::new(),
}
}
pub fn encode(&self, tile: &TileData) -> Result<Vec<u8>> {
let key = tile.coords();
if let Some(prev_data) = self.cache.get(&key) {
let delta = Self::compute_delta(&prev_data, &tile.data)?;
self.cache.insert(key, tile.data.clone());
Ok(delta)
} else {
self.cache.insert(key, tile.data.clone());
Ok(tile.data.to_vec())
}
}
fn compute_delta(old: &[u8], new: &[u8]) -> Result<Vec<u8>> {
let mut delta = Vec::new();
delta.extend_from_slice(&(new.len() as u32).to_le_bytes());
for (i, (&old_byte, &new_byte)) in old.iter().zip(new.iter()).enumerate() {
if old_byte != new_byte {
delta.extend_from_slice(&(i as u32).to_le_bytes());
delta.push(new_byte);
}
}
if new.len() > old.len() {
for (i, &byte) in new[old.len()..].iter().enumerate() {
let pos = old.len() + i;
delta.extend_from_slice(&(pos as u32).to_le_bytes());
delta.push(byte);
}
}
Ok(delta)
}
pub fn clear(&self) {
self.cache.clear();
}
pub fn cache_size(&self) -> usize {
self.cache.len()
}
}
impl Default for DeltaEncoder {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_message_stream() {
let (tx, rx) = mpsc::unbounded_channel();
let mut stream = MessageStream::new(rx);
let send_result = tx.send(Message::Ping { id: 1 });
assert!(send_result.is_ok());
let msg = stream.next_message().await;
assert!(msg.is_some());
if let Some(Message::Ping { id }) = msg {
assert_eq!(id, 1);
}
}
#[tokio::test]
async fn test_tile_stream() {
let (tx, rx) = mpsc::unbounded_channel();
let mut stream = TileStream::new(rx);
let tile = TileData::new(0, 0, 5, vec![1, 2, 3], "application/x-protobuf".to_string());
let send_result = tx.send(tile.clone());
assert!(send_result.is_ok());
let received = stream.next_tile().await;
assert!(received.is_some());
if let Some(tile) = received {
assert_eq!(tile.coords(), (0, 0, 5));
assert_eq!(tile.size(), 3);
}
}
#[test]
fn test_backpressure_controller() {
let mut controller = BackpressureController::new(100);
assert_eq!(controller.update(30), BackpressureState::Normal);
assert!(!controller.should_throttle());
assert_eq!(controller.update(75), BackpressureState::High);
assert!(controller.should_throttle());
assert_eq!(controller.update(95), BackpressureState::Critical);
assert!(controller.should_drop());
assert_eq!(controller.update(25), BackpressureState::Normal);
assert!(!controller.should_throttle());
}
#[test]
#[ignore]
fn test_delta_encoder() {
let encoder = DeltaEncoder::new();
let tile1 = TileData::new(
0,
0,
5,
vec![1, 2, 3, 4, 5],
"application/x-protobuf".to_string(),
);
let delta1 = encoder.encode(&tile1);
assert!(delta1.is_ok());
if let Ok(data) = delta1 {
assert_eq!(data.len(), 5);
}
let tile2 = TileData::new(
0,
0,
5,
vec![1, 2, 9, 4, 5],
"application/x-protobuf".to_string(),
);
let delta2 = encoder.encode(&tile2);
assert!(delta2.is_ok());
if let Ok(data) = delta2 {
assert!(data.len() < 5);
}
}
}