backup/wip

Author: gouslu

rust/otap-dataflow/crates/otap/src/experimental/azure_monitor_exporter/auth.rs

@@ -8,14 +8,16 @@ use azure_identity::{
     ManagedIdentityCredentialOptions, UserAssignedId,
 };
 use std::sync::Arc;
+use tokio::sync::RwLock;
 
 use crate::experimental::azure_monitor_exporter::config::{AuthConfig, AuthMethod};
 
 #[derive(Clone)]
 pub struct Auth {
     credential: Arc<dyn TokenCredential>,
     scope: String,
-    cached_token: AccessToken,
+    // Thread-safe shared token cache
+    cached_token: Arc<RwLock<Option<AccessToken>>>,
 }
 
 impl Auth {
@@ -25,31 +27,39 @@ impl Auth {
         Ok(Self {
             credential,
             scope: auth_config.scope.clone(),
-            cached_token: AccessToken {
-                token: "".into(),
-                expires_on: OffsetDateTime::now_utc(),
-            },
+            cached_token: Arc::new(RwLock::new(None)),
         })
     }
 
     pub fn from_credential(credential: Arc<dyn TokenCredential>, scope: String) -> Self {
         Self {
             credential,
             scope,
-            cached_token: AccessToken {
-                token: "".into(),
-                expires_on: OffsetDateTime::now_utc(),
-            },
+            cached_token: Arc::new(RwLock::new(None)),
         }
     }
 
-    pub async fn get_token(&mut self) -> Result<AccessToken, String> {
+    pub async fn get_token(&self) -> Result<AccessToken, String> {
         println!("[AzureMonitorExporter][Auth] Acquiring token");
 
-        if self.cached_token.expires_on
-            > OffsetDateTime::now_utc() + azure_core::time::Duration::minutes(5)
+        // Try to use cached token
         {
-            return Ok(self.cached_token.clone());
+            let cached = self.cached_token.read().await;
+            if let Some(token) = &*cached {
+                if token.expires_on > OffsetDateTime::now_utc() {
+                    return Ok(token.clone());
+                }
+            }
+        }
+
+        // Need to refresh - acquire write lock
+        let mut cached = self.cached_token.write().await;
+
+        // Double-check in case another thread refreshed while we waited
+        if let Some(token) = &*cached {
+            if token.expires_on > OffsetDateTime::now_utc() {
+                return Ok(token.clone());
+            }
         }
 
         let token_response = self
@@ -62,16 +72,14 @@ impl Auth {
             .map_err(|e| format!("Failed to get token: {e}"))?;
 
         // Update the cached token
-        self.cached_token = token_response.clone();
+        *cached = Some(token_response.clone());
 
         Ok(token_response)
     }
 
-    pub fn invalidate_token(&mut self) {
-        self.cached_token = AccessToken {
-            token: "".into(),
-            expires_on: OffsetDateTime::now_utc(),
-        };
+    pub async fn invalidate_token(&self) {
+        let mut cached = self.cached_token.write().await;
+        *cached = None;
     }
 
     #[allow(clippy::print_stdout)]
@@ -148,7 +156,7 @@ mod tests {
             call_count: call_count.clone(),
         });
 
-        let mut auth = Auth::from_credential(credential, "scope".to_string());
+        let auth = Auth::from_credential(credential, "scope".to_string());
 
         // First call should hit the credential
         let token1 = auth.get_token().await.unwrap();
@@ -171,7 +179,7 @@ mod tests {
             call_count: call_count.clone(),
         });
 
-        let mut auth = Auth::from_credential(credential, "scope".to_string());
+        let auth = Auth::from_credential(credential, "scope".to_string());
 
         // First call
         let _ = auth.get_token().await.unwrap();
@@ -191,22 +199,22 @@ mod tests {
             call_count: call_count.clone(),
         });
 
-        let mut auth = Auth::from_credential(credential, "scope".to_string());
+        let auth = Auth::from_credential(credential, "scope".to_string());
 
         // First call
         let _ = auth.get_token().await.unwrap();
         assert_eq!(*call_count.lock().unwrap(), 1);
 
         // Invalidate
-        auth.invalidate_token();
+        auth.invalidate_token().await;
 
         // Should refresh
         let _ = auth.get_token().await.unwrap();
         assert_eq!(*call_count.lock().unwrap(), 2);
     }
 
-    #[test]
-    fn test_new_with_managed_identity() {
+    #[tokio::test]
+    async fn test_new_with_managed_identity() {
         let auth_config = AuthConfig {
             method: AuthMethod::ManagedIdentity,
             client_id: Some("test-client-id".to_string()),
@@ -217,11 +225,12 @@ mod tests {
         assert!(auth.is_ok());
         let auth = auth.unwrap();
         assert_eq!(auth.scope, "https://test.scope");
-        assert_eq!(auth.cached_token.token.secret(), "");
+        // Check that cached_token is None initially
+        assert!(auth.cached_token.read().await.is_none());
     }
 
-    #[test]
-    fn test_new_with_system_assigned_managed_identity() {
+    #[tokio::test]
+    async fn test_new_with_system_assigned_managed_identity() {
         let auth_config = AuthConfig {
             method: AuthMethod::ManagedIdentity,
             client_id: None,
@@ -232,11 +241,12 @@ mod tests {
         assert!(auth.is_ok());
         let auth = auth.unwrap();
         assert_eq!(auth.scope, "https://test.scope");
-        assert_eq!(auth.cached_token.token.secret(), "");
+        // Check that cached_token is None initially
+        assert!(auth.cached_token.read().await.is_none());
     }
 
-    #[test]
-    fn test_new_with_development_auth() {
+    #[tokio::test]
+    async fn test_new_with_development_auth() {
         let auth_config = AuthConfig {
             method: AuthMethod::Development,
             client_id: None,
@@ -248,7 +258,8 @@ mod tests {
         match auth {
             Ok(auth) => {
                 assert_eq!(auth.scope, "https://test.scope");
-                assert_eq!(auth.cached_token.token.secret(), "");
+                // Check that cached_token is None initially
+                assert!(auth.cached_token.read().await.is_none());
             }
             Err(err) => {
                 // Expected if Azure CLI/Azure Developer CLI is not installed
@@ -257,6 +268,21 @@ mod tests {
         }
     }
 
+    #[tokio::test]
+    async fn test_from_credential_initializes_with_none() {
+        let credential = Arc::new(MockCredential {
+            token: "test_token".to_string(),
+            expires_in: azure_core::time::Duration::minutes(60),
+            call_count: Arc::new(Mutex::new(0)),
+        });
+
+        let auth = Auth::from_credential(credential, "test_scope".to_string());
+
+        // Check that initial cached token is None
+        assert!(auth.cached_token.read().await.is_none());
+        assert_eq!(auth.scope, "test_scope");
+    }
+
     #[tokio::test]
     async fn test_get_token_with_exactly_5_minute_buffer() {
         let call_count = Arc::new(Mutex::new(0));
@@ -267,7 +293,7 @@ mod tests {
             call_count: call_count.clone(),
         });
 
-        let mut auth = Auth::from_credential(credential, "scope".to_string());
+        let auth = Auth::from_credential(credential, "scope".to_string());
 
         // First call
         let _ = auth.get_token().await.unwrap();
@@ -298,7 +324,7 @@ mod tests {
         }
 
         let credential = Arc::new(FailingCredential);
-        let mut auth = Auth::from_credential(credential, "scope".to_string());
+        let auth = Auth::from_credential(credential, "scope".to_string());
 
         let result = auth.get_token().await;
         assert!(result.is_err());
@@ -314,7 +340,7 @@ mod tests {
             call_count: call_count.clone(),
         });
 
-        let mut auth = Auth::from_credential(credential, "scope".to_string());
+        let auth = Auth::from_credential(credential, "scope".to_string());
 
         // Get token twice and verify they are different instances (cloned)
         let token1 = auth.get_token().await.unwrap();
@@ -326,8 +352,8 @@ mod tests {
         assert_eq!(*call_count.lock().unwrap(), 1);
     }
 
-    #[test]
-    fn test_from_credential_initializes_with_expired_token() {
+    #[tokio::test]
+    async fn test_from_credential_initializes_with_expired_token() {
         let credential = Arc::new(MockCredential {
             token: "test_token".to_string(),
             expires_in: azure_core::time::Duration::minutes(60),
@@ -336,9 +362,9 @@ mod tests {
 
         let auth = Auth::from_credential(credential, "test_scope".to_string());
 
-        // Check that initial cached token is expired
-        assert_eq!(auth.cached_token.token.secret(), "");
-        assert!(auth.cached_token.expires_on <= OffsetDateTime::now_utc());
+        // Check that initial cached token is None
+        let cached = auth.cached_token.read().await;
+        assert!(cached.is_none(), "Initial cached token should be None");
         assert_eq!(auth.scope, "test_scope");
     }
 }

rust/otap-dataflow/crates/otap/src/experimental/azure_monitor_exporter/client.rs

@@ -3,6 +3,9 @@
 
 use azure_core::credentials::TokenCredential;
 use azure_core::time::OffsetDateTime;
+
+use bytes::Bytes;
+
 use reqwest::{
     Client,
     header::{AUTHORIZATION, CONTENT_ENCODING, CONTENT_TYPE, HeaderValue},
@@ -28,6 +31,8 @@ pub struct LogsIngestionClient {
 
     /// Token expiry time using monotonic clock for faster comparisons
     pub token_valid_until: Instant,
+
+    token_refresh_after: Instant,
 }
 
 impl LogsIngestionClient {
@@ -53,6 +58,7 @@ impl LogsIngestionClient {
             auth: Auth::from_credential(credential, scope),
             auth_header: HeaderValue::from_static("Bearer "),
             token_valid_until: Instant::now(),
+            token_refresh_after: Instant::now(),
         }
     }
 
@@ -79,14 +85,16 @@ impl LogsIngestionClient {
         );
 
         let auth =
-            Auth::new(&config.auth).map_err(|e| format!("Failed to create auth handler: {e}"))?;
+            Auth::new(&config.auth)
+                .map_err(|e| format!("Failed to create auth handler: {e}"))?;
 
         Ok(Self {
             http_client,
             endpoint,
             auth,
             auth_header: HeaderValue::from_static("Bearer "),
             token_valid_until: Instant::now(),
+            token_refresh_after: Instant::now(),
         })
     }
 
@@ -96,7 +104,7 @@ impl LogsIngestionClient {
         let now = Instant::now();
 
         // Fast path: token is still valid
-        if now < self.token_valid_until {
+        if now < self.token_refresh_after {
             return Ok(());
         }
 
@@ -114,10 +122,10 @@ impl LogsIngestionClient {
         // Calculate validity using Instant for faster comparisons
         // Refresh 5 minutes before expiry
         let valid_seconds = (token.expires_on - OffsetDateTime::now_utc())
-            .whole_seconds()
-            .saturating_sub(300); // 5 minutes = 300 seconds
+            .whole_seconds();
 
         self.token_valid_until = now + Duration::from_secs(valid_seconds.max(0) as u64);
+        self.token_refresh_after = self.token_valid_until - Duration::from_secs(300);
 
         println!("[AzureMonitorExporter] Acquired new token, valid for {} seconds, valid until {:?}, current time {:?}", valid_seconds, self.token_valid_until, now);
 
@@ -134,10 +142,7 @@ impl LogsIngestionClient {
     /// # Returns
     /// * `Ok(())` - If the request was successful
     /// * `Err(String)` - Error message if the request failed
-    pub async fn send(&mut self, body: Vec<u8>) -> Result<(), String> {
-        // Ensure we have a valid token (fast path when cached)
-        self.ensure_valid_token().await?;
-
+    pub async fn send(&mut self, body: Bytes) -> Result<(), String> {
         let start = Instant::now();
 
         // Send compressed body - avoid cloning headers by setting them individually
@@ -168,7 +173,9 @@ impl LogsIngestionClient {
             401 => {
                 // Invalidate token and force refresh on next call
                 self.token_valid_until = Instant::now();
-                self.auth.invalidate_token();
+                self.auth.invalidate_token().await;
+                self.ensure_valid_token().await?;
+
                 Err(format!("Authentication failed: {error}"))
             }
             403 => Err(format!("Authorization failed: {error}")),
@@ -183,6 +190,7 @@ impl LogsIngestionClient {
 mod tests {
     use super::*;
     use super::super::config::{ApiConfig, AuthConfig, AuthMethod};
+    use azure_core::Bytes;
     use azure_core::credentials::TokenRequestOptions;
     use azure_core::credentials::{AccessToken, TokenCredential};
     use std::sync::Mutex;
@@ -277,7 +285,7 @@ mod tests {
             "scope".to_string(),
         );
 
-        let result = client.send(vec![1, 2, 3]).await;
+        let result = client.send(Bytes::from(vec![1, 2, 3])).await;
         assert!(result.is_ok());
         assert_eq!(*call_count.lock().unwrap(), 1); // Token fetched once
     }
@@ -312,7 +320,7 @@ mod tests {
         );
 
         // This should fail with 401, but invalidate the token
-        let result = client.send(vec![1, 2, 3]).await;
+        let result = client.send(Bytes::from(vec![1, 2, 3])).await;
         assert!(result.is_err());
         assert!(result.unwrap_err().contains("Authentication failed"));
 
@@ -326,7 +334,7 @@ mod tests {
             .mount(&mock_server)
             .await;
 
-        let result = client.send(vec![1, 2, 3]).await;
+        let result = client.send(Bytes::from(vec![1, 2, 3])).await;
         assert!(result.is_ok());
         assert_eq!(*call_count.lock().unwrap(), 2); // Token fetched again
     }
@@ -357,7 +365,7 @@ mod tests {
             "scope".to_string(),
         );
 
-        let result = client.send(vec![1, 2, 3]).await;
+        let result = client.send(Bytes::from(vec![1, 2, 3])).await;
         assert!(result.is_err());
         assert!(result.unwrap_err().contains("Rate limited"));
     }