track cache handlers

crates/napi/Cargo.toml

@@ -73,6 +73,7 @@ tracing-chrome = "0.5.0"
 url = { workspace = true }
 urlencoding = { workspace = true }
 once_cell = { workspace = true }
+dashmap = "6.1.0"
 
 swc_core = { workspace = true, features = [
     "base_concurrent",

crates/napi/src/lib.rs

@@ -35,22 +35,18 @@ DEALINGS IN THE SOFTWARE.
 extern crate napi_derive;
 
 use std::{
-    env,
     panic::set_hook,
-    sync::{
-        atomic::{AtomicUsize, Ordering},
-        Arc, Once,
-    },
+    sync::{Arc, Once},
 };
 
 use backtrace::Backtrace;
+use dashmap::DashMap;
 use fxhash::FxHashSet;
 use napi::bindgen_prelude::*;
 use swc_core::{
     base::{Compiler, TransformOutput},
     common::{FilePathMapping, SourceMap},
 };
-
 #[cfg(not(target_arch = "wasm32"))]
 pub mod css;
 pub mod mdx;
@@ -100,7 +96,7 @@ pub fn complete_output(
     env: &Env,
     output: TransformOutput,
     eliminated_packages: FxHashSet<String>,
-    use_cache_directive_count: Arc<AtomicUsize>,
+    use_cache_telemetry_tracker: DashMap<String, usize>,
 ) -> napi::Result<Object> {
     let mut js_output = env.create_object()?;
     js_output.set_named_property("code", env.create_string_from_std(output.code)?)?;
@@ -113,10 +109,18 @@ pub fn complete_output(
             env.create_string_from_std(serde_json::to_string(&eliminated_packages)?)?,
         )?;
     }
-    js_output.set_named_property(
-        "useCacheCount",
-        env.create_int32(use_cache_directive_count.load(Ordering::SeqCst) as i32)?,
-    )?;
+    if !use_cache_telemetry_tracker.is_empty() {
+        js_output.set_named_property(
+            "useCacheTelemetryTracker",
+            env.create_string_from_std(serde_json::to_string(
+                &use_cache_telemetry_tracker
+                    .iter()
+                    .map(|entry| (entry.key().clone(), *entry.value()))
+                    .collect::<Vec<_>>(),
+            )?)?,
+        )?;
+    }
+
     Ok(js_output)
 }
 

crates/napi/src/transform.rs

@@ -31,10 +31,11 @@ use std::{
     fs::read_to_string,
     panic::{catch_unwind, AssertUnwindSafe},
     rc::Rc,
-    sync::{atomic::AtomicUsize, Arc},
+    sync::Arc,
 };
 
 use anyhow::{anyhow, bail, Context as _};
+use dashmap::DashMap;
 use fxhash::FxHashSet;
 use napi::bindgen_prelude::*;
 use next_custom_transforms::chain_transforms::{custom_before_pass, TransformOptions};
@@ -81,13 +82,15 @@ fn skip_filename() -> bool {
 }
 
 impl Task for TransformTask {
-    type Output = (TransformOutput, FxHashSet<String>, Arc<AtomicUsize>);
+    type Output = (TransformOutput, FxHashSet<String>, DashMap<String, usize>);
     type JsValue = Object;
 
     fn compute(&mut self) -> napi::Result<Self::Output> {
         GLOBALS.set(&Default::default(), || {
             let eliminated_packages: Rc<RefCell<fxhash::FxHashSet<String>>> = Default::default();
-            let use_cache_directive_count = Arc::new(AtomicUsize::new(0));
+            let use_cache_telemetry_tracker: Rc<RefCell<DashMap<String, usize>>> =
+                Default::default();
+
             let res = catch_unwind(AssertUnwindSafe(|| {
                 try_with_handler(
                     self.c.cm.clone(),
@@ -144,7 +147,7 @@ impl Task for TransformTask {
                                         comments.clone(),
                                         eliminated_packages.clone(),
                                         unresolved_mark,
-                                        use_cache_directive_count.clone(),
+                                        use_cache_telemetry_tracker.clone(),
                                     )
                                 },
                                 |_| noop_pass(),
@@ -167,7 +170,7 @@ impl Task for TransformTask {
                         (
                             o,
                             eliminated_packages.replace(Default::default()),
-                            use_cache_directive_count,
+                            use_cache_telemetry_tracker.replace(Default::default()),
                         )
                     })
                     .convert_err(),
@@ -182,9 +185,14 @@ impl Task for TransformTask {
     fn resolve(
         &mut self,
         env: Env,
-        (output, eliminated_packages, use_cache_directive_count): Self::Output,
+        (output, eliminated_packages, use_cache_telemetry_tracker): Self::Output,
     ) -> napi::Result<Self::JsValue> {
-        complete_output(&env, output, eliminated_packages, use_cache_directive_count)
+        complete_output(
+            &env,
+            output,
+            eliminated_packages,
+            use_cache_telemetry_tracker,
+        )
     }
 }
 

crates/next-core/src/next_shared/transforms/server_actions.rs

@@ -1,5 +1,3 @@
-use std::sync::{atomic::AtomicUsize, Arc};
-
 use anyhow::Result;
 use async_trait::async_trait;
 use next_custom_transforms::transforms::server_actions::{server_actions, Config};
@@ -63,7 +61,7 @@ impl CustomTransformer for NextServerActions {
                 cache_kinds: self.cache_kinds.await?.clone_value(),
             },
             ctx.comments.clone(),
-            Arc::new(AtomicUsize::new(0)),
+            Default::default(),
         );
         program.mutate(actions);
         Ok(())

crates/next-custom-transforms/Cargo.toml

@@ -30,6 +30,7 @@ sha1 = "0.10.1"
 tracing = { version = "0.1.37" }
 anyhow = { workspace = true }
 lazy_static = { workspace = true }
+dashmap = "6.1.0"
 
 swc_core = { workspace = true, features = [
   "base",

crates/next-custom-transforms/src/chain_transforms.rs

@@ -1,10 +1,6 @@
-use std::{
-    cell::RefCell,
-    path::PathBuf,
-    rc::Rc,
-    sync::{atomic::AtomicUsize, Arc},
-};
+use std::{cell::RefCell, path::PathBuf, rc::Rc, sync::Arc};
 
+use dashmap::DashMap;
 use either::Either;
 use fxhash::FxHashSet;
 use modularize_imports;
@@ -130,7 +126,7 @@ pub fn custom_before_pass<'a, C>(
     comments: C,
     eliminated_packages: Rc<RefCell<FxHashSet<String>>>,
     unresolved_mark: Mark,
-    use_cache_directive_count: Arc<AtomicUsize>,
+    use_cache_telemetry_tracker: Rc<RefCell<DashMap<String, usize>>>,
 ) -> impl Pass + 'a
 where
     C: Clone + Comments + 'a,
@@ -309,7 +305,7 @@ where
                     &file.name,
                     config.clone(),
                     comments.clone(),
-                    use_cache_directive_count,
+                    use_cache_telemetry_tracker,
                 )),
                 None => Either::Right(noop_pass()),
             },

crates/next-custom-transforms/src/transforms/server_actions.rs

@@ -1,13 +1,12 @@
 use std::{
+    cell::RefCell,
     collections::{BTreeMap, HashSet},
     convert::{TryFrom, TryInto},
     mem::{replace, take},
-    sync::{
-        atomic::{AtomicUsize, Ordering},
-        Arc,
-    },
+    rc::Rc,
 };
 
+use dashmap::DashMap;
 use hex::encode as hex_encode;
 use indoc::formatdoc;
 use rustc_hash::FxHashSet;
@@ -125,7 +124,7 @@ pub fn server_actions<C: Comments>(
     file_name: &FileName,
     config: Config,
     comments: C,
-    use_cache_directive_count: Arc<AtomicUsize>,
+    use_cache_telemetry_tracker: Rc<RefCell<DashMap<String, usize>>>,
 ) -> impl Pass {
     visit_mut_pass(ServerActions {
         config,
@@ -165,7 +164,7 @@ pub fn server_actions<C: Comments>(
         arrow_or_fn_expr_ident: None,
         exported_local_ids: HashSet::new(),
 
-        use_cache_directive_count,
+        use_cache_telemetry_tracker,
     })
 }
 
@@ -222,7 +221,7 @@ struct ServerActions<C: Comments> {
     arrow_or_fn_expr_ident: Option<Ident>,
     exported_local_ids: HashSet<Id>,
 
-    use_cache_directive_count: Arc<AtomicUsize>,
+    use_cache_telemetry_tracker: Rc<RefCell<DashMap<String, usize>>>,
 }
 
 impl<C: Comments> ServerActions<C> {
@@ -364,7 +363,7 @@ impl<C: Comments> ServerActions<C> {
                 has_file_directive: self.file_directive.is_some(),
                 is_allowed_position: true,
                 location: DirectiveLocation::FunctionBody,
-                use_cache_directive_count: Arc::clone(&self.use_cache_directive_count),
+                use_cache_telemetry_tracker: self.use_cache_telemetry_tracker.clone(),
             };
 
             body.stmts.retain(|stmt| {
@@ -391,7 +390,7 @@ impl<C: Comments> ServerActions<C> {
             has_file_directive: false,
             is_allowed_position: true,
             location: DirectiveLocation::Module,
-            use_cache_directive_count: Arc::clone(&self.use_cache_directive_count),
+            use_cache_telemetry_tracker: self.use_cache_telemetry_tracker.clone(),
         };
 
         stmts.retain(|item| {
@@ -2620,7 +2619,7 @@ struct DirectiveVisitor<'a> {
     directive: Option<Directive>,
     has_file_directive: bool,
     is_allowed_position: bool,
-    use_cache_directive_count: Arc<AtomicUsize>,
+    use_cache_telemetry_tracker: Rc<RefCell<DashMap<String, usize>>>,
 }
 
 impl DirectiveVisitor<'_> {
@@ -2675,8 +2674,6 @@ impl DirectiveVisitor<'_> {
                 // `use cache` or `use cache: foo`
                 if value == "use cache" || value.starts_with("use cache: ") {
                     // Increment telemetry counter tracking usage of "use cache" directives
-                    self.use_cache_directive_count
-                        .fetch_add(1, Ordering::SeqCst);
 
                     if in_fn_body && !allow_inline {
                         emit_error(ServerActionsErrorKind::InlineUseCacheInClientComponent {
@@ -2699,6 +2696,7 @@ impl DirectiveVisitor<'_> {
                             self.directive = Some(Directive::UseCache {
                                 cache_kind: RcStr::from("default"),
                             });
+                            self.increment_cache_usage_counter("default");
                         } else {
                             // Slice the value after "use cache: "
                             let cache_kind = RcStr::from(value.split_at("use cache: ".len()).1);
@@ -2710,6 +2708,7 @@ impl DirectiveVisitor<'_> {
                                 });
                             }
 
+                            self.increment_cache_usage_counter(&*cache_kind);
                             self.directive = Some(Directive::UseCache { cache_kind });
                         }
 
@@ -2778,6 +2777,28 @@ impl DirectiveVisitor<'_> {
 
         false
     }
+
+    // Increment telemetry counter tracking usage of "use cache" directives
+    fn increment_cache_usage_counter(&mut self, cache_kind: &str) {
+        // Retry up to 3 times to handle potential borrow conflicts
+        for _ in 0..3 {
+            match self.use_cache_telemetry_tracker.try_borrow_mut() {
+                Ok(tracker) => {
+                    if let Some(mut counter) = tracker.get_mut(cache_kind) {
+                        *counter += 1;
+                    } else {
+                        tracker.insert(cache_kind.to_string(), 1);
+                    }
+                    break;
+                }
+                Err(_) => {
+                    // If borrow fails, continue loop to retry
+                    println!("borrow failed, retrying...");
+                    continue;
+                }
+            }
+        }
+    }
 }
 
 pub(crate) struct ClosureReplacer<'a> {

crates/next-custom-transforms/tests/errors.rs

@@ -1,8 +1,4 @@
-use std::{
-    iter::FromIterator,
-    path::PathBuf,
-    sync::{atomic::AtomicUsize, Arc},
-};
+use std::{iter::FromIterator, path::PathBuf};
 
 use next_custom_transforms::transforms::{
     disallow_re_export_all_in_page::disallow_re_export_all_in_page,
@@ -195,7 +191,7 @@ fn react_server_actions_server_errors(input: PathBuf) {
                         cache_kinds: FxHashSet::default(),
                     },
                     tr.comments.as_ref().clone(),
-                    Arc::new(AtomicUsize::new(0)),
+                    Default::default(),
                 ),
             )
         },
@@ -236,7 +232,7 @@ fn react_server_actions_client_errors(input: PathBuf) {
                         cache_kinds: FxHashSet::default(),
                     },
                     tr.comments.as_ref().clone(),
-                    Arc::new(AtomicUsize::new(0)),
+                    Default::default(),
                 ),
             )
         },
@@ -295,7 +291,7 @@ fn use_cache_not_allowed(input: PathBuf) {
                         cache_kinds: FxHashSet::from_iter(["x".into()]),
                     },
                     tr.comments.as_ref().clone(),
-                    Arc::new(AtomicUsize::new(0)),
+                    Default::default(),
                 ),
             )
         },