use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use apache_avro::reader::datum::GenericDatumReader;
use apache_avro::writer::datum::GenericDatumWriter;
use apache_avro::{AvroSchema, Schema};
use bytes::BufMut;
use ruststream::BytesMut;
use ruststream::codec::{Codec, CodecError};
use serde::Serialize;
use serde::de::DeserializeOwned;
use crate::error::KafkaError;
use crate::schema_registry::{
MissingSubject, Registration, SchemaPrefetch, SchemaRegistry, SchemaType, WIRE_MAGIC,
parse_envelope, serde_name,
};
struct Prepared {
schema: &'static Schema,
writer: GenericDatumWriter<'static>,
reader: GenericDatumReader<'static>,
}
impl Prepared {
fn build(writer_schema: Schema, reader_schema: Option<Schema>) -> Result<Self, KafkaError> {
let schema: &'static Schema = Box::leak(Box::new(writer_schema));
let reader_schema: Option<&'static Schema> =
reader_schema.map(|reader| &*Box::leak(Box::new(reader)));
Ok(Self {
schema,
writer: GenericDatumWriter::builder(schema)
.build()
.map_err(KafkaError::wire_format)?,
reader: GenericDatumReader::builder(schema)
.maybe_reader_schema(reader_schema)
.build()
.map_err(KafkaError::wire_format)?,
})
}
}
#[derive(Clone)]
enum SchemaSource {
Local(&'static Prepared),
Registry {
registry: SchemaRegistry,
subjects: Arc<HashMap<&'static str, Registered>>,
prepared: Arc<Mutex<HashMap<u32, &'static Prepared>>>,
on_missing: MissingSubject,
},
}
struct Registered {
subject: String,
local: &'static Prepared,
}
impl Clone for Registered {
fn clone(&self) -> Self {
Self {
subject: self.subject.clone(),
local: self.local,
}
}
}
fn avro_record_name(schema: &Schema) -> Option<&'static str> {
match schema {
Schema::Record(record) => Some(String::leak(record.name.name().to_string())),
_ => None,
}
}
#[derive(Clone)]
pub struct AvroCodec {
source: SchemaSource,
reader_schema: Option<Arc<Schema>>,
prefetch: Option<SchemaPrefetch>,
}
impl std::fmt::Debug for AvroCodec {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let mut out = f.debug_struct("AvroCodec");
match &self.source {
SchemaSource::Local(_) => out.field("source", &"local"),
SchemaSource::Registry { subjects, .. } => out.field(
"subjects",
&subjects.values().map(|r| &r.subject).collect::<Vec<_>>(),
),
};
out.finish_non_exhaustive()
}
}
impl AvroCodec {
pub fn local(schema: Schema) -> Result<Self, KafkaError> {
Ok(Self {
source: SchemaSource::Local(Box::leak(Box::new(Prepared::build(schema, None)?))),
reader_schema: None,
prefetch: None,
})
}
#[must_use]
pub fn registry(prefetch: &SchemaPrefetch) -> Self {
Self {
source: SchemaSource::Registry {
registry: prefetch.registry().clone(),
subjects: Arc::new(HashMap::new()),
prepared: Arc::new(Mutex::new(HashMap::new())),
on_missing: prefetch.missing_subject(),
},
reader_schema: None,
prefetch: Some(prefetch.clone()),
}
}
#[must_use]
pub fn register<T: Serialize + AvroSchema>(mut self, subject: impl Into<String>) -> Self {
let subject = subject.into();
let schema = T::get_schema();
let name = avro_record_name(&schema).unwrap_or_else(|| {
panic!(
"the Avro schema of a registered type is not a named record, so nothing \
identifies its values on the publish path; register a record type under \
{subject:?}"
)
});
let definition = crate::avro::schema_json(&schema).unwrap_or_else(|err| {
panic!(
"the schema of the type registered under {subject:?} \
cannot be serialized: {err}"
)
});
let SchemaSource::Registry { subjects, .. } = &mut self.source else {
panic!("register is for a registry codec; a local one publishes a bare datum");
};
let table = Arc::get_mut(subjects).expect("not yet shared with a mount");
let local = Box::leak(Box::new(Prepared::build(schema, None).unwrap_or_else(
|err| panic!("the schema registered under {subject:?}: {err}"),
)));
if let Some(existing) = table.insert(
name,
Registered {
subject: subject.clone(),
local,
},
) && existing.subject != subject
{
panic!(
"two types registered on this codec are both called {name:?} to serde, so a \
publish could not tell {:?} from {subject:?}; give one of them a distinct \
`#[serde(rename = \"..\")]`",
existing.subject,
);
}
if let Some(prefetch) = &self.prefetch {
prefetch.record(Registration {
subject,
schema_type: SchemaType::Avro,
definition,
});
}
self
}
pub fn resolve_onto(mut self, schema: Schema) -> Result<Self, KafkaError> {
let schema = Arc::new(schema);
if let SchemaSource::Local(prepared) = &self.source {
self.source = SchemaSource::Local(Box::leak(Box::new(Prepared::build(
prepared.schema.clone(),
Some((*schema).clone()),
)?)));
}
self.reader_schema = Some(schema);
Ok(self)
}
fn prepared_for(&self, id: u32) -> Result<&'static Prepared, CodecError> {
let SchemaSource::Registry {
registry, prepared, ..
} = &self.source
else {
unreachable!("a local codec never looks a schema id up");
};
if let Some(found) = prepared
.lock()
.expect("prepared schema cache mutex poisoned")
.get(&id)
{
return Ok(found);
}
let schema = registry.cached_schema(id).ok_or_else(|| {
decode_error(format!(
"no schema is cached for id {id}: the delivery's writer schema was not resolved \
before it reached the codec. Attach a SchemaPrefetch to the broker \
(KafkaBroker::schema_prefetch), and check the registry is reachable - a \
synchronous codec cannot fetch it here"
))
})?;
let parsed = registry.parsed_avro(&schema).map_err(decode_source)?;
let built: &'static Prepared = Box::leak(Box::new(
Prepared::build((*parsed).clone(), self.reader_schema.as_deref().cloned())
.map_err(decode_source)?,
));
Ok(prepared
.lock()
.expect("prepared schema cache mutex poisoned")
.entry(id)
.or_insert(built))
}
}
impl Codec for AvroCodec {
fn encode<T: Serialize>(&self, value: &T) -> Result<BytesMut, CodecError> {
let mut buf = BytesMut::new();
let prepared = match &self.source {
SchemaSource::Local(prepared) => *prepared,
SchemaSource::Registry {
registry,
subjects,
on_missing,
..
} => {
let name = serde_name(value).ok_or_else(|| {
encode_error(
"this codec frames by the name serde gives a value's type, and this \
value has none - a bare number, a sequence or a map. Publish a named \
struct, or use a codec that needs no subject"
.to_owned(),
)
})?;
let registered = subjects.get(name).ok_or_else(|| {
encode_error(format!(
"no subject is registered for messages named {name:?}, so this codec has \
no id to frame them with. Add `.register::<{name}>(\"..\")` to it; the \
types it does carry are {:?}. A type that sets an Avro record name \
differing from its serde name registers under the Avro one and arrives \
here under the serde one, which also looks like this",
subjects.keys().collect::<Vec<_>>(),
))
})?;
match registry.cached_subject(®istered.subject) {
Some(schema) => {
buf.put_u8(WIRE_MAGIC);
buf.put_u32(schema.id());
self.prepared_for(schema.id())
.map_err(|err| encode_error(err.to_string()))?
}
None if *on_missing == MissingSubject::PublishUnframed => registered.local,
None => {
return Err(encode_error(format!(
"no schema is cached for subject {:?}, which messages named \
{name:?} publish under, so there is no id to frame them with. \
Attach the SchemaPrefetch this codec was built from to the broker \
(KafkaBroker::schema_prefetch) so connect resolves it; if it was \
resolved and has since gone, the subject was deleted from the \
registry and MissingSubject says what to do about that",
registered.subject,
)));
}
}
}
};
let value = apache_avro::to_value(value)
.map_err(encode_source)?
.resolve(prepared.schema)
.map_err(|err| {
encode_error(format!(
"the value does not fit the schema this codec writes: {err}"
))
})?;
prepared
.writer
.write_value(&mut (&mut buf).writer(), value)
.map_err(encode_source)?;
Ok(buf)
}
fn decode<T: DeserializeOwned>(&self, bytes: &[u8]) -> Result<T, CodecError> {
let (prepared, mut datum) = match &self.source {
SchemaSource::Local(prepared) => (*prepared, bytes),
SchemaSource::Registry { .. } => {
let (id, datum) = parse_envelope(bytes).ok_or_else(|| {
decode_error(format!(
"the delivery does not carry the Confluent wire format (a zero magic \
byte and a 4-byte schema id), so no writer schema can be named for it; \
its first bytes are {:02x?}",
&bytes[..bytes.len().min(8)],
))
})?;
(self.prepared_for(id)?, datum)
}
};
let value = prepared
.reader
.read_value(&mut datum)
.map_err(decode_source)?;
apache_avro::from_value(&value).map_err(decode_source)
}
}
fn encode_error(message: String) -> CodecError {
CodecError::Encode(Box::new(KafkaError::malformed(message)))
}
fn encode_source(err: impl std::error::Error + Send + Sync + 'static) -> CodecError {
CodecError::Encode(Box::new(err))
}
fn decode_error(message: String) -> CodecError {
CodecError::Decode(Box::new(KafkaError::malformed(message)))
}
fn decode_source(err: impl std::error::Error + Send + Sync + 'static) -> CodecError {
CodecError::Decode(Box::new(err))
}
#[cfg(test)]
mod tests {
use apache_avro::AvroSchema;
use serde::Deserialize;
use wiremock::matchers::{method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
use super::*;
use crate::schema_registry::SchemaType;
#[derive(Debug, PartialEq, Serialize, Deserialize, AvroSchema)]
#[serde(rename = "CodecOrder")]
struct Order {
id: i64,
item: String,
}
#[derive(Debug, PartialEq, Serialize, Deserialize, AvroSchema)]
#[serde(rename = "CodecOrder")]
struct OrderV2 {
id: i64,
item: String,
#[avro(default = r#""none""#)]
note: String,
}
fn order() -> Order {
Order {
id: 42,
item: "anvil".to_owned(),
}
}
#[test]
fn a_local_schema_round_trips_with_no_registry() {
let codec = AvroCodec::local(Order::get_schema()).expect("local");
let bytes = codec.encode(&order()).expect("encode");
assert_ne!(
&bytes[..],
br#"{"id":42,"item":"anvil"}"#,
"the wire is Avro"
);
assert_eq!(codec.decode::<Order>(&bytes).expect("decode"), order());
}
#[test]
fn a_local_codec_writes_no_envelope() {
let codec = AvroCodec::local(Order::get_schema()).expect("local");
let bytes = codec.encode(&order()).expect("encode");
assert!(
parse_envelope(&bytes).is_none() || bytes[0] != WIRE_MAGIC,
"a local codec puts a bare datum on the wire",
);
}
#[test]
fn a_value_the_schema_does_not_describe_is_refused() {
#[derive(Serialize)]
struct Other {
nope: String,
}
let codec = AvroCodec::local(Order::get_schema()).expect("local");
let err = codec
.encode(&Other {
nope: "x".to_owned(),
})
.expect_err("a value of another shape");
assert!(err.to_string().contains("does not fit the schema"));
}
#[test]
fn a_reader_schema_fills_what_the_writer_never_wrote() {
let writing = AvroCodec::local(Order::get_schema()).expect("writer");
let bytes = writing.encode(&order()).expect("encode");
let reading = AvroCodec::local(Order::get_schema()).expect("reader");
assert!(reading.decode::<OrderV2>(&bytes).is_err());
let resolving = AvroCodec::local(Order::get_schema())
.expect("reader")
.resolve_onto(OrderV2::get_schema())
.expect("resolve");
assert_eq!(
resolving.decode::<OrderV2>(&bytes).expect("resolve"),
OrderV2 {
id: 42,
item: "anvil".to_owned(),
note: "none".to_owned(),
},
);
}
async fn registry_with_order(id: u32) -> (MockServer, SchemaPrefetch) {
let server = MockServer::start().await;
let definition = crate::avro::schema_json(&Order::get_schema()).expect("json");
Mock::given(method("GET"))
.and(path("/subjects/orders-value/versions/latest"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"id": id,
"version": 1,
"schema": definition,
"schemaType": "AVRO",
})))
.mount(&server)
.await;
let prefetch = SchemaPrefetch::new(SchemaRegistry::new(server.uri()));
(server, prefetch)
}
#[tokio::test]
async fn a_registry_codec_frames_with_the_subjects_id() {
let (_server, prefetch) = registry_with_order(11).await;
let codec = AvroCodec::registry(&prefetch).register::<Order>("orders-value");
prefetch.warm_subjects().await.expect("warm");
let bytes = codec.encode(&order()).expect("encode");
let (id, datum) = parse_envelope(&bytes).expect("the wire format");
assert_eq!(id, 11);
assert_ne!(datum, b"");
assert_eq!(codec.decode::<Order>(&bytes).expect("decode"), order());
}
#[tokio::test]
async fn one_codec_frames_each_registered_type_under_its_own_subject() {
#[derive(Debug, PartialEq, Serialize, Deserialize, AvroSchema)]
#[serde(rename = "CodecShipment")]
struct Shipment {
id: i64,
}
let (server, prefetch) = registry_with_order(11).await;
Mock::given(method("GET"))
.and(path("/subjects/shipments-value/versions/latest"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"id": 12,
"version": 1,
"schema": crate::avro::schema_json(&Shipment::get_schema()).expect("json"),
"schemaType": "AVRO",
})))
.mount(&server)
.await;
let codec = AvroCodec::registry(&prefetch)
.register::<Order>("orders-value")
.register::<Shipment>("shipments-value");
prefetch.warm_subjects().await.expect("warm");
let (id, _) = parse_envelope(&codec.encode(&order()).expect("encode")).expect("framed");
assert_eq!(id, 11);
let shipped = codec.encode(&Shipment { id: 3 }).expect("encode");
let (id, _) = parse_envelope(&shipped).expect("framed");
assert_eq!(id, 12, "the second type framed under its own subject");
assert_eq!(
codec.decode::<Shipment>(&shipped).expect("decode"),
Shipment { id: 3 },
);
}
#[tokio::test]
async fn an_unregistered_type_names_itself_and_what_the_codec_carries() {
#[derive(Serialize)]
#[serde(rename = "Stranger")]
struct Stranger {
id: i64,
}
let (_server, prefetch) = registry_with_order(11).await;
let codec = AvroCodec::registry(&prefetch).register::<Order>("orders-value");
prefetch.warm_subjects().await.expect("warm");
let err = codec
.encode(&Stranger { id: 1 })
.expect_err("not registered");
assert!(err.to_string().contains("Stranger"));
assert!(err.to_string().contains("CodecOrder"), "{err}");
}
#[tokio::test]
async fn an_unnamed_value_is_refused_rather_than_guessed_at() {
let (_server, prefetch) = registry_with_order(11).await;
let codec = AvroCodec::registry(&prefetch).register::<Order>("orders-value");
let err = codec.encode(&7i64).expect_err("no name to key by");
assert!(err.to_string().contains("has none"));
}
#[tokio::test]
async fn an_unresolved_subject_names_the_attachment_it_is_missing() {
let (_server, prefetch) = registry_with_order(11).await;
let codec = AvroCodec::registry(&prefetch).register::<Order>("orders-value");
let err = codec.encode(&order()).expect_err("cold subject");
assert!(err.to_string().contains("schema_prefetch"));
}
#[tokio::test]
async fn an_unresolved_id_names_the_attachment_it_is_missing() {
let (_server, prefetch) = registry_with_order(11).await;
let codec = AvroCodec::registry(&prefetch).register::<Order>("orders-value");
let framed = [0u8, 0, 0, 0, 99, 1];
let err = codec.decode::<Order>(&framed).expect_err("cold id");
assert!(err.to_string().contains("no schema is cached for id 99"));
assert!(err.to_string().contains("SchemaPrefetch"));
}
#[tokio::test]
async fn an_unframed_delivery_on_a_registry_codec_is_refused() {
let (_server, prefetch) = registry_with_order(11).await;
let codec = AvroCodec::registry(&prefetch).register::<Order>("orders-value");
let err = codec
.decode::<Order>(br#"{"id":42}"#)
.expect_err("not framed");
assert!(err.to_string().contains("Confluent wire format"));
}
#[tokio::test]
async fn the_prefetch_resolves_a_deliverys_writer_schema() {
let (server, prefetch) = registry_with_order(11).await;
let definition = crate::avro::schema_json(&Order::get_schema()).expect("json");
Mock::given(method("GET"))
.and(path("/schemas/ids/11"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"schema": definition,
"schemaType": "AVRO",
})))
.expect(1)
.mount(&server)
.await;
let codec = AvroCodec::registry(&prefetch).register::<Order>("orders-value");
let framed = [0u8, 0, 0, 0, 11, 84, 10, 97, 110, 118, 105, 108];
assert!(codec.decode::<Order>(&framed).is_err());
prefetch.warm_delivery(&framed).await;
assert_eq!(codec.decode::<Order>(&framed).expect("decode"), order());
}
#[tokio::test]
async fn a_missing_subject_fails_startup_rather_than_the_first_publish() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/subjects/absent-value/versions/latest"))
.respond_with(ResponseTemplate::new(404))
.mount(&server)
.await;
let prefetch = SchemaPrefetch::new(SchemaRegistry::new(server.uri()));
let _codec = AvroCodec::registry(&prefetch).register::<Order>("absent-value");
let err = prefetch.warm_subjects().await.expect_err("absent subject");
assert!(err.to_string().contains("absent-value"));
}
#[tokio::test]
async fn an_unsigned_field_against_a_long_schema_is_reported() {
#[derive(Serialize)]
#[serde(rename = "CodecOrder")]
struct Wide {
id: u64,
item: String,
}
let (_server, prefetch) = registry_with_order(11).await;
let codec = AvroCodec::registry(&prefetch).register::<Order>("orders-value");
prefetch.warm_subjects().await.expect("warm");
let err = codec
.encode(&Wide {
id: 7,
item: "anvil".to_owned(),
})
.expect_err("u64 against a long schema");
assert!(err.to_string().contains("does not fit the schema"));
}
#[tokio::test]
async fn the_registered_flavor_is_checked_before_it_is_parsed() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/subjects/json-value/versions/latest"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"id": 3,
"version": 1,
"schema": "{\"type\":\"object\"}",
"schemaType": "JSON",
})))
.mount(&server)
.await;
let prefetch = SchemaPrefetch::new(SchemaRegistry::new(server.uri()));
let codec = AvroCodec::registry(&prefetch).register::<Order>("json-value");
prefetch.warm_subjects().await.expect("warm");
let err = codec.encode(&order()).expect_err("not an avro schema");
assert!(matches!(err, CodecError::Encode(_)));
let _ = SchemaType::Json;
}
}