Reduce Stripe queries

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<Option<MeteredPlanId>, HttpError> {
+    ) -> Result<Option<CustomerId>, 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<Self, HttpError> {
-        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(

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::<MeteredPlanId>().unwrap())
+        PlanKind::Metered("cus_test".into())
     }
 
     fn licensed_plan(level: PlanLevel) -> PlanKind {

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]

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,61 +434,72 @@ 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,
 }
 
 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<Option<MeteredPlanId>, ChannelError> {
-        match &self.metered_plan {
-            CachedMeteredPlan::Unknown => {
+    ) -> Result<Option<CustomerId>, ChannelError> {
+        match &self.customer {
+            CachedCustomer::Unknown => {
                 let plan_id: Option<Option<MeteredPlanId>> = 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())),
         }
     }