fix: improve row lock contention on concurrency counter

This commit is contained in:
Ruben Fiszel
2025-03-16 17:20:39 +01:00
parent 47ccbeab01
commit e8bb307502
6 changed files with 100 additions and 54 deletions
@@ -130,28 +130,28 @@
},
{
"ordinal": 25,
"name": "teams_command_script",
"type_info": "Text"
},
{
"ordinal": 26,
"name": "teams_team_id",
"type_info": "Text"
},
{
"ordinal": 27,
"name": "teams_team_name",
"type_info": "Text"
},
{
"ordinal": 28,
"name": "ai_models",
"type_info": "VarcharArray"
},
{
"ordinal": 29,
"ordinal": 26,
"name": "code_completion_model",
"type_info": "Varchar"
},
{
"ordinal": 27,
"name": "teams_command_script",
"type_info": "Text"
},
{
"ordinal": 28,
"name": "teams_team_id",
"type_info": "Text"
},
{
"ordinal": 29,
"name": "teams_team_name",
"type_info": "Text"
}
],
"parameters": {
@@ -185,10 +185,10 @@
true,
true,
true,
true,
true,
true,
false,
true,
true,
true,
true
]
},
@@ -130,28 +130,28 @@
},
{
"ordinal": 25,
"name": "teams_command_script",
"type_info": "Text"
},
{
"ordinal": 26,
"name": "teams_team_id",
"type_info": "Text"
},
{
"ordinal": 27,
"name": "teams_team_name",
"type_info": "Text"
},
{
"ordinal": 28,
"name": "ai_models",
"type_info": "VarcharArray"
},
{
"ordinal": 29,
"ordinal": 26,
"name": "code_completion_model",
"type_info": "Varchar"
},
{
"ordinal": 27,
"name": "teams_command_script",
"type_info": "Text"
},
{
"ordinal": 28,
"name": "teams_team_id",
"type_info": "Text"
},
{
"ordinal": 29,
"name": "teams_team_name",
"type_info": "Text"
}
],
"parameters": {
@@ -185,10 +185,10 @@
true,
true,
true,
true,
true,
true,
false,
true,
true,
true,
true
]
},
@@ -1,6 +1,6 @@
{
"db_name": "PostgreSQL",
"query": "INSERT INTO concurrency_counter(concurrency_id, job_uuids) VALUES ($1, $2)\n ON CONFLICT (concurrency_id) \n DO UPDATE SET job_uuids = jsonb_set(concurrency_counter.job_uuids, array[$3], '{}')\n RETURNING (SELECT COUNT(*) FROM jsonb_object_keys(job_uuids))",
"query": "INSERT INTO concurrency_counter(concurrency_id, job_uuids) \n VALUES ($1, $2)\n ON CONFLICT (concurrency_id) \n DO UPDATE SET job_uuids = jsonb_set(concurrency_counter.job_uuids, array[$3], '{}')\n RETURNING (SELECT COUNT(*) FROM jsonb_object_keys(job_uuids))",
"describe": {
"columns": [
{
@@ -20,5 +20,5 @@
null
]
},
"hash": "3e3d12a51cb524fbd3d6949e150cb608acfbe8c8eade1939e813086380c205e0"
"hash": "b6d657d98de40a25b2bf66471c376d8a2424dce6c4ab2a43cf863911c833fd30"
}
@@ -0,0 +1,22 @@
{
"db_name": "PostgreSQL",
"query": "SELECT job_uuids FROM concurrency_counter \n WHERE concurrency_id = $1 \n FOR UPDATE",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "job_uuids",
"type_info": "Jsonb"
}
],
"parameters": {
"Left": [
"Text"
]
},
"nullable": [
false
]
},
"hash": "c013880b8567bf5545b0c93ffb4b69a1e6798103df08344fbc9719e53f336c29"
}
@@ -15,7 +15,7 @@
]
},
"nullable": [
true
null
]
},
"hash": "ddf2eccb78a310ed00c7d8b9c3f05d394a7cbcf0038c72a78add5c7b02ef5927"
+37 -13
View File
@@ -1919,22 +1919,13 @@ pub async fn pull(
format!("{{\"{}\": {{}}}}", pulled_job.id.hyphenated().to_string()).as_str(),
)
.expect("Unable to serialize job_uuids column to proper JSON");
let running_job = sqlx::query_scalar!(
"INSERT INTO concurrency_counter(concurrency_id, job_uuids) VALUES ($1, $2)
ON CONFLICT (concurrency_id)
DO UPDATE SET job_uuids = jsonb_set(concurrency_counter.job_uuids, array[$3], '{}')
RETURNING (SELECT COUNT(*) FROM jsonb_object_keys(job_uuids))",
job_concurrency_key,
let (mut tx, running_job) = update_concurrency_counter(
tx,
job_concurrency_key.clone(),
jobs_uuids_init_json_value,
pulled_job.id.hyphenated().to_string(),
)
.fetch_one(&mut *tx)
.await
.map_err(|e| {
Error::internal_err(format!(
"Error getting concurrency count for script path {job_script_path}: {e:#}"
))
})?;
.await?;
tracing::debug!("running_job: {}", running_job.unwrap_or(0));
let completed_count = sqlx::query!(
@@ -2069,6 +2060,39 @@ pub async fn pull(
}
}
async fn update_concurrency_counter<'c>(
mut tx: Transaction<'c, sqlx::Postgres>,
job_concurrency_key: String,
jobs_uuids_init_json_value: serde_json::Value,
pulled_job_id: String,
) -> anyhow::Result<(Transaction<'c, sqlx::Postgres>, Option<i64>)> {
// 1. Try to lock the row first
let _ = sqlx::query!(
"SELECT job_uuids FROM concurrency_counter
WHERE concurrency_id = $1
FOR UPDATE",
job_concurrency_key
)
.fetch_optional(&mut *tx)
.await?;
// 2. Insert if missing, otherwise update
let running_job = sqlx::query_scalar!(
"INSERT INTO concurrency_counter(concurrency_id, job_uuids)
VALUES ($1, $2)
ON CONFLICT (concurrency_id)
DO UPDATE SET job_uuids = jsonb_set(concurrency_counter.job_uuids, array[$3], '{}')
RETURNING (SELECT COUNT(*) FROM jsonb_object_keys(job_uuids))",
job_concurrency_key,
jobs_uuids_init_json_value,
pulled_job_id
)
.fetch_one(&mut *tx)
.await?;
Ok((tx, running_job))
}
async fn pull_single_job_and_mark_as_running_no_concurrency_limit<'c>(
db: &Pool<Postgres>,
suspend_first: bool,