Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 39 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,7 @@ opentelemetry_sdk = { version = "0.32", features = ["trace", "rt-tokio
opentelemetry-otlp = { version = "0.32", default-features = false, features = ["trace", "grpc-tonic", "tls-ring"] }
metrics = "0.24"
metrics-exporter-prometheus = "0.18"
metrics-util = "0.20"

# Error handling
thiserror = "2"
Expand Down
88 changes: 86 additions & 2 deletions crates/buzz-db/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -52,8 +52,8 @@ pub use error::{DbError, Result};
pub use event::{EventQuery, ReactionEventInsertOutcome};

use chrono::{DateTime, Utc};
use sqlx::postgres::PgPoolOptions;
use sqlx::{PgPool, QueryBuilder, Row};
use sqlx::postgres::{PgConnection, PgPoolOptions};
use sqlx::{Connection, PgPool, QueryBuilder, Row};
use std::time::Duration;
use uuid::Uuid;

Expand Down Expand Up @@ -181,6 +181,27 @@ pub struct DbPoolStats {
pub max: u32,
}

/// Owns the detached Postgres session holding the relay usage-metrics advisory lock.
///
/// The connection deliberately does not return to the main pool: session advisory
/// locks must remain bound to this exact physical connection, and the poller
/// pings it before each leader-only collection tick.
pub struct UsageMetricsLeader {
connection: PgConnection,
}

impl UsageMetricsLeader {
/// Returns whether the lock-owning session is still reachable.
///
/// Bounded to 5 seconds — a blackholed connection (no RST) would otherwise
/// stall the entire poller tick until the OS TCP timeout.
pub async fn is_live(&mut self) -> bool {
tokio::time::timeout(std::time::Duration::from_secs(5), self.connection.ping())
.await
.is_ok_and(|r| r.is_ok())
}
}

/// Configuration for the Postgres connection pool.
#[derive(Debug, Clone)]
pub struct DbConfig {
Expand Down Expand Up @@ -343,6 +364,30 @@ impl Db {
}
}

/// Try to acquire the detached session advisory lock for relay usage metrics.
///
/// The returned guard owns the exact connection that acquired the lock. It is
/// detached from the shared pool so a stable leader neither returns a locked
/// session to other callers nor permanently consumes a pool slot. Dropping the
/// guard closes the connection and releases the session-scoped lock.
pub async fn try_lock_usage_metrics(
&self,
lock_key: i64,
) -> Result<Option<UsageMetricsLeader>> {
let mut connection = self.pool.acquire().await?;
let acquired = sqlx::query_scalar::<_, bool>("SELECT pg_try_advisory_lock($1)")
.bind(lock_key)
.fetch_one(&mut *connection)
.await?;
if acquired {
Ok(Some(UsageMetricsLeader {
connection: connection.detach(),
}))
} else {
Ok(None)
}
}

/// Return total number of communities on this relay.
pub async fn usage_community_count(&self) -> Result<i64> {
usage::community_count(&self.pool).await
Expand Down Expand Up @@ -3510,6 +3555,7 @@ mod tests {
//! channels, that fail-closed chain would go blind.
use super::*;
use buzz_core::CommunityId;
use sqlx::postgres::PgPoolOptions;
use sqlx::{Acquire, PgPool};
use uuid::Uuid;

Expand Down Expand Up @@ -4163,6 +4209,44 @@ mod tests {
assert_eq!(watermark.1, c.id.as_bytes().as_slice());
}

#[tokio::test]
#[ignore = "requires Postgres"]
async fn test_usage_metrics_lock_has_single_owner_and_releases_on_drop() {
let pool = PgPoolOptions::new()
.max_connections(2)
.connect(TEST_DB_URL)
.await
.expect("connect to test DB");
let first = Db::from_pool(pool.clone());
let second = Db::from_pool(pool);
let key = 0x4255_5A5A_4D45_5452;

let mut leader = first
.try_lock_usage_metrics(key)
.await
.expect("first lock attempt")
.expect("first database handle becomes leader");
assert!(leader.is_live().await, "lock owner remains reachable");
assert!(
second
.try_lock_usage_metrics(key)
.await
.expect("second lock attempt")
.is_none(),
"another session cannot become leader while the guard exists"
);

drop(leader);
assert!(
second
.try_lock_usage_metrics(key)
.await
.expect("lock attempt after leader drop")
.is_some(),
"dropping the detached session releases its advisory lock"
);
}

#[tokio::test]
#[ignore = "requires Postgres"]
async fn lookup_community_by_host_matches_case_insensitive_host_index() {
Expand Down
1 change: 1 addition & 0 deletions crates/buzz-relay/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,7 @@ url = { workspace = true }
moka = { workspace = true }
metrics = { workspace = true }
metrics-exporter-prometheus = { workspace = true }
metrics-util = { workspace = true }

[features]
dev = ["buzz-auth/dev"]
Expand Down
Loading
Loading