use crate::cache::{CompiledSchema, Lookup, SchemaCache};
use schema_registry_converter::async_impl::schema_registry::{self, SrSettings, SrSettingsBuilder};
use schema_registry_converter::error::SRCError;
use schema_registry_converter::schema_registry_common::{SchemaType, SubjectNameStrategy};
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::mpsc;
use tokio::task::JoinSet;
const FETCH_BACKOFF_INITIAL: Duration = Duration::from_millis(200);
const FETCH_BACKOFF_MAX: Duration = Duration::from_secs(5);
const MAX_CONCURRENT_FETCHES: usize = 4;
#[derive(Clone, Debug)]
pub(crate) struct RegistryHandle {
tx: mpsc::UnboundedSender<u32>,
pub(crate) cache: Arc<SchemaCache>,
}
impl RegistryHandle {
pub(crate) fn request(&self, id: u32) {
let _ = self.tx.send(id);
}
}
#[derive(Clone, Debug)]
pub(crate) struct RegistryConfig {
pub url: String,
pub basic_auth: Option<(String, Option<String>)>,
pub negative_cache_ttl: Duration,
}
fn sr_settings(cfg: &RegistryConfig) -> SrSettings {
let mut builder: SrSettingsBuilder = SrSettings::new_builder(cfg.url.clone());
if let Some((user, pass)) = &cfg.basic_auth {
builder.set_basic_authorization(user, pass.as_deref());
}
builder.build().expect("registry settings")
}
enum FetchOutcome {
Resolved,
Transient,
}
struct Backoff {
delay: Duration,
next_allowed: Instant,
}
pub(crate) fn spawn_fetcher(
cfg: RegistryConfig,
runtime: &tokio::runtime::Handle,
) -> RegistryHandle {
let cache = Arc::new(SchemaCache::new(cfg.negative_cache_ttl));
let (tx, mut rx) = mpsc::unbounded_channel::<u32>();
let task_cache = Arc::clone(&cache);
let settings = Arc::new(sr_settings(&cfg));
runtime.spawn(async move {
let mut in_flight: HashSet<u32> = HashSet::new();
let mut backoff: HashMap<u32, Backoff> = HashMap::new();
let mut tasks: JoinSet<(u32, FetchOutcome)> = JoinSet::new();
loop {
tokio::select! {
biased;
Some(joined) = tasks.join_next(), if !tasks.is_empty() => {
let Ok((id, outcome)) = joined else {
tracing::error!("registry fetch task panicked");
continue;
};
in_flight.remove(&id);
match outcome {
FetchOutcome::Resolved => {
backoff.remove(&id);
}
FetchOutcome::Transient => match backoff.get_mut(&id) {
Some(b) => {
b.delay = (b.delay * 2).min(FETCH_BACKOFF_MAX);
b.next_allowed = Instant::now() + b.delay;
}
None => {
backoff.insert(
id,
Backoff {
delay: FETCH_BACKOFF_INITIAL,
next_allowed: Instant::now() + FETCH_BACKOFF_INITIAL,
},
);
}
},
}
}
maybe_id = rx.recv(), if tasks.len() < MAX_CONCURRENT_FETCHES => {
let Some(id) = maybe_id else {
break;
};
if in_flight.contains(&id) {
continue;
}
if !matches!(task_cache.get(id), Lookup::Missing) {
continue;
}
if backoff.get(&id).is_some_and(|b| Instant::now() < b.next_allowed) {
continue;
}
in_flight.insert(id);
let cache = Arc::clone(&task_cache);
let settings = Arc::clone(&settings);
tasks.spawn(async move {
let outcome = fetch_one(id, &settings, &cache).await;
(id, outcome)
});
}
}
}
});
RegistryHandle { tx, cache }
}
async fn fetch_one(id: u32, settings: &SrSettings, cache: &SchemaCache) -> FetchOutcome {
match schema_registry::get_schema_by_id_and_type(id, settings, SchemaType::Avro).await {
Ok(registered) => {
if !registered.references.is_empty() {
cache.insert_failed(
id,
format!(
"schema {id} uses {} registry reference(s), which spate-avro \
does not support yet",
registered.references.len()
),
);
return FetchOutcome::Resolved;
}
let compiled = CompiledSchema::compile(id, ®istered.schema);
match compiled.unusable_reason() {
None => {
tracing::info!(schema_id = id, "schema fetched and compiled");
cache.insert_ready(compiled);
}
Some(reason) => {
cache.insert_failed(id, reason);
}
}
FetchOutcome::Resolved
}
Err(e) if is_permanent(&e) => {
tracing::warn!(schema_id = id, error = %e, "registry reports schema id unknown");
cache.insert_failed(id, format!("registry fetch for schema {id} failed: {e}"));
FetchOutcome::Resolved
}
Err(e) => {
tracing::warn!(schema_id = id, error = %e, "registry fetch failed transiently; will retry");
FetchOutcome::Transient
}
}
}
fn is_permanent(e: &SRCError) -> bool {
e.error.contains("status 404")
}
pub(crate) async fn prewarm(cfg: &RegistryConfig, subjects: &[String], cache: &SchemaCache) {
let settings = sr_settings(cfg);
for subject in subjects {
let strategy = SubjectNameStrategy::RecordNameStrategy(subject.clone());
match schema_registry::get_schema_by_subject(&settings, &strategy).await {
Ok(registered) if registered.references.is_empty() => {
let compiled = CompiledSchema::compile(registered.id, ®istered.schema);
match compiled.unusable_reason() {
None => {
tracing::info!(subject, schema_id = registered.id, "pre-warmed schema");
cache.insert_ready(compiled);
}
Some(reason) => {
tracing::warn!(subject, %reason, "pre-warm parse failed; skipping subject");
}
}
}
Ok(_) => {
tracing::warn!(subject, "pre-warm skipped: schema references unsupported");
}
Err(e) => tracing::warn!(subject, error = %e, "pre-warm fetch failed"),
}
}
}