mod unified;
use proc_macro::TokenStream;
use proc_macro2::TokenStream as TokenStream2;
use quote::{ToTokens, quote};
use syn::{Error, Expr, FnArg, Ident, ItemFn, Pat, PatType, ReturnType, Type, TypePath};
use crate::parse::{
FailurePolicyArg, SubscriberArgs, WorkersArg, doc_description, position_type, source_tokens,
vec_element,
};
pub(crate) fn subscriber(args: &SubscriberArgs, func: &ItemFn) -> syn::Result<TokenStream> {
let parts = handler_parts(args, func)?;
Ok(unified::expand(args, &parts, func)?.into())
}
struct HandlerParts<'a> {
vis: &'a syn::Visibility,
name: &'a Ident,
block: &'a syn::Block,
pat: &'a Pat,
input_ty: &'a Type,
pair: Option<(&'a Type, &'a Type)>,
description: TokenStream2,
source_expr: TokenStream2,
ctx_param: TokenStream2,
ctx_ty: TokenStream2,
state_ty: Option<TokenStream2>,
extractors: Vec<(&'a Pat, &'a Type)>,
outs: Vec<OutParam<'a>>,
shape: Shape,
settings_chain: TokenStream2,
settings_source_ty: TokenStream2,
settings_state_ty: TokenStream2,
}
fn context_type(func: &ItemFn) -> TokenStream2 {
let Some(FnArg::Typed(PatType { ty, .. })) = func.sig.inputs.get(1) else {
return quote!(());
};
if !is_context_param(ty) {
return inferred_context_type(func);
}
if let Type::Reference(reference) = &**ty
&& let Type::Path(path) = &*reference.elem
&& let Some(segment) = path.path.segments.last()
&& let syn::PathArguments::AngleBracketed(args) = &segment.arguments
{
for arg in &args.args {
if let syn::GenericArgument::Type(context_ty) = arg {
return quote!(#context_ty);
}
}
}
quote!(())
}
fn inferred_context_type(func: &ItemFn) -> TokenStream2 {
for arg in func.sig.inputs.iter().skip(1) {
if let FnArg::Typed(PatType { ty, .. }) = arg
&& let Some(key) = ctx_extractor_key(ty)
{
return quote!(<#key as ::ruststream::ContextField>::Context);
}
}
quote!(())
}
struct OutParam<'a> {
pat: &'a Pat,
ty: &'a Type,
bounds: &'a syn::punctuated::Punctuated<syn::TypeParamBound, syn::Token![+]>,
marker: TokenStream2,
bodies: Option<BodyDecl<'a>>,
}
enum BodyDecl<'a> {
List(Vec<&'a Type>),
Set(&'a Type),
}
fn out_param_args(ty: &Type) -> Option<Vec<&Type>> {
let Type::Path(path) = ty else {
return None;
};
let segment = path.path.segments.last()?;
if segment.ident != "Out" {
return None;
}
let syn::PathArguments::AngleBracketed(args) = &segment.arguments else {
return None;
};
let types: Vec<&Type> = args
.args
.iter()
.filter_map(|arg| match arg {
syn::GenericArgument::Type(ty) => Some(ty),
_ => None,
})
.collect();
if types.is_empty() || types.len() > 3 {
return None;
}
Some(types)
}
fn from_headers_ty(ty: &Type) -> Option<&Type> {
let Type::Path(path) = ty else {
return None;
};
let segment = path.path.segments.last()?;
if segment.ident != "Headers" {
return None;
}
let syn::PathArguments::AngleBracketed(args) = &segment.arguments else {
return None;
};
let mut types = args.args.iter().filter_map(|arg| match arg {
syn::GenericArgument::Type(ty) => Some(ty),
_ => None,
});
let contract = types.next()?;
if types.next().is_some() {
return None;
}
Some(contract)
}
fn message_pair_args(ty: &Type) -> Option<(&Type, &Type)> {
let Type::Path(path) = ty else {
return None;
};
let segment = path.path.segments.last()?;
if segment.ident != "Message" {
return None;
}
let syn::PathArguments::AngleBracketed(args) = &segment.arguments else {
return None;
};
let mut types = args.args.iter().filter_map(|arg| match arg {
syn::GenericArgument::Type(ty) => Some(ty),
_ => None,
});
let headers = types.next()?;
let payload = types.next()?;
if types.next().is_some() {
return None;
}
Some((headers, payload))
}
fn ctx_extractor_key(ty: &Type) -> Option<&Type> {
let Type::Path(path) = ty else {
return None;
};
let segment = path.path.segments.last()?;
if segment.ident != "Ctx" {
return None;
}
let syn::PathArguments::AngleBracketed(args) = &segment.arguments else {
return None;
};
let mut types = args.args.iter().filter_map(|arg| match arg {
syn::GenericArgument::Type(ty) => Some(ty),
_ => None,
});
let key = types.next()?;
if types.next().is_some() {
return None;
}
Some(key)
}
fn state_type(func: &ItemFn) -> Option<TokenStream2> {
let FnArg::Typed(PatType { ty, .. }) = func.sig.inputs.get(1)? else {
return None;
};
if !is_context_param(ty) {
return None;
}
if let Type::Reference(reference) = &**ty
&& let Type::Path(path) = &*reference.elem
&& let Some(segment) = path.path.segments.last()
&& let syn::PathArguments::AngleBracketed(args) = &segment.arguments
{
let mut types = args.args.iter().filter_map(|arg| match arg {
syn::GenericArgument::Type(ty) => Some(ty),
_ => None,
});
let _context_ty = types.next();
if let Some(state_ty) = types.next() {
return Some(quote!(#state_ty));
}
}
None
}
fn is_context_param(ty: &Type) -> bool {
if let Type::Reference(reference) = ty
&& reference.mutability.is_some()
&& let Type::Path(path) = &*reference.elem
&& let Some(segment) = path.path.segments.last()
{
return segment.ident == "Context";
}
false
}
fn collect_extractors(func: &ItemFn, ctx_present: bool) -> syn::Result<Vec<(&Pat, &Type)>> {
let start = if ctx_present { 2 } else { 1 };
let mut extractors = Vec::new();
for arg in func.sig.inputs.iter().skip(start) {
let FnArg::Typed(PatType { pat, ty, .. }) = arg else {
return Err(Error::new_spanned(
arg,
"a #[subscriber] handler cannot take `self`",
));
};
if is_context_param(ty) {
return Err(Error::new_spanned(
ty,
"the `&mut Context` parameter must come immediately after the message, before any \
extractor parameters",
));
}
extractors.push((&**pat, &**ty));
}
Ok(extractors)
}
fn extractor_preds(
extractors: &[(&Pat, &Type)],
ctx_ty: &TokenStream2,
state: &TokenStream2,
) -> Vec<TokenStream2> {
extractors
.iter()
.map(|(_, ty)| quote!(#ty: ::ruststream::runtime::FromContext<#ctx_ty, #state>))
.collect()
}
fn where_clause(preds: &[TokenStream2]) -> TokenStream2 {
if preds.is_empty() {
return quote!();
}
quote!(where #(#preds),*)
}
fn extractor_prelude(
extractors: &[(&Pat, &Type)],
ctx_param: &TokenStream2,
ctx_ty: &TokenStream2,
state: &TokenStream2,
reject: &TokenStream2,
) -> TokenStream2 {
let binds = extractors.iter().map(|(pat, ty)| {
if from_headers_ty(ty).is_some() {
return quote! {
let #pat = {
let __rs_policy = #ctx_param.decode_policy();
match <#ty>::extract(&mut *#ctx_param, __rs_policy) {
::core::result::Result::Ok(__rs_value) => __rs_value,
::core::result::Result::Err(__rs_err) => { #reject; }
}
};
};
}
quote! {
let #pat = match <#ty as ::ruststream::runtime::FromContext<#ctx_ty, #state>>::from_context(
&mut *#ctx_param,
)
.await
{
::core::result::Result::Ok(__rs_value) => __rs_value,
::core::result::Result::Err(__rs_err) => { #reject; }
};
}
});
quote!(#(#binds)*)
}
fn failure_step(args: &SubscriberArgs) -> TokenStream2 {
let Some(failure) = &args.on_failure else {
return quote!();
};
let panic = failure
.panic
.as_ref()
.map(failure_policy_tokens)
.map(|policy| quote!(.with_panic(#policy)));
let decode = failure
.decode
.as_ref()
.map(failure_policy_tokens)
.map(|policy| quote!(.with_decode(#policy)));
quote! {
.on_failure(::ruststream::runtime::FailurePolicies::default() #panic #decode)
}
}
fn reject_batch_headers(shape: Shape, extractors: &[(&Pat, &Type)]) -> syn::Result<()> {
if shape == Shape::Batch
&& let Some((_, ty)) = extractors
.iter()
.find(|(_, ty)| from_headers_ty(ty).is_some())
{
return Err(Error::new_spanned(
ty,
"headers are per-delivery: a batch pairs each element with its contract - take the \
batch as `&[Message<H, T>]` and read `element.headers` (`Headers<..>` extraction \
stays on the single-message forms)",
));
}
Ok(())
}
fn headers_schema_expr(
extractors: &[(&Pat, &Type)],
input_static: &TokenStream2,
pair: Option<(&Type, &Type)>,
) -> TokenStream2 {
if let Some((headers, _)) = pair {
return quote! {{
#[allow(unused_imports)]
use ::ruststream::__private::NoSchemaProbe as _;
::ruststream::__private::Probe::<#headers>::new().schema_json()
}};
}
extractors
.iter()
.find_map(|(_, ty)| from_headers_ty(ty))
.map_or_else(
|| {
quote! {{
#[allow(unused_imports)]
use ::ruststream::__private::NoHeadersSchemaProbe as _;
::ruststream::__private::Probe::<#input_static>::new().headers_schema_json()
}}
},
|contract| {
quote! {{
#[allow(unused_imports)]
use ::ruststream::__private::NoSchemaProbe as _;
::ruststream::__private::Probe::<#contract>::new().schema_json()
}}
},
)
}
fn outgoing_entry(channel: &TokenStream2, message_ty: &TokenStream2) -> TokenStream2 {
quote! {
__rs_outgoing.push(
::ruststream::runtime::OutgoingMessageMetadata::new(
{ let __rs_channel: &'static str = #channel; __rs_channel },
::core::any::type_name::<#message_ty>(),
)
.with_message_name({
#[allow(unused_imports)]
use ::ruststream::__private::NoMessageProbe as _;
::ruststream::__private::Probe::<#message_ty>::new().message_name()
})
.with_message_description({
#[allow(unused_imports)]
use ::ruststream::__private::NoMessageProbe as _;
::ruststream::__private::Probe::<#message_ty>::new().message_description()
})
.with_payload_schema({
#[allow(unused_imports)]
use ::ruststream::__private::NoSchemaProbe as _;
::ruststream::__private::Probe::<#message_ty>::new().schema_json()
})
.with_headers_schema({
#[allow(unused_imports)]
use ::ruststream::__private::NoHeadersSchemaProbe as _;
::ruststream::__private::Probe::<#message_ty>::new().headers_schema_json()
})
.with_serialized({
#[allow(unused_imports)]
use ::ruststream::__private::NoSerializedProbe as _;
::ruststream::__private::Probe::<#message_ty>::new().serialized_wire()
}),
);
}
}
fn failure_policy_tokens(policy: &FailurePolicyArg) -> TokenStream2 {
match policy {
FailurePolicyArg::FailFast => quote!(::ruststream::runtime::FailurePolicy::FailFast),
FailurePolicyArg::Drop => quote!(::ruststream::runtime::FailurePolicy::Drop),
FailurePolicyArg::Retry => quote!(::ruststream::runtime::FailurePolicy::Retry),
FailurePolicyArg::RetryAfter(expr) => {
quote!(::ruststream::runtime::FailurePolicy::RetryAfter(#expr))
}
FailurePolicyArg::Skip => quote!(::ruststream::runtime::FailurePolicy::Skip),
FailurePolicyArg::Value(expr) => quote!(#expr),
}
}
fn workers_step(args: &SubscriberArgs, handler: &Ident, shape: Shape) -> syn::Result<TokenStream2> {
let Some(WorkersArg { count, by_key }) = &args.workers else {
return Ok(quote!());
};
let count = workers_count(count, handler)?;
if let Some(marker) = by_key {
if shape == Shape::Batch {
return Err(Error::new(
marker.span(),
"by_key lanes order single messages per key; they do not apply to a batch \
handler",
));
}
return Ok(quote!(.workers_by_key(#count)));
}
Ok(quote!(.workers(#count)))
}
fn workers_count(count: &Expr, handler: &Ident) -> syn::Result<TokenStream2> {
if let Expr::Lit(syn::ExprLit {
lit: syn::Lit::Int(literal),
..
}) = count
{
if literal.base10_parse::<usize>()? == 0 {
return Err(Error::new(
literal.span(),
"workers(0) is not a policy; the minimum is 1",
));
}
return Ok(quote! {
match ::core::num::NonZeroUsize::new(#literal) {
::core::option::Option::Some(count) => count,
::core::option::Option::None => ::core::num::NonZeroUsize::MIN,
}
});
}
let misconfigured = format!(
"workers(..) on subscriber `{handler}` needs a non-zero count; the configured value is 0",
);
Ok(quote! {
::core::num::NonZeroUsize::new(#count).expect(#misconfigured)
})
}
fn split_outs<'a>(extractors: &mut Vec<(&'a Pat, &'a Type)>) -> syn::Result<Vec<OutParam<'a>>> {
let mut outs = Vec::new();
let mut kept = Vec::new();
for (pat, ty) in extractors.drain(..) {
let Some(args) = out_param_args(ty) else {
kept.push((pat, ty));
continue;
};
let Type::ImplTrait(capability) = args[0] else {
return Err(Error::new_spanned(
args[0],
"an Out parameter names the capability it needs, not a publisher type: write \
`Out<impl Publisher>` (or a capability like `impl OwnedTransactions`); the \
concrete publisher is inferred from the policy attached at the include site",
));
};
let marker = args.get(1).map_or_else(
|| quote!(::ruststream::runtime::DefaultSlot),
|marker| quote!(#marker),
);
let bodies = match args.get(2) {
Some(body) => body_decl(body)?,
None => None,
};
outs.push(OutParam {
pat,
ty,
bounds: &capability.bounds,
marker,
bodies,
});
}
*extractors = kept;
if outs.len() > 3 {
return Err(Error::new_spanned(
outs[3].pat,
"a #[subscriber] handler takes at most three Out parameters",
));
}
for (index, out) in outs.iter().enumerate() {
let name = out.marker.to_string();
if outs
.iter()
.skip(index + 1)
.any(|other| other.marker.to_string() == name)
{
return Err(Error::new_spanned(
out.pat,
"every Out parameter needs its own slot marker: two parameters bind the same \
slot (name each with `Out<impl Publisher, MyMarker>` and derive the markers \
with #[derive(OutSlot)])",
));
}
}
Ok(outs)
}
fn body_decl(body: &Type) -> syn::Result<Option<BodyDecl<'_>>> {
let Type::Tuple(tuple) = body else {
return Ok(Some(BodyDecl::Set(body)));
};
let bodies: Vec<&Type> = tuple.elems.iter().collect();
if bodies.is_empty() {
return Ok(None);
}
if bodies.len() > 4 {
return Err(Error::new_spanned(
bodies[4],
"an Out parameter lists at most four message types inline; declare a \
#[derive(OutMessages)] set enum for a larger one",
));
}
for (index, body) in bodies.iter().enumerate() {
let name = body.to_token_stream().to_string();
if bodies
.iter()
.skip(index + 1)
.any(|other| other.to_token_stream().to_string() == name)
{
return Err(Error::new_spanned(
body,
"this type is already in the Out parameter's message list",
));
}
}
Ok(Some(BodyDecl::List(bodies)))
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum Shape {
Single,
Batch,
}
fn resolve_shape(reference: &syn::TypeReference) -> syn::Result<(Shape, &Type)> {
let elem = &*reference.elem;
if is_u8_slice(elem) {
return Err(Error::new_spanned(
elem,
"a payload's bytes take a named type: #[derive(Deserialized)] on a newtype over \
the payload view (struct Frame<'a>(&'a [u8]);) and take `&Frame<'_>`",
));
}
match elem {
Type::Slice(slice) => {
if is_u8_slice(&slice.elem)
|| matches!(&*slice.elem, Type::Reference(r) if is_u8_slice(&r.elem))
{
return Err(Error::new_spanned(
&slice.elem,
"a batch of payloads takes a named element type: #[derive(Deserialized)] on \
a newtype over the payload view (struct Frame<'a>(&'a [u8]);) and take \
`&[Frame<'_>]`",
));
}
Ok((Shape::Batch, &slice.elem))
}
other => Ok((Shape::Single, other)),
}
}
fn is_u8_slice(ty: &Type) -> bool {
if let Type::Slice(slice) = ty
&& let Type::Path(TypePath {
qself: None, path, ..
}) = &*slice.elem
{
return path.is_ident("u8");
}
false
}
fn subst_elided_lifetime(ty: &Type, lifetime: &str) -> (Type, bool) {
use syn::visit_mut::VisitMut;
struct Subst {
lifetime: syn::Lifetime,
found: bool,
}
impl VisitMut for Subst {
fn visit_lifetime_mut(&mut self, node: &mut syn::Lifetime) {
if node.ident == "_" {
*node = self.lifetime.clone();
self.found = true;
}
}
}
let mut subst = Subst {
lifetime: syn::Lifetime::new(lifetime, proc_macro2::Span::call_site()),
found: false,
};
let mut rewritten = ty.clone();
subst.visit_type_mut(&mut rewritten);
(rewritten, subst.found)
}
fn message_param(func: &ItemFn) -> syn::Result<MessageParam<'_>> {
let first = func.sig.inputs.first().ok_or_else(|| {
Error::new_spanned(
&func.sig,
"a #[subscriber] handler must take exactly one message parameter",
)
})?;
let FnArg::Typed(PatType { pat, ty, .. }) = first else {
return Err(Error::new_spanned(
first,
"a #[subscriber] handler cannot take `self`",
));
};
let Type::Reference(reference) = &**ty else {
return Err(Error::new_spanned(
ty,
"the message parameter must be a reference `&T`",
));
};
let (shape, input_ty) = resolve_shape(reference)?;
Ok((pat, shape, input_ty, message_pair_args(input_ty)))
}
type MessageParam<'a> = (&'a Pat, Shape, &'a Type, Option<(&'a Type, &'a Type)>);
fn handler_parts<'a>(args: &SubscriberArgs, func: &'a ItemFn) -> syn::Result<HandlerParts<'a>> {
let (pat, shape, input_ty, pair) = message_param(func)?;
let description = doc_description(&func.attrs);
let (source_ty, source_expr) = source_tokens(&args.source)?;
let (settings_source_ty, start_at_step, position_state) = match &args.start_at {
Some(position) => {
let position_ty = position_type(position)?;
(
quote!(::ruststream::StartAt<#source_ty, #position_ty>),
quote!(.start_at(#position)),
quote!(::ruststream::runtime::Fixed),
)
}
None => (
source_ty.clone(),
quote!(),
quote!(::ruststream::runtime::Open),
),
};
let ctx_arg = func
.sig
.inputs
.get(1)
.filter(|arg| matches!(arg, FnArg::Typed(pt) if is_context_param(&pt.ty)));
let ctx_param = if let Some(FnArg::Typed(PatType { pat, .. })) = ctx_arg {
quote!(#pat)
} else {
quote!(_ctx)
};
let mut extractors = collect_extractors(func, ctx_arg.is_some())?;
let outs = split_outs(&mut extractors)?;
let ctx_ty = context_type(func);
let state_ty = state_type(func);
reject_batch_headers(shape, &extractors)?;
let workers_step = workers_step(args, &func.sig.ident, shape)?;
let failure_step = failure_step(args);
let workers_state = state_marker(args.workers.is_some());
let failure_state = state_marker(args.on_failure.is_some());
let settings_chain = quote!(#workers_step #failure_step #start_at_step);
let settings_state_ty = quote!((
#workers_state,
#failure_state,
#position_state,
::ruststream::runtime::Open
));
Ok(HandlerParts {
vis: &func.vis,
name: &func.sig.ident,
block: &func.block,
pat,
input_ty,
pair,
description,
source_expr,
ctx_param,
ctx_ty,
state_ty,
extractors,
outs,
shape,
settings_chain,
settings_source_ty,
settings_state_ty,
})
}
fn state_marker(fixed: bool) -> TokenStream2 {
if fixed {
quote!(::ruststream::runtime::Fixed)
} else {
quote!(::ruststream::runtime::Open)
}
}
fn batch_reply_body<'a>(
func: &'a ItemFn,
block: &syn::Block,
) -> syn::Result<(&'a Type, TokenStream2)> {
let declared_ty = match &func.sig.output {
ReturnType::Type(_, ty) => &**ty,
ReturnType::Default => {
return Err(Error::new_spanned(
&func.sig,
"a batch publishing handler must return the replies: Vec<Reply>, or \
Result<Vec<Reply>, HandlerOutcome>",
));
}
};
if let Some(ok_ty) = crate::parse::publish_result_reply(declared_ty) {
let Some(elem) = vec_element(ok_ty) else {
return Err(Error::new_spanned(
ok_ty,
"a batch publishing handler replies with a Vec: Result<Vec<Reply>, HandlerOutcome>",
));
};
Ok((elem, quote!((async move #block).await)))
} else {
let Some(elem) = vec_element(declared_ty) else {
return Err(Error::new_spanned(
declared_ty,
"a batch publishing handler returns the replies: Vec<Reply>, or \
Result<Vec<Reply>, HandlerOutcome>",
));
};
Ok((
elem,
quote!(::core::result::Result::Ok((async move #block).await)),
))
}
}
fn publishing_reply<'a>(
func: &'a ItemFn,
block: &syn::Block,
) -> syn::Result<(&'a Type, TokenStream2)> {
let declared_ty = match &func.sig.output {
ReturnType::Type(_, ty) => &**ty,
ReturnType::Default => {
return Err(Error::new_spanned(
&func.sig,
"a publishing handler must return the reply value",
));
}
};
Ok(match crate::parse::publish_result_reply(declared_ty) {
Some(reply_ty) => (reply_ty, quote!((async move #block).await)),
None => (
declared_ty,
quote!(::core::result::Result::Ok((async move #block).await)),
),
})
}