git.iko.soy/reshape
on main help sign in
Add automatic retry of queries
40fb397fabianlindfors
Cargo.toml+2−1
compare with main
unifiedsplit
@@ -19,4 +19,5 @@ anyhow = { version = "1.0.44", features = ["backtrace"] }
19clap = { version = "3.0.0", features = ["derive"] }
19clap = { version = "3.0.0", features = ["derive"] }
20toml = "0.5"
20toml = "0.5"
21version = "3.0.0"
21version = "3.0.0"
22−colored = "2"
22+colored = "2"
23+rand = "0.8"
src/db.rs+60−3
@@ -1,5 +1,8 @@
1+use std::{cmp::min, time::Duration};
2+
1use anyhow::{anyhow, Context};
3use anyhow::{anyhow, Context};
2use postgres::{types::ToSql, NoTls, Row};
4use postgres::{types::ToSql, NoTls, Row};
5+use rand::prelude::*;
3
6
4// DbLocker wraps a regular DbConn, only allowing access using the
7// DbLocker wraps a regular DbConn, only allowing access using the
5// `lock` method. This method will acquire the advisory lock before
8// `lock` method. This method will acquire the advisory lock before
@@ -103,12 +106,12 @@ impl DbConn {
103
106
104impl Conn for DbConn {
107impl Conn for DbConn {
105 fn run(&mut self, query: &str) -> anyhow::Result<()> {
108 fn run(&mut self, query: &str) -> anyhow::Result<()> {
106− self.client.batch_execute(query)?;
109+ retry_automatically(|| self.client.batch_execute(query))?;
107 Ok(())
110 Ok(())
108 }
111 }
109
112
110 fn query(&mut self, query: &str) -> anyhow::Result<Vec<Row>> {
113 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, &[]))?;
112 Ok(rows)
115 Ok(rows)
113 }
116 }
114
117
@@ -117,7 +120,7 @@ impl Conn for DbConn {
117 query: &str,
120 query: &str,
118 params: &[&(dyn ToSql + Sync)],
121 params: &[&(dyn ToSql + Sync)],
119 ) -> anyhow::Result<Vec<Row>> {
122 ) -> anyhow::Result<Vec<Row>> {
120− let rows = self.client.query(query, params)?;
123+ let rows = retry_automatically(|| self.client.query(query, params))?;
121 Ok(rows)
124 Ok(rows)
122 }
125 }
123
126
@@ -168,3 +171,57 @@ impl Conn for Transaction<'_> {
168 Ok(Transaction { transaction })
171 Ok(Transaction { transaction })
169 }
172 }
170}
173}
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+}