feat: Add timeouts for tokenserver database calls. (#1561)

Closes: SYNC-4270

syncserver/src/tokenserver/handlers.rs

@@ -13,6 +13,7 @@ use tokenserver_db::{
     params::{GetNodeId, PostUser, PutUser, ReplaceUsers},
     Db,
 };
+use tokio::time::timeout;
 
 use super::{
     extractors::{DbWrapper, TokenserverRequest},
@@ -120,6 +121,24 @@ struct UserUpdates {
     uid: i64,
 }
 
+/// Optionally apply a timeout when running a future
+async fn apply_timeout<O, E>(
+    duration: Option<Duration>,
+    fut: impl std::future::Future<Output = Result<O, E>>,
+) -> Result<O, TokenserverError>
+where
+    E: Into<TokenserverError>,
+{
+    match duration {
+        Some(duration) => match timeout(duration, fut).await {
+            Ok(Ok(result)) => Ok(result),
+            Ok(Err(e)) => Err(e.into()),
+            Err(_) => Err(TokenserverError::elapsed()),
+        },
+        None => fut.await.map_err(Into::into),
+    }
+}
+
 async fn update_user(
     req: &TokenserverRequest,
     db: Box<dyn Db>,
@@ -197,25 +216,32 @@ async fn update_user(
             email: req.auth_data.email.clone(),
             generation,
             client_state: req.auth_data.client_state.clone(),
-            node_id: db
-                .get_node_id(GetNodeId {
+            node_id: apply_timeout(
+                db.timeout(),
+                db.get_node_id(GetNodeId {
                     service_id: req.service_id,
                     node: req.user.node.clone(),
-                })
-                .await?
-                .id,
+                }),
+            )
+            .await?
+            .id,
             keys_changed_at,
             created_at: timestamp,
         };
-        let uid = db.post_user(post_user_params).await?.id;
+        let uid = apply_timeout(db.timeout(), db.post_user(post_user_params))
+            .await?
+            .id;
 
         // Make sure each old row is marked as replaced (they might not be, due to races in row
         // creation)
-        db.replace_users(ReplaceUsers {
-            email: req.auth_data.email.clone(),
-            service_id: req.service_id,
-            replaced_at: timestamp,
-        })
+        apply_timeout(
+            db.timeout(),
+            db.replace_users(ReplaceUsers {
+                email: req.auth_data.email.clone(),
+                service_id: req.service_id,
+                replaced_at: timestamp,
+            }),
+        )
         .await?;
 
         Ok(UserUpdates {
@@ -232,7 +258,7 @@ async fn update_user(
                 keys_changed_at,
             };
 
-            db.put_user(params).await?;
+            apply_timeout(db.timeout(), db.put_user(params)).await?;
         }
 
         Ok(UserUpdates {
@@ -250,7 +276,7 @@ pub async fn heartbeat(DbWrapper(db): DbWrapper) -> Result<HttpResponse, Error>
         Value::String(env!("CARGO_PKG_VERSION").to_owned()),
     );
 
-    match db.check().await {
+    match apply_timeout(db.timeout(), db.check()).await {
         Ok(result) => {
             if result {
                 checklist.insert("database".to_owned(), Value::from("Ok"));

tokenserver-common/Cargo.toml

@@ -12,5 +12,6 @@ serde.workspace = true
 serde_json.workspace = true
 jsonwebtoken.workspace = true
 thiserror.workspace = true
+tokio.workspace = true
 
 syncserver-common = { path = "../syncserver-common" }

tokenserver-common/src/error.rs

@@ -131,6 +131,17 @@ impl TokenserverError {
         }
     }
 
+    pub fn elapsed() -> Self {
+        Self {
+            status: "elapsed",
+            location: ErrorLocation::Body,
+            description: "Elapsed".to_owned(),
+            http_status: StatusCode::GATEWAY_TIMEOUT,
+            context: "Elapsed".to_owned(),
+            ..Self::default()
+        }
+    }
+
     pub fn resource_unavailable() -> Self {
         Self {
             location: ErrorLocation::Body,

tokenserver-db/src/models.rs

@@ -12,7 +12,7 @@ use syncserver_db_common::{sync_db_method, DbFuture};
 
 use std::{
     sync::Arc,
-    time::{SystemTime, UNIX_EPOCH},
+    time::{Duration, SystemTime, UNIX_EPOCH},
 };
 
 use super::{
@@ -40,6 +40,7 @@ pub struct TokenserverDb {
     service_id: Option<i32>,
     spanner_node_id: Option<i32>,
     blocking_threadpool: Arc<BlockingThreadpool>,
+    pub timeout: Option<Duration>,
 }
 
 /// Despite the db conn structs being !Sync (see Arc<MysqlDbInner> above) we
@@ -68,6 +69,7 @@ impl TokenserverDb {
         service_id: Option<i32>,
         spanner_node_id: Option<i32>,
         blocking_threadpool: Arc<BlockingThreadpool>,
+        timeout: Option<Duration>,
     ) -> Self {
         let inner = DbInner {
             #[cfg(not(test))]
@@ -84,6 +86,7 @@ impl TokenserverDb {
             service_id,
             spanner_node_id,
             blocking_threadpool,
+            timeout,
         }
     }
 
@@ -683,6 +686,10 @@ impl Db for TokenserverDb {
     sync_db_method!(get_or_create_user, get_or_create_user_sync, GetOrCreateUser);
     sync_db_method!(get_service_id, get_service_id_sync, GetServiceId);
 
+    fn timeout(&self) -> Option<Duration> {
+        self.timeout
+    }
+
     #[cfg(test)]
     sync_db_method!(get_user, get_user_sync, GetUser);
 
@@ -722,6 +729,10 @@ impl Db for TokenserverDb {
 }
 
 pub trait Db {
+    fn timeout(&self) -> Option<Duration> {
+        None
+    }
+
     fn replace_user(
         &self,
         params: params::ReplaceUser,

tokenserver-db/src/pool.rs

@@ -41,6 +41,7 @@ pub struct TokenserverPool {
     pub service_id: Option<i32>,
     spanner_node_id: Option<i32>,
     blocking_threadpool: Arc<BlockingThreadpool>,
+    pub timeout: Option<Duration>,
 }
 
 impl TokenserverPool {
@@ -68,13 +69,17 @@ impl TokenserverPool {
         } else {
             builder
         };
+        let timeout = settings
+            .database_request_timeout
+            .map(|v| Duration::from_secs(v as u64));
 
         Ok(Self {
             inner: builder.build(manager)?,
             metrics: metrics.clone(),
             spanner_node_id: settings.spanner_node_id,
             service_id: None,
             blocking_threadpool,
+            timeout,
         })
     }
 
@@ -87,6 +92,7 @@ impl TokenserverPool {
             self.service_id,
             self.spanner_node_id,
             self.blocking_threadpool.clone(),
+            self.timeout,
         ))
     }
 
@@ -104,6 +110,7 @@ impl TokenserverPool {
             self.service_id,
             self.spanner_node_id,
             self.blocking_threadpool.clone(),
+            self.timeout,
         ))
     }
 }
@@ -126,6 +133,7 @@ impl DbPool for TokenserverPool {
             self.service_id,
             self.spanner_node_id,
             self.blocking_threadpool.clone(),
+            self.timeout,
         )) as Box<dyn Db>)
     }
 

tokenserver-settings/src/lib.rs

@@ -14,6 +14,8 @@ pub struct Settings {
     pub database_pool_min_idle: Option<u32>,
     /// Pool timeout when waiting for a slot to become available, in seconds
     pub database_pool_connection_timeout: Option<u32>,
+    /// Database request timeout, in seconds
+    pub database_request_timeout: Option<u32>,
     // XXX: This is a temporary setting used to enable Tokenserver-related features. In
     // the future, Tokenserver will always be enabled, and this setting will be
     // removed.
@@ -78,6 +80,7 @@ impl Default for Settings {
             database_pool_max_size: 10,
             database_pool_min_idle: None,
             database_pool_connection_timeout: Some(30),
+            database_request_timeout: None,
             enabled: false,
             fxa_email_domain: "api-accounts.stage.mozaws.net".to_owned(),
             fxa_metrics_hash_secret: "secret".to_owned(),