diff --git a/lib/bencher_schema/src/model/organization/plan.rs b/lib/bencher_schema/src/model/organization/plan.rs index 1b4cc0c252..670308e2ed 100644 --- a/lib/bencher_schema/src/model/organization/plan.rs +++ b/lib/bencher_schema/src/model/organization/plan.rs @@ -1,6 +1,6 @@ #![cfg(feature = "plus")] -use bencher_billing::Biller; +use bencher_billing::{Biller, CustomerId}; use bencher_json::{ DateTime, Entitlements, JsonPlan, Jwt, LicensedPlanId, MeteredPlanId, OrganizationUuid, PlanLevel, project::Visibility, @@ -93,7 +93,7 @@ impl QueryPlan { biller: Option<&Biller>, public_user: &PublicUser, query_organization: &QueryOrganization, - ) -> Result, HttpError> { + ) -> Result, HttpError> { let Some(biller) = biller else { return Ok(None); }; @@ -108,13 +108,13 @@ impl QueryPlan { return Ok(None); }; - let plan_status = biller + let (plan_status, customer_id) = biller .get_metered_plan_status(&metered_plan_id) .await .map_err(not_found_error)?; if plan_status.is_active() { - Ok(Some(metered_plan_id)) + Ok(Some(customer_id)) } else { Err(payment_required_error(PlanKindError::InactiveMeteredPlan { organization: query_organization.clone(), @@ -214,7 +214,7 @@ impl InsertPlan { } pub enum PlanKind { - Metered(MeteredPlanId), + Metered(CustomerId), Licensed(LicenseUsage), None, } @@ -254,11 +254,11 @@ impl PlanKind { query_organization: &QueryOrganization, visibility: Visibility, ) -> Result { - if let Some(metered_plan_id) = + if let Some(customer_id) = QueryPlan::get_active_metered_plan(context, biller, public_user, query_organization) .await? { - Ok(Self::Metered(metered_plan_id)) + Ok(Self::Metered(customer_id)) } else if let Some(license_usage) = LicenseUsage::get( public_conn!(context, public_user), licensor, @@ -376,7 +376,7 @@ impl PlanKind { usage: u32, ) -> Result<(), HttpError> { match self { - Self::Metered(metered_plan_id) => { + Self::Metered(customer_id) => { let Some(biller) = biller else { return Err(issue_error( "No Biller when checking usage", @@ -385,7 +385,7 @@ impl PlanKind { )); }; biller - .record_metrics_usage(&metered_plan_id, usage) + .record_metrics_usage(&customer_id, usage) .await .map_err(|e| { issue_error( diff --git a/lib/bencher_schema/src/model/runner/job.rs b/lib/bencher_schema/src/model/runner/job.rs index 78ee33f58c..e1ed1d1945 100644 --- a/lib/bencher_schema/src/model/runner/job.rs +++ b/lib/bencher_schema/src/model/runner/job.rs @@ -475,7 +475,7 @@ fn insert_job_duration( #[cfg(test)] mod tests { - use bencher_json::{DateTime, Entitlements, MeteredPlanId}; + use bencher_json::{DateTime, Entitlements}; use diesel::{Connection as _, QueryDsl as _}; use pretty_assertions::assert_eq; @@ -490,7 +490,7 @@ mod tests { }; fn metered_plan() -> PlanKind { - PlanKind::Metered("test_plan".parse::().unwrap()) + PlanKind::Metered("cus_test".into()) } fn licensed_plan(level: PlanLevel) -> PlanKind { diff --git a/plus/api_runners/Cargo.toml b/plus/api_runners/Cargo.toml index 666c0d3e3d..45e005ff1e 100644 --- a/plus/api_runners/Cargo.toml +++ b/plus/api_runners/Cargo.toml @@ -15,6 +15,7 @@ plus = [ "bencher_otel?/plus", "bencher_schema/plus", "bencher_token/plus", + "dep:bencher_billing", "dep:bencher_token", "dep:diesel", "dep:dropshot", @@ -30,7 +31,7 @@ plus = [ "dep:tokio-tungstenite", "dep:uuid", ] -sentry = ["bencher_schema/sentry", "dep:bencher_billing", "dep:sentry"] +sentry = ["bencher_schema/sentry", "dep:sentry"] otel = ["bencher_endpoint/otel", "bencher_schema/otel", "dep:bencher_otel"] [dependencies] diff --git a/plus/api_runners/src/channel.rs b/plus/api_runners/src/channel.rs index d6bcc33e9f..5eb82f17ab 100644 --- a/plus/api_runners/src/channel.rs +++ b/plus/api_runners/src/channel.rs @@ -5,6 +5,7 @@ use std::time::Duration; +use bencher_billing::CustomerId; use bencher_json::{ DEFAULT_POLL_TIMEOUT, JobStatus, JobUuid, JsonClaimedJob, JsonSpec, MeteredPlanId, Priority, RunnerResourceId, @@ -77,6 +78,9 @@ enum ChannelError { #[error("{0}")] Json(#[from] serde_json::Error), + #[error("{0}")] + Billing(#[from] bencher_billing::BillingError), + /// Job is in an unexpected state for the requested transition. #[error("Invalid state transition to {target:?} for job {job_id:?}, found {current:?}")] InvalidStateTransition { @@ -403,14 +407,11 @@ async fn bill_stripe_best_effort( organization_id: OrganizationId, billing_state: &mut BillingState, ) { - let metered_plan_id = match billing_state - .metered_plan_id(context, organization_id) - .await - { + let customer_id = match billing_state.customer_id(context, organization_id).await { Ok(Some(id)) => id, Ok(None) => return, Err(e) => { - slog::warn!(log, "Failed to look up metered plan for billing"; "job_id" => ?job_id, "error" => %e); + slog::warn!(log, "Failed to look up customer for billing"; "job_id" => ?job_id, "error" => %e); #[cfg(feature = "sentry")] sentry::capture_error(&e); return; @@ -421,7 +422,7 @@ async fn bill_stripe_best_effort( return; }; - if let Err(e) = biller.record_runner_usage(&metered_plan_id, delta).await { + if let Err(e) = biller.record_runner_usage(&customer_id, delta).await { slog::warn!(log, "Failed to record runner billing"; "job_id" => ?job_id, "delta" => delta, "error" => %e); #[cfg(feature = "otel")] bencher_otel::ApiMeter::increment(bencher_otel::ApiCounter::RunnerMinutesBillingFailed); @@ -433,25 +434,25 @@ async fn bill_stripe_best_effort( } } -/// Cached result of the metered plan lookup for an organization. +/// Cached result of the customer ID lookup for an organization. /// -/// Avoids querying `schema::plan::table` on every heartbeat, since the plan +/// Avoids querying the DB and Stripe on every heartbeat, since the customer /// will not change mid-job. -enum CachedMeteredPlan { +enum CachedCustomer { /// Not yet looked up. Unknown, /// Looked up and no metered plan exists for the organization. None, - /// Looked up and found a metered plan. - Some(MeteredPlanId), + /// Looked up and resolved the Stripe customer ID. + Some(CustomerId), } -/// Per-job billing state: caches the metered plan lookup and tracks Sentry +/// Per-job billing state: caches the customer ID lookup and tracks Sentry /// reporting so only the first billing failure per job is sent. /// /// Created at the start of each job execution and dropped when the job finishes. struct BillingState { - metered_plan: CachedMeteredPlan, + customer: CachedCustomer, #[cfg(feature = "sentry")] reported: bool, } @@ -459,35 +460,46 @@ struct BillingState { impl BillingState { fn new() -> Self { Self { - metered_plan: CachedMeteredPlan::Unknown, + customer: CachedCustomer::Unknown, #[cfg(feature = "sentry")] reported: false, } } - /// Return the cached metered plan ID, querying the DB on first call. - async fn metered_plan_id( + /// Return the cached customer ID, querying the DB and Stripe on first call. + async fn customer_id( &mut self, context: &ApiContext, organization_id: OrganizationId, - ) -> Result, ChannelError> { - match &self.metered_plan { - CachedMeteredPlan::Unknown => { + ) -> Result, ChannelError> { + match &self.customer { + CachedCustomer::Unknown => { let plan_id: Option> = schema::plan::table .filter(schema::plan::organization_id.eq(organization_id)) .select(schema::plan::metered_plan) .first(auth_conn!(context)) .optional()?; - if let Some(id) = plan_id.flatten() { - self.metered_plan = CachedMeteredPlan::Some(id.clone()); - Ok(Some(id)) + if let Some(metered_plan_id) = plan_id.flatten() { + let Some(biller) = context.biller.as_ref() else { + self.customer = CachedCustomer::None; + return Ok(None); + }; + let (status, customer_id) = + biller.get_metered_plan_status(&metered_plan_id).await?; + if status.is_active() { + self.customer = CachedCustomer::Some(customer_id.clone()); + Ok(Some(customer_id)) + } else { + self.customer = CachedCustomer::None; + Ok(None) + } } else { - self.metered_plan = CachedMeteredPlan::None; + self.customer = CachedCustomer::None; Ok(None) } }, - CachedMeteredPlan::None => Ok(None), - CachedMeteredPlan::Some(id) => Ok(Some(id.clone())), + CachedCustomer::None => Ok(None), + CachedCustomer::Some(id) => Ok(Some(id.clone())), } } diff --git a/plus/bencher_billing/src/biller.rs b/plus/bencher_billing/src/biller.rs index 4c0e6bc7eb..5071834d9d 100644 --- a/plus/bencher_billing/src/biller.rs +++ b/plus/bencher_billing/src/biller.rs @@ -625,10 +625,13 @@ impl Biller { pub async fn get_metered_plan_status( &self, metered_plan_id: &MeteredPlanId, - ) -> Result { + ) -> Result<(PlanStatus, CustomerId), BillingError> { let subscription_id: SubscriptionId = metered_plan_id.as_ref().into(); let subscription = self.get_subscription(&subscription_id).await?; - Ok(Self::map_status(&subscription.status)) + Ok(( + Self::map_status(&subscription.status), + subscription.customer.id().clone(), + )) } pub async fn get_licensed_plan_status( @@ -655,36 +658,32 @@ impl Biller { pub async fn record_metrics_usage( &self, - metered_plan_id: &MeteredPlanId, + customer_id: &CustomerId, quantity: u32, ) -> Result { - self.record_metered_usage(METRICS_METER_NAME, metered_plan_id, quantity) + self.record_metered_usage(METRICS_METER_NAME, customer_id, quantity) .await } pub async fn record_runner_usage( &self, - metered_plan_id: &MeteredPlanId, + customer_id: &CustomerId, minutes: u32, ) -> Result { - self.record_metered_usage(RUNNER_MINUTES_METER_NAME, metered_plan_id, minutes) + self.record_metered_usage(RUNNER_MINUTES_METER_NAME, customer_id, minutes) .await } async fn record_metered_usage( &self, meter_name: &str, - metered_plan_id: &MeteredPlanId, + customer_id: &CustomerId, quantity: u32, ) -> Result { - let subscription_id: SubscriptionId = metered_plan_id.as_ref().into(); - let subscription = self.get_subscription(&subscription_id).await?; - let customer_id = subscription.customer.id().to_string(); - CreateBillingMeterEvent::new( meter_name, HashMap::from([ - (METER_CUSTOMER_KEY.to_owned(), customer_id), + (METER_CUSTOMER_KEY.to_owned(), customer_id.to_string()), (METER_VALUE_KEY.to_owned(), quantity.to_string()), ]), ) @@ -740,8 +739,8 @@ mod tests { use std::collections::HashSet; use bencher_json::{ - Entitlements, MeteredPlanId, OrganizationUuid, PlanLevel, PlanStatus, UserUuid, - organization::plan::{DEFAULT_PRICE_NAME, METRICS_METER_NAME}, + Entitlements, OrganizationUuid, PlanLevel, PlanStatus, UserUuid, + organization::plan::{DEFAULT_PRICE_NAME, METRICS_METER_NAME, RUNNER_MINUTES_METER_NAME}, system::{ config::{JsonBilling, JsonProduct, JsonProducts}, payment::{JsonCard, JsonCustomer}, @@ -756,6 +755,8 @@ mod tests { use crate::Biller; + use super::{METER_CUSTOMER_KEY, METER_VALUE_KEY}; + const TEST_BILLING_KEY: &str = "TEST_BILLING_KEY"; fn billing_key() -> Option { @@ -787,6 +788,7 @@ mod tests { } } + #[expect(clippy::too_many_arguments)] async fn metered_subscription( biller: &Biller, organization: OrganizationUuid, @@ -794,7 +796,8 @@ mod tests { payment_method_id: PaymentMethodId, plan_level: PlanLevel, price_name: String, - usage_count: usize, + runner_minutes: usize, + metrics_quantity: u32, ) { let create_subscription = biller .create_metered_subscription( @@ -814,19 +817,20 @@ mod tests { let metered_plan_id = &subscription_id.as_ref().parse().unwrap(); biller.get_metered_plan(metered_plan_id).await.unwrap(); - let plan_status = biller + let (plan_status, customer_id) = biller .get_metered_plan_status(metered_plan_id) .await .unwrap(); assert_eq!(plan_status, PlanStatus::Active); - record_metrics_usage(biller, metered_plan_id, usage_count).await; + record_runner_usage(biller, &customer_id, runner_minutes).await; + record_metrics_usage(biller, &customer_id, metrics_quantity).await; biller .cancel_metered_subscription(&subscription_id.parse().unwrap()) .await .unwrap(); - let plan_status = biller + let (plan_status, _) = biller .get_metered_plan_status(metered_plan_id) .await .unwrap(); @@ -879,20 +883,34 @@ mod tests { assert_eq!(plan_status, PlanStatus::Canceled); } - async fn record_metrics_usage( - biller: &Biller, - metered_plan_id: &MeteredPlanId, - usage_count: usize, - ) { - for _ in 0..usage_count { - let quantity = u32::from(rand::random::()); - biller - .record_metrics_usage(metered_plan_id, quantity) - .await - .unwrap(); + async fn record_runner_usage(biller: &Biller, customer_id: &CustomerId, runner_minutes: usize) { + for _ in 0..runner_minutes { + let event = biller.record_runner_usage(customer_id, 1).await.unwrap(); + assert_eq!(event.event_name, RUNNER_MINUTES_METER_NAME); + assert_eq!( + event.payload.get(METER_CUSTOMER_KEY), + Some(&customer_id.to_string()), + ); + assert_eq!(event.payload.get(METER_VALUE_KEY), Some(&"1".to_owned()),); } } + async fn record_metrics_usage(biller: &Biller, customer_id: &CustomerId, quantity: u32) { + let event = biller + .record_metrics_usage(customer_id, quantity) + .await + .unwrap(); + assert_eq!(event.event_name, METRICS_METER_NAME); + assert_eq!( + event.payload.get(METER_CUSTOMER_KEY), + Some(&customer_id.to_string()), + ); + assert_eq!( + event.payload.get(METER_VALUE_KEY), + Some(&quantity.to_string()), + ); + } + fn make_price(price_id: &str) -> stripe_shared::Price { stripe_shared::Price { active: false, @@ -1125,6 +1143,7 @@ mod tests { payment_method_id.clone(), PlanLevel::Team, METRICS_METER_NAME.into(), + 5, 10, ) .await; @@ -1151,6 +1170,7 @@ mod tests { payment_method_id.clone(), PlanLevel::Enterprise, METRICS_METER_NAME.into(), + 10, 25, ) .await;