Merge pull request #547 from theseus-rs/refactor-callsitecache

refactor: update CallSiteCache to use a DashMap instead of RwLock<HashMap>

Author: brianheineman

Cargo.lock

@@ -68,7 +68,7 @@ test-log = "0.2.18"
 thiserror = "2.0.12"
 thread-priority = "2.1.0"
 tokio = { version = "1.47.1", default-features = false, features = ["macros", "rt", "sync"] }
-tracing = { version = "0.1.41", default-features = false, features = ["release_max_level_info", "std"] }
+tracing = "0.1.41"
 tracing-subscriber = "0.3.19"
 walkdir = "2.5.0"
 whoami = "1.6.0"

Cargo.toml

@@ -13,6 +13,8 @@ pub(crate) fn initialize() -> Result<()> {
 
     let format = tracing_subscriber::fmt::format()
         .with_level(true)
+        .with_target(false)
+        .with_thread_ids(false)
         .with_thread_names(true)
         .with_timer(fmt::time::uptime())
         .compact();

ristretto_cli/src/logging.rs

@@ -3,11 +3,10 @@
 //! This module provides a thread-safe cache that tracks the resolution state of invokedynamic call
 //! sites.
 
-use crate::Error::{InternalError, PoisonedLock};
+use crate::Error::InternalError;
 use crate::Result;
+use dashmap::DashMap;
 use ristretto_classloader::Value;
-use std::collections::HashMap;
-use std::sync::RwLock;
 
 /// Unique identifier for an invokedynamic call site
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
@@ -41,14 +40,14 @@ pub enum CallSiteState {
 #[derive(Debug)]
 pub struct CallSiteCache {
     /// Maps call site keys to their resolution states
-    states: RwLock<HashMap<CallSiteKey, CallSiteState>>,
+    states: DashMap<CallSiteKey, CallSiteState>,
 }
 
 impl CallSiteCache {
     /// Create a new empty call site cache
     pub fn new() -> Self {
         Self {
-            states: RwLock::new(HashMap::new()),
+            states: DashMap::new(),
         }
     }
 
@@ -83,40 +82,29 @@ impl CallSiteCache {
         debug!("CallSiteCache: Checking cache for key: {key:?}");
 
         // Check current state
-        {
-            let map = self
-                .states
-                .read()
-                .map_err(|error| PoisonedLock(format!("Failed to acquire cache lock: {error}")))?;
-
-            match map.get(&key) {
-                Some(CallSiteState::InProgress) => {
-                    debug!("CallSiteCache: RECURSION DETECTED for key: {key:?}",);
+        match self.states.get(&key) {
+            Some(ref state) => match &**state {
+                CallSiteState::InProgress => {
+                    debug!("CallSiteCache: RECURSION DETECTED for key: {key:?}");
                     return Err(InternalError(format!(
                         "Recursive invokedynamic call site resolution detected for class '{}' at index {}",
                         key.class_name, key.instruction_index
                     )));
                 }
-                Some(CallSiteState::Resolved(value)) => {
+                CallSiteState::Resolved(value) => {
                     debug!("CallSiteCache: Returning cached result for key: {key:?}");
                     return Ok(value.clone());
                 }
-                None => {
-                    debug!("CallSiteCache: Key not found in cache, will resolve: {key:?}");
-                    // Call site not yet resolved, continue to resolution
-                }
+            },
+            None => {
+                debug!("CallSiteCache: Key not found in cache, will resolve: {key:?}");
+                // Call site not yet resolved, continue to resolution
             }
         }
 
         // Mark as in progress
-        {
-            let mut map = self
-                .states
-                .write()
-                .map_err(|error| PoisonedLock(format!("Failed to acquire cache lock: {error}")))?;
-            debug!("CallSiteCache: Marking as InProgress: {key:?}");
-            map.insert(key.clone(), CallSiteState::InProgress);
-        }
+        debug!("CallSiteCache: Marking as InProgress: {key:?}");
+        self.states.insert(key.clone(), CallSiteState::InProgress);
 
         // Perform resolution
         debug!("CallSiteCache: Starting resolution for key: {key:?}");
@@ -129,51 +117,32 @@ impl CallSiteCache {
         // Update cache based on result
         if let Ok(value) = &result {
             // Store successful resolution
-            let mut map = self
-                .states
-                .write()
-                .map_err(|error| PoisonedLock(format!("Failed to acquire cache lock: {error}")))?;
             debug!("CallSiteCache: Caching successful result for key: {key:?}",);
-            map.insert(key, CallSiteState::Resolved(value.clone()));
+            self.states
+                .insert(key, CallSiteState::Resolved(value.clone()));
         } else {
             // Remove in-progress marker on failure to allow retry
-            let mut map = self
-                .states
-                .write()
-                .map_err(|error| PoisonedLock(format!("Failed to acquire cache lock: {error}")))?;
             debug!("CallSiteCache: Removing failed resolution from cache for key: {key:?}");
-            map.remove(&key);
+            self.states
+                .remove_if(&key, |_, state| matches!(state, CallSiteState::InProgress));
         }
 
         result
     }
 
     /// Clear all cached call sites
-    pub fn clear(&self) -> Result<()> {
-        let mut map = self
-            .states
-            .write()
-            .map_err(|error| PoisonedLock(format!("Failed to acquire cache lock: {error}")))?;
-        map.clear();
-        Ok(())
+    pub fn clear(&self) {
+        self.states.clear();
     }
 
     /// Get the number of cached call sites
-    pub fn len(&self) -> Result<usize> {
-        let map = self
-            .states
-            .read()
-            .map_err(|error| PoisonedLock(format!("Failed to acquire cache lock: {error}")))?;
-        Ok(map.len())
+    pub fn len(&self) -> usize {
+        self.states.len()
     }
 
     /// Check if the cache is empty
-    pub fn is_empty(&self) -> Result<bool> {
-        let map = self
-            .states
-            .read()
-            .map_err(|error| PoisonedLock(format!("Failed to acquire cache lock: {error}")))?;
-        Ok(map.is_empty())
+    pub fn is_empty(&self) -> bool {
+        self.states.is_empty()
     }
 }
 

ristretto_vm/src/call_site_cache.rs

@@ -8,7 +8,9 @@ use bitflags::bitflags;
 use ristretto_classfile::VersionSpecification::{
     Any, Between, GreaterThan, GreaterThanOrEqual, LessThanOrEqual,
 };
-use ristretto_classfile::{FieldAccessFlags, JAVA_8, JAVA_11, JAVA_17, JAVA_21, MethodAccessFlags};
+use ristretto_classfile::{
+    FieldAccessFlags, JAVA_8, JAVA_11, JAVA_17, JAVA_21, MethodAccessFlags, ReferenceKind,
+};
 use ristretto_classloader::Error::IllegalAccessError;
 use ristretto_classloader::{Class, Method, Value};
 use ristretto_macros::intrinsic_method;
@@ -220,6 +222,7 @@ pub(crate) async fn resolve(
         class_object.value("name")?.as_string()?
     };
     let class = thread.class(class_name.clone()).await?;
+    let _reference_kind = get_reference_kind(flags)?;
     let member_name_flags = MemberNameFlags::from_bits_truncate(flags);
 
     // Handle methods/constructors
@@ -322,6 +325,15 @@ pub(crate) async fn resolve(
     }
 }
 
+/// Extracts the reference kind from the flags of a member name.
+fn get_reference_kind(flags: i32) -> Result<ReferenceKind> {
+    let flags = flags as u32;
+    let shift = MemberNameFlags::REFERENCE_KIND_SHIFT.bits();
+    let mask = MemberNameFlags::REFERENCE_KIND_MASK.bits() as u32;
+    let reference_kind = ((flags >> shift) & mask) as u8;
+    ReferenceKind::try_from(reference_kind).map_err(Into::into)
+}
+
 /// Returns true if `caller` is permitted to access a method of `declaring` with the given access
 /// flags.
 ///
@@ -682,6 +694,15 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_get_reference_kind() -> Result<()> {
+        assert_eq!(
+            get_reference_kind(0x0601_0000)?,
+            ReferenceKind::InvokeStatic
+        );
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_set_call_site_target_normal() -> Result<()> {
         let (_vm, thread) = crate::test::thread().await.expect("thread");