Add exceptionPredicate in Tracer for customized exception filtering l…

…ogic (#1496)
alibaba · May 21, 2020 · 5523fd0 · 5523fd0
1 parent 46076b3
commit 5523fd0
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 0 deletions.
diff --git a/sentinel-core/src/main/java/com/alibaba/csp/sentinel/Tracer.java b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/Tracer.java
@@ -20,6 +20,7 @@
 import com.alibaba.csp.sentinel.context.NullContext;
 import com.alibaba.csp.sentinel.slots.block.BlockException;
 import com.alibaba.csp.sentinel.util.AssertUtil;
+import com.alibaba.csp.sentinel.util.function.Predicate;
 
 /**
  * This class is used to record other exceptions except block exception.
@@ -32,6 +33,8 @@ public class Tracer {
     protected static Class<? extends Throwable>[] traceClasses;
     protected static Class<? extends Throwable>[] ignoreClasses;
 
+    protected static Predicate<Throwable> exceptionPredicate;
+
     protected Tracer() {}
 
     /**
@@ -164,6 +167,24 @@ public static Class<? extends Throwable>[] getExceptionsToIgnore() {
         return ignoreClasses;
     }
 
+    /**
+     * Get exception predicate
+     * @return the exception predicate.
+     */
+    public static Predicate<? extends Throwable> getExceptionPredicate() {
+        return exceptionPredicate;
+    }
+
+    /**
+     * set an exception predicate which indicates the exception should be traced(return true) or ignored(return false)
+     * except for {@link BlockException}
+     * @param exceptionPredicate the exception predicate
+     */
+    public static void setExceptionPredicate(Predicate<Throwable> exceptionPredicate) {
+        AssertUtil.notNull(exceptionPredicate, "exception predicate must not be null");
+        Tracer.exceptionPredicate = exceptionPredicate;
+    }
+
     private static void checkNotNull(Class<? extends Throwable>[] classes) {
         AssertUtil.notNull(classes, "trace or ignore classes must not be null");
         for (Class<? extends Throwable> clazz : classes) {
@@ -181,6 +202,10 @@ protected static boolean shouldTrace(Throwable t) {
         if (t == null || t instanceof BlockException) {
             return false;
         }
+        if (exceptionPredicate != null) {
+            return exceptionPredicate.test(t);
+        }
+
         if (ignoreClasses != null) {
             for (Class<? extends Throwable> clazz : ignoreClasses) {
                 if (clazz != null && clazz.isAssignableFrom(t.getClass())) {

diff --git a/sentinel-core/src/test/java/com/alibaba/csp/sentinel/TracerTest.java b/sentinel-core/src/test/java/com/alibaba/csp/sentinel/TracerTest.java
@@ -2,6 +2,7 @@
 
 import com.alibaba.csp.sentinel.context.ContextTestUtil;
 import com.alibaba.csp.sentinel.context.ContextUtil;
+import com.alibaba.csp.sentinel.util.function.Predicate;
 import org.junit.After;
 import org.junit.Assert;
 import org.junit.Before;
@@ -52,6 +53,25 @@ public void setExceptionsToTrace() {
         Assert.assertFalse(Tracer.shouldTrace(new Exception()));
     }
 
+    @Test
+    public void setExceptionPredicate() {
+
+        Predicate<Throwable> throwablePredicate = new Predicate<Throwable>() {
+            @Override
+            public boolean test(Throwable throwable) {
+                if (throwable instanceof TraceException) {
+                    return true;
+                } else if (throwable instanceof IgnoreException) {
+                    return false;
+                }
+                return false;
+            }
+        };
+        Tracer.setExceptionPredicate(throwablePredicate);
+        Assert.assertTrue(Tracer.shouldTrace(new TraceException()));
+        Assert.assertFalse(Tracer.shouldTrace(new IgnoreException()));
+    }
+
     @Test
     public void setExceptionsToIgnore() {
         Tracer.ignoreClasses = null;
@@ -88,6 +108,11 @@ public void testNull2() {
         Tracer.setExceptionsToIgnore(IgnoreException.class, null);
     }
 
+    @Test(expected = IllegalArgumentException.class)
+    public void testNull3() {
+        Tracer.setExceptionPredicate(null);
+    }
+
     private class TraceException extends Exception {}
 
     private class TraceException2 extends Exception {}