From 5523fd0d427ba0f66c3e110c2de6347edadcd125 Mon Sep 17 00:00:00 2001
From: pleasecheckhere2016 <707748808@qq.com>
Date: Thu, 21 May 2020 15:17:19 +0800
Subject: [PATCH] Add exceptionPredicate in Tracer for customized exception
 filtering logic (#1496)

---
 .../java/com/alibaba/csp/sentinel/Tracer.java | 25 +++++++++++++++++++
 .../com/alibaba/csp/sentinel/TracerTest.java  | 25 +++++++++++++++++++
 2 files changed, 50 insertions(+)

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
index f29d6c5c..299c80ee 100755
--- 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.ContextUtil;
 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 class Tracer {
         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 @@ public class Tracer {
         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
index 2a0e21a6..7acb3d7e 100644
--- 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 @@ package com.alibaba.csp.sentinel;
 
 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 class TracerTest extends Tracer {
         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 class TracerTest extends Tracer {
         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 {}