kjuulh/noleader
1
0
Fork 0
Code Issues 1 Pull Requests Actions Packages Projects Releases 2 Wiki Activity

feat: add postgres

Browse Source
This commit is contained in:
kjuulh
2025-09-24 21:16:35 +02:00
parent 8b27219af5
commit fc190a12d4
13 changed files with 1152 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
crates/noleader/Cargo.toml
View File

@@ -19,6 +19,13 @@ tokio.workspace = true
tokio-util = "0.7"
rand = "0.9.1"
async-trait = "0.1.89"
# fork until dangerous set migrate table name is stable. Should be any version after 8.6
sqlx = { git = "https://github.com/launchbadge/sqlx", features = [
"uuid",
"postgres",
"runtime-tokio",
"tls-rustls",
], rev = "064d649abdfd1742e5fdcc20176a6b415b9c25d3" }
[dev-dependencies]
tracing-subscriber = { version = "0.3", features = ["env-filter"] }

6
crates/noleader/examples/lots_of_candidates/main.rs
View File

@@ -1,4 +1,3 @@
use tokio_util::sync::CancellationToken;
use tracing_subscriber::EnvFilter;
#[tokio::main]
@@ -32,10 +31,7 @@ async fn main() -> anyhow::Result<()> {
async move {
tracing::debug!(leader_id, "starting leader");
leader
.start(CancellationToken::default())
.await
.expect("to succeed");
leader.start().await.expect("to succeed");
}
});

96
crates/noleader/examples/lots_of_postgres/main.rs Normal file
View File

@@ -0,0 +1,96 @@
use anyhow::Context;
use tokio_util::sync::CancellationToken;
use tracing_subscriber::EnvFilter;
#[tokio::main]
async fn main() -> anyhow::Result<()> {
// Set up logger
tracing_subscriber::fmt()
.with_env_filter(
EnvFilter::from_default_env()
.add_directive("noleader=debug".parse().unwrap())
.add_directive("lots_of_candidates=debug".parse().unwrap())
.add_directive("info".parse().unwrap()),
)
.init();
let mykey = "myleaderkey";
let mut handles = Vec::new();
let db_url = &std::env::var("DATABASE_URL").context("DATABASE_URL is missing")?;
let pool = sqlx::PgPool::connect_lazy(db_url)?;
let cancel = CancellationToken::new();
let mut cancelled_resp = Vec::new();
tokio::spawn({
let cancel = cancel.clone();
async move {
tokio::signal::ctrl_c().await.expect("to receive shutdown");
cancel.cancel();
}
});
for _ in 0..100 {
let pool = pool.clone();
let cancel = cancel.child_token();
let item_cancellation = CancellationToken::new();
cancelled_resp.push(item_cancellation.child_token());
let handle = tokio::spawn(async move {
let mut leader = noleader::Leader::new_postgres_pool(mykey, pool);
leader.with_cancellation(cancel);
let leader_id = leader.leader_id().await.to_string();
tokio::spawn({
let leader = leader.clone();
let leader_id = leader_id.clone();
async move {
tracing::debug!(leader_id, "starting leader");
let res = leader.start().await;
tracing::warn!("shutting down");
item_cancellation.cancel();
if let Err(e) = res {
tracing::error!("lots failed: {e:?}");
}
}
});
loop {
tokio::time::sleep(std::time::Duration::from_millis(10000)).await;
match leader.is_leader().await {
noleader::Status::Leader => {
tracing::info!(leader_id, "is leader");
}
noleader::Status::Candidate => {
//tracing::debug!("is candiate");
}
}
}
#[allow(unreachable_code)]
Ok::<(), anyhow::Error>(())
});
handles.push(handle);
}
for cancel in cancelled_resp {
cancel.cancelled().await;
}
for handle in handles {
handle.abort();
}
Ok(())
}

49
crates/noleader/examples/postgres/main.rs Normal file
View File

@@ -0,0 +1,49 @@
use anyhow::Context;
use tokio::signal;
use tracing_subscriber::EnvFilter;
#[tokio::main]
async fn main() -> anyhow::Result<()> {
// Set up logger
tracing_subscriber::fmt()
.with_env_filter(
EnvFilter::from_default_env()
.add_directive("noleader=debug".parse().unwrap())
.add_directive("lots_of_candidates=debug".parse().unwrap())
.add_directive("info".parse().unwrap()),
)
.init();
let mykey = "postgres";
let mut leader = noleader::Leader::new_postgres(
mykey,
&std::env::var("DATABASE_URL").context("DATABASE_URL is missing")?,
);
leader.with_cancel_task(async move {
signal::ctrl_c().await.unwrap();
});
let leader_id = leader.leader_id().await.to_string();
leader
.acquire_and_run({
move |token| {
let leader_id = leader_id.clone();
async move {
loop {
if token.is_cancelled() {
return Ok(());
}
tracing::info!(leader_id, "do work as leader");
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
}
}
}
})
.await?;
Ok(())
}

8
crates/noleader/migrations/postgres/20250924114540_initial.sql Normal file
View File

@@ -0,0 +1,8 @@
-- Add migration script here
CREATE TABLE IF NOT EXISTS noleader_leaders (
key TEXT PRIMARY KEY NOT NULL,
value TEXT NOT NULL,
revision BIGINT NOT NULL,
heartbeat TIMESTAMPTZ NOT NULL DEFAULT now()
);

16
crates/noleader/src/backend.rs
View File

@@ -1,8 +1,9 @@
use std::{ops::Deref, sync::Arc};
use crate::backend::nats::NatsBackend;
use crate::backend::{nats::NatsBackend, postgres::PostgresBackend};
mod nats;
mod postgres;
pub struct Backend {
inner: Arc<dyn BackendEdge + Send + Sync + 'static>,
@@ -20,6 +21,18 @@ impl Backend {
inner: Arc::new(NatsBackend::new(client, bucket)),
}
}
pub fn postgres(database_url: &str) -> Self {
Self {
inner: Arc::new(PostgresBackend::new(database_url)),
}
}
pub fn postgres_with_pool(pool: sqlx::PgPool) -> Self {
Self {
inner: Arc::new(PostgresBackend::new_with_pool("bogus", pool)),
}
}
}
impl Deref for Backend {
@@ -35,6 +48,7 @@ pub trait BackendEdge {
async fn setup(&self) -> anyhow::Result<()>;
async fn get(&self, key: &Key) -> anyhow::Result<LeaderValue>;
async fn update(&self, key: &Key, val: &LeaderId) -> anyhow::Result<()>;
async fn release(&self, key: &Key, val: &LeaderId) -> anyhow::Result<()>;
}
pub enum LeaderValue {

6
crates/noleader/src/backend/nats.rs
View File

@@ -122,4 +122,10 @@ impl BackendEdge for NatsBackend {
Ok(())
}
async fn release(&self, _key: &Key, _val: &LeaderId) -> anyhow::Result<()> {
// TODO: implement release for nats
Ok(())
}
}

211
crates/noleader/src/backend/postgres.rs Normal file
View File

@@ -0,0 +1,211 @@
use std::{
sync::atomic::{AtomicU64, Ordering},
time::Duration,
};
use anyhow::Context;
use sqlx::{PgPool, postgres::PgPoolOptions};
use tokio::sync::OnceCell;
use crate::backend::{BackendEdge, Key, LeaderId, LeaderValue};
pub struct PostgresBackend {
database_url: String,
revision: AtomicU64,
pool: OnceCell<PgPool>,
migrated: OnceCell<()>,
}
impl PostgresBackend {
pub fn new(database_url: &str) -> Self {
Self {
database_url: database_url.into(),
revision: AtomicU64::new(0),
pool: OnceCell::new(),
migrated: OnceCell::new(),
}
}
pub fn new_with_pool(database_url: &str, pool: PgPool) -> Self {
Self {
database_url: database_url.into(),
revision: AtomicU64::new(0),
pool: OnceCell::new_with(Some(pool)),
migrated: OnceCell::new(),
}
}
async fn db(&self) -> anyhow::Result<PgPool> {
let pool = self
.pool
.get_or_try_init(|| async move {
PgPoolOptions::new()
.max_connections(1)
.min_connections(0)
.idle_timeout(Some(Duration::from_secs(5)))
.connect_lazy(&self.database_url)
.context("connect postgres noleader")
})
.await?;
Ok(pool.clone())
}
async fn migrate(&self) -> anyhow::Result<()> {
self.migrated
.get_or_try_init(|| async move {
let db = self.db().await?;
let mut migrate = sqlx::migrate!("./migrations/postgres/");
migrate
.set_locking(false)
.dangerous_set_table_name("_sqlx_noleader_migrations")
.run(&db)
.await
.context("migrate noleader")?;
Ok::<_, anyhow::Error>(())
})
.await?;
Ok(())
}
}
#[async_trait::async_trait]
impl BackendEdge for PostgresBackend {
async fn setup(&self) -> anyhow::Result<()> {
self.migrate().await?;
Ok(())
}
async fn get(&self, key: &Key) -> anyhow::Result<LeaderValue> {
let rec = sqlx::query!(
"
SELECT value, revision
FROM noleader_leaders
WHERE
key = $1
AND heartbeat >= now() - interval '60 seconds'
",
key.0
)
.fetch_optional(&self.db().await?)
.await
.context("get noleader key")?;
let Some(val) = rec else {
anyhow::bail!("key doesn't exist, we've lost leadership status")
};
// Update our local revision to match what's in the database
self.revision.store(val.revision as u64, Ordering::Relaxed);
let Ok(id) = uuid::Uuid::parse_str(&val.value) else {
tracing::warn!("value is not a valid uuid: {}", val.value);
return Ok(LeaderValue::Unknown);
};
Ok(LeaderValue::Found { id: id.into() })
}
async fn update(&self, key: &Key, val: &LeaderId) -> anyhow::Result<()> {
let current_rev = self.revision.load(Ordering::Relaxed);
let new_rev = current_rev + 1;
let res = sqlx::query!(
r#"
INSERT INTO noleader_leaders (key, value, revision, heartbeat)
VALUES ($1, $2, $3, now())
ON CONFLICT (key)
DO UPDATE SET
value = EXCLUDED.value,
revision = EXCLUDED.revision,
heartbeat = now()
WHERE
(
-- Normal case: revision matches (we're the current leader updating)
noleader_leaders.revision = $4
OR
-- Override case: heartbeat is old (stale leader)
noleader_leaders.heartbeat < now() - INTERVAL '60 seconds'
)
RETURNING value, revision
"#,
key.0,
val.0.to_string(),
new_rev as i64, // new revision
current_rev as i64, // expected current revision
)
.fetch_optional(&self.db().await?)
.await;
let res = match res {
Ok(res) => res,
Err(e) => match &e {
sqlx::Error::Database(database_error) => {
if database_error.is_unique_violation() {
anyhow::bail!("update conflict: another leader holds lock")
} else {
anyhow::bail!(e);
}
}
_ => {
anyhow::bail!(e);
}
},
};
match res {
Some(rec) => {
if rec.value == val.0.to_string() && rec.revision == new_rev as i64 {
tracing::debug!(
val = val.0.to_string(),
revision = rec.revision,
"successfully updated leader"
);
// Only update our local revision if the update succeeded with our expected value
self.revision.store(rec.revision as u64, Ordering::Relaxed);
} else {
anyhow::bail!(
"update conflict: expected value={}, revision={}, got value={}, revision={}",
val.0.to_string(),
new_rev,
rec.value,
rec.revision
);
}
}
None => {
anyhow::bail!(
"update rejected: another leader is holding the lock or revision mismatch"
)
}
}
Ok(())
}
async fn release(&self, key: &Key, val: &LeaderId) -> anyhow::Result<()> {
let rev = self.revision.load(Ordering::Relaxed);
sqlx::query!(
"
DELETE FROM noleader_leaders
WHERE
key = $1
AND value = $2
AND revision = $3
",
key.0,
val.0.to_string(),
rev as i64, // new revision
)
.execute(&self.db().await?)
.await
.context("failed to release lock, it will expire naturally")?;
Ok(())
}
}

71
crates/noleader/src/lib.rs
View File

@@ -20,6 +20,8 @@ pub struct Leader {
shutting_down: Arc<AtomicBool>,
is_leader: Arc<AtomicBool>,
inner: Arc<RwLock<InnerLeader>>,
cancellation: CancellationToken,
}
const DEFAULT_INTERVAL: Duration = std::time::Duration::from_secs(10);
@@ -31,6 +33,7 @@ impl Leader {
shutting_down: Arc::new(AtomicBool::new(false)),
is_leader: Arc::new(AtomicBool::new(false)),
inner: Arc::new(RwLock::new(InnerLeader::new(backend, key))),
cancellation: CancellationToken::new(),
}
}
@@ -38,21 +41,48 @@ impl Leader {
Self::new(key, Backend::nats(client, bucket))
}
pub fn new_postgres(key: &str, database_url: &str) -> Self {
Self::new(key, Backend::postgres(database_url))
}
pub fn new_postgres_pool(key: &str, pool: sqlx::PgPool) -> Self {
Self::new(key, Backend::postgres_with_pool(pool))
}
pub fn with_cancellation(&mut self, cancellation: CancellationToken) -> &mut Self {
self.cancellation = cancellation;
self
}
pub fn with_cancel_task<T>(&mut self, f: T) -> &mut Self
where
T: Future<Output = ()> + Send + 'static,
{
let cancel = self.cancellation.clone();
tokio::spawn(async move {
f.await;
cancel.cancel();
});
self
}
pub async fn acquire_and_run<F, Fut>(&self, f: F) -> anyhow::Result<()>
where
F: Fn(CancellationToken) -> Fut,
Fut: Future<Output = anyhow::Result<()>> + Send + 'static,
{
let parent_token = CancellationToken::default();
let parent_token = self.cancellation.clone();
let s = self.clone();
let server_token = parent_token.child_token();
// Start the server election process in another task, this is because start is blocking
let handle = tokio::spawn({
let server_token = server_token.child_token();
async move {
match s.start(server_token).await {
match s.start().await {
Ok(_) => {}
Err(e) => tracing::error!("leader election process failed: {}", e),
}
@@ -72,6 +102,11 @@ impl Leader {
server_token.cancel();
// Close down the task as well, it should already be stopped, but this forces the task to close
handle.abort();
{
self.inner.write().await.cleanup().await?;
}
res?;
Ok(())
@@ -96,11 +131,21 @@ impl Leader {
Fut: Future<Output = anyhow::Result<()>> + Send + 'static,
{
loop {
if cancellation_token.is_cancelled() {
return Ok(());
}
let cancellation_token = cancellation_token.child_token();
let is_leader = self.is_leader.clone();
if !is_leader.load(Ordering::Relaxed) {
tokio::time::sleep(Duration::from_secs(1)).await;
tokio::select! {
_ = tokio::time::sleep(Duration::from_secs(1)) => {}
_ = cancellation_token.cancelled() => {
return Ok(());
}
}
continue;
}
@@ -111,7 +156,7 @@ impl Leader {
tokio::select! {
_ = tokio::time::sleep(std::time::Duration::from_millis(500)) => {}
_ = cancellation_token.cancelled() => {
break;
return;
}
}
@@ -123,6 +168,7 @@ impl Leader {
});
let res = f(child_token).await;
guard.abort();
res?;
}
@@ -133,7 +179,7 @@ impl Leader {
inner.leader_id.clone().into()
}
pub async fn start(&self, cancellation_token: CancellationToken) -> anyhow::Result<()> {
pub async fn start(&self) -> anyhow::Result<()> {
let mut attempts = 1;
{
@@ -153,7 +199,7 @@ impl Leader {
tokio::select! {
_ = sleep_fut => {},
_ = cancellation_token.cancelled() => {
_ = self.cancellation.cancelled() => {
self.shutting_down.store(true, std::sync::atomic::Ordering::Relaxed); // Ordering can be relaxed, because our operation is an atomic update
return Ok(())
}
@@ -214,7 +260,6 @@ struct InnerLeader {
key: Key,
leader_id: LeaderId,
revision: u64,
}
#[derive(Default, Clone)]
@@ -230,7 +275,6 @@ impl InnerLeader {
Self {
backend,
leader_id: LeaderId::new(),
revision: u64::MIN,
key: key.into(),
@@ -275,6 +319,15 @@ impl InnerLeader {
Ok(())
}
pub async fn cleanup(&self) -> anyhow::Result<()> {
self.backend
.release(&self.key, &self.leader_id)
.await
.context("cleanup")?;
Ok(())
}
async fn update_leadership(&mut self) -> anyhow::Result<()> {
let val = self
.backend