meta_srv/election/rds/
mysql.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::sync::Arc;
16use std::sync::atomic::{AtomicBool, Ordering};
17use std::time::Duration;
18
19use common_meta::key::{CANDIDATES_ROOT, ELECTION_KEY};
20use common_telemetry::{error, warn};
21use common_time::Timestamp;
22use snafu::{OptionExt, ResultExt, ensure};
23use sqlx::mysql::{MySqlArguments, MySqlRow};
24use sqlx::pool::PoolConnection;
25use sqlx::query::Query;
26use sqlx::{MySql, MySqlPool, MySqlTransaction, Row};
27use tokio::sync::{Mutex, MutexGuard, broadcast};
28use tokio::time::MissedTickBehavior;
29
30use crate::election::rds::{LEASE_SEP, Lease, RdsLeaderKey, parse_value_and_expire_time};
31use crate::election::{
32    Election, LeaderChangeMessage, listen_leader_change, send_leader_change_and_set_flags,
33};
34use crate::error::{
35    AcquireMySqlClientSnafu, DecodeSqlValueSnafu, DeserializeFromJsonSnafu,
36    LeaderLeaseChangedSnafu, LeaderLeaseExpiredSnafu, MySqlExecutionSnafu, NoLeaderSnafu, Result,
37    SerializeToJsonSnafu, SqlExecutionTimeoutSnafu, UnexpectedSnafu,
38};
39use crate::metasrv::{ElectionRef, LeaderValue, MetasrvNodeInfo};
40
41struct ElectionSqlFactory<'a> {
42    table_name: &'a str,
43}
44
45struct ElectionSqlSet {
46    campaign: String,
47    // SQL to put a value with expire time.
48    //
49    // Parameters for the query:
50    // `$1`: key,
51    // `$2`: value,
52    // `$3`: lease time in seconds
53    //
54    // Returns:
55    // If the key already exists, return the previous value.
56    put_value_with_lease: String,
57    // SQL to update a value with expire time.
58    //
59    // Parameters for the query:
60    // `$1`: updated value,
61    // `$2`: lease time in seconds
62    // `$3`: key,
63    // `$4`: previous value,
64    update_value_with_lease: String,
65    // SQL to get a value with expire time.
66    //
67    // Parameters:
68    // `$1`: key
69    get_value_with_lease: String,
70    // SQL to get all values with expire time with the given key prefix.
71    //
72    // Parameters:
73    // `$1`: key prefix like 'prefix%'
74    //
75    // Returns:
76    // column 0: value,
77    // column 1: current timestamp
78    get_value_with_lease_by_prefix: String,
79    // SQL to delete a value.
80    //
81    // Parameters:
82    // `?`: key
83    //
84    // Returns:
85    // Rows affected
86    delete_value: String,
87}
88
89impl<'a> ElectionSqlFactory<'a> {
90    fn new(table_name: &'a str) -> Self {
91        Self { table_name }
92    }
93
94    fn build(self) -> ElectionSqlSet {
95        ElectionSqlSet {
96            campaign: self.campaign_sql(),
97            put_value_with_lease: self.put_value_with_lease_sql(),
98            update_value_with_lease: self.update_value_with_lease_sql(),
99            get_value_with_lease: self.get_value_with_lease_sql(),
100            get_value_with_lease_by_prefix: self.get_value_with_lease_by_prefix_sql(),
101            delete_value: self.delete_value_sql(),
102        }
103    }
104
105    /// Use `SELECT FOR UPDATE` to lock for compatibility with other MySQL-compatible databases
106    /// instead of directly using `GET_LOCK`.
107    fn campaign_sql(&self) -> String {
108        format!("SELECT * FROM `{}` FOR UPDATE;", self.table_name)
109    }
110
111    fn put_value_with_lease_sql(&self) -> String {
112        format!(
113            r#"
114            INSERT INTO `{}` (k, v) VALUES (
115                ?,
116                CONCAT(
117                    ?,
118                    '{}',
119                    DATE_FORMAT(DATE_ADD(NOW(4), INTERVAL ? SECOND), '%Y-%m-%d %T.%f')
120                )
121            )
122            ON DUPLICATE KEY UPDATE v = VALUES(v);
123            "#,
124            self.table_name, LEASE_SEP
125        )
126    }
127
128    fn update_value_with_lease_sql(&self) -> String {
129        format!(
130            r#"UPDATE `{}`
131               SET v = CONCAT(?, '{}', DATE_FORMAT(DATE_ADD(NOW(4), INTERVAL ? SECOND), '%Y-%m-%d %T.%f'))
132               WHERE k = ? AND v = ?"#,
133            self.table_name, LEASE_SEP
134        )
135    }
136
137    fn get_value_with_lease_sql(&self) -> String {
138        format!(
139            r#"SELECT v, DATE_FORMAT(NOW(4), '%Y-%m-%d %T.%f') FROM `{}` WHERE k = ?"#,
140            self.table_name
141        )
142    }
143
144    fn get_value_with_lease_by_prefix_sql(&self) -> String {
145        format!(
146            r#"SELECT v, DATE_FORMAT(NOW(4), '%Y-%m-%d %T.%f') FROM `{}` WHERE k LIKE ?"#,
147            self.table_name
148        )
149    }
150
151    fn delete_value_sql(&self) -> String {
152        format!("DELETE FROM {} WHERE k = ?;", self.table_name)
153    }
154}
155
156enum Executor<'a> {
157    Default(MutexGuard<'a, ElectionMysqlClient>),
158    Txn(TransactionWithExecutionTimeout<'a>),
159}
160
161impl Executor<'_> {
162    async fn query(
163        &mut self,
164        query: Query<'_, MySql, MySqlArguments>,
165        sql: &str,
166    ) -> Result<Vec<MySqlRow>> {
167        match self {
168            Executor::Default(client) => client.query(query, sql).await,
169            Executor::Txn(txn) => txn.query(query, sql).await,
170        }
171    }
172
173    async fn execute(&mut self, query: Query<'_, MySql, MySqlArguments>, sql: &str) -> Result<u64> {
174        match self {
175            Executor::Default(client) => client.execute(query, sql).await,
176            Executor::Txn(txn) => txn.execute(query, sql).await,
177        }
178    }
179
180    async fn commit(self) -> Result<()> {
181        match self {
182            Executor::Txn(txn) => txn.commit().await,
183            _ => Ok(()),
184        }
185    }
186}
187
188/// MySQL client for election.
189pub struct ElectionMysqlClient {
190    current: Option<PoolConnection<MySql>>,
191    pool: MySqlPool,
192
193    /// The client-side timeout for statement execution.
194    ///
195    /// This timeout is enforced by the client application and is independent of any server-side timeouts.
196    /// If a statement takes longer than this duration to execute, the client will abort the operation.
197    execution_timeout: Duration,
198
199    /// The maximum execution time for the statement.
200    ///
201    /// This timeout is enforced by the server and is independent of any client-side timeouts.
202    /// If a statement takes longer than this duration to execute, the server will abort the operation.
203    max_execution_time: Duration,
204
205    /// The lock wait timeout for the session.
206    ///
207    /// This timeout determines how long the server waits for a lock to be acquired before timing out.
208    /// If a lock cannot be acquired within this duration, the server will abort the operation.
209    innode_lock_wait_timeout: Duration,
210
211    /// The wait timeout for the session.
212    ///
213    /// This timeout determines how long the server waits for activity on a noninteractive connection
214    /// before closing it. If a connection is idle for longer than this duration, the server will
215    /// terminate it.
216    wait_timeout: Duration,
217
218    /// The table name for election.
219    table_for_election: String,
220}
221
222impl ElectionMysqlClient {
223    pub fn new(
224        pool: MySqlPool,
225        execution_timeout: Duration,
226        max_execution_time: Duration,
227        innode_lock_wait_timeout: Duration,
228        wait_timeout: Duration,
229        table_for_election: &str,
230    ) -> Self {
231        Self {
232            current: None,
233            pool,
234            execution_timeout,
235            max_execution_time,
236            innode_lock_wait_timeout,
237            wait_timeout,
238            table_for_election: table_for_election.to_string(),
239        }
240    }
241
242    fn create_table_sql(&self) -> String {
243        format!(
244            r#"
245            CREATE TABLE IF NOT EXISTS `{}` (
246                k VARBINARY(3072) PRIMARY KEY,
247                v BLOB
248            );"#,
249            self.table_for_election
250        )
251    }
252
253    fn insert_once_sql(&self) -> String {
254        format!(
255            "INSERT IGNORE INTO `{}` (k, v) VALUES ('__place_holder_for_lock', '');",
256            self.table_for_election
257        )
258    }
259
260    fn check_version_sql(&self) -> String {
261        "SELECT @@version;".to_string()
262    }
263
264    async fn reset_client(&mut self) -> Result<()> {
265        self.current = None;
266        self.maybe_init_client().await
267    }
268
269    async fn ensure_table_exists(&mut self) -> Result<()> {
270        let create_table_sql = self.create_table_sql();
271        let query = sqlx::query(&create_table_sql);
272        self.execute(query, &create_table_sql).await?;
273        // Insert at least one row for `SELECT * FOR UPDATE` to work.
274        let insert_once_sql = self.insert_once_sql();
275        let query = sqlx::query(&insert_once_sql);
276        self.execute(query, &insert_once_sql).await?;
277        Ok(())
278    }
279
280    async fn maybe_init_client(&mut self) -> Result<()> {
281        if self.current.is_none() {
282            let client = self.pool.acquire().await.context(AcquireMySqlClientSnafu)?;
283
284            self.current = Some(client);
285            let (query, sql) = if !self.wait_timeout.is_zero() {
286                let sql = "SET SESSION wait_timeout = ?, innodb_lock_wait_timeout = ?, max_execution_time = ?;";
287                (
288                    sqlx::query(sql)
289                        .bind(self.wait_timeout.as_secs())
290                        .bind(self.innode_lock_wait_timeout.as_secs())
291                        .bind(self.max_execution_time.as_millis() as u64),
292                    sql,
293                )
294            } else {
295                let sql = "SET SESSION innodb_lock_wait_timeout = ?, max_execution_time = ?;";
296                (
297                    sqlx::query(sql)
298                        .bind(self.innode_lock_wait_timeout.as_secs())
299                        .bind(self.max_execution_time.as_millis() as u64),
300                    sql,
301                )
302            };
303            self.set_session_isolation_level().await?;
304            self.execute(query, sql).await?;
305            self.check_version(&self.check_version_sql()).await?;
306        }
307        Ok(())
308    }
309
310    /// Set session isolation level to serializable.
311    ///
312    /// # Panics
313    /// if `current` is `None`.
314    async fn set_session_isolation_level(&mut self) -> Result<()> {
315        let sql = "SET SESSION TRANSACTION ISOLATION LEVEL SERIALIZABLE";
316        let query = sqlx::query(sql);
317        self.execute(query, sql).await?;
318        Ok(())
319    }
320
321    /// Check if the MySQL version is supported.
322    ///
323    /// # Panics
324    /// if `current` is `None`.
325    async fn check_version(&mut self, sql: &str) -> Result<()> {
326        // Check if the MySQL version is supported.
327        let query = sqlx::query(sql);
328        // Safety: `maybe_init_client` ensures `current` is not `None`.
329        let client = self.current.as_mut().unwrap();
330        let row = tokio::time::timeout(self.execution_timeout, query.fetch_one(&mut **client))
331            .await
332            .map_err(|_| {
333                SqlExecutionTimeoutSnafu {
334                    sql,
335                    duration: self.execution_timeout,
336                }
337                .build()
338            })?;
339        match row {
340            Ok(row) => {
341                let version: String = row.try_get(0).context(DecodeSqlValueSnafu {})?;
342                if !version.starts_with("8.0") && !version.starts_with("5.7") {
343                    warn!(
344                        "Unsupported MySQL version: {}, expected: [5.7, 8.0]",
345                        version
346                    );
347                }
348            }
349            Err(e) => {
350                warn!(e; "Failed to check MySQL version through sql: {}", sql);
351            }
352        }
353        Ok(())
354    }
355
356    /// Returns the result of the query.
357    ///
358    /// # Panics
359    /// if `current` is `None`.
360    async fn execute(&mut self, query: Query<'_, MySql, MySqlArguments>, sql: &str) -> Result<u64> {
361        // Safety: `maybe_init_client` ensures `current` is not `None`.
362        let client = self.current.as_mut().unwrap();
363        let future = query.execute(&mut **client);
364        let rows_affected = tokio::time::timeout(self.execution_timeout, future)
365            .await
366            .map_err(|_| {
367                SqlExecutionTimeoutSnafu {
368                    sql,
369                    duration: self.execution_timeout,
370                }
371                .build()
372            })?
373            .context(MySqlExecutionSnafu { sql })?
374            .rows_affected();
375        Ok(rows_affected)
376    }
377
378    /// Returns the result of the query.
379    ///
380    /// # Panics
381    /// if `current` is `None`.
382    async fn query(
383        &mut self,
384        query: Query<'_, MySql, MySqlArguments>,
385        sql: &str,
386    ) -> Result<Vec<MySqlRow>> {
387        // Safety: `maybe_init_client` ensures `current` is not `None`.
388        let client = self.current.as_mut().unwrap();
389        let future = query.fetch_all(&mut **client);
390        tokio::time::timeout(self.execution_timeout, future)
391            .await
392            .map_err(|_| {
393                SqlExecutionTimeoutSnafu {
394                    sql,
395                    duration: self.execution_timeout,
396                }
397                .build()
398            })?
399            .context(MySqlExecutionSnafu { sql })
400    }
401
402    async fn transaction(&mut self) -> Result<TransactionWithExecutionTimeout<'_>> {
403        use sqlx::Acquire;
404        let client = self.current.as_mut().unwrap();
405        let transaction = client
406            .begin()
407            .await
408            .context(MySqlExecutionSnafu { sql: "BEGIN" })?;
409
410        Ok(TransactionWithExecutionTimeout {
411            transaction,
412            execution_timeout: self.execution_timeout,
413        })
414    }
415}
416
417struct TransactionWithExecutionTimeout<'a> {
418    transaction: MySqlTransaction<'a>,
419    execution_timeout: Duration,
420}
421
422impl TransactionWithExecutionTimeout<'_> {
423    async fn query(
424        &mut self,
425        query: Query<'_, MySql, MySqlArguments>,
426        sql: &str,
427    ) -> Result<Vec<MySqlRow>> {
428        let res = tokio::time::timeout(
429            self.execution_timeout,
430            query.fetch_all(&mut *self.transaction),
431        )
432        .await
433        .map_err(|_| {
434            SqlExecutionTimeoutSnafu {
435                sql,
436                duration: self.execution_timeout,
437            }
438            .build()
439        })?
440        .context(MySqlExecutionSnafu { sql })?;
441        Ok(res)
442    }
443
444    async fn execute(&mut self, query: Query<'_, MySql, MySqlArguments>, sql: &str) -> Result<u64> {
445        let res = tokio::time::timeout(
446            self.execution_timeout,
447            query.execute(&mut *self.transaction),
448        )
449        .await
450        .map_err(|_| {
451            SqlExecutionTimeoutSnafu {
452                sql,
453                duration: self.execution_timeout,
454            }
455            .build()
456        })?
457        .context(MySqlExecutionSnafu { sql })?;
458        Ok(res.rows_affected())
459    }
460
461    async fn commit(self) -> Result<()> {
462        tokio::time::timeout(self.execution_timeout, self.transaction.commit())
463            .await
464            .map_err(|_| {
465                SqlExecutionTimeoutSnafu {
466                    sql: "COMMIT",
467                    duration: self.execution_timeout,
468                }
469                .build()
470            })?
471            .context(MySqlExecutionSnafu { sql: "COMMIT" })?;
472        Ok(())
473    }
474}
475
476/// MySQL implementation of Election.
477pub struct MySqlElection {
478    leader_value: String,
479    client: Mutex<ElectionMysqlClient>,
480    is_leader: AtomicBool,
481    leader_infancy: AtomicBool,
482    leader_watcher: broadcast::Sender<LeaderChangeMessage>,
483    store_key_prefix: String,
484    candidate_lease_ttl: Duration,
485    meta_lease_ttl: Duration,
486    sql_set: ElectionSqlSet,
487}
488
489impl MySqlElection {
490    pub async fn with_mysql_client(
491        leader_value: String,
492        mut client: ElectionMysqlClient,
493        store_key_prefix: String,
494        candidate_lease_ttl: Duration,
495        meta_lease_ttl: Duration,
496        table_name: &str,
497    ) -> Result<ElectionRef> {
498        let sql_factory = ElectionSqlFactory::new(table_name);
499        client.maybe_init_client().await?;
500        client.ensure_table_exists().await?;
501        let tx = listen_leader_change(leader_value.clone());
502        Ok(Arc::new(Self {
503            leader_value,
504            client: Mutex::new(client),
505            is_leader: AtomicBool::new(false),
506            leader_infancy: AtomicBool::new(false),
507            leader_watcher: tx,
508            store_key_prefix,
509            candidate_lease_ttl,
510            meta_lease_ttl,
511            sql_set: sql_factory.build(),
512        }))
513    }
514
515    fn election_key(&self) -> String {
516        format!("{}{}", self.store_key_prefix, ELECTION_KEY)
517    }
518
519    fn candidate_root(&self) -> String {
520        format!("{}{}", self.store_key_prefix, CANDIDATES_ROOT)
521    }
522
523    fn candidate_key(&self) -> String {
524        format!("{}{}", self.candidate_root(), self.leader_value)
525    }
526
527    async fn maybe_init_client(&self) -> Result<()> {
528        let mut client = self.client.lock().await;
529        client.maybe_init_client().await?;
530        Ok(())
531    }
532}
533
534#[async_trait::async_trait]
535impl Election for MySqlElection {
536    type Leader = LeaderValue;
537
538    fn is_leader(&self) -> bool {
539        self.is_leader.load(Ordering::Relaxed)
540    }
541
542    fn in_leader_infancy(&self) -> bool {
543        self.leader_infancy
544            .compare_exchange(true, false, Ordering::AcqRel, Ordering::Acquire)
545            .is_ok()
546    }
547
548    async fn register_candidate(&self, node_info: &MetasrvNodeInfo) -> Result<()> {
549        self.maybe_init_client().await?;
550        let key = self.candidate_key();
551        let node_info =
552            serde_json::to_string(node_info).with_context(|_| SerializeToJsonSnafu {
553                input: format!("{node_info:?}"),
554            })?;
555
556        {
557            let client = self.client.lock().await;
558            let mut executor = Executor::Default(client);
559            let res = self
560                .put_value_with_lease(
561                    &key,
562                    &node_info,
563                    self.candidate_lease_ttl.as_secs(),
564                    &mut executor,
565                )
566                .await?;
567            // May registered before, just update the lease.
568            if !res {
569                warn!("Candidate already registered, update the lease");
570                self.delete_value(&key, &mut executor).await?;
571                self.put_value_with_lease(
572                    &key,
573                    &node_info,
574                    self.candidate_lease_ttl.as_secs(),
575                    &mut executor,
576                )
577                .await?;
578            }
579        }
580
581        // Check if the current lease has expired and renew the lease.
582        let mut keep_alive_interval = tokio::time::interval(self.candidate_lease_ttl / 2);
583        loop {
584            let _ = keep_alive_interval.tick().await;
585            let client = self.client.lock().await;
586            let mut executor = Executor::Default(client);
587            let lease = self
588                .get_value_with_lease(&key, &mut executor)
589                .await?
590                .unwrap_or_default();
591
592            ensure!(
593                lease.expire_time > lease.current,
594                UnexpectedSnafu {
595                    violated: format!(
596                        "Candidate lease expired at {:?} (current time: {:?}), key: {:?}",
597                        lease.expire_time,
598                        lease.current,
599                        String::from_utf8_lossy(&key.into_bytes())
600                    ),
601                }
602            );
603
604            self.update_value_with_lease(
605                &key,
606                &lease.origin,
607                &node_info,
608                self.candidate_lease_ttl.as_secs(),
609                &mut executor,
610            )
611            .await?;
612            std::mem::drop(executor);
613        }
614    }
615
616    async fn all_candidates(&self) -> Result<Vec<MetasrvNodeInfo>> {
617        self.maybe_init_client().await?;
618        let key_prefix = self.candidate_root();
619        let client = self.client.lock().await;
620        let mut executor = Executor::Default(client);
621        let (mut candidates, current) = self
622            .get_value_with_lease_by_prefix(&key_prefix, &mut executor)
623            .await?;
624        // Remove expired candidates
625        candidates.retain(|c| c.1 > current);
626        let mut valid_candidates = Vec::with_capacity(candidates.len());
627        for (c, _) in candidates {
628            let node_info: MetasrvNodeInfo =
629                serde_json::from_str(&c).with_context(|_| DeserializeFromJsonSnafu {
630                    input: format!("{:?}", c),
631                })?;
632            valid_candidates.push(node_info);
633        }
634        Ok(valid_candidates)
635    }
636
637    async fn campaign(&self) -> Result<()> {
638        self.maybe_init_client().await?;
639        let mut keep_alive_interval = tokio::time::interval(self.meta_lease_ttl / 2);
640        keep_alive_interval.set_missed_tick_behavior(MissedTickBehavior::Delay);
641        loop {
642            self.do_campaign().await?;
643            keep_alive_interval.tick().await;
644        }
645    }
646
647    async fn reset_campaign(&self) {
648        if let Err(err) = self.client.lock().await.reset_client().await {
649            error!(err; "Failed to reset client");
650        }
651    }
652
653    async fn leader(&self) -> Result<Self::Leader> {
654        self.maybe_init_client().await?;
655        if self.is_leader.load(Ordering::Relaxed) {
656            Ok(self.leader_value.as_bytes().into())
657        } else {
658            let key = self.election_key();
659
660            let client = self.client.lock().await;
661            let mut executor = Executor::Default(client);
662            if let Some(lease) = self.get_value_with_lease(&key, &mut executor).await? {
663                ensure!(lease.expire_time > lease.current, NoLeaderSnafu);
664                Ok(lease.leader_value.as_bytes().into())
665            } else {
666                NoLeaderSnafu.fail()
667            }
668        }
669    }
670
671    async fn resign(&self) -> Result<()> {
672        todo!()
673    }
674
675    fn subscribe_leader_change(&self) -> broadcast::Receiver<LeaderChangeMessage> {
676        self.leader_watcher.subscribe()
677    }
678}
679
680impl MySqlElection {
681    /// Returns value, expire time and current time.
682    async fn get_value_with_lease(
683        &self,
684        key: &str,
685        executor: &mut Executor<'_>,
686    ) -> Result<Option<Lease>> {
687        let key = key.as_bytes();
688        let query = sqlx::query(&self.sql_set.get_value_with_lease).bind(key);
689        let res = executor
690            .query(query, &self.sql_set.get_value_with_lease)
691            .await?;
692
693        if res.is_empty() {
694            return Ok(None);
695        }
696        // Safety: Checked if res is empty above.
697        let current_time_str = String::from_utf8_lossy(res[0].try_get(1).unwrap());
698        let current_time = match Timestamp::from_str(&current_time_str, None) {
699            Ok(ts) => ts,
700            Err(_) => UnexpectedSnafu {
701                violated: format!("Invalid timestamp: {}", current_time_str),
702            }
703            .fail()?,
704        };
705        // Safety: Checked if res is empty above.
706        let value_and_expire_time = String::from_utf8_lossy(res[0].try_get(0).unwrap_or_default());
707        let (value, expire_time) = parse_value_and_expire_time(&value_and_expire_time)?;
708
709        Ok(Some(Lease {
710            leader_value: value,
711            expire_time,
712            current: current_time,
713            origin: value_and_expire_time.to_string(),
714        }))
715    }
716
717    /// Returns all values and expire time with the given key prefix. Also returns the current time.
718    async fn get_value_with_lease_by_prefix(
719        &self,
720        key_prefix: &str,
721        executor: &mut Executor<'_>,
722    ) -> Result<(Vec<(String, Timestamp)>, Timestamp)> {
723        let key_prefix = format!("{}%", key_prefix).as_bytes().to_vec();
724        let query = sqlx::query(&self.sql_set.get_value_with_lease_by_prefix).bind(key_prefix);
725        let res = executor
726            .query(query, &self.sql_set.get_value_with_lease_by_prefix)
727            .await?;
728
729        let mut values_with_leases = vec![];
730        let mut current = Timestamp::default();
731        for row in res {
732            let current_time_str = row.try_get(1).unwrap_or_default();
733            current = match Timestamp::from_str(current_time_str, None) {
734                Ok(ts) => ts,
735                Err(_) => UnexpectedSnafu {
736                    violated: format!("Invalid timestamp: {}", current_time_str),
737                }
738                .fail()?,
739            };
740
741            let value_and_expire_time = String::from_utf8_lossy(row.try_get(0).unwrap_or_default());
742            let (value, expire_time) = parse_value_and_expire_time(&value_and_expire_time)?;
743
744            values_with_leases.push((value, expire_time));
745        }
746        Ok((values_with_leases, current))
747    }
748
749    async fn update_value_with_lease(
750        &self,
751        key: &str,
752        prev: &str,
753        updated: &str,
754        lease_ttl: u64,
755        executor: &mut Executor<'_>,
756    ) -> Result<()> {
757        let key = key.as_bytes();
758        let prev = prev.as_bytes();
759        let updated = updated.as_bytes();
760
761        let query = sqlx::query(&self.sql_set.update_value_with_lease)
762            .bind(updated)
763            .bind(lease_ttl as f64)
764            .bind(key)
765            .bind(prev);
766        let res = executor
767            .execute(query, &self.sql_set.update_value_with_lease)
768            .await?;
769
770        ensure!(
771            res == 1,
772            UnexpectedSnafu {
773                violated: format!("Failed to update key: {}", String::from_utf8_lossy(key)),
774            }
775        );
776
777        Ok(())
778    }
779
780    /// Returns `true` if the insertion is successful
781    async fn put_value_with_lease(
782        &self,
783        key: &str,
784        value: &str,
785        lease_ttl_secs: u64,
786        executor: &mut Executor<'_>,
787    ) -> Result<bool> {
788        let key = key.as_bytes();
789        let lease_ttl_secs = lease_ttl_secs as f64;
790        let query = sqlx::query(&self.sql_set.put_value_with_lease)
791            .bind(key)
792            .bind(value)
793            .bind(lease_ttl_secs);
794        let res = executor
795            .execute(query, &self.sql_set.put_value_with_lease)
796            .await?;
797        Ok(res == 1)
798    }
799
800    /// Returns `true` if the deletion is successful.
801    /// Caution: Should only delete the key if the lease is expired.
802    async fn delete_value(&self, key: &str, executor: &mut Executor<'_>) -> Result<bool> {
803        let key = key.as_bytes();
804        let query = sqlx::query(&self.sql_set.delete_value).bind(key);
805        let res = executor.execute(query, &self.sql_set.delete_value).await?;
806
807        Ok(res == 1)
808    }
809
810    /// Attempts to acquire leadership by executing a campaign. This function continuously checks
811    /// if the current lease is still valid.
812    async fn do_campaign(&self) -> Result<()> {
813        let lease = {
814            let client = self.client.lock().await;
815            let mut executor = Executor::Default(client);
816            self.get_value_with_lease(&self.election_key(), &mut executor)
817                .await?
818        };
819
820        let is_leader = self.is_leader();
821        // If current leader value is the same as the leader value in the remote lease,
822        // it means the current leader is still valid.
823        let is_current_leader = lease
824            .as_ref()
825            .map(|lease| lease.leader_value == self.leader_value)
826            .unwrap_or(false);
827        match (self.lease_check(&lease), is_leader, is_current_leader) {
828            // If the leader lease is valid and I'm the leader, renew the lease.
829            (Ok(_), true, true) => {
830                let mut client = self.client.lock().await;
831                let txn = client.transaction().await?;
832                let mut executor = Executor::Txn(txn);
833                let query = sqlx::query(&self.sql_set.campaign);
834                executor.query(query, &self.sql_set.campaign).await?;
835                // Safety: Checked if lease is not None above.
836                self.renew_lease(executor, lease.unwrap()).await?;
837            }
838            // If the leader lease expires and I'm the leader, notify the leader watcher and step down.
839            // Another instance should be elected as the leader in this case.
840            (Err(err), true, _) => {
841                warn!(err; "Leader lease expired, step down...");
842                self.step_down_without_lock().await?;
843            }
844            (Ok(_), true, false) => {
845                warn!("Leader lease expired, step down...");
846                self.step_down_without_lock().await?;
847            }
848            // If the leader lease expires and I'm not the leader, elect myself.
849            (Err(err), false, _) => {
850                warn!(err; "Leader lease expired, elect myself.");
851                let mut client = self.client.lock().await;
852                let txn = client.transaction().await?;
853                let mut executor = Executor::Txn(txn);
854                let query = sqlx::query(&self.sql_set.campaign);
855                executor.query(query, &self.sql_set.campaign).await?;
856                self.elected(executor, lease).await?;
857            }
858            // If the leader lease is valid and I'm the leader, but I don't think I'm the leader.
859            // Just re-elect myself.
860            (Ok(_), false, true) => {
861                warn!("I should be the leader, but I don't think so. Something went wrong.");
862                let mut client = self.client.lock().await;
863                let txn = client.transaction().await?;
864                let mut executor = Executor::Txn(txn);
865                let query = sqlx::query(&self.sql_set.campaign);
866                executor.query(query, &self.sql_set.campaign).await?;
867                // Safety: Checked if lease is not None above.
868                self.renew_lease(executor, lease.unwrap()).await?;
869            }
870            // If the leader lease is valid and I'm not the leader, do nothing.
871            (Ok(_), false, false) => {}
872        }
873        Ok(())
874    }
875
876    /// Renew the lease
877    async fn renew_lease(&self, mut executor: Executor<'_>, lease: Lease) -> Result<()> {
878        let key = self.election_key();
879        self.update_value_with_lease(
880            &key,
881            &lease.origin,
882            &self.leader_value,
883            self.meta_lease_ttl.as_secs(),
884            &mut executor,
885        )
886        .await?;
887        executor.commit().await?;
888
889        if !self.is_leader() {
890            let key = self.election_key();
891            let leader_key = RdsLeaderKey {
892                name: self.leader_value.clone().into_bytes(),
893                key: key.clone().into_bytes(),
894                ..Default::default()
895            };
896            send_leader_change_and_set_flags(
897                &self.is_leader,
898                &self.leader_infancy,
899                &self.leader_watcher,
900                LeaderChangeMessage::Elected(Arc::new(leader_key)),
901            );
902        }
903
904        Ok(())
905    }
906
907    /// Performs a lease check during the election process.
908    ///
909    /// This function performs the following checks and actions:
910    ///
911    /// - **Case 1**: If the current instance is not the leader but the lease has expired, it raises an error
912    ///   to re-initiate the campaign. If the leader failed to renew the lease, its session will expire and the lock
913    ///   will be released.
914    /// - **Case 2**: If all checks pass, the function returns without performing any actions.
915    fn lease_check(&self, lease: &Option<Lease>) -> Result<Lease> {
916        let lease = lease.as_ref().context(NoLeaderSnafu)?;
917        // Case 1: Lease expired
918        ensure!(lease.expire_time > lease.current, LeaderLeaseExpiredSnafu);
919        // Case 2: Everything is fine
920        Ok(lease.clone())
921    }
922
923    /// Still consider itself as the leader locally but failed to acquire the lock. Step down without deleting the key.
924    async fn step_down_without_lock(&self) -> Result<()> {
925        let key = self.election_key().into_bytes();
926        let leader_key = RdsLeaderKey {
927            name: self.leader_value.clone().into_bytes(),
928            key: key.clone(),
929            ..Default::default()
930        };
931        send_leader_change_and_set_flags(
932            &self.is_leader,
933            &self.leader_infancy,
934            &self.leader_watcher,
935            LeaderChangeMessage::StepDown(Arc::new(leader_key)),
936        );
937        Ok(())
938    }
939
940    /// Elected as leader. The leader should put the key and notify the leader watcher.
941    /// Caution: Should only elected while holding the lock.
942    async fn elected(
943        &self,
944        mut executor: Executor<'_>,
945        expected_lease: Option<Lease>,
946    ) -> Result<()> {
947        let key = self.election_key();
948        let leader_key = RdsLeaderKey {
949            name: self.leader_value.clone().into_bytes(),
950            key: key.clone().into_bytes(),
951            ..Default::default()
952        };
953        let remote_lease = self.get_value_with_lease(&key, &mut executor).await?;
954        ensure!(
955            expected_lease.map(|lease| lease.origin) == remote_lease.map(|lease| lease.origin),
956            LeaderLeaseChangedSnafu
957        );
958        self.delete_value(&key, &mut executor).await?;
959        self.put_value_with_lease(
960            &key,
961            &self.leader_value,
962            self.meta_lease_ttl.as_secs(),
963            &mut executor,
964        )
965        .await?;
966        executor.commit().await?;
967
968        send_leader_change_and_set_flags(
969            &self.is_leader,
970            &self.leader_infancy,
971            &self.leader_watcher,
972            LeaderChangeMessage::Elected(Arc::new(leader_key)),
973        );
974        Ok(())
975    }
976}
977
978#[cfg(test)]
979mod tests {
980    use std::assert_matches::assert_matches;
981    use std::env;
982
983    use common_meta::maybe_skip_mysql_integration_test;
984    use common_telemetry::init_default_ut_logging;
985
986    use super::*;
987    use crate::error;
988    use crate::utils::mysql::create_mysql_pool;
989
990    async fn create_mysql_client(
991        table_name: Option<&str>,
992        execution_timeout: Duration,
993        wait_timeout: Duration,
994    ) -> Result<Mutex<ElectionMysqlClient>> {
995        init_default_ut_logging();
996        let endpoint = env::var("GT_MYSQL_ENDPOINTS").unwrap_or_default();
997        if endpoint.is_empty() {
998            return UnexpectedSnafu {
999                violated: "MySQL endpoint is empty".to_string(),
1000            }
1001            .fail();
1002        }
1003        let pool = create_mysql_pool(&[endpoint], None).await.unwrap();
1004        let mut client = ElectionMysqlClient::new(
1005            pool,
1006            execution_timeout,
1007            execution_timeout,
1008            Duration::from_secs(1),
1009            wait_timeout,
1010            table_name.unwrap_or("default_greptime_metakv_election"),
1011        );
1012        client.maybe_init_client().await?;
1013        if table_name.is_some() {
1014            client.ensure_table_exists().await?;
1015        }
1016        Ok(Mutex::new(client))
1017    }
1018
1019    async fn drop_table(client: &Mutex<ElectionMysqlClient>, table_name: &str) {
1020        let mut client = client.lock().await;
1021        let sql = format!("DROP TABLE IF EXISTS {};", table_name);
1022        client.execute(sqlx::query(&sql), &sql).await.unwrap();
1023    }
1024
1025    #[tokio::test]
1026    async fn test_mysql_crud() {
1027        maybe_skip_mysql_integration_test!();
1028        let key = "test_key".to_string();
1029        let value = "test_value".to_string();
1030
1031        let uuid = uuid::Uuid::new_v4().to_string();
1032        let table_name = "test_mysql_crud_greptime_metakv";
1033        let candidate_lease_ttl = Duration::from_secs(10);
1034        let meta_lease_ttl = Duration::from_secs(2);
1035
1036        let execution_timeout = Duration::from_secs(10);
1037        let idle_session_timeout = Duration::from_secs(0);
1038        let client = create_mysql_client(Some(table_name), execution_timeout, idle_session_timeout)
1039            .await
1040            .unwrap();
1041
1042        {
1043            let mut a = client.lock().await;
1044            let txn = a.transaction().await.unwrap();
1045            let mut executor = Executor::Txn(txn);
1046            let raw_query = format!("SELECT * FROM {} FOR UPDATE;", table_name);
1047            let query = sqlx::query(&raw_query);
1048            let _ = executor.query(query, &raw_query).await.unwrap();
1049        }
1050
1051        let (tx, _) = broadcast::channel(100);
1052        let mysql_election = MySqlElection {
1053            leader_value: "test_leader".to_string(),
1054            client,
1055            is_leader: AtomicBool::new(false),
1056            leader_infancy: AtomicBool::new(true),
1057            leader_watcher: tx,
1058            store_key_prefix: uuid,
1059            candidate_lease_ttl,
1060            meta_lease_ttl,
1061            sql_set: ElectionSqlFactory::new(table_name).build(),
1062        };
1063        let client = mysql_election.client.lock().await;
1064        let mut executor = Executor::Default(client);
1065        let res = mysql_election
1066            .put_value_with_lease(&key, &value, 10, &mut executor)
1067            .await
1068            .unwrap();
1069        assert!(res);
1070
1071        let lease = mysql_election
1072            .get_value_with_lease(&key, &mut executor)
1073            .await
1074            .unwrap()
1075            .unwrap();
1076        assert_eq!(lease.leader_value, value);
1077
1078        mysql_election
1079            .update_value_with_lease(&key, &lease.origin, &value, 10, &mut executor)
1080            .await
1081            .unwrap();
1082
1083        let res = mysql_election
1084            .delete_value(&key, &mut executor)
1085            .await
1086            .unwrap();
1087        assert!(res);
1088
1089        let res = mysql_election
1090            .get_value_with_lease(&key, &mut executor)
1091            .await
1092            .unwrap();
1093        assert!(res.is_none());
1094
1095        for i in 0..10 {
1096            let key = format!("test_key_{}", i);
1097            let value = format!("test_value_{}", i);
1098            mysql_election
1099                .put_value_with_lease(&key, &value, 10, &mut executor)
1100                .await
1101                .unwrap();
1102        }
1103
1104        let key_prefix = "test_key".to_string();
1105        let (res, _) = mysql_election
1106            .get_value_with_lease_by_prefix(&key_prefix, &mut executor)
1107            .await
1108            .unwrap();
1109        assert_eq!(res.len(), 10);
1110
1111        for i in 0..10 {
1112            let key = format!("test_key_{}", i);
1113            let res = mysql_election
1114                .delete_value(&key, &mut executor)
1115                .await
1116                .unwrap();
1117            assert!(res);
1118        }
1119
1120        let (res, current) = mysql_election
1121            .get_value_with_lease_by_prefix(&key_prefix, &mut executor)
1122            .await
1123            .unwrap();
1124        assert!(res.is_empty());
1125        assert!(current == Timestamp::default());
1126
1127        // Should drop manually.
1128        std::mem::drop(executor);
1129        drop_table(&mysql_election.client, table_name).await;
1130    }
1131
1132    async fn candidate(
1133        leader_value: String,
1134        candidate_lease_ttl: Duration,
1135        store_key_prefix: String,
1136        table_name: String,
1137    ) {
1138        let meta_lease_ttl = Duration::from_secs(2);
1139        let execution_timeout = Duration::from_secs(10);
1140        let idle_session_timeout = Duration::from_secs(0);
1141        let client =
1142            create_mysql_client(Some(&table_name), execution_timeout, idle_session_timeout)
1143                .await
1144                .unwrap();
1145
1146        let (tx, _) = broadcast::channel(100);
1147        let mysql_election = MySqlElection {
1148            leader_value,
1149            client,
1150            is_leader: AtomicBool::new(false),
1151            leader_infancy: AtomicBool::new(true),
1152            leader_watcher: tx,
1153            store_key_prefix,
1154            candidate_lease_ttl,
1155            meta_lease_ttl,
1156            sql_set: ElectionSqlFactory::new(&table_name).build(),
1157        };
1158
1159        let node_info = MetasrvNodeInfo {
1160            addr: "test_addr".to_string(),
1161            version: "test_version".to_string(),
1162            git_commit: "test_git_commit".to_string(),
1163            start_time_ms: 0,
1164            total_cpu_millicores: 0,
1165            total_memory_bytes: 0,
1166            cpu_usage_millicores: 0,
1167            memory_usage_bytes: 0,
1168            hostname: "test_hostname".to_string(),
1169        };
1170        mysql_election.register_candidate(&node_info).await.unwrap();
1171    }
1172
1173    #[tokio::test]
1174    async fn test_candidate_registration() {
1175        maybe_skip_mysql_integration_test!();
1176        let leader_value_prefix = "test_leader".to_string();
1177        let candidate_lease_ttl = Duration::from_secs(2);
1178        let execution_timeout = Duration::from_secs(10);
1179        let meta_lease_ttl = Duration::from_secs(2);
1180        let idle_session_timeout = Duration::from_secs(0);
1181        let uuid = uuid::Uuid::new_v4().to_string();
1182        let table_name = "test_candidate_registration_greptime_metakv";
1183        let mut handles = vec![];
1184        let client = create_mysql_client(Some(table_name), execution_timeout, idle_session_timeout)
1185            .await
1186            .unwrap();
1187
1188        for i in 0..10 {
1189            let leader_value = format!("{}{}", leader_value_prefix, i);
1190            let handle = tokio::spawn(candidate(
1191                leader_value,
1192                candidate_lease_ttl,
1193                uuid.clone(),
1194                table_name.to_string(),
1195            ));
1196            handles.push(handle);
1197        }
1198        // Wait for candidates to register themselves and renew their leases at least once.
1199        tokio::time::sleep(candidate_lease_ttl / 2 + Duration::from_secs(1)).await;
1200
1201        let (tx, _) = broadcast::channel(100);
1202        let leader_value = "test_leader".to_string();
1203        let mysql_election = MySqlElection {
1204            leader_value,
1205            client,
1206            is_leader: AtomicBool::new(false),
1207            leader_infancy: AtomicBool::new(true),
1208            leader_watcher: tx,
1209            store_key_prefix: uuid.clone(),
1210            candidate_lease_ttl,
1211            meta_lease_ttl,
1212            sql_set: ElectionSqlFactory::new(table_name).build(),
1213        };
1214
1215        let candidates = mysql_election.all_candidates().await.unwrap();
1216        assert_eq!(candidates.len(), 10);
1217
1218        for handle in handles {
1219            handle.abort();
1220        }
1221
1222        // Wait for the candidate leases to expire.
1223        tokio::time::sleep(candidate_lease_ttl + Duration::from_secs(1)).await;
1224        let candidates = mysql_election.all_candidates().await.unwrap();
1225        assert!(candidates.is_empty());
1226
1227        // Garbage collection
1228        let client = mysql_election.client.lock().await;
1229        let mut executor = Executor::Default(client);
1230        for i in 0..10 {
1231            let key = format!("{}{}{}{}", uuid, CANDIDATES_ROOT, leader_value_prefix, i);
1232            let res = mysql_election
1233                .delete_value(&key, &mut executor)
1234                .await
1235                .unwrap();
1236            assert!(res);
1237        }
1238
1239        // Should drop manually.
1240        std::mem::drop(executor);
1241        drop_table(&mysql_election.client, table_name).await;
1242    }
1243
1244    async fn elected(
1245        election: &MySqlElection,
1246        table_name: &str,
1247        expected_lease: Option<Lease>,
1248    ) -> Result<()> {
1249        let mut client = election.client.lock().await;
1250        let txn = client.transaction().await.unwrap();
1251        let mut executor = Executor::Txn(txn);
1252        let raw_query = format!("SELECT * FROM {} FOR UPDATE;", table_name);
1253        let query = sqlx::query(&raw_query);
1254        let _ = executor.query(query, &raw_query).await.unwrap();
1255        election.elected(executor, expected_lease).await
1256    }
1257
1258    async fn get_lease(election: &MySqlElection) -> Option<Lease> {
1259        let client = election.client.lock().await;
1260        let mut executor = Executor::Default(client);
1261        election
1262            .get_value_with_lease(&election.election_key(), &mut executor)
1263            .await
1264            .unwrap()
1265    }
1266
1267    #[tokio::test]
1268    async fn test_elected_with_incorrect_lease_fails() {
1269        maybe_skip_mysql_integration_test!();
1270        let leader_value = "test_leader".to_string();
1271        let candidate_lease_ttl = Duration::from_secs(5);
1272        let meta_lease_ttl = Duration::from_secs(2);
1273        let execution_timeout = Duration::from_secs(10);
1274        let idle_session_timeout = Duration::from_secs(0);
1275        let uuid = uuid::Uuid::new_v4().to_string();
1276        let table_name = "test_elected_failed_greptime_metakv";
1277        let client = create_mysql_client(Some(table_name), execution_timeout, idle_session_timeout)
1278            .await
1279            .unwrap();
1280
1281        let (tx, _) = broadcast::channel(100);
1282        let leader_mysql_election = MySqlElection {
1283            leader_value: leader_value.clone(),
1284            client,
1285            is_leader: AtomicBool::new(false),
1286            leader_infancy: AtomicBool::new(true),
1287            leader_watcher: tx,
1288            store_key_prefix: uuid,
1289            candidate_lease_ttl,
1290            meta_lease_ttl,
1291            sql_set: ElectionSqlFactory::new(table_name).build(),
1292        };
1293
1294        let incorrect_lease = Lease::default();
1295        let err = elected(&leader_mysql_election, table_name, Some(incorrect_lease))
1296            .await
1297            .unwrap_err();
1298        assert_matches!(err, error::Error::LeaderLeaseChanged { .. });
1299        let lease = get_lease(&leader_mysql_election).await;
1300        assert!(lease.is_none());
1301        drop_table(&leader_mysql_election.client, table_name).await;
1302    }
1303
1304    #[tokio::test]
1305    async fn test_reelection_with_idle_session_timeout() {
1306        maybe_skip_mysql_integration_test!();
1307        let leader_value = "test_leader".to_string();
1308        let uuid = uuid::Uuid::new_v4().to_string();
1309        let table_name = "test_reelection_greptime_metakv";
1310        let candidate_lease_ttl = Duration::from_secs(5);
1311        let meta_lease_ttl = Duration::from_secs(5);
1312        let execution_timeout = Duration::from_secs(10);
1313        let idle_session_timeout = Duration::from_secs(2);
1314        let client = create_mysql_client(Some(table_name), execution_timeout, idle_session_timeout)
1315            .await
1316            .unwrap();
1317
1318        let (tx, _) = broadcast::channel(100);
1319        let leader_mysql_election = MySqlElection {
1320            leader_value: leader_value.clone(),
1321            client,
1322            is_leader: AtomicBool::new(false),
1323            leader_infancy: AtomicBool::new(true),
1324            leader_watcher: tx,
1325            store_key_prefix: uuid,
1326            candidate_lease_ttl,
1327            meta_lease_ttl,
1328            sql_set: ElectionSqlFactory::new(table_name).build(),
1329        };
1330
1331        elected(&leader_mysql_election, table_name, None)
1332            .await
1333            .unwrap();
1334        let lease = get_lease(&leader_mysql_election).await.unwrap();
1335        assert_eq!(lease.leader_value, leader_value);
1336        assert!(lease.expire_time > lease.current);
1337        assert!(leader_mysql_election.is_leader());
1338        // Wait for mysql server close the inactive connection.
1339        tokio::time::sleep(Duration::from_millis(2100)).await;
1340        // Should be failed.
1341        leader_mysql_election
1342            .client
1343            .lock()
1344            .await
1345            .query(sqlx::query("SELECT 1"), "SELECT 1")
1346            .await
1347            .unwrap_err();
1348        // Reset the client.
1349        leader_mysql_election
1350            .client
1351            .lock()
1352            .await
1353            .reset_client()
1354            .await
1355            .unwrap();
1356
1357        // Should able to re-elected.
1358        elected(&leader_mysql_election, table_name, Some(lease.clone()))
1359            .await
1360            .unwrap();
1361        let lease = get_lease(&leader_mysql_election).await.unwrap();
1362        assert_eq!(lease.leader_value, leader_value);
1363        assert!(lease.expire_time > lease.current);
1364        assert!(leader_mysql_election.is_leader());
1365        drop_table(&leader_mysql_election.client, table_name).await;
1366    }
1367
1368    #[tokio::test]
1369    async fn test_elected_and_step_down() {
1370        maybe_skip_mysql_integration_test!();
1371        let leader_value = "test_leader".to_string();
1372        let candidate_lease_ttl = Duration::from_secs(5);
1373        let meta_lease_ttl = Duration::from_secs(2);
1374        let execution_timeout = Duration::from_secs(10);
1375        let idle_session_timeout = Duration::from_secs(0);
1376        let uuid = uuid::Uuid::new_v4().to_string();
1377        let table_name = "test_elected_and_step_down_greptime_metakv";
1378        let client = create_mysql_client(Some(table_name), execution_timeout, idle_session_timeout)
1379            .await
1380            .unwrap();
1381
1382        let (tx, mut rx) = broadcast::channel(100);
1383        let leader_mysql_election = MySqlElection {
1384            leader_value: leader_value.clone(),
1385            client,
1386            is_leader: AtomicBool::new(false),
1387            leader_infancy: AtomicBool::new(true),
1388            leader_watcher: tx,
1389            store_key_prefix: uuid,
1390            candidate_lease_ttl,
1391            meta_lease_ttl,
1392            sql_set: ElectionSqlFactory::new(table_name).build(),
1393        };
1394
1395        elected(&leader_mysql_election, table_name, None)
1396            .await
1397            .unwrap();
1398        let lease = get_lease(&leader_mysql_election).await.unwrap();
1399        assert_eq!(lease.leader_value, leader_value);
1400        assert!(lease.expire_time > lease.current);
1401        assert!(leader_mysql_election.is_leader());
1402
1403        match rx.recv().await {
1404            Ok(LeaderChangeMessage::Elected(key)) => {
1405                assert_eq!(String::from_utf8_lossy(key.name()), leader_value);
1406                assert_eq!(
1407                    String::from_utf8_lossy(key.key()),
1408                    leader_mysql_election.election_key()
1409                );
1410                assert_eq!(key.lease_id(), i64::default());
1411                assert_eq!(key.revision(), i64::default());
1412            }
1413            _ => panic!("Expected LeaderChangeMessage::Elected"),
1414        }
1415
1416        leader_mysql_election
1417            .step_down_without_lock()
1418            .await
1419            .unwrap();
1420        let lease = get_lease(&leader_mysql_election).await.unwrap();
1421        assert_eq!(lease.leader_value, leader_value);
1422        assert!(!leader_mysql_election.is_leader());
1423
1424        match rx.recv().await {
1425            Ok(LeaderChangeMessage::StepDown(key)) => {
1426                assert_eq!(String::from_utf8_lossy(key.name()), leader_value);
1427                assert_eq!(
1428                    String::from_utf8_lossy(key.key()),
1429                    leader_mysql_election.election_key()
1430                );
1431                assert_eq!(key.lease_id(), i64::default());
1432                assert_eq!(key.revision(), i64::default());
1433            }
1434            _ => panic!("Expected LeaderChangeMessage::StepDown"),
1435        }
1436
1437        elected(&leader_mysql_election, table_name, Some(lease.clone()))
1438            .await
1439            .unwrap();
1440        let lease = get_lease(&leader_mysql_election).await.unwrap();
1441        assert_eq!(lease.leader_value, leader_value);
1442        assert!(lease.expire_time > lease.current);
1443        assert!(leader_mysql_election.is_leader());
1444
1445        match rx.recv().await {
1446            Ok(LeaderChangeMessage::Elected(key)) => {
1447                assert_eq!(String::from_utf8_lossy(key.name()), leader_value);
1448                assert_eq!(
1449                    String::from_utf8_lossy(key.key()),
1450                    leader_mysql_election.election_key()
1451                );
1452                assert_eq!(key.lease_id(), i64::default());
1453                assert_eq!(key.revision(), i64::default());
1454            }
1455            _ => panic!("Expected LeaderChangeMessage::Elected"),
1456        }
1457
1458        drop_table(&leader_mysql_election.client, table_name).await;
1459    }
1460
1461    #[tokio::test]
1462    async fn test_campaign() {
1463        maybe_skip_mysql_integration_test!();
1464        let leader_value = "test_leader".to_string();
1465        let uuid = uuid::Uuid::new_v4().to_string();
1466        let table_name = "test_leader_action_greptime_metakv";
1467        let candidate_lease_ttl = Duration::from_secs(5);
1468        let meta_lease_ttl = Duration::from_secs(2);
1469        let execution_timeout = Duration::from_secs(10);
1470        let idle_session_timeout = Duration::from_secs(0);
1471        let client = create_mysql_client(Some(table_name), execution_timeout, idle_session_timeout)
1472            .await
1473            .unwrap();
1474
1475        let (tx, mut rx) = broadcast::channel(100);
1476        let leader_mysql_election = MySqlElection {
1477            leader_value: leader_value.clone(),
1478            client,
1479            is_leader: AtomicBool::new(false),
1480            leader_infancy: AtomicBool::new(true),
1481            leader_watcher: tx,
1482            store_key_prefix: uuid,
1483            candidate_lease_ttl,
1484            meta_lease_ttl,
1485            sql_set: ElectionSqlFactory::new(table_name).build(),
1486        };
1487
1488        // Step 1: No leader exists, campaign and elected.
1489        leader_mysql_election.do_campaign().await.unwrap();
1490        let lease = get_lease(&leader_mysql_election).await.unwrap();
1491        assert_eq!(lease.leader_value, leader_value);
1492        assert!(lease.expire_time > lease.current);
1493        assert!(leader_mysql_election.is_leader());
1494
1495        match rx.recv().await {
1496            Ok(LeaderChangeMessage::Elected(key)) => {
1497                assert_eq!(String::from_utf8_lossy(key.name()), leader_value);
1498                assert_eq!(
1499                    String::from_utf8_lossy(key.key()),
1500                    leader_mysql_election.election_key()
1501                );
1502                assert_eq!(key.lease_id(), i64::default());
1503                assert_eq!(key.revision(), i64::default());
1504            }
1505            _ => panic!("Expected LeaderChangeMessage::Elected"),
1506        }
1507
1508        // Step 2: As a leader, renew the lease.
1509        leader_mysql_election.do_campaign().await.unwrap();
1510        let new_lease = get_lease(&leader_mysql_election).await.unwrap();
1511        assert_eq!(lease.leader_value, leader_value);
1512        // The lease should be renewed.
1513        assert!(new_lease.expire_time > lease.expire_time);
1514        assert!(new_lease.expire_time > new_lease.current);
1515        assert!(leader_mysql_election.is_leader());
1516
1517        // Step 3: Something wrong, the leader lease expired.
1518        tokio::time::sleep(meta_lease_ttl + Duration::from_secs(1)).await;
1519        leader_mysql_election.do_campaign().await.unwrap();
1520        let lease = get_lease(&leader_mysql_election).await.unwrap();
1521        assert_eq!(lease.leader_value, leader_value);
1522        assert!(lease.expire_time <= lease.current);
1523        assert!(!leader_mysql_election.is_leader());
1524
1525        match rx.recv().await {
1526            Ok(LeaderChangeMessage::StepDown(key)) => {
1527                assert_eq!(String::from_utf8_lossy(key.name()), leader_value);
1528                assert_eq!(
1529                    String::from_utf8_lossy(key.key()),
1530                    leader_mysql_election.election_key()
1531                );
1532                assert_eq!(key.lease_id(), i64::default());
1533                assert_eq!(key.revision(), i64::default());
1534            }
1535            _ => panic!("Expected LeaderChangeMessage::StepDown"),
1536        }
1537
1538        // Step 4: Re-elect itself.
1539        leader_mysql_election.do_campaign().await.unwrap();
1540        let lease = get_lease(&leader_mysql_election).await.unwrap();
1541        assert_eq!(lease.leader_value, leader_value);
1542        assert!(lease.expire_time > lease.current);
1543        assert!(leader_mysql_election.is_leader());
1544
1545        match rx.recv().await {
1546            Ok(LeaderChangeMessage::Elected(key)) => {
1547                assert_eq!(String::from_utf8_lossy(key.name()), leader_value);
1548                assert_eq!(
1549                    String::from_utf8_lossy(key.key()),
1550                    leader_mysql_election.election_key()
1551                );
1552                assert_eq!(key.lease_id(), i64::default());
1553                assert_eq!(key.revision(), i64::default());
1554            }
1555            _ => panic!("Expected LeaderChangeMessage::Elected"),
1556        }
1557
1558        // Step 5: Something wrong, the leader key is deleted by other followers.
1559        {
1560            let client = leader_mysql_election.client.lock().await;
1561            let mut executor = Executor::Default(client);
1562            leader_mysql_election
1563                .delete_value(&leader_mysql_election.election_key(), &mut executor)
1564                .await
1565                .unwrap();
1566        }
1567        leader_mysql_election.do_campaign().await.unwrap();
1568        let res = get_lease(&leader_mysql_election).await;
1569        assert!(res.is_none());
1570        assert!(!leader_mysql_election.is_leader());
1571
1572        match rx.recv().await {
1573            Ok(LeaderChangeMessage::StepDown(key)) => {
1574                assert_eq!(String::from_utf8_lossy(key.name()), leader_value);
1575                assert_eq!(
1576                    String::from_utf8_lossy(key.key()),
1577                    leader_mysql_election.election_key()
1578                );
1579                assert_eq!(key.lease_id(), i64::default());
1580                assert_eq!(key.revision(), i64::default());
1581            }
1582            _ => panic!("Expected LeaderChangeMessage::StepDown"),
1583        }
1584
1585        // Step 6: Re-elect itself.
1586        leader_mysql_election.do_campaign().await.unwrap();
1587        let lease = get_lease(&leader_mysql_election).await.unwrap();
1588        assert_eq!(lease.leader_value, leader_value);
1589        assert!(lease.expire_time > lease.current);
1590        assert!(leader_mysql_election.is_leader());
1591
1592        match rx.recv().await {
1593            Ok(LeaderChangeMessage::Elected(key)) => {
1594                assert_eq!(String::from_utf8_lossy(key.name()), leader_value);
1595                assert_eq!(
1596                    String::from_utf8_lossy(key.key()),
1597                    leader_mysql_election.election_key()
1598                );
1599                assert_eq!(key.lease_id(), i64::default());
1600                assert_eq!(key.revision(), i64::default());
1601            }
1602            _ => panic!("Expected LeaderChangeMessage::Elected"),
1603        }
1604
1605        // Step 7: Something wrong, the leader key changed by others.
1606        let another_leader_key = "another_leader";
1607        {
1608            let client = leader_mysql_election.client.lock().await;
1609            let mut executor = Executor::Default(client);
1610            leader_mysql_election
1611                .delete_value(&leader_mysql_election.election_key(), &mut executor)
1612                .await
1613                .unwrap();
1614            leader_mysql_election
1615                .put_value_with_lease(
1616                    &leader_mysql_election.election_key(),
1617                    another_leader_key,
1618                    10,
1619                    &mut executor,
1620                )
1621                .await
1622                .unwrap();
1623        }
1624        leader_mysql_election.do_campaign().await.unwrap();
1625        let lease = get_lease(&leader_mysql_election).await.unwrap();
1626        // Different from pg, mysql will not delete the key, just step down.
1627        assert_eq!(lease.leader_value, another_leader_key);
1628        assert!(lease.expire_time > lease.current);
1629        assert!(!leader_mysql_election.is_leader());
1630
1631        match rx.recv().await {
1632            Ok(LeaderChangeMessage::StepDown(key)) => {
1633                assert_eq!(String::from_utf8_lossy(key.name()), leader_value);
1634                assert_eq!(
1635                    String::from_utf8_lossy(key.key()),
1636                    leader_mysql_election.election_key()
1637                );
1638                assert_eq!(key.lease_id(), i64::default());
1639                assert_eq!(key.revision(), i64::default());
1640            }
1641            _ => panic!("Expected LeaderChangeMessage::StepDown"),
1642        }
1643
1644        drop_table(&leader_mysql_election.client, table_name).await;
1645    }
1646
1647    #[tokio::test]
1648    async fn test_follower_action() {
1649        maybe_skip_mysql_integration_test!();
1650        common_telemetry::init_default_ut_logging();
1651        let candidate_lease_ttl = Duration::from_secs(5);
1652        let meta_lease_ttl = Duration::from_secs(1);
1653        let execution_timeout = Duration::from_secs(10);
1654        let idle_session_timeout = Duration::from_secs(0);
1655        let uuid = uuid::Uuid::new_v4().to_string();
1656        let table_name = "test_follower_action_greptime_metakv";
1657
1658        let follower_client =
1659            create_mysql_client(Some(table_name), execution_timeout, idle_session_timeout)
1660                .await
1661                .unwrap();
1662        let (tx, mut rx) = broadcast::channel(100);
1663        let follower_mysql_election = MySqlElection {
1664            leader_value: "test_follower".to_string(),
1665            client: follower_client,
1666            is_leader: AtomicBool::new(false),
1667            leader_infancy: AtomicBool::new(true),
1668            leader_watcher: tx,
1669            store_key_prefix: uuid.clone(),
1670            candidate_lease_ttl,
1671            meta_lease_ttl,
1672            sql_set: ElectionSqlFactory::new(table_name).build(),
1673        };
1674
1675        let leader_client =
1676            create_mysql_client(Some(table_name), execution_timeout, idle_session_timeout)
1677                .await
1678                .unwrap();
1679        let (tx, _) = broadcast::channel(100);
1680        let leader_mysql_election = MySqlElection {
1681            leader_value: "test_leader".to_string(),
1682            client: leader_client,
1683            is_leader: AtomicBool::new(false),
1684            leader_infancy: AtomicBool::new(true),
1685            leader_watcher: tx,
1686            store_key_prefix: uuid,
1687            candidate_lease_ttl,
1688            meta_lease_ttl,
1689            sql_set: ElectionSqlFactory::new(table_name).build(),
1690        };
1691
1692        leader_mysql_election.do_campaign().await.unwrap();
1693
1694        // Step 1: As a follower, the leader exists and the lease is not expired. Do nothing.
1695        follower_mysql_election.do_campaign().await.unwrap();
1696
1697        // Step 2: As a follower, the leader exists but the lease expired. Re-elect itself.
1698        tokio::time::sleep(meta_lease_ttl + Duration::from_secs(1)).await;
1699        follower_mysql_election.do_campaign().await.unwrap();
1700        assert!(follower_mysql_election.is_leader());
1701
1702        match rx.recv().await {
1703            Ok(LeaderChangeMessage::Elected(key)) => {
1704                assert_eq!(String::from_utf8_lossy(key.name()), "test_follower");
1705                assert_eq!(
1706                    String::from_utf8_lossy(key.key()),
1707                    follower_mysql_election.election_key()
1708                );
1709                assert_eq!(key.lease_id(), i64::default());
1710                assert_eq!(key.revision(), i64::default());
1711            }
1712            _ => panic!("Expected LeaderChangeMessage::Elected"),
1713        }
1714
1715        drop_table(&follower_mysql_election.client, table_name).await;
1716    }
1717
1718    #[tokio::test]
1719    async fn test_wait_timeout() {
1720        maybe_skip_mysql_integration_test!();
1721        common_telemetry::init_default_ut_logging();
1722        let execution_timeout = Duration::from_secs(10);
1723        let idle_session_timeout = Duration::from_secs(1);
1724
1725        let client = create_mysql_client(None, execution_timeout, idle_session_timeout)
1726            .await
1727            .unwrap();
1728        tokio::time::sleep(Duration::from_millis(1100)).await;
1729        // Wait for the idle session timeout.
1730        let err = client
1731            .lock()
1732            .await
1733            .query(sqlx::query("SELECT 1"), "SELECT 1")
1734            .await
1735            .unwrap_err();
1736        assert_matches!(err, error::Error::MySqlExecution { .. });
1737        // Reset the client and try again.
1738        client.lock().await.reset_client().await.unwrap();
1739        let _ = client
1740            .lock()
1741            .await
1742            .query(sqlx::query("SELECT 1"), "SELECT 1")
1743            .await
1744            .unwrap();
1745    }
1746}
meta_srv/election/rds/mysql.rs

meta_srv/election/rds/
mysql.rs
