Prevent deadlock when create a new meter/price on Stripe (#19196)

This also puts the entire state of `StripeBilling` behind a `RwLock`.
When fetching the existing prices and meters, or when inserting new
ones, we acquire a write lock and hold it until the Stripe request
completes. This prevents two concurrent calls to `get_or_insert_price`
from inserting the same data twice.

Creating a new meter/price is unusual, so in practice we'll acquire a
read lock most of the time.

/cc @rtfeldman @maxdeviant 

Release Notes:

- N/A

crates/collab/src/stripe_billing.rs

@@ -8,11 +8,16 @@ use serde::{Deserialize, Serialize};
 use tokio::sync::RwLock;
 
 pub struct StripeBilling {
-    meters_by_event_name: RwLock<HashMap<String, StripeMeter>>,
-    price_ids_by_meter_id: RwLock<HashMap<String, stripe::PriceId>>,
+    state: RwLock<StripeBillingState>,
     client: Arc<stripe::Client>,
 }
 
+#[derive(Default)]
+struct StripeBillingState {
+    meters_by_event_name: HashMap<String, StripeMeter>,
+    price_ids_by_meter_id: HashMap<String, stripe::PriceId>,
+}
+
 pub struct StripeModel {
     input_tokens_price: StripeBillingPrice,
     input_cache_creation_tokens_price: StripeBillingPrice,
@@ -29,36 +34,36 @@ impl StripeBilling {
     pub fn new(client: Arc<stripe::Client>) -> Self {
         Self {
             client,
-            meters_by_event_name: RwLock::new(HashMap::default()),
-            price_ids_by_meter_id: RwLock::new(HashMap::default()),
+            state: RwLock::default(),
         }
     }
 
     pub async fn initialize(&self) -> Result<()> {
-        log::info!("initializing StripeBilling");
+        log::info!("StripeBilling: initializing");
 
-        {
-            let meters = StripeMeter::list(&self.client).await?.data;
-            let mut meters_by_event_name = self.meters_by_event_name.write().await;
-            for meter in meters {
-                meters_by_event_name.insert(meter.event_name.clone(), meter);
-            }
+        let mut state = self.state.write().await;
+
+        let (meters, prices) = futures::try_join!(
+            StripeMeter::list(&self.client),
+            stripe::Price::list(&self.client, &stripe::ListPrices::default())
+        )?;
+
+        for meter in meters.data {
+            state
+                .meters_by_event_name
+                .insert(meter.event_name.clone(), meter);
         }
 
-        {
-            let prices = stripe::Price::list(&self.client, &stripe::ListPrices::default())
-                .await?
-                .data;
-            let mut price_ids_by_meter_id = self.price_ids_by_meter_id.write().await;
-            for price in prices {
+        for price in prices.data {
             if let Some(recurring) = price.recurring {
                 if let Some(meter) = recurring.meter {
-                        price_ids_by_meter_id.insert(meter, price.id);
-                    }
+                    state.price_ids_by_meter_id.insert(meter, price.id);
                 }
             }
         }
 
+        log::info!("StripeBilling: initialized");
+
         Ok(())
     }
 
@@ -105,8 +110,21 @@ impl StripeBilling {
         price_description: &str,
         price_per_million_tokens: Cents,
     ) -> Result<StripeBillingPrice> {
-        let meter =
-            if let Some(meter) = self.meters_by_event_name.read().await.get(meter_event_name) {
+        // Fast code path when the meter and the price already exist.
+        {
+            let state = self.state.read().await;
+            if let Some(meter) = state.meters_by_event_name.get(meter_event_name) {
+                if let Some(price_id) = state.price_ids_by_meter_id.get(&meter.id) {
+                    return Ok(StripeBillingPrice {
+                        id: price_id.clone(),
+                        meter_event_name: meter_event_name.to_string(),
+                    });
+                }
+            }
+        }
+
+        let mut state = self.state.write().await;
+        let meter = if let Some(meter) = state.meters_by_event_name.get(meter_event_name) {
             meter.clone()
         } else {
             let meter = StripeMeter::create(
@@ -118,15 +136,13 @@ impl StripeBilling {
                 },
             )
             .await?;
-                self.meters_by_event_name
-                    .write()
-                    .await
+            state
+                .meters_by_event_name
                 .insert(meter_event_name.to_string(), meter.clone());
             meter
         };
 
-        let price_id =
-            if let Some(price_id) = self.price_ids_by_meter_id.read().await.get(&meter.id) {
+        let price_id = if let Some(price_id) = state.price_ids_by_meter_id.get(&meter.id) {
             price_id.clone()
         } else {
             let price = stripe::Price::create(
@@ -172,9 +188,8 @@ impl StripeBilling {
                 },
             )
             .await?;
-                self.price_ids_by_meter_id
-                    .write()
-                    .await
+            state
+                .price_ids_by_meter_id
                 .insert(meter.id, price.id.clone());
             price.id
         };