@@ -1,5 +1,8 @@
1+use std::{cmp::min, time::Duration};
2+
13use anyhow::{anyhow, Context};
24use postgres::{types::ToSql, NoTls, Row};
5+use rand::prelude::*;
36
47// DbLocker wraps a regular DbConn, only allowing access using the
58// `lock` method. This method will acquire the advisory lock before
@@ -103,12 +106,12 @@ impl DbConn {
103106
104107impl Conn for DbConn {
105108 fn run(&mut self, query: &str) -> anyhow::Result<()> {
106− self.client.batch_execute(query)?;
109+ retry_automatically(|| self.client.batch_execute(query))?;
107110 Ok(())
108111 }
109112
110113 fn query(&mut self, query: &str) -> anyhow::Result<Vec<Row>> {
111− let rows = self.client.query(query, &[])?;
114+ let rows = retry_automatically(|| self.client.query(query, &[]))?;
112115 Ok(rows)
113116 }
114117
@@ -117,7 +120,7 @@ impl Conn for DbConn {
117120 query: &str,
118121 params: &[&(dyn ToSql + Sync)],
119122 ) -> anyhow::Result<Vec<Row>> {
120− let rows = self.client.query(query, params)?;
123+ let rows = retry_automatically(|| self.client.query(query, params))?;
121124 Ok(rows)
122125 }
123126
@@ -168,3 +171,57 @@ impl Conn for Transaction<'_> {
168171 Ok(Transaction { transaction })
169172 }
170173}
174+
175+// Retry a database operation with exponential backoff and jitter
176+fn retry_automatically<T>(
177+ mut f: impl FnMut() -> Result<T, postgres::Error>,
178+) -> Result<T, postgres::Error> {
179+ const STARTING_WAIT_TIME: u64 = 100;
180+ const MAX_WAIT_TIME: u64 = 3_200;
181+ const MAX_ATTEMPTS: u32 = 10;
182+
183+ let mut rng = rand::thread_rng();
184+ let mut attempts = 0;
185+ loop {
186+ let result = f();
187+
188+ let error = match result {
189+ Ok(_) => return result,
190+ Err(err) => err,
191+ };
192+
193+ // If we got a database error, we check if it's retryable.
194+ // If we didn't get a database error, then it's most likely some kind of connection
195+ // error which should also be retried.
196+ if let Some(db_error) = error.as_db_error() {
197+ if !error_retryable(db_error) {
198+ return Err(error);
199+ }
200+ }
201+
202+ attempts += 1;
203+ if attempts >= MAX_ATTEMPTS {
204+ return Err(error);
205+ }
206+
207+ // The wait time increases exponentially, starting at 100ms and doubling up to a max of 3.2s.
208+ let wait_time = min(
209+ MAX_WAIT_TIME,
210+ STARTING_WAIT_TIME * u64::pow(2, attempts - 1),
211+ );
212+
213+ // The jitter is up to half the wait time
214+ let jitter: u64 = rng.gen_range(0..wait_time / 2);
215+
216+ std::thread::sleep(Duration::from_millis(wait_time + jitter));
217+ }
218+}
219+
220+// Check if a database error can be retried
221+fn error_retryable(error: &postgres::error::DbError) -> bool {
222+ match error.code() {
223+ // Caused by lock_timeout being exceeded
224+ &postgres::error::SqlState::LOCK_NOT_AVAILABLE => true,
225+ _ => false,
226+ }
227+}