Skip to content

Commit c06ab00

Browse files
authored
refactor: simplify caches (#448)
1 parent 27b6e7f commit c06ab00

File tree

8 files changed

+147
-70
lines changed

8 files changed

+147
-70
lines changed

Cargo.lock

Lines changed: 21 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

crates/pgt_workspace/Cargo.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,9 +13,9 @@ version = "0.0.0"
1313

1414
[dependencies]
1515
biome_deserialize = "0.6.0"
16-
dashmap = "5.5.3"
1716
futures = "0.3.31"
1817
globset = "0.4.16"
18+
lru = "0.12"
1919

2020
ignore = { workspace = true }
2121
pgt_analyse = { workspace = true, features = ["serde"] }

crates/pgt_workspace/src/workspace/server.rs

Lines changed: 23 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
use std::{
2+
collections::HashMap,
23
fs,
34
panic::RefUnwindSafe,
45
path::{Path, PathBuf},
@@ -8,7 +9,6 @@ use std::{
89
use analyser::AnalyserVisitorBuilder;
910
use async_helper::run_async;
1011
use connection_manager::ConnectionManager;
11-
use dashmap::DashMap;
1212
use document::{
1313
AsyncDiagnosticsMapper, CursorPositionFilter, DefaultMapper, Document, ExecuteStatementMapper,
1414
SyncDiagnosticsMapper,
@@ -67,7 +67,7 @@ pub(super) struct WorkspaceServer {
6767
/// Stores the schema cache for this workspace
6868
schema_cache: SchemaCacheManager,
6969

70-
documents: DashMap<PgTPath, Document>,
70+
documents: RwLock<HashMap<PgTPath, Document>>,
7171

7272
connection: ConnectionManager,
7373
}
@@ -89,7 +89,7 @@ impl WorkspaceServer {
8989
pub(crate) fn new() -> Self {
9090
Self {
9191
settings: RwLock::default(),
92-
documents: DashMap::default(),
92+
documents: RwLock::new(HashMap::new()),
9393
schema_cache: SchemaCacheManager::new(),
9494
connection: ConnectionManager::new(),
9595
}
@@ -262,7 +262,8 @@ impl Workspace for WorkspaceServer {
262262
/// Add a new file to the workspace
263263
#[tracing::instrument(level = "info", skip_all, fields(path = params.path.as_path().as_os_str().to_str()), err)]
264264
fn open_file(&self, params: OpenFileParams) -> Result<(), WorkspaceError> {
265-
self.documents
265+
let mut documents = self.documents.write().unwrap();
266+
documents
266267
.entry(params.path.clone())
267268
.or_insert_with(|| Document::new(params.content, params.version));
268269

@@ -275,7 +276,8 @@ impl Workspace for WorkspaceServer {
275276

276277
/// Remove a file from the workspace
277278
fn close_file(&self, params: super::CloseFileParams) -> Result<(), WorkspaceError> {
278-
self.documents
279+
let mut documents = self.documents.write().unwrap();
280+
documents
279281
.remove(&params.path)
280282
.ok_or_else(WorkspaceError::not_found)?;
281283

@@ -288,13 +290,15 @@ impl Workspace for WorkspaceServer {
288290
version = params.version
289291
), err)]
290292
fn change_file(&self, params: super::ChangeFileParams) -> Result<(), WorkspaceError> {
291-
match self.documents.entry(params.path.clone()) {
292-
dashmap::mapref::entry::Entry::Occupied(mut entry) => {
293+
let mut documents = self.documents.write().unwrap();
294+
295+
match documents.entry(params.path.clone()) {
296+
std::collections::hash_map::Entry::Occupied(mut entry) => {
293297
entry
294298
.get_mut()
295299
.update_content(params.content, params.version);
296300
}
297-
dashmap::mapref::entry::Entry::Vacant(entry) => {
301+
std::collections::hash_map::Entry::Vacant(entry) => {
298302
entry.insert(Document::new(params.content, params.version));
299303
}
300304
}
@@ -307,8 +311,8 @@ impl Workspace for WorkspaceServer {
307311
}
308312

309313
fn get_file_content(&self, params: GetFileContentParams) -> Result<String, WorkspaceError> {
310-
let document = self
311-
.documents
314+
let documents = self.documents.read().unwrap();
315+
let document = documents
312316
.get(&params.path)
313317
.ok_or(WorkspaceError::not_found())?;
314318
Ok(document.get_document_content().to_string())
@@ -322,8 +326,8 @@ impl Workspace for WorkspaceServer {
322326
&self,
323327
params: code_actions::CodeActionsParams,
324328
) -> Result<code_actions::CodeActionsResult, WorkspaceError> {
325-
let parser = self
326-
.documents
329+
let documents = self.documents.read().unwrap();
330+
let parser = documents
327331
.get(&params.path)
328332
.ok_or(WorkspaceError::not_found())?;
329333

@@ -364,8 +368,8 @@ impl Workspace for WorkspaceServer {
364368
&self,
365369
params: ExecuteStatementParams,
366370
) -> Result<ExecuteStatementResult, WorkspaceError> {
367-
let parser = self
368-
.documents
371+
let documents = self.documents.read().unwrap();
372+
let parser = documents
369373
.get(&params.path)
370374
.ok_or(WorkspaceError::not_found())?;
371375

@@ -422,8 +426,8 @@ impl Workspace for WorkspaceServer {
422426
}
423427
};
424428

425-
let doc = self
426-
.documents
429+
let documents = self.documents.read().unwrap();
430+
let doc = documents
427431
.get(&params.path)
428432
.ok_or(WorkspaceError::not_found())?;
429433

@@ -607,8 +611,8 @@ impl Workspace for WorkspaceServer {
607611
&self,
608612
params: GetCompletionsParams,
609613
) -> Result<CompletionsResult, WorkspaceError> {
610-
let parsed_doc = self
611-
.documents
614+
let documents = self.documents.read().unwrap();
615+
let parsed_doc = documents
612616
.get(&params.path)
613617
.ok_or(WorkspaceError::not_found())?;
614618

@@ -621,7 +625,7 @@ impl Workspace for WorkspaceServer {
621625

622626
let schema_cache = self.schema_cache.load(pool)?;
623627

624-
match get_statement_for_completions(&parsed_doc, params.position) {
628+
match get_statement_for_completions(parsed_doc, params.position) {
625629
None => {
626630
tracing::debug!("No statement found.");
627631
Ok(CompletionsResult::default())

crates/pgt_workspace/src/workspace/server/annotation.rs

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,20 @@
1-
use std::sync::Arc;
1+
use std::num::NonZeroUsize;
2+
use std::sync::{Arc, Mutex};
23

3-
use dashmap::DashMap;
4+
use lru::LruCache;
45
use pgt_lexer::SyntaxKind;
56

67
use super::statement_identifier::StatementId;
78

9+
const DEFAULT_CACHE_SIZE: usize = 1000;
10+
811
#[derive(Debug, Clone, PartialEq, Eq)]
912
pub struct StatementAnnotations {
1013
ends_with_semicolon: bool,
1114
}
1215

1316
pub struct AnnotationStore {
14-
db: DashMap<StatementId, Arc<StatementAnnotations>>,
17+
db: Mutex<LruCache<StatementId, Arc<StatementAnnotations>>>,
1518
}
1619

1720
const WHITESPACE_TOKENS: [SyntaxKind; 6] = [
@@ -25,7 +28,11 @@ const WHITESPACE_TOKENS: [SyntaxKind; 6] = [
2528

2629
impl AnnotationStore {
2730
pub fn new() -> AnnotationStore {
28-
AnnotationStore { db: DashMap::new() }
31+
AnnotationStore {
32+
db: Mutex::new(LruCache::new(
33+
NonZeroUsize::new(DEFAULT_CACHE_SIZE).unwrap(),
34+
)),
35+
}
2936
}
3037

3138
#[allow(unused)]
@@ -34,8 +41,10 @@ impl AnnotationStore {
3441
statement_id: &StatementId,
3542
content: &str,
3643
) -> Arc<StatementAnnotations> {
37-
if let Some(existing) = self.db.get(statement_id).map(|x| x.clone()) {
38-
return existing;
44+
let mut cache = self.db.lock().unwrap();
45+
46+
if let Some(existing) = cache.get(statement_id) {
47+
return existing.clone();
3948
}
4049

4150
let lexed = pgt_lexer::lex(content);
@@ -51,7 +60,7 @@ impl AnnotationStore {
5160
ends_with_semicolon,
5261
});
5362

54-
self.db.insert(statement_id.clone(), annotations.clone());
63+
cache.put(statement_id.clone(), annotations.clone());
5564

5665
annotations
5766
}

crates/pgt_workspace/src/workspace/server/connection_manager.rs

Lines changed: 21 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
1+
use std::collections::HashMap;
2+
use std::sync::RwLock;
13
use std::time::{Duration, Instant};
24

3-
use dashmap::DashMap;
45
use sqlx::{PgPool, Postgres, pool::PoolOptions, postgres::PgConnectOptions};
56

67
use crate::settings::DatabaseSettings;
@@ -16,13 +17,13 @@ struct CachedPool {
1617

1718
#[derive(Default)]
1819
pub struct ConnectionManager {
19-
pools: DashMap<ConnectionKey, CachedPool>,
20+
pools: RwLock<HashMap<ConnectionKey, CachedPool>>,
2021
}
2122

2223
impl ConnectionManager {
2324
pub fn new() -> Self {
2425
Self {
25-
pools: DashMap::new(),
26+
pools: RwLock::new(HashMap::new()),
2627
}
2728
}
2829

@@ -41,8 +42,19 @@ impl ConnectionManager {
4142
return None;
4243
}
4344

44-
// If we have a cached pool, update its last_accessed time and return it
45-
if let Some(mut cached_pool) = self.pools.get_mut(&key) {
45+
// Try read lock first for cache hit
46+
if let Ok(pools) = self.pools.read() {
47+
if let Some(cached_pool) = pools.get(&key) {
48+
// Can't update last_accessed with read lock, but that's okay for occasional misses
49+
return Some(cached_pool.pool.clone());
50+
}
51+
}
52+
53+
// Cache miss or need to update timestamp - use write lock
54+
let mut pools = self.pools.write().unwrap();
55+
56+
// Double-check after acquiring write lock
57+
if let Some(cached_pool) = pools.get_mut(&key) {
4658
cached_pool.last_accessed = Instant::now();
4759
return Some(cached_pool.pool.clone());
4860
}
@@ -69,7 +81,7 @@ impl ConnectionManager {
6981
idle_timeout: Duration::from_secs(60 * 5),
7082
};
7183

72-
self.pools.insert(key, cached_pool);
84+
pools.insert(key, cached_pool);
7385

7486
Some(pool)
7587
}
@@ -78,8 +90,10 @@ impl ConnectionManager {
7890
fn cleanup_idle_pools(&self, ignore_key: &ConnectionKey) {
7991
let now = Instant::now();
8092

93+
let mut pools = self.pools.write().unwrap();
94+
8195
// Use retain to keep only non-idle connections
82-
self.pools.retain(|key, cached_pool| {
96+
pools.retain(|key, cached_pool| {
8397
let idle_duration = now.duration_since(cached_pool.last_accessed);
8498
if idle_duration > cached_pool.idle_timeout && key != ignore_key {
8599
tracing::debug!(
Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,29 +1,38 @@
1-
use std::sync::Arc;
1+
use std::num::NonZeroUsize;
2+
use std::sync::{Arc, Mutex};
23

3-
use dashmap::DashMap;
4+
use lru::LruCache;
45
use pgt_query_ext::diagnostics::*;
56

67
use super::statement_identifier::StatementId;
78

9+
const DEFAULT_CACHE_SIZE: usize = 1000;
10+
811
pub struct PgQueryStore {
9-
db: DashMap<StatementId, Arc<Result<pgt_query_ext::NodeEnum, SyntaxDiagnostic>>>,
12+
db: Mutex<LruCache<StatementId, Arc<Result<pgt_query_ext::NodeEnum, SyntaxDiagnostic>>>>,
1013
}
1114

1215
impl PgQueryStore {
1316
pub fn new() -> PgQueryStore {
14-
PgQueryStore { db: DashMap::new() }
17+
PgQueryStore {
18+
db: Mutex::new(LruCache::new(
19+
NonZeroUsize::new(DEFAULT_CACHE_SIZE).unwrap(),
20+
)),
21+
}
1522
}
1623

1724
pub fn get_or_cache_ast(
1825
&self,
1926
statement: &StatementId,
2027
) -> Arc<Result<pgt_query_ext::NodeEnum, SyntaxDiagnostic>> {
21-
if let Some(existing) = self.db.get(statement).map(|x| x.clone()) {
22-
return existing;
28+
let mut cache = self.db.lock().unwrap();
29+
30+
if let Some(existing) = cache.get(statement) {
31+
return existing.clone();
2332
}
2433

2534
let r = Arc::new(pgt_query_ext::parse(statement.content()).map_err(SyntaxDiagnostic::from));
26-
self.db.insert(statement.clone(), r.clone());
35+
cache.put(statement.clone(), r.clone());
2736
r
2837
}
2938
}

0 commit comments

Comments
 (0)