fix(state): keep failed and stale reads from poisoning the chat cache

Three related races in ChatStore. A read that errored (busy/IO) was indistinguishable from an absent row, so the default got cached — and the next update would write that default back over the chat's real settings (forward channel, templates, formats). A cache-miss backfill inserted unconditionally, letting a stale DB snapshot overwrite the value a concurrent set had just written. And sweep eviction removed per-chat locks unconditionally, so a lock pulled out from under an in-flight update let a second writer create a fresh one and enter the critical section concurrently; eviction now keeps any lock with a holder (the same rule rate_limit's prune applies), which still bounds the map because an uncontended leftover is caught by a later sweep. Two regression tests pin the failed-read and contended-lock cases.
This commit is contained in:
2026-09-24 02:47:49 +08:00
parent d774f0b37c
commit eecd4320f8
+103 -22
View File
@@ -69,27 +69,39 @@ impl ChatStore {
let chat_key = chat_id.to_string();
let payload = self
.pool
.with_conn_or(
log::Level::Warn,
"chat_state read failed",
None,
move |conn| {
// Concurrent handler tasks (batch-forwards) may write
// chat_state while this read runs; the shared busy timeout
// handles the write-lock collision instead of failing the
// query.
conn.query_row(
"SELECT payload FROM chat_state WHERE chat_id = ?1",
params![chat_key],
|row| row.get::<_, String>(0),
)
.optional()
},
)
.await
.unwrap_or_default();
.with_conn(move |conn| {
// Concurrent handler tasks (batch-forwards) may write
// chat_state while this read runs; the shared busy timeout
// handles the write-lock collision instead of failing the
// query.
conn.query_row(
"SELECT payload FROM chat_state WHERE chat_id = ?1",
params![chat_key],
|row| row.get::<_, String>(0),
)
.optional()
})
.await;
let payload = match payload {
// A row — or its documented absence — is a real answer and may be
// cached.
Ok(payload) => payload.unwrap_or_default(),
Err(e) => {
// A failed read must not be cached: the default would become
// what every later get returns, and the next update would
// write it back over the chat's real settings. The next get
// simply retries the DB.
log::warn!("chat_state read failed: {e}");
return ChatData::default();
}
};
let data: ChatData = serde_json::from_str(&payload).unwrap_or_default();
self.cache.lock().insert(chat_id, data.clone());
// Only fill a miss: an unconditional insert would let this (possibly
// stale) snapshot overwrite what a concurrent set just wrote.
self.cache
.lock()
.entry(chat_id)
.or_insert_with(|| data.clone());
data
}
@@ -191,12 +203,20 @@ impl ChatStore {
}
if !evicted_chats.is_empty() {
let mut cache = self.cache.lock();
let mut locks = self.locks.lock();
for chat_id in &evicted_chats {
cache.remove(chat_id);
locks.remove(chat_id);
}
}
// Per-chat locks go only while uncontended (the same rule as
// rate_limit's prune): pulling a lock out from under an in-flight
// update — between its `lock_for` clone and its `lock().await` —
// would let a second writer `lock_for` a fresh one and enter the
// critical section concurrently. A contended lock stays until a later
// sweep, and dropping the uncontended ones also catches chats an
// earlier sweep had to skip, so the map stays bounded.
self.locks
.lock()
.retain(|_, lock| Arc::strong_count(lock) > 1);
if !removed.is_empty() {
log::info!(
"pruned {} expired edit-before-forward record(s)",
@@ -360,4 +380,65 @@ mod tests {
"eviction dropped state the DB never received"
);
}
#[tokio::test]
async fn a_failed_read_is_not_cached() {
// A read that errors (busy, IO, a missing table) answers the default;
// caching that answer would make the next get return it blind and the
// next update write it back over the chat's real settings.
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("f.db");
let pool = crate::db::open_store(path.to_str().unwrap()).unwrap();
let raw = rusqlite::Connection::open(&path).unwrap();
raw.execute_batch("DROP TABLE chat_state").unwrap();
let store = ChatStore::new(pool);
let first = store.get(7).await;
assert!(first.forward_channel_id.is_none());
assert!(
!store.cache.lock().contains_key(&7),
"a failed read must not poison the cache"
);
// The next get retries the DB and sees the real row.
raw.execute_batch(
"CREATE TABLE chat_state (chat_id TEXT PRIMARY KEY, payload TEXT NOT NULL)",
)
.unwrap();
let real = ChatData {
forward_channel_id: Some(42),
..ChatData::default()
};
raw.execute(
"INSERT INTO chat_state (chat_id, payload) VALUES ('7', ?1)",
rusqlite::params![serde_json::to_string(&real).unwrap()],
)
.unwrap();
assert_eq!(store.get(7).await.forward_channel_id, Some(42));
}
#[tokio::test]
async fn prune_spares_a_lock_someone_still_holds() {
// The sweep evicts uncontended locks only: removing one an update
// still holds (its `lock_for` clone alive) would let a second writer
// create a fresh lock and enter the critical section concurrently.
let dir = tempfile::tempdir().unwrap();
let pool = crate::db::open_store(dir.path().join("l.db").to_str().unwrap()).unwrap();
let store = ChatStore::new(pool);
store.set(1, &ChatData::default()).await;
let held = store.lock_for(1); // an update between lock_for and lock().await
store.prune_expired(Duration::from_secs(60)).await;
assert!(
store.locks.lock().contains_key(&1),
"a contended lock must survive the sweep"
);
drop(held);
store.prune_expired(Duration::from_secs(60)).await;
assert!(
!store.locks.lock().contains_key(&1),
"the next sweep drops it once uncontended"
);
}
}