mirror of
https://github.com/fluxerapp/fluxer
synced 2026-10-07 19:22:14 +09:00
fix(kv): align rust and typescript schema migration (#2328)
This commit is contained in:
Hampus
+101
-4
@@ -10,15 +10,22 @@ use rustls::{
|
||||
};
|
||||
use serde_json::{Map, Number, Value};
|
||||
use std::str::FromStr;
|
||||
use std::time::Duration;
|
||||
use tokio_postgres::{
|
||||
Config as PgConfig, Row,
|
||||
config::SslMode,
|
||||
error::SqlState,
|
||||
types::{ToSql, Type},
|
||||
};
|
||||
use tokio_postgres_rustls::MakeRustlsConnect;
|
||||
|
||||
const POSTGRES_KV_SCHEMA_LOCK_NAMESPACE: i32 = 0x4658_4b56;
|
||||
const POSTGRES_KV_SCHEMA_LOCK_TIMEOUT: &str = "120s";
|
||||
const POSTGRES_KV_SCHEMA_MIGRATION_TIMEOUT: &str = "30min";
|
||||
const POSTGRES_KV_MIGRATION_TABLE: &str = "__fluxer_schema_migrations";
|
||||
const POSTGRES_KV_MESSAGES_PARTITION_MIGRATION: &str = "messages_partition_key_v1";
|
||||
const POSTGRES_KV_SCHEMA_ATTEMPTS: u32 = 3;
|
||||
const POSTGRES_KV_SCHEMA_RETRY_DELAY: Duration = Duration::from_millis(250);
|
||||
const CACHED_JSON_FIELDS: &[&str] = &["message_id"];
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
@@ -156,7 +163,41 @@ WHERE att.attrelid = to_regclass($1)
|
||||
Ok(row.and_then(|row| row.get::<_, Option<bool>>("c_collated")) == Some(true))
|
||||
}
|
||||
|
||||
fn is_concurrent_ddl_conflict(error: &anyhow::Error) -> bool {
|
||||
error
|
||||
.chain()
|
||||
.filter_map(|cause| cause.downcast_ref::<tokio_postgres::Error>())
|
||||
.filter_map(tokio_postgres::Error::code)
|
||||
.any(|code| {
|
||||
*code == SqlState::UNIQUE_VIOLATION
|
||||
|| *code == SqlState::DUPLICATE_TABLE
|
||||
|| *code == SqlState::DUPLICATE_OBJECT
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn ensure_kv_schema(pool: &Pool, kv_table: &str) -> anyhow::Result<()> {
|
||||
let mut attempt = 1;
|
||||
loop {
|
||||
match ensure_kv_schema_once(pool, kv_table).await {
|
||||
Ok(()) => return Ok(()),
|
||||
Err(err) => {
|
||||
if attempt >= POSTGRES_KV_SCHEMA_ATTEMPTS || !is_concurrent_ddl_conflict(&err) {
|
||||
return Err(err);
|
||||
}
|
||||
tracing::warn!(
|
||||
error = ?err,
|
||||
kv_table,
|
||||
attempt,
|
||||
"retrying Postgres KV schema after a concurrent DDL conflict"
|
||||
);
|
||||
attempt += 1;
|
||||
tokio::time::sleep(POSTGRES_KV_SCHEMA_RETRY_DELAY).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn ensure_kv_schema_once(pool: &Pool, kv_table: &str) -> anyhow::Result<()> {
|
||||
let table = quote_identifier(kv_table)?;
|
||||
let old_partition_index = quote_identifier(&format!("{kv_table}_partition_idx"))?;
|
||||
let partition_row_index = quote_identifier(&format!("{kv_table}_partition_row_idx"))?;
|
||||
@@ -185,9 +226,12 @@ pub async fn ensure_kv_schema(pool: &Pool, kv_table: &str) -> anyhow::Result<()>
|
||||
.await
|
||||
.context("failed to acquire Postgres KV schema lock")?;
|
||||
transaction
|
||||
.query_one("SELECT set_config('statement_timeout', '0', true)", &[])
|
||||
.query_one(
|
||||
"SELECT set_config('statement_timeout', $1, true)",
|
||||
&[&POSTGRES_KV_SCHEMA_MIGRATION_TIMEOUT],
|
||||
)
|
||||
.await
|
||||
.context("failed to clear Postgres KV schema lock timeout")?;
|
||||
.context("failed to configure Postgres KV schema migration timeout")?;
|
||||
transaction
|
||||
.batch_execute(&format!(
|
||||
r#"
|
||||
@@ -219,14 +263,67 @@ CREATE INDEX IF NOT EXISTS {partition_row_index} ON {table} (table_name, partiti
|
||||
CREATE INDEX IF NOT EXISTS {expires_index} ON {table} (expires_at) WHERE expires_at IS NOT NULL;
|
||||
CREATE INDEX IF NOT EXISTS {messages_message_index} ON {table} (partition_key, ((CASE WHEN row_data -> 'message_id' ->> 'value' ~ '^-?[0-9]+$' THEN (row_data -> 'message_id' ->> 'value')::bigint END))) WHERE table_name = 'messages';
|
||||
CREATE INDEX IF NOT EXISTS {message_reactions_message_index} ON {table} (partition_key, ((CASE WHEN row_data -> 'message_id' ->> 'value' ~ '^-?[0-9]+$' THEN (row_data -> 'message_id' ->> 'value')::bigint END))) WHERE table_name = 'message_reactions';
|
||||
"#
|
||||
))
|
||||
.await
|
||||
.context("failed to ensure Postgres KV schema")?;
|
||||
let migrated = transaction
|
||||
.query_typed_opt(
|
||||
&format!("SELECT 1 FROM {table} WHERE table_name = $1 AND row_key = $2 LIMIT 1"),
|
||||
&[
|
||||
(&POSTGRES_KV_MIGRATION_TABLE, Type::TEXT),
|
||||
(&POSTGRES_KV_MESSAGES_PARTITION_MIGRATION, Type::TEXT),
|
||||
],
|
||||
)
|
||||
.await
|
||||
.context("failed to read the Postgres KV migration marker")?;
|
||||
if migrated.is_none() {
|
||||
let pending = transaction
|
||||
.query_typed_opt(
|
||||
&format!(
|
||||
r#"
|
||||
SELECT 1
|
||||
FROM {table}
|
||||
WHERE table_name = 'messages'
|
||||
AND partition_key = row_key
|
||||
AND split_part(row_key, chr(31), 3) <> ''
|
||||
LIMIT 1"#
|
||||
),
|
||||
&[],
|
||||
)
|
||||
.await
|
||||
.context("failed to probe the Postgres KV messages partition backfill")?;
|
||||
if pending.is_some() {
|
||||
transaction
|
||||
.batch_execute(&format!(
|
||||
r#"
|
||||
UPDATE {table}
|
||||
SET partition_key = split_part(row_key, chr(31), 1) || chr(31) || split_part(row_key, chr(31), 2)
|
||||
WHERE table_name = 'messages'
|
||||
AND partition_key = row_key
|
||||
AND split_part(row_key, chr(31), 3) <> '';
|
||||
DROP INDEX IF EXISTS {old_partition_index};
|
||||
"#
|
||||
))
|
||||
))
|
||||
.await
|
||||
.context("failed to backfill Postgres KV message partition keys")?;
|
||||
}
|
||||
transaction
|
||||
.execute_typed(
|
||||
&format!(
|
||||
r#"INSERT INTO {table} (table_name, partition_key, row_key, row_data)
|
||||
VALUES ($1, $2, $2, jsonb_build_object('applied_at', now()))
|
||||
ON CONFLICT (table_name, row_key) DO NOTHING"#
|
||||
),
|
||||
&[
|
||||
(&POSTGRES_KV_MIGRATION_TABLE, Type::TEXT),
|
||||
(&POSTGRES_KV_MESSAGES_PARTITION_MIGRATION, Type::TEXT),
|
||||
],
|
||||
)
|
||||
.await
|
||||
.context("failed to record the Postgres KV migration marker")?;
|
||||
}
|
||||
transaction
|
||||
.batch_execute(&format!("DROP INDEX IF EXISTS {old_partition_index};"))
|
||||
.await
|
||||
.context("failed to ensure Postgres KV schema")?;
|
||||
transaction
|
||||
|
||||
@@ -0,0 +1,279 @@
|
||||
// SPDX-License-Identifier: AGPL-3.0-or-later
|
||||
|
||||
use deadpool_postgres::Pool;
|
||||
use fluxer_svc::postgres::{PostgresConfig, connect, ensure_kv_schema};
|
||||
use std::net::TcpListener;
|
||||
use std::process::{Command, Stdio};
|
||||
use std::time::Duration;
|
||||
|
||||
const KV_TABLE_WARMUP: &str = "svc_schema_warmup";
|
||||
const KV_TABLE_BOOT: &str = "svc_schema_boot";
|
||||
const KV_TABLE_HELPER: &str = "svc_schema_helper";
|
||||
const KV_TABLE_RACE: &str = "svc_schema_race";
|
||||
const KV_TABLE_BACKFILL: &str = "svc_schema_backfill";
|
||||
const SEPARATOR: char = '\u{1f}';
|
||||
|
||||
fn docker_available() -> bool {
|
||||
Command::new("docker")
|
||||
.arg("version")
|
||||
.stdin(Stdio::null())
|
||||
.stdout(Stdio::null())
|
||||
.stderr(Stdio::null())
|
||||
.status()
|
||||
.is_ok_and(|status| status.success())
|
||||
}
|
||||
|
||||
fn docker(args: &[&str]) -> anyhow::Result<()> {
|
||||
let output = Command::new("docker").args(args).output()?;
|
||||
if output.status.success() {
|
||||
return Ok(());
|
||||
}
|
||||
anyhow::bail!(
|
||||
"docker {} failed: {}",
|
||||
args.join(" "),
|
||||
String::from_utf8_lossy(&output.stderr).trim().to_owned()
|
||||
)
|
||||
}
|
||||
|
||||
fn free_port() -> anyhow::Result<u16> {
|
||||
let listener = TcpListener::bind("127.0.0.1:0")?;
|
||||
Ok(listener.local_addr()?.port())
|
||||
}
|
||||
|
||||
fn postgres_config(port: u16, kv_table: &str, max_connections: usize) -> PostgresConfig {
|
||||
PostgresConfig {
|
||||
url: None,
|
||||
host: "127.0.0.1".to_owned(),
|
||||
port,
|
||||
database: "fluxer".to_owned(),
|
||||
username: "fluxer".to_owned(),
|
||||
password: Some("fluxer".to_owned()),
|
||||
ssl: false,
|
||||
ssl_ca: None,
|
||||
max_connections,
|
||||
kv_table: kv_table.to_owned(),
|
||||
prepared_statements: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn start_postgres(container: &str, port: u16) -> anyhow::Result<()> {
|
||||
docker(&[
|
||||
"run",
|
||||
"-d",
|
||||
"--name",
|
||||
container,
|
||||
"-e",
|
||||
"POSTGRES_USER=fluxer",
|
||||
"-e",
|
||||
"POSTGRES_PASSWORD=fluxer",
|
||||
"-e",
|
||||
"POSTGRES_DB=fluxer",
|
||||
"-p",
|
||||
&format!("127.0.0.1:{port}:5432"),
|
||||
"postgres:16-alpine",
|
||||
"-c",
|
||||
"fsync=off",
|
||||
"-c",
|
||||
"synchronous_commit=off",
|
||||
])
|
||||
}
|
||||
|
||||
async fn wait_for_postgres(container: &str, port: u16) -> anyhow::Result<()> {
|
||||
for _ in 0..180 {
|
||||
tokio::time::sleep(Duration::from_millis(500)).await;
|
||||
if docker(&[
|
||||
"exec",
|
||||
container,
|
||||
"pg_isready",
|
||||
"-U",
|
||||
"fluxer",
|
||||
"-d",
|
||||
"fluxer",
|
||||
])
|
||||
.is_err()
|
||||
{
|
||||
continue;
|
||||
}
|
||||
if let Ok(pool) = connect(&postgres_config(port, KV_TABLE_WARMUP, 1)).await {
|
||||
pool.close();
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
anyhow::bail!("postgres never came up")
|
||||
}
|
||||
|
||||
fn unlocked_create_table_sql(kv_table: &str) -> String {
|
||||
format!(
|
||||
r#"
|
||||
CREATE TABLE IF NOT EXISTS "{kv_table}" (
|
||||
table_name text NOT NULL,
|
||||
partition_key text COLLATE "C" NOT NULL,
|
||||
row_key text COLLATE "C" NOT NULL,
|
||||
row_data jsonb NOT NULL,
|
||||
expires_at timestamptz,
|
||||
updated_at timestamptz NOT NULL DEFAULT now(),
|
||||
PRIMARY KEY (table_name, row_key)
|
||||
);
|
||||
"#
|
||||
)
|
||||
}
|
||||
|
||||
fn legacy_row_key(id: &str) -> String {
|
||||
format!("\"c\"{SEPARATOR}\"b\"{SEPARATOR}\"{id}\"")
|
||||
}
|
||||
|
||||
fn migrated_partition_key() -> String {
|
||||
format!("\"c\"{SEPARATOR}\"b\"")
|
||||
}
|
||||
|
||||
async fn insert_legacy_message(pool: &Pool, kv_table: &str, id: &str) -> anyhow::Result<()> {
|
||||
let client = pool.get().await?;
|
||||
client
|
||||
.execute(
|
||||
&format!(
|
||||
"INSERT INTO \"{kv_table}\" (table_name, partition_key, row_key, row_data) VALUES ('messages', $1, $1, '{{}}'::jsonb)"
|
||||
),
|
||||
&[&legacy_row_key(id)],
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn partition_key_of(pool: &Pool, kv_table: &str, id: &str) -> anyhow::Result<Option<String>> {
|
||||
let client = pool.get().await?;
|
||||
let row = client
|
||||
.query_opt(
|
||||
&format!(
|
||||
"SELECT partition_key FROM \"{kv_table}\" WHERE table_name = 'messages' AND row_key = $1"
|
||||
),
|
||||
&[&legacy_row_key(id)],
|
||||
)
|
||||
.await?;
|
||||
Ok(row.map(|row| row.get::<_, String>("partition_key")))
|
||||
}
|
||||
|
||||
async fn migration_marker_count(pool: &Pool, kv_table: &str) -> anyhow::Result<i64> {
|
||||
let client = pool.get().await?;
|
||||
let row = client
|
||||
.query_one(
|
||||
&format!(
|
||||
"SELECT count(*) AS n FROM \"{kv_table}\" WHERE table_name = '__fluxer_schema_migrations' AND row_key = 'messages_partition_key_v1'"
|
||||
),
|
||||
&[],
|
||||
)
|
||||
.await?;
|
||||
Ok(row.get::<_, i64>("n"))
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn concurrent_boots_all_ensure_the_kv_schema() -> anyhow::Result<()> {
|
||||
if !docker_available() {
|
||||
eprintln!("skipping: docker is not available");
|
||||
return Ok(());
|
||||
}
|
||||
let port = free_port()?;
|
||||
let container = format!("fluxer-kvschema-boot-{}-{port}", std::process::id());
|
||||
start_postgres(&container, port)?;
|
||||
let result = async {
|
||||
wait_for_postgres(&container, port).await?;
|
||||
let mut handles = Vec::new();
|
||||
for _ in 0..4 {
|
||||
handles.push(tokio::spawn(async move {
|
||||
connect(&postgres_config(port, KV_TABLE_BOOT, 2))
|
||||
.await
|
||||
.map(|pool| pool.close())
|
||||
}));
|
||||
}
|
||||
let mut outcomes = Vec::new();
|
||||
for handle in handles {
|
||||
outcomes.push(match handle.await? {
|
||||
Ok(()) => "ok".to_owned(),
|
||||
Err(err) => format!("{err:#}"),
|
||||
});
|
||||
}
|
||||
anyhow::Ok(outcomes)
|
||||
}
|
||||
.await;
|
||||
let _ = docker(&["rm", "-f", &container]);
|
||||
let outcomes = result?;
|
||||
assert_eq!(outcomes, vec!["ok".to_owned(); 4], "{outcomes:#?}");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn boot_survives_a_peer_that_creates_the_table_without_the_schema_lock() -> anyhow::Result<()>
|
||||
{
|
||||
if !docker_available() {
|
||||
eprintln!("skipping: docker is not available");
|
||||
return Ok(());
|
||||
}
|
||||
let port = free_port()?;
|
||||
let container = format!("fluxer-kvschema-race-{}-{port}", std::process::id());
|
||||
start_postgres(&container, port)?;
|
||||
let result = async {
|
||||
wait_for_postgres(&container, port).await?;
|
||||
let helper = connect(&postgres_config(port, KV_TABLE_HELPER, 4)).await?;
|
||||
let mut peer = helper.get().await?;
|
||||
let peer_transaction = peer.transaction().await?;
|
||||
peer_transaction
|
||||
.batch_execute(&unlocked_create_table_sql(KV_TABLE_RACE))
|
||||
.await?;
|
||||
let booting = tokio::spawn(async move {
|
||||
connect(&postgres_config(port, KV_TABLE_RACE, 2))
|
||||
.await
|
||||
.map(|pool| pool.close())
|
||||
});
|
||||
tokio::time::sleep(Duration::from_millis(500)).await;
|
||||
peer_transaction.commit().await?;
|
||||
let booted = match booting.await? {
|
||||
Ok(()) => "ok".to_owned(),
|
||||
Err(err) => format!("{err:#}"),
|
||||
};
|
||||
drop(peer);
|
||||
helper.close();
|
||||
anyhow::Ok(booted)
|
||||
}
|
||||
.await;
|
||||
let _ = docker(&["rm", "-f", &container]);
|
||||
let booted = result?;
|
||||
assert_eq!(booted, "ok");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn backfills_message_partition_keys_once_and_never_scans_again() -> anyhow::Result<()> {
|
||||
if !docker_available() {
|
||||
eprintln!("skipping: docker is not available");
|
||||
return Ok(());
|
||||
}
|
||||
let port = free_port()?;
|
||||
let container = format!("fluxer-kvschema-backfill-{}-{port}", std::process::id());
|
||||
start_postgres(&container, port)?;
|
||||
let result = async {
|
||||
wait_for_postgres(&container, port).await?;
|
||||
let helper = connect(&postgres_config(port, KV_TABLE_HELPER, 4)).await?;
|
||||
helper
|
||||
.get()
|
||||
.await?
|
||||
.batch_execute(&unlocked_create_table_sql(KV_TABLE_BACKFILL))
|
||||
.await?;
|
||||
insert_legacy_message(&helper, KV_TABLE_BACKFILL, "m1").await?;
|
||||
ensure_kv_schema(&helper, KV_TABLE_BACKFILL).await?;
|
||||
let first = partition_key_of(&helper, KV_TABLE_BACKFILL, "m1").await?;
|
||||
let markers = migration_marker_count(&helper, KV_TABLE_BACKFILL).await?;
|
||||
insert_legacy_message(&helper, KV_TABLE_BACKFILL, "m2").await?;
|
||||
ensure_kv_schema(&helper, KV_TABLE_BACKFILL).await?;
|
||||
let untouched = partition_key_of(&helper, KV_TABLE_BACKFILL, "m2").await?;
|
||||
let still_migrated = partition_key_of(&helper, KV_TABLE_BACKFILL, "m1").await?;
|
||||
helper.close();
|
||||
anyhow::Ok((first, markers, untouched, still_migrated))
|
||||
}
|
||||
.await;
|
||||
let _ = docker(&["rm", "-f", &container]);
|
||||
let (first, markers, untouched, still_migrated) = result?;
|
||||
assert_eq!(first, Some(migrated_partition_key()));
|
||||
assert_eq!(markers, 1);
|
||||
assert_eq!(untouched, Some(legacy_row_key("m2")));
|
||||
assert_eq!(still_migrated, Some(migrated_partition_key()));
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user