fix(gcp_pubsub source): Fix handling of auth token (#12757)

* Regenerate the GCP auth token only *after* one interval

* Break the receive loop immediately after a receive error

* Fix handling of credentials generation in gcp_pubsub source

The gcp_pubsub source would take the credentials on the first time
through the streaming pull loop, in order to start the token regenerate
task, leaving them empty if it needed to restart the loop. This drops
the `take` and moves the token regenerator into a stream that can be
polled in the main request loop, exiting the loop and restarting the
stream when the token is regenerated.

Author: bruceg

src/gcp.rs

@@ -1,7 +1,7 @@
 use std::sync::{Arc, RwLock};
 use std::time::Duration;
 
-use futures::StreamExt;
+use futures::{Stream, StreamExt};
 pub use goauth::scopes::Scope;
 use goauth::{
     auth::{JwtClaims, Token, TokenErr},
@@ -13,6 +13,7 @@ use once_cell::sync::Lazy;
 use serde::{Deserialize, Serialize};
 use smpl_jwt::Jwt;
 use snafu::{ResultExt, Snafu};
+use tokio::time::Instant;
 use tokio_stream::wrappers::IntervalStream;
 
 use crate::{config::ProxyConfig, http::HttpClient, http::HttpError};
@@ -70,10 +71,13 @@ impl GcpAuthConfig {
 }
 
 #[derive(Clone, Debug)]
-pub struct GcpCredentials {
+pub struct GcpCredentials(Arc<Inner>);
+
+#[derive(Debug)]
+struct Inner {
     creds: Option<Credentials>,
     scope: Scope,
-    token: Arc<RwLock<Token>>,
+    token: RwLock<Token>,
 }
 
 impl GcpCredentials {
@@ -83,24 +87,24 @@ impl GcpCredentials {
         let token = goauth::get_token(&jwt, &creds)
             .await
             .context(GetTokenSnafu)?;
-        Ok(Self {
+        Ok(Self(Arc::new(Inner {
             creds: Some(creds),
             scope,
-            token: Arc::new(RwLock::new(token)),
-        })
+            token: RwLock::new(token),
+        })))
     }
 
     async fn new_implicit(scope: Scope) -> crate::Result<Self> {
         let token = get_token_implicit().await?;
-        Ok(Self {
+        Ok(Self(Arc::new(Inner {
             creds: None,
             scope,
-            token: Arc::new(RwLock::new(token)),
-        })
+            token: RwLock::new(token),
+        })))
     }
 
     pub fn make_token(&self) -> String {
-        let token = self.token.read().unwrap();
+        let token = self.0.token.read().unwrap();
         format!("{} {}", token.token_type(), token.access_token())
     }
 
@@ -111,35 +115,39 @@ impl GcpCredentials {
     }
 
     async fn regenerate_token(&self) -> crate::Result<()> {
-        let token = match &self.creds {
+        let token = match &self.0.creds {
             Some(creds) => {
-                let jwt = make_jwt(creds, &self.scope).unwrap(); // Errors caught above
+                let jwt = make_jwt(creds, &self.0.scope).unwrap(); // Errors caught above
                 goauth::get_token(&jwt, creds).await?
             }
             None => get_token_implicit().await?,
         };
-        *self.token.write().unwrap() = token;
+        *self.0.token.write().unwrap() = token;
         Ok(())
     }
 
     pub fn spawn_regenerate_token(&self) {
         let this = self.clone();
+        tokio::spawn(async move { this.token_regenerator().for_each(|_| async {}).await });
+    }
 
-        let period = this.token.read().unwrap().expires_in() as u64 / 2;
-        let interval = IntervalStream::new(tokio::time::interval(Duration::from_secs(period)));
-        let task = interval.for_each(move |_| {
-            let this = this.clone();
-            async move {
-                debug!("Renewing GCP authentication token.");
-                if let Err(error) = this.regenerate_token().await {
-                    error!(
-                        message = "Failed to update GCP authentication token.",
-                        %error
-                    );
+    pub fn token_regenerator(&self) -> impl Stream<Item = ()> + 'static {
+        let period = Duration::from_secs(self.0.token.read().unwrap().expires_in() as u64 / 2);
+        let this = self.clone();
+        IntervalStream::new(tokio::time::interval_at(Instant::now() + period, period)).then(
+            move |_| {
+                let this = this.clone();
+                async move {
+                    debug!("Renewing GCP authentication token.");
+                    if let Err(error) = this.regenerate_token().await {
+                        error!(
+                            message = "Failed to update GCP authentication token.",
+                            %error
+                        );
+                    }
                 }
-            }
-        });
-        tokio::spawn(task);
+            },
+        )
     }
 }
 

src/sources/gcp_pubsub.rs

@@ -1,4 +1,4 @@
-use std::{sync::Arc, time::Duration};
+use std::{pin::Pin, sync::Arc, time::Duration};
 
 use chrono::{DateTime, NaiveDateTime, Utc};
 use codecs::decoding::{DeserializerConfig, FramingConfig};
@@ -27,7 +27,8 @@ use crate::{
     },
     serde::{bool_or_struct, default_decoding, default_framing_message_based},
     shutdown::ShutdownSignal,
-    sources::util::{self, finalizer::UnorderedFinalizer},
+    sources::util,
+    sources::util::finalizer::{EmptyStream, UnorderedFinalizer},
     tls::{TlsConfig, TlsSettings},
     SourceSender,
 };
@@ -215,15 +216,24 @@ impl PubsubSource {
                 .context(EndpointTlsSnafu)?;
         }
 
-        while self.run_once(&endpoint).await {
-            info!(timeout_secs = 1, "Retrying after timeout");
+        let mut token_generator = match &self.credentials {
+            Some(credentials) => credentials.clone().token_regenerator().boxed(),
+            None => EmptyStream::default().boxed(),
+        };
+
+        while self.run_once(&endpoint, &mut token_generator).await {
+            info!(timeout_secs = 1, "Retrying after timeout.");
             tokio::time::sleep(self.retry_delay).await;
         }
 
         Ok(())
     }
 
-    async fn run_once(&mut self, endpoint: &Endpoint) -> bool {
+    async fn run_once(
+        &mut self,
+        endpoint: &Endpoint,
+        token_generator: &mut Pin<Box<dyn Stream<Item = ()> + Send>>,
+    ) -> bool {
         let connection = match endpoint.connect().await {
             Ok(connection) => connection,
             Err(error) => {
@@ -252,6 +262,7 @@ impl PubsubSource {
         // Handle shutdown during startup, the streaming pull doesn't
         // start if there is no data in the subscription.
         let request_stream = self.request_stream();
+        debug!("Starting streaming pull.");
         let stream = tokio::select! {
             _ = &mut self.shutdown => return false,
             result = client.streaming_pull(request_stream) => match result {
@@ -264,24 +275,27 @@ impl PubsubSource {
         };
         let mut stream = stream.into_inner();
 
-        if let Some(credentials) = self.credentials.take() {
-            credentials.spawn_regenerate_token();
-        }
-
         let (finalizer, mut ack_stream) =
             Finalizer::maybe_new(self.acknowledgements, self.shutdown.clone());
 
         loop {
             tokio::select! {
                 _ = &mut self.shutdown => return false,
+                _ = &mut token_generator.next() => {
+                    debug!("New authentication token generated, restarting stream.");
+                    return true;
+                },
                 receipts = ack_stream.next() => if let Some((status, receipts)) = receipts {
                     if status == BatchStatus::Delivered {
                         self.ack_ids.lock().await.extend(receipts);
                     }
                 },
                 response = stream.next() => match response {
                     Some(Ok(response)) => self.handle_response(response, &finalizer).await,
-                    Some(Err(error)) => emit!(GcpPubsubReceiveError { error }),
+                    Some(Err(error)) => {
+                        emit!(GcpPubsubReceiveError { error });
+                        break;
+                    }
                     None => break,
                 },
             }

src/sources/util/finalizer.rs

@@ -82,7 +82,7 @@ where
             let (finalizer, stream) = Self::new(shutdown);
             (Some(finalizer), stream.boxed())
         } else {
-            (None, EmptyStream(Default::default()).boxed())
+            (None, EmptyStream::default().boxed())
         }
     }
 
@@ -194,8 +194,9 @@ impl<T> Future for FinalizerFuture<T> {
     }
 }
 
-#[derive(Clone, Copy)]
-struct EmptyStream<T>(PhantomData<T>);
+#[derive(Clone, Copy, Derivative)]
+#[derivative(Default(bound = ""))]
+pub struct EmptyStream<T>(PhantomData<T>);
 
 impl<T> Stream for EmptyStream<T> {
     type Item = T;