use async_trait::async_trait;
use reqwest::Client;
use std::time::Duration;
use tokio::sync::Mutex;
use tokio::time::Instant;
use url::Url;
use uuid::Uuid;
use crate::{
transport::{Transport, TransportError},
EventData,
};
#[derive(Debug, Clone)]
pub struct BatchConfig {
pub max_batch_size: usize,
pub max_batch_age: Duration,
}
impl Default for BatchConfig {
fn default() -> Self {
Self {
max_batch_size: 100,
max_batch_age: Duration::from_secs(5),
}
}
}
impl BatchConfig {
pub fn new(max_batch_size: usize, max_batch_age: Duration) -> Self {
Self {
max_batch_size,
max_batch_age,
}
}
}
pub struct BatchingHttpTransport {
client: Client,
batch_url: Url,
config: BatchConfig,
auth_token: Option<String>,
buffer: Mutex<BatchBuffer>,
}
struct BatchBuffer {
events: Vec<EventData>,
oldest_event_time: Option<Instant>,
}
impl BatchBuffer {
fn new() -> Self {
Self {
events: Vec::new(),
oldest_event_time: None,
}
}
fn push(&mut self, event: EventData) {
if self.events.is_empty() {
self.oldest_event_time = Some(Instant::now());
}
self.events.push(event);
}
fn take(&mut self) -> Vec<EventData> {
self.oldest_event_time = None;
std::mem::take(&mut self.events)
}
fn is_empty(&self) -> bool {
self.events.is_empty()
}
fn len(&self) -> usize {
self.events.len()
}
fn age(&self) -> Option<Duration> {
self.oldest_event_time.map(|t| t.elapsed())
}
}
impl BatchingHttpTransport {
pub fn new(
base_url: Url,
org_id: Uuid,
app_id: Uuid,
config: BatchConfig,
auth_token: Option<String>,
) -> Result<Self, TransportError> {
let batch_url = base_url
.join(&format!(
"/api/orgs/{}/apps/{}/events/batch",
org_id, app_id
))
.map_err(|e| TransportError::Configuration(format!("Invalid URL: {}", e)))?;
Ok(Self {
client: Client::new(),
batch_url,
config,
auth_token,
buffer: Mutex::new(BatchBuffer::new()),
})
}
pub fn with_default_config(
base_url: Url,
org_id: Uuid,
app_id: Uuid,
auth_token: Option<String>,
) -> Result<Self, TransportError> {
Self::new(base_url, org_id, app_id, BatchConfig::default(), auth_token)
}
async fn flush(&self) -> Result<(), TransportError> {
let events = {
let mut buffer = self.buffer.lock().await;
if buffer.is_empty() {
return Ok(());
}
buffer.take()
};
self.send_batch(events).await
}
async fn send_batch(&self, events: Vec<EventData>) -> Result<(), TransportError> {
if events.is_empty() {
return Ok(());
}
let mut request = self.client.post(self.batch_url.clone()).json(&events);
if let Some(token) = &self.auth_token {
request = request.bearer_auth(token);
}
let response = request
.send()
.await
.map_err(|e| TransportError::Send(format!("HTTP batch request failed: {}", e)))?;
if !response.status().is_success() {
return Err(TransportError::Send(format!(
"Server returned error status for batch: {}",
response.status()
)));
}
Ok(())
}
async fn should_flush(&self) -> bool {
let buffer = self.buffer.lock().await;
if buffer.len() >= self.config.max_batch_size {
return true;
}
if let Some(age) = buffer.age() {
if age >= self.config.max_batch_age {
return true;
}
}
false
}
}
#[async_trait]
impl Transport for BatchingHttpTransport {
async fn connect(&mut self) -> Result<(), TransportError> {
Ok(())
}
async fn send(&mut self, event: EventData) -> Result<(), TransportError> {
{
let mut buffer = self.buffer.lock().await;
buffer.push(event);
}
if self.should_flush().await {
self.flush().await?;
}
Ok(())
}
async fn close(&mut self) -> Result<(), TransportError> {
self.flush().await
}
}
#[cfg(test)]
mod tests {
use super::*;
use chrono::Utc;
#[test]
fn test_batch_config_default() {
let config = BatchConfig::default();
assert_eq!(config.max_batch_size, 100);
assert_eq!(config.max_batch_age, Duration::from_secs(5));
}
#[test]
fn test_batch_config_custom() {
let config = BatchConfig::new(50, Duration::from_millis(500));
assert_eq!(config.max_batch_size, 50);
assert_eq!(config.max_batch_age, Duration::from_millis(500));
}
#[test]
fn test_batch_buffer_operations() {
let mut buffer = BatchBuffer::new();
assert!(buffer.is_empty());
assert_eq!(buffer.len(), 0);
assert!(buffer.age().is_none());
let event = EventData {
event_type: "test".to_string(),
event_data: serde_json::json!({}),
event_timestamp: Utc::now(),
process_instance_id: None,
};
buffer.push(event);
assert!(!buffer.is_empty());
assert_eq!(buffer.len(), 1);
assert!(buffer.age().is_some());
let events = buffer.take();
assert_eq!(events.len(), 1);
assert!(buffer.is_empty());
assert!(buffer.age().is_none());
}
#[test]
fn test_batching_transport_creation() {
let base_url = Url::parse("http://localhost:4318").unwrap();
let org_id = Uuid::new_v4();
let app_id = Uuid::new_v4();
let transport =
BatchingHttpTransport::with_default_config(base_url.clone(), org_id, app_id, None);
assert!(transport.is_ok());
let custom_config = BatchConfig::new(50, Duration::from_millis(100));
let transport = BatchingHttpTransport::new(base_url, org_id, app_id, custom_config, None);
assert!(transport.is_ok());
}
#[tokio::test]
async fn test_batching_transport_url() {
let base_url = Url::parse("http://localhost:4318").unwrap();
let org_id = Uuid::parse_str("12345678-1234-1234-1234-123456789012").unwrap();
let app_id = Uuid::parse_str("87654321-4321-4321-4321-210987654321").unwrap();
let transport = BatchingHttpTransport::with_default_config(base_url, org_id, app_id, None)
.expect("should create transport");
assert_eq!(
transport.batch_url.as_str(),
"http://localhost:4318/api/orgs/12345678-1234-1234-1234-123456789012/apps/87654321-4321-4321-4321-210987654321/events/batch"
);
}
#[tokio::test]
async fn test_should_flush_by_size() {
let base_url = Url::parse("http://localhost:4318").unwrap();
let org_id = Uuid::new_v4();
let app_id = Uuid::new_v4();
let config = BatchConfig::new(2, Duration::from_secs(60));
let transport = BatchingHttpTransport::new(base_url, org_id, app_id, config, None)
.expect("should create transport");
{
let mut buffer = transport.buffer.lock().await;
buffer.push(EventData {
event_type: "test".to_string(),
event_data: serde_json::json!({}),
event_timestamp: Utc::now(),
process_instance_id: None,
});
}
assert!(!transport.should_flush().await);
{
let mut buffer = transport.buffer.lock().await;
buffer.push(EventData {
event_type: "test".to_string(),
event_data: serde_json::json!({}),
event_timestamp: Utc::now(),
process_instance_id: None,
});
}
assert!(transport.should_flush().await);
}
#[tokio::test]
async fn test_should_flush_by_age() {
let base_url = Url::parse("http://localhost:4318").unwrap();
let org_id = Uuid::new_v4();
let app_id = Uuid::new_v4();
let config = BatchConfig::new(1000, Duration::from_millis(10));
let transport = BatchingHttpTransport::new(base_url, org_id, app_id, config, None)
.expect("should create transport");
{
let mut buffer = transport.buffer.lock().await;
buffer.push(EventData {
event_type: "test".to_string(),
event_data: serde_json::json!({}),
event_timestamp: Utc::now(),
process_instance_id: None,
});
}
assert!(!transport.should_flush().await);
tokio::time::sleep(Duration::from_millis(15)).await;
assert!(transport.should_flush().await);
}
}