diff --git a/README.md b/README.md index fe941aa..b9a9f54 100644 --- a/README.md +++ b/README.md @@ -52,6 +52,7 @@ On the wire to Kafka, `traits` and `metadata` are produced as JSON-encoded strin | Env var | Required | Default | |---|---|---| +| `SINK` | no | `kafka` (the only sink so far) | | `KAFKA_BOOTSTRAP_SERVERS` | yes | — | | `KAFKA_TOPIC` | yes | — | | `KAFKA_EXTERNAL_TOPIC` | yes | — | diff --git a/src/config.rs b/src/config.rs index d7582af..4e0e38f 100644 --- a/src/config.rs +++ b/src/config.rs @@ -5,28 +5,50 @@ use crate::sink::KafkaAuth; #[derive(Debug, Clone)] pub struct Config { pub listen_addr: String, - pub kafka_bootstrap_servers: String, - pub kafka_topic: String, - pub kafka_external_topic: String, - pub kafka_auth: KafkaAuth, + pub sink: SinkConfig, pub database_url: String, } +#[derive(Debug, Clone)] +pub enum SinkConfig { + Kafka(KafkaConfig), +} + +#[derive(Debug, Clone)] +pub struct KafkaConfig { + pub bootstrap_servers: String, + pub topic: String, + pub external_topic: String, + pub auth: KafkaAuth, +} + impl Config { pub fn from_env() -> Result { - let kafka_auth = parse_auth( + let sink = match parse_sink(env::var("SINK").ok().as_deref())? { + SinkKind::Kafka => SinkConfig::Kafka(KafkaConfig::from_env()?), + }; + + Ok(Config { + listen_addr: env::var("LISTEN_ADDR").unwrap_or_else(|_| "0.0.0.0:8080".to_string()), + sink, + database_url: required("DATABASE_URL")?, + }) + } +} + +impl KafkaConfig { + fn from_env() -> Result { + let auth = parse_auth( &env::var("KAFKA_AUTH").unwrap_or_else(|_| "scram".to_string()), env::var("KAFKA_USERNAME").ok(), env::var("KAFKA_PASSWORD").ok(), )?; - Ok(Config { - listen_addr: env::var("LISTEN_ADDR").unwrap_or_else(|_| "0.0.0.0:8080".to_string()), - kafka_bootstrap_servers: required("KAFKA_BOOTSTRAP_SERVERS")?, - kafka_topic: required("KAFKA_TOPIC")?, - kafka_external_topic: required("KAFKA_EXTERNAL_TOPIC")?, - kafka_auth, - database_url: required("DATABASE_URL")?, + Ok(KafkaConfig { + bootstrap_servers: required("KAFKA_BOOTSTRAP_SERVERS")?, + topic: required("KAFKA_TOPIC")?, + external_topic: required("KAFKA_EXTERNAL_TOPIC")?, + auth, }) } } @@ -35,6 +57,19 @@ fn required(name: &'static str) -> Result { env::var(name).map_err(|_| ConfigError::MissingRequired(name)) } +#[derive(Debug, PartialEq)] +enum SinkKind { + Kafka, +} + +/// Which sink `SINK` selects; unset means Kafka. +fn parse_sink(value: Option<&str>) -> Result { + match value.unwrap_or("kafka") { + "kafka" => Ok(SinkKind::Kafka), + other => Err(ConfigError::InvalidSink(other.to_string())), + } +} + fn parse_auth( mode: &str, username: Option, @@ -56,6 +91,8 @@ pub enum ConfigError { MissingRequired(&'static str), #[error("Invalid value for {0}: {1:?} (expected scram or none)")] Invalid(&'static str, String), + #[error("Invalid value for SINK: {0:?} (expected kafka)")] + InvalidSink(String), } #[cfg(test)] @@ -98,4 +135,18 @@ mod tests { Err(ConfigError::Invalid("KAFKA_AUTH", v)) if v == "iam" )); } + + #[test] + fn parse_sink_defaults_to_kafka() { + assert_eq!(parse_sink(None).unwrap(), SinkKind::Kafka); + assert_eq!(parse_sink(Some("kafka")).unwrap(), SinkKind::Kafka); + } + + #[test] + fn parse_sink_rejects_unknown_value() { + assert!(matches!( + parse_sink(Some("clickhouse")), + Err(ConfigError::InvalidSink(v)) if v == "clickhouse" + )); + } } diff --git a/src/main.rs b/src/main.rs index e9b507f..d5185e6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -16,9 +16,9 @@ use tokio::signal; use tracing::info; use auth::ContextState; -use config::Config; +use config::{Config, SinkConfig}; use environment_keys::EnvironmentKeys; -use sink::{KafkaSink, KafkaSinkConfig}; +use sink::{KafkaSink, KafkaSinkConfig, Sink}; const ENVIRONMENT_KEY_REFRESH_AFTER: Duration = Duration::from_secs(300); @@ -33,19 +33,23 @@ async fn main() -> anyhow::Result<()> { .init(); let config = Config::from_env()?; - info!( - listen_addr = %config.listen_addr, - kafka_bootstrap_servers = %config.kafka_bootstrap_servers, - kafka_topic = %config.kafka_topic, - kafka_auth = ?config.kafka_auth, - "Starting flagsmith-analytics server" - ); - - let sink = KafkaSink::new(KafkaSinkConfig::new( - config.kafka_bootstrap_servers.clone(), - config.kafka_topic.clone(), - config.kafka_auth.clone(), - ))?; + let (sink, external_topic) = match &config.sink { + SinkConfig::Kafka(kafka) => { + info!( + listen_addr = %config.listen_addr, + kafka_bootstrap_servers = %kafka.bootstrap_servers, + kafka_topic = %kafka.topic, + kafka_auth = ?kafka.auth, + "Starting flagsmith-analytics server" + ); + let sink = KafkaSink::new(KafkaSinkConfig::new( + kafka.bootstrap_servers.clone(), + kafka.topic.clone(), + kafka.auth.clone(), + ))?; + (Sink::Kafka(sink), kafka.external_topic.as_str()) + } + }; let pool = PgPoolOptions::new() .max_connections(8) @@ -53,7 +57,7 @@ async fn main() -> anyhow::Result<()> { .connect_lazy(&config.database_url)?; let context_state = ContextState { environment_keys: EnvironmentKeys::new(pool, ENVIRONMENT_KEY_REFRESH_AFTER), - external_topic: Arc::from(config.kafka_external_topic.as_str()), + external_topic: Arc::from(external_topic), }; let app = routes::app(sink.clone(), context_state); diff --git a/src/models/event.rs b/src/models/event.rs index d9d2d6b..2f44b77 100644 --- a/src/models/event.rs +++ b/src/models/event.rs @@ -1,4 +1,4 @@ -use serde::{Deserialize, Serialize}; +use serde::{Deserialize, Serialize, Serializer}; /// Reserved system event for flag-exposure tracking. `feature_name` is required /// for events of this type. @@ -73,9 +73,26 @@ pub struct EventRecord<'a> { pub identifier: Option<&'a str>, #[serde(skip_serializing_if = "serde_json::Value::is_null")] pub value: &'a serde_json::Value, - #[serde(skip_serializing_if = "Option::is_none")] - pub traits: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub metadata: Option, + #[serde( + skip_serializing_if = "Option::is_none", + serialize_with = "json_string" + )] + pub traits: Option<&'a serde_json::Value>, + #[serde( + skip_serializing_if = "Option::is_none", + serialize_with = "json_string" + )] + pub metadata: Option<&'a serde_json::Value>, pub timestamp: i64, } + +/// Serializes a JSON value as its encoded text, in a string. +fn json_string( + value: &Option<&serde_json::Value>, + serializer: S, +) -> Result { + match value { + Some(value) => serializer.collect_str(value), + None => serializer.serialize_none(), + } +} diff --git a/src/routes/health.rs b/src/routes/health.rs index 942b003..d696d36 100644 --- a/src/routes/health.rs +++ b/src/routes/health.rs @@ -1,7 +1,7 @@ use axum::{extract::State, http::StatusCode, Json}; use serde::Serialize; -use crate::sink::KafkaSink; +use crate::sink::Sink; #[derive(Serialize)] pub struct HealthResponse { @@ -12,7 +12,7 @@ pub struct HealthResponse { /// Liveness plus one local check: a fatal producer error is permanent, so /// report it and let ECS replace the task. No downstream call, so this cannot hang. -pub async fn health_check(State(sink): State) -> (StatusCode, Json) { +pub async fn health_check(State(sink): State) -> (StatusCode, Json) { match sink.fatal_error() { None => ( StatusCode::OK, diff --git a/src/routes/ingest.rs b/src/routes/ingest.rs index 1cea1ab..4a039e0 100644 --- a/src/routes/ingest.rs +++ b/src/routes/ingest.rs @@ -12,7 +12,7 @@ use crate::lookup::mask_key; use crate::models::{ Event, EventBatch, EventRecord, IngestResponse, RejectedEvent, FLAG_EXPOSURE_EVENT, }; -use crate::sink::KafkaSink; +use crate::sink::Sink; /// Header-derived fields are copied into every record of the batch, so an /// unbounded header would inflate every record. @@ -177,15 +177,15 @@ fn partition_events(events: &[Event]) -> (Vec<&Event>, Vec) { } /// POST /v1/events -/// Validates each event independently: valid events are serialized and -/// forwarded to the sink, invalid ones are reported back as `rejected` -/// entries without blocking the rest of the batch. +/// Validates each event independently: valid events are forwarded to the +/// sink, invalid ones are reported back as `rejected` entries without +/// blocking the rest of the batch. /// Returns 202 when at least one event was accepted and forwarded, 400 when /// none were (all invalid, or empty batch), 503 on sink failure /// (delivery not confirmed; retrying is safe but may duplicate events from /// a partially delivered batch). pub async fn ingest_batch( - State(sink): State, + State(sink): State, Extension(ctx): Extension, headers: HeaderMap, Json(batch): Json, @@ -225,28 +225,26 @@ pub async fn ingest_batch( .map(parse_sdk_user_agent) .unwrap_or_else(|| ("unknown".to_string(), "unknown".to_string())); - let records: Vec> = valid + // TODO: move this to a method. + let records: Vec = valid .iter() - .map(|event| { - let record = EventRecord { - environment_key: &ctx.client_api_key, - collected_at, - sdk_language: &sdk_language, - sdk_version: &sdk_version, - event: &event.event, - feature_name: event.feature_name.as_deref(), - identifier: event.identifier.as_deref(), - value: &event.value, - traits: event.traits.as_ref().map(ToString::to_string), - metadata: event.metadata.as_ref().map(ToString::to_string), - timestamp: event.timestamp, - }; - serde_json::to_vec(&record).expect("EventRecord serialization cannot fail") + .map(|event| EventRecord { + environment_key: &ctx.client_api_key, + collected_at, + sdk_language: &sdk_language, + sdk_version: &sdk_version, + event: &event.event, + feature_name: event.feature_name.as_deref(), + identifier: event.identifier.as_deref(), + value: &event.value, + traits: event.traits.as_ref(), + metadata: event.metadata.as_ref(), + timestamp: event.timestamp, }) .collect(); let record_count = records.len(); - match sink.put(ctx.destination.as_deref(), records).await { + match sink.put(ctx.destination.as_deref(), &records).await { Ok(()) => { debug!( client_api_key = %mask_key(&ctx.client_api_key), @@ -288,7 +286,7 @@ mod tests { use rdkafka::types::RDKafkaRespErr; use tower::ServiceExt; - use crate::sink::{KafkaAuth, KafkaSinkConfig}; + use crate::sink::{KafkaAuth, KafkaSink, KafkaSinkConfig}; use crate::test_helpers::{ assert_topic_empty, mock_kafka, read_records, FLAGSMITH_WAREHOUSE_TOPIC, }; @@ -514,8 +512,8 @@ mod tests { feature_name: event.feature_name.as_deref(), identifier: event.identifier.as_deref(), value: &event.value, - traits: event.traits.as_ref().map(ToString::to_string), - metadata: event.metadata.as_ref().map(ToString::to_string), + traits: event.traits.as_ref(), + metadata: event.metadata.as_ref(), timestamp: event.timestamp, }; @@ -534,17 +532,17 @@ mod tests { } /// A sink whose bootstrap server never answers, so every put times out. - fn unreachable_sink() -> KafkaSink { + fn unreachable_sink() -> Sink { let mut config = KafkaSinkConfig::new( "127.0.0.1:1".to_string(), FLAGSMITH_WAREHOUSE_TOPIC.to_string(), KafkaAuth::None, ); config.message_timeout = Duration::from_secs(2); - KafkaSink::new(config).expect("sink") + Sink::Kafka(KafkaSink::new(config).expect("sink")) } - fn app_with_destination(sink: &KafkaSink, destination: Option<&str>) -> Router { + fn app_with_destination(sink: &Sink, destination: Option<&str>) -> Router { let ctx = EnvironmentContext { client_api_key: Arc::from("test-client-key"), destination: destination.map(Arc::from), @@ -555,7 +553,7 @@ mod tests { .with_state(sink.clone()) } - fn app(sink: &KafkaSink) -> Router { + fn app(sink: &Sink) -> Router { app_with_destination(sink, None) } diff --git a/src/routes/mod.rs b/src/routes/mod.rs index fd2b277..7e12d2d 100644 --- a/src/routes/mod.rs +++ b/src/routes/mod.rs @@ -6,12 +6,12 @@ use tower_http::cors::{Any, CorsLayer}; use tower_http::decompression::RequestDecompressionLayer; use crate::auth::{extract_environment_context, ContextState}; -use crate::sink::KafkaSink; +use crate::sink::Sink; pub use health::health_check; pub use ingest::ingest_batch; -pub fn app(sink: KafkaSink, context_state: ContextState) -> Router { +pub fn app(sink: Sink, context_state: ContextState) -> Router { Router::new() .route("/health", get(health_check)) .route( diff --git a/src/sink/kafka.rs b/src/sink/kafka.rs index 33ba54b..d41fb13 100644 --- a/src/sink/kafka.rs +++ b/src/sink/kafka.rs @@ -10,6 +10,8 @@ use rdkafka::statistics::Statistics; use rdkafka::util::Timeout; use tracing::{debug, error, info, warn}; +use crate::models::EventRecord; + const CLIENT_ID: &str = "events-ingestion-api"; /// How long `put` waits for room in the local queue before failing the batch. const ENQUEUE_TIMEOUT: Duration = Duration::from_millis(500); @@ -146,11 +148,17 @@ impl KafkaSink { pub async fn put( &self, destination: Option<&str>, - records: Vec>, - ) -> Result<(), SinkError> { + records: &[EventRecord<'_>], + ) -> Result<(), DeliveryError> { let topic = destination.unwrap_or(&self.default_topic); let total = records.len(); - let sends = records.iter().map(|bytes| { + let payloads: Vec> = records + .iter() + .map(|record| { + serde_json::to_vec(record).expect("EventRecord serialization cannot fail") + }) + .collect(); + let sends = payloads.iter().map(|bytes| { self.producer.send( FutureRecord::<(), _>::to(topic).payload(bytes), Timeout::After(ENQUEUE_TIMEOUT), @@ -159,7 +167,7 @@ impl KafkaSink { let mut failures = join_all(sends).await.into_iter().filter_map(Result::err); match failures.next() { None => Ok(()), - Some((error, _record)) => Err(SinkError::Delivery { + Some((error, _record)) => Err(DeliveryError { failed: 1 + failures.count(), total, error, @@ -190,13 +198,11 @@ impl KafkaSink { } #[derive(Debug, thiserror::Error)] -pub enum SinkError { - #[error("Kafka delivery failed for {failed} of {total} records: {error}")] - Delivery { - failed: usize, - total: usize, - error: KafkaError, - }, +#[error("Kafka delivery failed for {failed} of {total} records: {error}")] +pub struct DeliveryError { + failed: usize, + total: usize, + error: KafkaError, } #[cfg(test)] diff --git a/src/sink/mod.rs b/src/sink/mod.rs index 2a73a9b..4369129 100644 --- a/src/sink/mod.rs +++ b/src/sink/mod.rs @@ -1,3 +1,47 @@ mod kafka; +use std::time::Duration; + pub use kafka::{KafkaAuth, KafkaSink, KafkaSinkConfig}; + +use crate::models::EventRecord; + +/// Where accepted events are written. +#[derive(Clone)] +pub enum Sink { + Kafka(KafkaSink), +} + +impl Sink { + /// Write every record to the destination (the sink's default when `None`). + /// Returns once the sink has accepted all of them. + pub async fn put( + &self, + destination: Option<&str>, + records: &[EventRecord<'_>], + ) -> Result<(), SinkError> { + match self { + Self::Kafka(sink) => Ok(sink.put(destination, records).await?), + } + } + + /// A permanent sink failure that should fail the health check. + pub fn fatal_error(&self) -> Option { + match self { + Self::Kafka(sink) => sink.fatal_error(), + } + } + + /// Drain anything still buffered, bounded by `timeout`. + pub async fn shutdown(self, timeout: Duration) { + match self { + Self::Kafka(sink) => sink.shutdown(timeout).await, + } + } +} + +#[derive(Debug, thiserror::Error)] +pub enum SinkError { + #[error(transparent)] + Kafka(#[from] kafka::DeliveryError), +} diff --git a/src/test_helpers.rs b/src/test_helpers.rs index 35c85b3..9545c21 100644 --- a/src/test_helpers.rs +++ b/src/test_helpers.rs @@ -8,7 +8,7 @@ use rdkafka::producer::DefaultProducerContext; use rdkafka::Message; use sqlx::PgPool; -use crate::sink::{KafkaAuth, KafkaSink, KafkaSinkConfig}; +use crate::sink::{KafkaAuth, KafkaSink, KafkaSinkConfig, Sink}; pub(crate) const FLAGSMITH_WAREHOUSE_TOPIC: &str = "events-default"; @@ -18,7 +18,7 @@ pub(crate) type MockKafkaCluster = MockCluster<'static, DefaultProducerContext>; /// test. pub(crate) struct MockKafka { pub(crate) cluster: MockKafkaCluster, - pub(crate) sink: KafkaSink, + pub(crate) sink: Sink, } pub(crate) fn mock_kafka(topics: &[&str]) -> MockKafka { @@ -32,7 +32,7 @@ pub(crate) fn mock_kafka(topics: &[&str]) -> MockKafka { KafkaAuth::None, ); config.message_timeout = Duration::from_secs(5); - let sink = KafkaSink::new(config).expect("sink"); + let sink = Sink::Kafka(KafkaSink::new(config).expect("sink")); MockKafka { cluster, sink } }