fix(postgres): prevent custom type stream deadlock
This commit is contained in:
1 parent
0448d5bd2c
commit
2324d84c4e
2 files changed
+106
No files matched your search
@@ -1413,12 +1413,34 @@ fn postgres_select_stream_outcome(stream: tokio_postgres::RowStream) -> Postgres
|
|||||||
PostgresSelectStreamOutcome::Binary { stream, metadata }
|
PostgresSelectStreamOutcome::Binary { stream, metadata }
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn prepare_unnamed_select_metadata(
|
||||||
|
client: &deadpool_postgres::Client,
|
||||||
|
sql: &str,
|
||||||
|
) -> Result<PreparedSelectMetadata, tokio_postgres::Error> {
|
||||||
|
// query_typed_raw sends Describe and Execute together. If its result has an
|
||||||
|
// unknown user-defined type, tokio-postgres then resolves that type with a
|
||||||
|
// second query on the same connection before returning the RowStream. A
|
||||||
|
// large result can fill the connection buffers and block that type lookup.
|
||||||
|
// Describe once without executing so custom types are cached before the
|
||||||
|
// actual unnamed stream starts. The stream remains unnamed to tolerate
|
||||||
|
// server-side prepared statement loss in snapshot/proxy environments.
|
||||||
|
let stmt = client.prepare(sql).await?;
|
||||||
|
Ok(prepared_select_metadata(stmt.columns()))
|
||||||
|
}
|
||||||
|
|
||||||
async fn start_postgres_select_stream(
|
async fn start_postgres_select_stream(
|
||||||
client: &deadpool_postgres::Client,
|
client: &deadpool_postgres::Client,
|
||||||
sql: &str,
|
sql: &str,
|
||||||
force_unnamed: bool,
|
force_unnamed: bool,
|
||||||
) -> Result<PostgresSelectStreamOutcome, tokio_postgres::Error> {
|
) -> Result<PostgresSelectStreamOutcome, tokio_postgres::Error> {
|
||||||
if force_unnamed || postgres_client_uses_unnamed_statements(client) {
|
if force_unnamed || postgres_client_uses_unnamed_statements(client) {
|
||||||
|
let metadata = prepare_unnamed_select_metadata(client, sql).await?;
|
||||||
|
if let Some(unsupported_type) = metadata.unsupported_type {
|
||||||
|
return Ok(PostgresSelectStreamOutcome::TextFallback {
|
||||||
|
column_types: metadata.column_types,
|
||||||
|
unsupported_type,
|
||||||
|
});
|
||||||
|
}
|
||||||
return postgres_query_unnamed(client, sql).await.map(postgres_select_stream_outcome);
|
return postgres_query_unnamed(client, sql).await.map(postgres_select_stream_outcome);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -13243,6 +13265,12 @@ mod tests {
|
|||||||
.expect("stream unnamed query inside backup snapshot");
|
.expect("stream unnamed query inside backup snapshot");
|
||||||
assert_eq!(columns, vec!["value"]);
|
assert_eq!(columns, vec!["value"]);
|
||||||
assert_eq!(rows[0][0], serde_json::json!(45));
|
assert_eq!(rows[0][0], serde_json::json!(45));
|
||||||
|
let prepared_count = client
|
||||||
|
.query_typed_one("SELECT count(*)::int8 FROM pg_prepared_statements", &[])
|
||||||
|
.await
|
||||||
|
.expect("count server-side prepared statements")
|
||||||
|
.get::<_, i64>(0);
|
||||||
|
assert_eq!(prepared_count, 0, "snapshot stream metadata preparation must not retain a named statement");
|
||||||
client.execute_typed("SELECT 1", &[]).await.expect("snapshot remains usable");
|
client.execute_typed("SELECT 1", &[]).await.expect("snapshot remains usable");
|
||||||
client.execute_typed("ROLLBACK", &[]).await.expect("rollback backup snapshot");
|
client.execute_typed("ROLLBACK", &[]).await.expect("rollback backup snapshot");
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -153,6 +153,84 @@ async fn live_postgres_backup_snapshot_streams_after_server_deallocates_statemen
|
|||||||
let _ = std::fs::remove_file(db_path);
|
let _ = std::fs::remove_file(db_path);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[ignore = "requires DBX_LIVE_MANUAL_TXN_POSTGRES_* env vars pointing at writable PostgreSQL with pgvector"]
|
||||||
|
fn live_postgres_backup_snapshot_streams_wide_pgvector_result() {
|
||||||
|
std::thread::Builder::new()
|
||||||
|
.name("dbx-live-pg-wide-vector".to_string())
|
||||||
|
.stack_size(16 * 1024 * 1024)
|
||||||
|
.spawn(|| {
|
||||||
|
tokio::runtime::Builder::new_multi_thread()
|
||||||
|
.worker_threads(2)
|
||||||
|
.thread_stack_size(16 * 1024 * 1024)
|
||||||
|
.enable_all()
|
||||||
|
.build()
|
||||||
|
.expect("build wide pgvector test runtime")
|
||||||
|
.block_on(live_postgres_backup_snapshot_streams_wide_pgvector_result_inner())
|
||||||
|
})
|
||||||
|
.expect("spawn wide pgvector test thread")
|
||||||
|
.join()
|
||||||
|
.expect("wide pgvector test thread should not panic");
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn live_postgres_backup_snapshot_streams_wide_pgvector_result_inner() {
|
||||||
|
let mut config = live_config("DBX_LIVE_MANUAL_TXN_POSTGRES", DatabaseType::Postgres, 5432);
|
||||||
|
config.query_timeout_secs = 2;
|
||||||
|
let database = config.database.clone().expect("database");
|
||||||
|
let connection_id = config.id.clone();
|
||||||
|
let (state, db_path) = app_state_with_config(config).await;
|
||||||
|
|
||||||
|
if let Err(error) =
|
||||||
|
execute_sql_statement(&state, &connection_id, &database, "CREATE EXTENSION IF NOT EXISTS vector", None, None)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
eprintln!("skipping wide pgvector stream regression because the vector extension is unavailable: {error}");
|
||||||
|
let _ = std::fs::remove_file(db_path);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
let suffix = uuid::Uuid::new_v4().simple().to_string();
|
||||||
|
let schema = format!("dbx_7986_{}", &suffix[..12]);
|
||||||
|
let setup_sql = [
|
||||||
|
format!(r#"CREATE SCHEMA "{schema}""#),
|
||||||
|
format!(r#"CREATE TABLE "{schema}".wide_vector_probe (id int PRIMARY KEY, embedding vector(1536) NOT NULL)"#),
|
||||||
|
format!(
|
||||||
|
r#"INSERT INTO "{schema}".wide_vector_probe (id, embedding)
|
||||||
|
SELECT i, ('[' || array_to_string(array_fill((i % 10)::text, ARRAY[1536]), ',') || ']')::vector
|
||||||
|
FROM generate_series(1, 5000) AS source(i)"#
|
||||||
|
),
|
||||||
|
];
|
||||||
|
for sql in setup_sql {
|
||||||
|
execute_sql_statement(&state, &connection_id, &database, &sql, None, None)
|
||||||
|
.await
|
||||||
|
.expect("create wide pgvector stream fixture");
|
||||||
|
}
|
||||||
|
|
||||||
|
let snapshot =
|
||||||
|
begin_database_backup_snapshot_core(&state, &connection_id, &database).await.expect("begin snapshot");
|
||||||
|
let mut rows_seen = 0_u64;
|
||||||
|
let select_sql = format!(r#"SELECT id, embedding FROM "{schema}".wide_vector_probe ORDER BY id"#);
|
||||||
|
let streamed = stream_rows_in_manual_transaction(&state, &snapshot.session_id, &select_sql, 100, |batch| {
|
||||||
|
for row in &batch {
|
||||||
|
assert_eq!(row.len(), 2);
|
||||||
|
assert_eq!(row[1].as_array().map(Vec::len), Some(1536));
|
||||||
|
}
|
||||||
|
rows_seen += batch.len() as u64;
|
||||||
|
Ok(())
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
.expect("stream wide pgvector result through backup snapshot");
|
||||||
|
|
||||||
|
rollback_manual_transaction(&state, &snapshot.session_id).await.expect("rollback snapshot");
|
||||||
|
assert_eq!(streamed, 5000);
|
||||||
|
assert_eq!(rows_seen, 5000);
|
||||||
|
|
||||||
|
execute_sql_statement(&state, &connection_id, &database, &format!(r#"DROP SCHEMA "{schema}" CASCADE"#), None, None)
|
||||||
|
.await
|
||||||
|
.expect("drop wide pgvector stream fixture");
|
||||||
|
let _ = std::fs::remove_file(db_path);
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
#[ignore = "requires DBX_LIVE_MANUAL_TXN_POSTGRES_* env vars pointing at writable PostgreSQL"]
|
#[ignore = "requires DBX_LIVE_MANUAL_TXN_POSTGRES_* env vars pointing at writable PostgreSQL"]
|
||||||
async fn live_postgres_backup_snapshot_times_out_when_row_stream_stalls() {
|
async fn live_postgres_backup_snapshot_times_out_when_row_stream_stalls() {
|
||||||
|
|||||||
Reference in new issue
Block a user