mirror of
https://github.com/TheFunny/TelegramTwitterMediaBot.git
synced 2026-09-25 23:52:04 +00:00
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:
+103
-22
@@ -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"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user