1use std::sync::Arc;
5use std::time::Instant;
6
7use runlight::BoxFuture;
8use runlight::store::{
9 Db, DbError, Dialect, Hold, Param, Row, TablesBusyError, fill_placeholders, lock_wait, mysql_literal,
10};
11use sqlx::pool::PoolConnection;
12use sqlx::{AssertSqlSafe, MySql, MySqlPool};
13use tokio::sync::Mutex;
14
15use crate::rows::{db_error, mysql_row};
16
17const LOCK: &str = "CONCAT('runlight:', LEFT(SHA2(COALESCE(DATABASE(), ''), 256), 48))";
20
21#[derive(Clone)]
25pub struct MysqlDb {
26 pool: MySqlPool,
27}
28
29impl MysqlDb {
30 pub fn new(pool: MySqlPool) -> MysqlDb {
32 MysqlDb { pool }
33 }
34}
35
36const NO_BACKSLASH_ESCAPES_OFF: &str = "SET SESSION sql_mode = TRIM(BOTH ',' FROM REPLACE(CONCAT(',', @@SESSION.sql_mode, ','), ',NO_BACKSLASH_ESCAPES,', ','))";
38
39fn text(sql: &str, params: &[Param]) -> Result<AssertSqlSafe<String>, DbError> {
43 Ok(AssertSqlSafe(format!("{NO_BACKSLASH_ESCAPES_OFF}; {}", fill_placeholders(sql, params, true, mysql_literal)?)))
44}
45
46async fn exec(conn: &mut PoolConnection<MySql>, sql: &str) -> Result<Vec<Row>, DbError> {
47 let rows = sqlx::raw_sql(AssertSqlSafe(sql.to_string())).fetch_all(&mut **conn).await.map_err(db_error)?;
48 rows.iter().map(mysql_row).collect()
49}
50
51impl Db for MysqlDb {
52 fn dialect(&self) -> Dialect {
53 Dialect::Mysql
54 }
55
56 fn all<'a>(&'a self, sql: &'a str, params: Vec<Param>) -> BoxFuture<'a, Result<Vec<Row>, DbError>> {
57 Box::pin(async move {
58 let rows = sqlx::raw_sql(text(sql, ¶ms)?).fetch_all(&self.pool).await.map_err(db_error)?;
59 rows.iter().map(mysql_row).collect()
60 })
61 }
62
63 fn run<'a>(&'a self, sql: &'a str, params: Vec<Param>) -> BoxFuture<'a, Result<u64, DbError>> {
64 Box::pin(async move {
65 Ok(sqlx::raw_sql(text(sql, ¶ms)?).execute(&self.pool).await.map_err(db_error)?.rows_affected())
66 })
67 }
68
69 fn hold(&self, kind: Hold) -> BoxFuture<'_, Result<Arc<dyn Db>, DbError>> {
70 Box::pin(async move {
71 let mut slot = Slot { conn: Some(self.pool.acquire().await.map_err(db_error)?), nested: false };
74 let conn = slot.conn.as_mut().ok_or_else(ended)?;
75 match kind {
76 Hold::Transaction => begin(conn).await?,
77 Hold::Exclusive => {
78 let wait = lock_wait();
80 let until = Instant::now() + wait;
81 loop {
82 let left = until.saturating_duration_since(Instant::now()).as_millis().div_ceil(1000);
83 let seconds = left.clamp(1, 5);
84 let rows = exec(conn, &format!("SELECT GET_LOCK({LOCK}, {seconds}) AS ok")).await?;
85 let row = rows.into_iter().next().unwrap_or_default();
86 if row.is_null("ok") {
87 return Err(DbError("Runlight: MySQL refused the lock for creating tables".into()));
88 }
89 if row.num("ok") == 1.0 {
90 break;
91 }
92 if Instant::now() >= until {
94 return Err(TablesBusyError { waited: wait }.into());
95 }
96 }
97 }
98 }
99 Ok(Arc::new(Held { conn: Arc::new(Mutex::new(slot)), kind, role: Role::Owner }) as Arc<dyn Db>)
100 })
101 }
102}
103
104async fn begin(conn: &mut PoolConnection<MySql>) -> Result<(), DbError> {
107 exec(conn, "SET TRANSACTION ISOLATION LEVEL READ COMMITTED").await?;
108 exec(conn, "START TRANSACTION").await?;
109 Ok(())
110}
111
112struct Slot {
115 conn: Option<PoolConnection<MySql>>,
116 nested: bool,
118}
119
120impl Slot {
121 fn close(&mut self) {
122 if let Some(conn) = self.conn.take() {
123 drop(conn.detach());
124 }
125 }
126}
127
128impl Drop for Slot {
129 fn drop(&mut self) {
130 self.close();
131 }
132}
133
134struct Held {
136 conn: Arc<Mutex<Slot>>,
137 kind: Hold,
138 role: Role,
139}
140
141impl Drop for Held {
142 fn drop(&mut self) {
143 if self.role == Role::Owner
145 && let Ok(mut slot) = self.conn.try_lock()
146 {
147 slot.close();
148 }
149 }
150}
151
152#[derive(Clone, Copy, PartialEq, Eq)]
154enum Role {
155 Owner,
157 Nested,
159 Joined,
161}
162
163fn ended() -> DbError {
164 DbError("Runlight: the held connection has been let go".into())
165}
166
167impl Db for Held {
168 fn dialect(&self) -> Dialect {
169 Dialect::Mysql
170 }
171
172 fn all<'a>(&'a self, sql: &'a str, params: Vec<Param>) -> BoxFuture<'a, Result<Vec<Row>, DbError>> {
173 Box::pin(async move {
174 let mut slot = self.conn.lock().await;
175 let conn = slot.conn.as_mut().ok_or_else(ended)?;
176 let rows = sqlx::raw_sql(text(sql, ¶ms)?).fetch_all(&mut **conn).await.map_err(db_error)?;
177 rows.iter().map(mysql_row).collect()
178 })
179 }
180
181 fn run<'a>(&'a self, sql: &'a str, params: Vec<Param>) -> BoxFuture<'a, Result<u64, DbError>> {
182 Box::pin(async move {
183 let mut slot = self.conn.lock().await;
184 let conn = slot.conn.as_mut().ok_or_else(ended)?;
185 Ok(sqlx::raw_sql(text(sql, ¶ms)?).execute(&mut **conn).await.map_err(db_error)?.rows_affected())
186 })
187 }
188
189 fn hold(&self, kind: Hold) -> BoxFuture<'_, Result<Arc<dyn Db>, DbError>> {
190 Box::pin(async move {
191 if kind == Hold::Transaction && self.kind == Hold::Exclusive {
192 let mut slot = self.conn.lock().await;
193 let slot = &mut *slot;
194 let conn = slot.conn.as_mut().ok_or_else(ended)?;
195 slot.nested = true;
197 begin(conn).await?;
198 return Ok(Arc::new(Held { conn: self.conn.clone(), kind, role: Role::Nested }) as Arc<dyn Db>);
199 }
200 Ok(Arc::new(Held { conn: self.conn.clone(), kind: self.kind, role: Role::Joined }) as Arc<dyn Db>)
201 })
202 }
203
204 fn finish(&self, ok: bool) -> BoxFuture<'_, Result<(), DbError>> {
205 Box::pin(async move {
206 let mut slot = self.conn.lock().await;
207 match self.role {
208 Role::Joined => Ok(()),
209 Role::Nested => {
210 let slot = &mut *slot;
211 let Some(conn) = slot.conn.as_mut() else { return Ok(()) };
212 exec(conn, if ok { "COMMIT" } else { "ROLLBACK" }).await?;
213 slot.nested = false;
214 Ok(())
215 }
216 Role::Owner if slot.nested => {
219 slot.close();
220 Ok(())
221 }
222 Role::Owner => {
223 let Some(mut conn) = slot.conn.take() else { return Ok(()) };
224 let result = match self.kind {
225 Hold::Transaction => exec(&mut conn, if ok { "COMMIT" } else { "ROLLBACK" }).await.map(|_| ()),
226 Hold::Exclusive => exec(&mut conn, &format!("DO RELEASE_LOCK({LOCK})")).await.map(|_| ()),
229 };
230 if result.is_err() {
231 drop(conn.detach());
234 }
235 if self.kind == Hold::Exclusive { Ok(()) } else { result }
236 }
237 }
238 })
239 }
240}