Micro-optimize GCS token access (#1353)

Avoid `format!`-ing a String for each request, and remove an `Arc` indirection.

crates/symbolicator-service/src/download/gcs.rs

@@ -10,7 +10,7 @@ use crate::utils::gcs::{self, GcsToken};
 use crate::utils::http::DownloadTimeouts;
 
 /// An LRU cache for GCS OAuth tokens.
-type GcsTokenCache = moka::future::Cache<Arc<GcsSourceKey>, CacheEntry<Arc<GcsToken>>>;
+type GcsTokenCache = moka::future::Cache<Arc<GcsSourceKey>, CacheEntry<GcsToken>>;
 
 /// Downloader implementation that supports the GCS source.
 #[derive(Debug)]
@@ -35,16 +35,16 @@ impl GcsDownloader {
     ///
     /// If the cache contains a valid token, then this token is returned. Otherwise, a new token is
     /// requested from GCS and stored in the cache.
-    async fn get_token(&self, source_key: &Arc<GcsSourceKey>) -> CacheEntry<Arc<GcsToken>> {
+    async fn get_token(&self, source_key: &Arc<GcsSourceKey>) -> CacheEntry<GcsToken> {
         metric!(counter("source.gcs.token.access") += 1);
 
         let init = Box::pin(async {
             metric!(counter("source.gcs.token.computation") += 1);
             let token = gcs::request_new_token(&self.client, source_key).await;
-            token.map(Arc::new).map_err(CacheError::from)
+            token.map_err(CacheError::from)
         });
         let replace_if =
-            |entry: &CacheEntry<Arc<GcsToken>>| entry.as_ref().map_or(true, |t| t.is_expired());
+            |entry: &CacheEntry<GcsToken>| entry.as_ref().map_or(true, |t| t.is_expired());
 
         self.token_cache
             .entry_by_ref(source_key)

crates/symbolicator-service/src/utils/gcs.rs

@@ -1,6 +1,9 @@
 //! Access to Google Cloud Storeage
 
+use std::sync::Arc;
+
 use chrono::{DateTime, Duration, Utc};
+use jsonwebtoken::errors::Error as JwtError;
 use jsonwebtoken::EncodingKey;
 use reqwest::Client;
 use serde::{Deserialize, Serialize};
@@ -10,9 +13,9 @@ use url::Url;
 use symbolicator_sources::GcsSourceKey;
 
 /// A JWT token usable for GCS.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct GcsToken {
-    access_token: String,
+    bearer_token: Arc<str>,
     expires_at: DateTime<Utc>,
 }
 
@@ -23,18 +26,18 @@ impl GcsToken {
     }
 
     /// The token in the HTTP Bearer-header format, header value only.
-    pub fn bearer_token(&self) -> String {
-        format!("Bearer {}", self.access_token)
+    pub fn bearer_token(&self) -> &str {
+        &self.bearer_token
     }
 }
 
 #[derive(Serialize)]
-struct JwtClaims {
+struct JwtClaims<'s> {
     #[serde(rename = "iss")]
-    issuer: String,
-    scope: String,
+    issuer: &'s str,
+    scope: &'s str,
     #[serde(rename = "aud")]
-    audience: String,
+    audience: &'s str,
     #[serde(rename = "exp")]
     expiration: i64,
     #[serde(rename = "iat")]
@@ -57,22 +60,11 @@ pub enum GcsError {
     #[error("failed to construct URL")]
     InvalidUrl,
     #[error("failed encoding JWT")]
-    Jwt(#[from] jsonwebtoken::errors::Error),
+    Jwt(#[from] JwtError),
     #[error("failed to send authentication request")]
     Auth(#[source] reqwest::Error),
 }
 
-/// Returns the JWT key parsed from a string.
-///
-/// Because Google provides this key in JSON format a lot of users just copy-paste this key
-/// directly, leaving the escaped newlines from the JSON-encoding in place.  In normal
-/// base64 this should not occur so we pre-process the key to convert these back to real
-/// newlines, ensuring they are in the correct PEM format.
-fn key_from_string(key: &str) -> Result<EncodingKey, jsonwebtoken::errors::Error> {
-    let buffer = key.replace("\\n", "\n");
-    EncodingKey::from_rsa_pem(buffer.as_bytes())
-}
-
 /// Returns the URL for an object.
 ///
 /// This can be used to e.g. fetch the objects's metadata.
@@ -92,21 +84,32 @@ pub fn download_url(bucket: &str, object: &str) -> Result<Url, GcsError> {
     Ok(url)
 }
 
+/// Returns the JWT key parsed from a string.
+///
+/// Because Google provides this key in JSON format a lot of users just copy-paste this key
+/// directly, leaving the escaped newlines from the JSON-encoding in place.  In normal
+/// base64 this should not occur so we pre-process the key to convert these back to real
+/// newlines, ensuring they are in the correct PEM format.
+fn key_from_string(key: &str) -> Result<EncodingKey, JwtError> {
+    let buffer = key.replace("\\n", "\n");
+    EncodingKey::from_rsa_pem(buffer.as_bytes())
+}
+
 /// Computes a JWT authentication assertion for the given GCS bucket.
-fn get_auth_jwt(source_key: &GcsSourceKey, expiration: i64) -> Result<String, GcsError> {
+fn get_auth_jwt(source_key: &GcsSourceKey, expiration: i64) -> Result<String, JwtError> {
     let header = jsonwebtoken::Header::new(jsonwebtoken::Algorithm::RS256);
 
     let jwt_claims = JwtClaims {
-        issuer: source_key.client_email.clone(),
-        scope: "https://www.googleapis.com/auth/devstorage.read_only".into(),
-        audience: "https://www.googleapis.com/oauth2/v4/token".into(),
+        issuer: &source_key.client_email,
+        scope: "https://www.googleapis.com/auth/devstorage.read_only",
+        audience: "https://www.googleapis.com/oauth2/v4/token",
         expiration,
         issued_at: Utc::now().timestamp(),
     };
 
     let key = key_from_string(&source_key.private_key)?;
 
-    Ok(jsonwebtoken::encode(&header, &jwt_claims, &key)?)
+    jsonwebtoken::encode(&header, &jwt_claims, &key)
 }
 
 /// Requests a new GCS OAuth token.
@@ -133,9 +136,10 @@ pub async fn request_new_token(
         .json::<GcsTokenResponse>()
         .await
         .map_err(GcsError::Auth)?;
+    let bearer_token = format!("Bearer {}", token.access_token).into();
 
     Ok(GcsToken {
-        access_token: token.access_token,
+        bearer_token,
         expires_at,
     })
 }