add callback for record with exception

task-executor-kit/src/main/java/com/mushiny/task/TaskDefinition.java

@@ -2,6 +2,7 @@ package com.mushiny.task;
 
 import java.util.Collection;
 import java.util.function.BiFunction;
+import java.util.function.Consumer;
 
 /**
  * @author lihao
@@ -18,6 +19,8 @@ public class TaskDefinition<T, U, R> {
      */
     private BiFunction<T, U, R> handler;
 
+    private Consumer<T> exceptionHandler;
+
     /**
      * 共有变量，推荐只用于查询，不推荐更新数据
      */
@@ -29,6 +32,13 @@ public class TaskDefinition<T, U, R> {
         this.context = context;
     }
 
+    public TaskDefinition(Collection<T> data, BiFunction<T, U, R> handler, U context, Consumer<T> exceptionHandler) {
+        this.data = data;
+        this.handler = handler;
+        this.exceptionHandler = exceptionHandler;
+        this.context = context;
+    }
+
     public Collection<T> getData() {
         return data;
     }
@@ -48,4 +58,12 @@ public class TaskDefinition<T, U, R> {
     public U getContext() {
         return context;
     }
+
+    public Consumer<T> getExceptionHandler() {
+        return exceptionHandler;
+    }
+
+    public void setExceptionHandler(Consumer<T> exceptionHandler) {
+        this.exceptionHandler = exceptionHandler;
+    }
 }

task-executor-kit/src/main/java/com/mushiny/task/TaskExecutor.java

 package com.mushiny.task;
 
 import java.util.Collection;
+import java.util.Objects;
 import java.util.concurrent.*;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.BiFunction;
+import java.util.function.Consumer;
 
 /**
  * @author lihao
@@ -22,9 +24,9 @@ public class TaskExecutor<T, U, R> {
 
     private CountDownLatch countDownLatch;
 
-    private Long countDownTimeOut;
+    private final Long countDownTimeOut;
 
-    private TimeUnit timeUnit;
+    private final TimeUnit timeUnit;
 
     public TaskExecutor(ThreadPoolExecutor pool, Long countDownTimeOut, TimeUnit timeUnit) {
         this.pool = pool;
@@ -44,6 +46,10 @@ public class TaskExecutor<T, U, R> {
         init(new TaskDefinition<>(data, handler, context));
     }
 
+    public void init(BiFunction<T, U, R> handler, Collection<T> data, U context, Consumer<T> exceptionHandler) {
+        init(new TaskDefinition<>(data, handler, context, exceptionHandler));
+    }
+
     /**
      * 执行函数
      */
@@ -56,6 +62,9 @@ public class TaskExecutor<T, U, R> {
             try {
                 Object apply = t.getHandler().apply(d, t.getContext());
             } catch (Exception ex) {
+                if (Objects.nonNull(t.getExceptionHandler())) {
+                    t.getExceptionHandler().accept(d);
+                }
                 failCount.addAndGet(1);
             } finally {
                 countDownLatch.countDown();

task-executor-kit/src/test/java/com/mushiny/task/TaskTests.java

@@ -98,4 +98,28 @@ public class TaskTests {
         boolean result = executor.execute();
         Assert.assertFalse(result);
     }
+
+
+    @Test
+    public void errorRecordCallbackTest() {
+        TaskExecutor<User, User, String> executor = TaskPoolFactory.createExecutor();
+        List<User> errors = new ArrayList<>();
+        executor.init((data, context) -> {
+            data.setName(data.getName() + "handle");
+            try {
+                Thread.sleep(SLEEP_TIME);
+            } catch (InterruptedException e) {
+                e.printStackTrace();
+            }
+            context.getAge();
+            return "success";
+        }, list, null, errors::add);
+
+        boolean result = executor.execute();
+        System.out.println("========= start ===========");
+        errors.forEach(System.out::println);
+        System.out.println("=========  end  ===========");
+        Assert.assertFalse(result);
+
+    }
 }