Merge pull request #19 from mattnworb/mattbrown/move-allowances

move allowances Map from Feline class to avoid possible deadlock
spotify · Aug 5, 2022 · 44bef6e · 44bef6e
2 parents 926016a + a67ddb9
commit 44bef6e
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 18 deletions.
diff --git a/feline/src/main/java/com/spotify/feline/AllowancesTransformer.java b/feline/src/main/java/com/spotify/feline/AllowancesTransformer.java
@@ -17,24 +17,44 @@
 
 package com.spotify.feline;
 
-import java.util.Map;
+import java.util.HashSet;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 import net.bytebuddy.agent.builder.AgentBuilder;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.dynamic.DynamicType;
+import net.bytebuddy.matcher.ElementMatcher;
 import net.bytebuddy.utility.JavaModule;
 
 /**
  * This transformer applies {@link AllowAdvice} to every method registered with {@link
  * Feline#allowBlockingCallsInside(String, String)}.
  */
-class AllowancesTransformer implements AgentBuilder.Transformer {
+class AllowancesTransformer implements AgentBuilder.Transformer, ElementMatcher<TypeDescription> {
 
-  private final Map<String, Set<String>> allowances;
+  private final ConcurrentMap<String, Set<String>> allowances;
 
-  AllowancesTransformer(Map<String, Set<String>> allowances) {
-    this.allowances = allowances;
+  AllowancesTransformer() {
+    this.allowances = new ConcurrentHashMap<>();
+  }
+
+  void allow(final String className, final String methodName) {
+    allowances.compute(
+        className,
+        (key, set) -> {
+          if (set == null) {
+            set = new HashSet<>();
+          }
+          set.add(methodName);
+          return set;
+        });
+  }
+
+  @Override
+  public boolean matches(final TypeDescription typeDescription) {
+    return allowances.containsKey(typeDescription.getName());
   }
 
   @Override

diff --git a/feline/src/main/java/com/spotify/feline/Feline.java b/feline/src/main/java/com/spotify/feline/Feline.java
@@ -18,7 +18,6 @@
 
 import java.io.IOException;
 import java.lang.instrument.Instrumentation;
-import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
@@ -32,7 +31,7 @@
 /** Detects blocking calls to @link CompletableFuture and notifies registered consumers. */
 public class Feline {
 
-  private static final Map<String, Set<String>> allowances = new ConcurrentHashMap<>();
+  private static final AllowancesTransformer allowancesTransformer = new AllowancesTransformer();
 
   /**
    * Registers a consumer that will be invoked when blocking calls are detected. Consumers can throw
@@ -152,15 +151,7 @@ public static boolean removeOnExitConsumer(
    * @param methodName a method name
    */
   public static void allowBlockingCallsInside(final String className, final String methodName) {
-    allowances.compute(
-        className,
-        (key, allowances) -> {
-          if (allowances == null) {
-            allowances = new HashSet<>();
-          }
-          allowances.add(methodName);
-          return allowances;
-        });
+    allowancesTransformer.allow(className, methodName);
   }
 
   /**
@@ -214,8 +205,8 @@ public static boolean removeThreadLocalInitialValueConsumer(final Runnable consu
         .asTerminalTransformation()
 
         // Instrument allowed/disallowed methods
-        .type(it -> allowances.containsKey(it.getName()))
-        .transform(new AllowancesTransformer(allowances))
+        .type(allowancesTransformer)
+        .transform(allowancesTransformer)
         .asTerminalTransformation()
 
         // instrument ThreadLocal