git.iko.soy/reshape
on main help sign in
reshape/src/db.rs
224 lines 7.1 kB history raw
1use std::{cmp::min, time::Duration};
3use anyhow::{anyhow, Context};
4use postgres::{types::ToSql, NoTls, Row};
5use rand::prelude::*;
7// DbLocker wraps a regular DbConn, only allowing access using the
8// `lock` method. This method will acquire the advisory lock before
9// allowing access to the database, and then release it afterwards.
10//
11// We use advisory locks to avoid multiple Reshape instances working
12// on the same database as the same time. DbLocker is the only way to
13// get a DbConn which ensures that all database access is protected by
14// a lock.
15//
16// Postgres docs on advisory locks:
17// https://www.postgresql.org/docs/current/explicit-locking.html#ADVISORY-LOCKS
18pub struct DbLocker {
19 client: DbConn,
22impl DbLocker {
23 // Advisory lock keys in Postgres are 64-bit integers.
24 // The key we use was chosen randomly.
25 const LOCK_KEY: i64 = 4036779288569897133;
27 pub fn connect(config: &postgres::Config) -> anyhow::Result<Self> {
28 let mut pg = config.connect(NoTls)?;
30 // When running DDL queries that acquire locks, we risk causing a "lock queue".
31 // When attempting to acquire a lock, Postgres will wait for any long running queries to complete.
32 // At the same time, it will block other queries until the lock has been acquired and released.
33 // This has the bad effect of the long-running query blocking other queries because of us, forming
34 // a queue of other queries until we release our lock.
35 //
36 // We set the lock_timeout setting to avoid this. This puts an upper bound for how long Postgres will
37 // wait to acquire locks and also the maximum amount of time a long-running query can block other queries.
38 // We should also add automatic retries to handle these timeouts gracefully.
39 //
40 // Reference: https://medium.com/paypal-tech/postgresql-at-scale-database-schema-changes-without-downtime-20d3749ed680
41 //
42 // TODO: Make lock_timeout configurable
43 pg.simple_query("SET lock_timeout = '1s'")
44 .context("failed to set lock_timeout")?;
46 Ok(Self {
47 client: DbConn::new(pg),
48 })
49 }
51 pub fn lock(
52 &mut self,
53 f: impl FnOnce(&mut DbConn) -> anyhow::Result<()>,
54 ) -> anyhow::Result<()> {
55 self.acquire_lock()?;
56 let result = f(&mut self.client);
57 self.release_lock()?;
59 result
60 }
62 fn acquire_lock(&mut self) -> anyhow::Result<()> {
63 let success = self
64 .client
65 .query(&format!("SELECT pg_try_advisory_lock({})", Self::LOCK_KEY))?
66 .first()
67 .ok_or_else(|| anyhow!("unexpectedly failed when acquiring advisory lock"))
68 .map(|row| row.get::<'_, _, bool>(0))?;
70 if success {
71 Ok(())
72 } else {
73 Err(anyhow!("another instance of Reshape is already running"))
74 }
75 }
77 fn release_lock(&mut self) -> anyhow::Result<()> {
78 self.client
79 .query(&format!("SELECT pg_advisory_unlock({})", Self::LOCK_KEY))?
80 .first()
81 .ok_or_else(|| anyhow!("unexpectedly failed when releasing advisory lock"))?;
82 Ok(())
83 }
86pub trait Conn {
87 fn run(&mut self, query: &str) -> anyhow::Result<()>;
88 fn query(&mut self, query: &str) -> anyhow::Result<Vec<Row>>;
89 fn query_with_params(
90 &mut self,
91 query: &str,
92 params: &[&(dyn ToSql + Sync)],
93 ) -> anyhow::Result<Vec<Row>>;
94 fn transaction(&mut self) -> anyhow::Result<Transaction>;
97pub struct DbConn {
98 client: postgres::Client,
101impl DbConn {
102 fn new(client: postgres::Client) -> Self {
103 DbConn { client }
104 }
107impl Conn for DbConn {
108 fn run(&mut self, query: &str) -> anyhow::Result<()> {
109 retry_automatically(|| self.client.batch_execute(query))?;
110 Ok(())
111 }
113 fn query(&mut self, query: &str) -> anyhow::Result<Vec<Row>> {
114 let rows = retry_automatically(|| self.client.query(query, &[]))?;
115 Ok(rows)
116 }
118 fn query_with_params(
119 &mut self,
120 query: &str,
121 params: &[&(dyn ToSql + Sync)],
122 ) -> anyhow::Result<Vec<Row>> {
123 let rows = retry_automatically(|| self.client.query(query, params))?;
124 Ok(rows)
125 }
127 fn transaction(&mut self) -> anyhow::Result<Transaction> {
128 let transaction = self.client.transaction()?;
129 Ok(Transaction { transaction })
130 }
133pub struct Transaction<'a> {
134 transaction: postgres::Transaction<'a>,
137impl Transaction<'_> {
138 pub fn commit(self) -> anyhow::Result<()> {
139 self.transaction.commit()?;
140 Ok(())
141 }
143 pub fn rollback(self) -> anyhow::Result<()> {
144 self.transaction.rollback()?;
145 Ok(())
146 }
149impl Conn for Transaction<'_> {
150 fn run(&mut self, query: &str) -> anyhow::Result<()> {
151 self.transaction.batch_execute(query)?;
152 Ok(())
153 }
155 fn query(&mut self, query: &str) -> anyhow::Result<Vec<Row>> {
156 let rows = self.transaction.query(query, &[])?;
157 Ok(rows)
158 }
160 fn query_with_params(
161 &mut self,
162 query: &str,
163 params: &[&(dyn ToSql + Sync)],
164 ) -> anyhow::Result<Vec<Row>> {
165 let rows = self.transaction.query(query, params)?;
166 Ok(rows)
167 }
169 fn transaction(&mut self) -> anyhow::Result<Transaction> {
170 let transaction = self.transaction.transaction()?;
171 Ok(Transaction { transaction })
172 }
175// Retry a database operation with exponential backoff and jitter
176fn 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;
183 let mut rng = rand::thread_rng();
184 let mut attempts = 0;
185 loop {
186 let result = f();
188 let error = match result {
189 Ok(_) => return result,
190 Err(err) => err,
191 };
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 }
202 attempts += 1;
203 if attempts >= MAX_ATTEMPTS {
204 return Err(error);
205 }
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 );
213 // The jitter is up to half the wait time
214 let jitter: u64 = rng.gen_range(0..wait_time / 2);
216 std::thread::sleep(Duration::from_millis(wait_time + jitter));
217 }
220// Check if a database error can be retried
221fn error_retryable(error: &postgres::error::DbError) -> bool {
222 // LOCK_NOT_AVAILABLE is caused by lock_timeout being exceeded
223 matches!(error.code(), &postgres::error::SqlState::LOCK_NOT_AVAILABLE)