Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Sign up
Appearance settings

Commit c06ab00

Browse files
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)

AltStyle によって変換されたページ (->オリジナル) /