@rivetkit/engine 1.0.0

package/packages/gasoline-macros/src/lib.rs

@@ -0,0 +1,527 @@
+use proc_macro::TokenStream;
+use quote::{ToTokens, quote};
+use syn::{
+	GenericArgument, Ident, Item, ItemFn, ItemStruct, LitStr, PathArguments, ReturnType, Type,
+	parse::{Parse, ParseStream},
+	parse_macro_input,
+	spanned::Spanned,
+};
+
+#[derive(Default)]
+struct Config {
+	max_retries: Option<syn::Expr>,
+	timeout: Option<syn::Expr>,
+}
+
+impl Config {
+	fn max_retries(&self) -> proc_macro2::TokenStream {
+		self.max_retries
+			.as_ref()
+			.map(|e| e.to_token_stream())
+			.unwrap_or_else(|| quote! { 5 })
+	}
+
+	fn timeout(&self) -> proc_macro2::TokenStream {
+		self.timeout
+			.as_ref()
+			.map(|e| e.to_token_stream())
+			.unwrap_or_else(|| quote! { 30 })
+	}
+}
+
+struct MessageConfig {
+	tail_ttl: u64,
+}
+
+impl Default for MessageConfig {
+	fn default() -> Self {
+		MessageConfig { tail_ttl: 90 }
+	}
+}
+
+#[proc_macro_attribute]
+pub fn workflow(attr: TokenStream, item: TokenStream) -> TokenStream {
+	let name = parse_macro_input!(attr as OptionalIdent)
+		.ident
+		.map(|x| x.to_string())
+		.unwrap_or_else(|| "Workflow".to_string());
+	let item_fn = parse_macro_input!(item as ItemFn);
+
+	if let Err(err) = parse_empty_config(&item_fn.attrs) {
+		return err.into_compile_error().into();
+	}
+
+	let ctx_ty = syn::parse_str("&mut WorkflowCtx").unwrap();
+	let TraitFnOutput {
+		ctx_ident,
+		input_ident,
+		input_type,
+		output_type,
+	} = match parse_trait_fn(&ctx_ty, "Workflow", &item_fn) {
+		Ok(x) => x,
+		Err(err) => return err,
+	};
+
+	let struct_ident = Ident::new(&name, proc_macro2::Span::call_site());
+	let fn_name = item_fn.sig.ident.to_string();
+	if !fn_name
+		.chars()
+		.all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '_')
+	{
+		return error(
+			item_fn.sig.ident.span(),
+			"invalid workflow name, must be [a-z0-9_]",
+		);
+	}
+
+	let fn_body = item_fn.block;
+	let vis = item_fn.vis;
+
+	let expanded = quote! {
+		#vis struct #struct_ident;
+
+		impl gas::workflow::WorkflowInput for #input_type {
+			type Workflow = #struct_ident;
+		}
+
+		#[async_trait::async_trait]
+		impl gas::workflow::Workflow for #struct_ident {
+			type Input = #input_type;
+			type Output = #output_type;
+
+			const NAME: &'static str = #fn_name;
+
+			async fn run(#ctx_ident: #ctx_ty, #input_ident: &Self::Input) -> Result<Self::Output> {
+				#fn_body
+			}
+		}
+	};
+
+	// eprintln!("\n\n{expanded}\n");
+
+	TokenStream::from(expanded)
+}
+
+#[proc_macro_attribute]
+pub fn activity(attr: TokenStream, item: TokenStream) -> TokenStream {
+	let name = parse_macro_input!(attr as Ident).to_string();
+	let item_fn = parse_macro_input!(item as ItemFn);
+
+	let config = match parse_config(&item_fn.attrs) {
+		Ok(x) => x,
+		Err(err) => return err.into_compile_error().into(),
+	};
+
+	let ctx_ty = syn::parse_str("&ActivityCtx").unwrap();
+	let TraitFnOutput {
+		ctx_ident,
+		input_ident,
+		input_type,
+		output_type,
+	} = match parse_trait_fn(&ctx_ty, "Activity", &item_fn) {
+		Ok(x) => x,
+		Err(err) => return err,
+	};
+
+	let struct_ident = Ident::new(&name, proc_macro2::Span::call_site());
+	let fn_name = item_fn.sig.ident.to_string();
+	if !fn_name
+		.chars()
+		.all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '_')
+	{
+		return error(
+			item_fn.sig.ident.span(),
+			"invalid activity name, must be [a-z0-9_]",
+		);
+	}
+
+	let fn_body = item_fn.block;
+	let vis = item_fn.vis;
+
+	let max_retries = config.max_retries();
+	let timeout = config.timeout();
+
+	let expanded = quote! {
+		#vis struct #struct_ident;
+
+		impl gas::activity::ActivityInput for #input_type {
+			type Activity = #struct_ident;
+		}
+
+		#[async_trait::async_trait]
+		impl gas::activity::Activity for #struct_ident {
+			type Input = #input_type;
+			type Output = #output_type;
+
+			const NAME: &'static str = #fn_name;
+			const MAX_RETRIES: usize = #max_retries;
+			const TIMEOUT: std::time::Duration = std::time::Duration::from_secs(#timeout);
+
+			async fn run(#ctx_ident: #ctx_ty, #input_ident: &Self::Input) -> Result<Self::Output> {
+				#fn_body
+			}
+		}
+	};
+
+	// eprintln!("\n\n{expanded}\n");
+
+	TokenStream::from(expanded)
+}
+
+#[proc_macro_attribute]
+pub fn operation(attr: TokenStream, item: TokenStream) -> TokenStream {
+	let name = parse_macro_input!(attr as OptionalIdent)
+		.ident
+		.map(|x| x.to_string())
+		.unwrap_or_else(|| "Operation".to_string());
+
+	let item_fn = parse_macro_input!(item as ItemFn);
+
+	let config = match parse_config(&item_fn.attrs) {
+		Ok(x) => x,
+		Err(err) => return err.into_compile_error().into(),
+	};
+
+	let ctx_ty = syn::parse_str("&OperationCtx").unwrap();
+	let TraitFnOutput {
+		ctx_ident,
+		input_ident,
+		input_type,
+		output_type,
+	} = match parse_trait_fn(&ctx_ty, "Operation", &item_fn) {
+		Ok(x) => x,
+		Err(err) => return err,
+	};
+
+	let struct_ident = Ident::new(&name, proc_macro2::Span::call_site());
+	let fn_name = item_fn.sig.ident.to_string();
+	if !fn_name
+		.chars()
+		.all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '_')
+	{
+		return error(
+			item_fn.sig.ident.span(),
+			"invalid operation name, must be [a-z0-9_]",
+		);
+	}
+
+	let generics = &item_fn.sig.generics;
+	let fn_body = item_fn.block;
+	let vis = item_fn.vis;
+
+	let timeout = config.timeout();
+
+	let expanded = quote! {
+		#vis struct #struct_ident;
+
+		impl gas::operation::OperationInput for #input_type {
+			type Operation = #struct_ident;
+		}
+
+		#[async_trait::async_trait]
+		impl gas::operation::Operation for #struct_ident {
+			type Input = #input_type;
+			type Output = #output_type;
+
+			const NAME: &'static str = #fn_name;
+			const TIMEOUT: std::time::Duration = std::time::Duration::from_secs(#timeout);
+
+			async fn run #generics (#ctx_ident: #ctx_ty, #input_ident: &Self::Input) -> Result<Self::Output> {
+				#fn_body
+			}
+		}
+	};
+
+	// eprintln!("\n\n{expanded}\n");
+
+	TokenStream::from(expanded)
+}
+
+struct TraitFnOutput {
+	ctx_ident: syn::Ident,
+	input_ident: syn::Ident,
+	input_type: Box<syn::Type>,
+	output_type: syn::Type,
+}
+
+fn parse_trait_fn(
+	ctx_ty: &syn::Type,
+	trait_name: &str,
+	item_fn: &syn::ItemFn,
+) -> Result<TraitFnOutput, TokenStream> {
+	// Check if is async
+	if item_fn.sig.asyncness.is_none() {
+		return Err(error(item_fn.sig.span(), "function must be async"));
+	}
+
+	let mut arg_names = vec![];
+	let mut arg_types = vec![];
+
+	for input in &item_fn.sig.inputs {
+		if let syn::FnArg::Typed(arg) = input {
+			match arg.pat.as_ref() {
+				syn::Pat::Ident(ident) => {
+					let arg_name = ident.ident.to_string();
+					arg_names.push(arg_name);
+					arg_types.push((*arg.ty).clone());
+				}
+				_ => {
+					return Err(error(arg.pat.span(), "unsupported input parameter pattern"));
+				}
+			}
+		} else {
+			return Err(error(input.span(), "unsupported input parameter type"));
+		}
+	}
+
+	if arg_types.len() != 2 || &arg_types[0] != ctx_ty {
+		return Err(error(
+			item_fn.sig.span(),
+			&format!(
+				"{} function must have exactly two parameters: ctx: {} and input: &YourInputType",
+				trait_name,
+				ctx_ty.to_token_stream()
+			),
+		));
+	}
+
+	let input_type = if let syn::Type::Reference(syn::TypeReference { elem, .. }) = &arg_types[1] {
+		elem.clone()
+	} else {
+		return Err(error(arg_types[1].span(), "input type must be a reference"));
+	};
+
+	let output_type = match &item_fn.sig.output {
+		ReturnType::Type(_, ty) => match ty.as_ref() {
+			Type::Path(path) => {
+				let segment = path.path.segments.last().unwrap();
+				if segment.ident == "Result" {
+					match &segment.arguments {
+						PathArguments::AngleBracketed(args) => {
+							if let Some(GenericArgument::Type(ty)) = args.args.first() {
+								ty.clone()
+							} else {
+								return Err(error(args.span(), "unsupported Result type"));
+							}
+						}
+						_ => {
+							return Err(error(segment.arguments.span(), "unsupported Result type"));
+						}
+					}
+				} else {
+					return Err(error(
+						path.span(),
+						&format!("{} function must return a Result type", trait_name),
+					));
+				}
+			}
+			_ => return Err(error(ty.span(), "unsupported output type")),
+		},
+		_ => {
+			return Err(error(
+				item_fn.sig.output.span(),
+				&format!("{} function must have a return type", trait_name),
+			));
+		}
+	};
+
+	Ok(TraitFnOutput {
+		ctx_ident: Ident::new(&arg_names[0], proc_macro2::Span::call_site()),
+		input_ident: Ident::new(&arg_names[1], proc_macro2::Span::call_site()),
+		input_type,
+		output_type,
+	})
+}
+
+#[proc_macro_attribute]
+pub fn signal(attr: TokenStream, item: TokenStream) -> TokenStream {
+	let name = parse_macro_input!(attr as LitStr);
+	if !name
+		.value()
+		.chars()
+		.all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '_')
+	{
+		return error(name.span(), "invalid signal name, must be [a-z0-9_]");
+	}
+
+	let item = parse_macro_input!(item as Item);
+	let (ident, attrs) = match item {
+		Item::Struct(ref item_struct) => (&item_struct.ident, &item_struct.attrs),
+		Item::Enum(ref item_enum) => (&item_enum.ident, &item_enum.attrs),
+		_ => return error(item.span(), "expected struct or enum"),
+	};
+
+	// If also a message, don't derive serde traits
+	let also_message = attrs
+		.iter()
+		.filter_map(|attr| attr.path().segments.last())
+		.any(|seg| seg.ident == "message");
+	let serde_derive = if also_message {
+		quote! {}
+	} else {
+		quote! { #[derive(serde::Serialize, serde::Deserialize)] }
+	};
+
+	let expanded = quote! {
+		#serde_derive
+		#item
+
+		impl gas::signal::Signal for #ident {
+			const NAME: &'static str = #name;
+		}
+
+		#[async_trait::async_trait]
+		impl gas::listen::Listen for #ident {
+			async fn listen(ctx: &mut gas::prelude::ListenCtx) -> gas::prelude::WorkflowResult<Self> {
+				let row = ctx.listen_any(&[<Self as gas::signal::Signal>::NAME]).await?;
+				Self::parse(&row.signal_name, &row.body)
+			}
+
+			fn parse(_name: &str, body: &serde_json::value::RawValue) -> gas::prelude::WorkflowResult<Self> {
+				serde_json::from_str(body.get()).map_err(WorkflowError::DeserializeSignalBody)
+			}
+		}
+	};
+
+	TokenStream::from(expanded)
+}
+
+#[proc_macro_attribute]
+pub fn message(attr: TokenStream, item: TokenStream) -> TokenStream {
+	let name = parse_macro_input!(attr as LitStr);
+	if !name
+		.value()
+		.chars()
+		.all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '_')
+	{
+		return error(name.span(), "invalid message name, must be [a-z0-9_]");
+	}
+
+	let item_struct = parse_macro_input!(item as ItemStruct);
+
+	// If also a signal, don't derive serde traits
+	let also_signal = item_struct
+		.attrs
+		.iter()
+		.filter_map(|attr| attr.path().segments.last())
+		.any(|seg| seg.ident == "signal");
+	let serde_derive = if also_signal {
+		quote! {}
+	} else {
+		quote! { #[derive(serde::Serialize, serde::Deserialize)] }
+	};
+
+	let config = match parse_msg_config(&item_struct.attrs) {
+		Ok(x) => x,
+		Err(err) => return err.into_compile_error().into(),
+	};
+
+	let struct_ident = &item_struct.ident;
+	let tail_ttl = config.tail_ttl;
+
+	let expanded = quote! {
+		#serde_derive
+		#[derive(Debug)]
+		#item_struct
+
+		impl gas::message::Message for #struct_ident {
+			const NAME: &'static str = #name;
+			const TAIL_TTL: std::time::Duration = std::time::Duration::from_secs(#tail_ttl);
+		}
+	};
+
+	TokenStream::from(expanded)
+}
+
+fn error(span: proc_macro2::Span, msg: &str) -> TokenStream {
+	syn::Error::new(span, msg).to_compile_error().into()
+}
+
+fn parse_config(attrs: &[syn::Attribute]) -> syn::Result<Config> {
+	let mut config = Config::default();
+
+	for attr in attrs {
+		let syn::Meta::NameValue(name_value) = &attr.meta else {
+			continue;
+		};
+
+		let ident = name_value.path.require_ident()?;
+
+		// Verify config property
+		if ident == "max_retries" {
+			config.max_retries = Some(syn::parse::<syn::Expr>(
+				name_value.value.to_token_stream().into(),
+			)?);
+		} else if ident == "timeout" {
+			config.timeout = Some(syn::parse::<syn::Expr>(
+				name_value.value.to_token_stream().into(),
+			)?);
+		} else if ident != "doc" {
+			return Err(syn::Error::new(
+				name_value.span(),
+				format!("Unknown config property `{ident}`"),
+			));
+		}
+	}
+
+	Ok(config)
+}
+
+fn parse_msg_config(attrs: &[syn::Attribute]) -> syn::Result<MessageConfig> {
+	let mut config = MessageConfig::default();
+
+	for attr in attrs {
+		let syn::Meta::NameValue(name_value) = &attr.meta else {
+			continue;
+		};
+
+		let ident = name_value.path.require_ident()?;
+
+		// Verify config property
+		if ident == "tail_ttl" {
+			config.tail_ttl = syn::parse::<syn::LitInt>(name_value.value.to_token_stream().into())?
+				.base10_parse()?;
+		} else if ident != "doc" {
+			return Err(syn::Error::new(
+				name_value.span(),
+				format!("Unknown config property `{ident}`"),
+			));
+		}
+	}
+
+	Ok(config)
+}
+
+fn parse_empty_config(attrs: &[syn::Attribute]) -> syn::Result<()> {
+	for attr in attrs {
+		let syn::Meta::NameValue(name_value) = &attr.meta else {
+			continue;
+		};
+
+		let ident = name_value.path.require_ident()?;
+
+		if ident != "doc" {
+			return Err(syn::Error::new(
+				name_value.span(),
+				format!("Unknown config property `{ident}`"),
+			));
+		}
+	}
+
+	Ok(())
+}
+
+struct OptionalIdent {
+	ident: Option<Ident>,
+}
+
+impl Parse for OptionalIdent {
+	fn parse(input: ParseStream) -> syn::Result<Self> {
+		if input.is_empty() {
+			Ok(OptionalIdent { ident: None })
+		} else {
+			let ident: Ident = input.parse()?;
+			Ok(OptionalIdent { ident: Some(ident) })
+		}
+	}
+}

package/packages/guard/Cargo.toml

@@ -0,0 +1,56 @@
+[package]
+name = "rivet-guard"
+version.workspace = true
+edition.workspace = true
+authors.workspace = true
+license.workspace = true
+
+[lib]
+name = "rivet_guard"
+path = "src/lib.rs"
+
+[dependencies]
+anyhow.workspace = true
+axum.workspace = true
+bytes.workspace = true
+gas.workspace = true
+http-body.workspace = true
+http-body-util.workspace = true
+hyper-tungstenite.workspace = true
+tower.workspace = true
+futures.workspace = true
+# TODO: Make this use workspace version
+hyper = "1.6.0"
+indoc.workspace = true
+once_cell.workspace = true
+pegboard-gateway.workspace = true
+pegboard.workspace = true
+pegboard-runner.workspace = true
+regex.workspace = true
+rivet-api-builder.workspace = true
+rivet-api-public.workspace = true
+rivet-cache.workspace = true
+rivet-config.workspace = true
+rivet-data.workspace = true
+rivet-error.workspace = true
+rivet-guard-core.workspace = true
+rivet-logs.workspace = true
+rivet-metrics.workspace = true
+rivet-pools.workspace = true
+rivet-runner-protocol.workspace = true
+rivet-runtime.workspace = true
+rustls-pemfile.workspace = true
+rustls.workspace = true
+serde_json.workspace = true
+serde.workspace = true
+tokio.workspace = true
+tokio-tungstenite.workspace = true
+tracing.workspace = true
+universaldb.workspace = true
+universalpubsub.workspace = true
+url.workspace = true
+urlencoding.workspace = true
+uuid.workspace = true
+
+[dev-dependencies]
+uuid.workspace = true

package/packages/guard/src/cache/actor.rs

@@ -0,0 +1,43 @@
+use std::{
+	collections::hash_map::DefaultHasher,
+	hash::{Hash, Hasher},
+};
+
+use anyhow::Result;
+use gas::prelude::*;
+
+use crate::routing::pegboard_gateway::X_RIVET_ACTOR;
+
+#[tracing::instrument(skip_all)]
+pub fn build_cache_key(
+	target: &str,
+	path: &str,
+	method: &hyper::Method,
+	headers: &hyper::HeaderMap,
+) -> Result<u64> {
+	// Check target
+	ensure!(target == "actor", "wrong target");
+
+	// Find actor to route to
+	let actor_id_str = headers
+		.get(X_RIVET_ACTOR)
+		.ok_or_else(|| {
+			crate::errors::MissingHeader {
+				header: X_RIVET_ACTOR.to_string(),
+			}
+			.build()
+		})?
+		.to_str()
+		.context("invalid x-rivet-actor header")?;
+	let actor_id = Id::parse(actor_id_str).context("invalid x-rivet-actor header")?;
+
+	// Create a hash using target, actor_id, path, and method
+	let mut hasher = DefaultHasher::new();
+	target.hash(&mut hasher);
+	actor_id.hash(&mut hasher);
+	path.hash(&mut hasher);
+	method.as_str().hash(&mut hasher);
+	let hash = hasher.finish();
+
+	Ok(hash)
+}

package/packages/guard/src/cache/mod.rs

@@ -0,0 +1,69 @@
+use std::{
+	collections::hash_map::DefaultHasher,
+	hash::{Hash, Hasher},
+	sync::Arc,
+};
+
+use anyhow::Result;
+use gas::prelude::*;
+use rivet_guard_core::CacheKeyFn;
+
+pub mod actor;
+
+use crate::routing::X_RIVET_TARGET;
+
+/// Creates the main cache key function that handles all incoming requests
+#[tracing::instrument(skip_all)]
+pub fn create_cache_key_function(_ctx: StandaloneCtx) -> CacheKeyFn {
+	Arc::new(move |hostname, path, method, _port_type, headers| {
+		tracing::debug!("building cache key");
+
+		let target = match read_target(headers) {
+			Ok(target) => target,
+			Err(err) => {
+				tracing::debug!(?err, "failed parsing target for cache key");
+
+				return Ok(host_path_method_cache_key(hostname, path, method));
+			}
+		};
+
+		let cache_key = match actor::build_cache_key(target, path, method, headers) {
+			Ok(key) => Some(key),
+			Err(err) => {
+				tracing::debug!(?err, "failed to create actor cache key");
+
+				None
+			}
+		};
+
+		// Fallback to hostname + path + method hash if actor did not work
+		if let Some(cache_key) = cache_key {
+			Ok(cache_key)
+		} else {
+			Ok(host_path_method_cache_key(hostname, path, method))
+		}
+	})
+}
+
+fn read_target(headers: &hyper::HeaderMap) -> Result<&str> {
+	// Read target
+	let target = headers.get(X_RIVET_TARGET).ok_or_else(|| {
+		crate::errors::MissingHeader {
+			header: X_RIVET_TARGET.to_string(),
+		}
+		.build()
+	})?;
+
+	Ok(target.to_str()?)
+}
+
+fn host_path_method_cache_key(hostname: &str, path: &str, method: &hyper::Method) -> u64 {
+	// Extract just the hostname, stripping the port if present
+	let hostname_only = hostname.split(':').next().unwrap_or(hostname);
+
+	let mut hasher = DefaultHasher::new();
+	hostname_only.hash(&mut hasher);
+	path.hash(&mut hasher);
+	method.as_str().hash(&mut hasher);
+	hasher.finish()
+}