Fix cacheInvalidateCallback does not work

core/src/test/scala/com/pingcap/tispark/accumulator/AccumulatorSuite.scala

@@ -0,0 +1,72 @@
+/*
+ *
+ * Copyright 2022 PingCAP, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.pingcap.tispark.accumulator;
+
+import org.apache.log4j.spi.LoggingEvent
+import org.apache.log4j.{AppenderSkeleton, Logger}
+import org.apache.spark.sql.BaseTiSparkTest
+
+import java.util
+import java.util.stream.Collectors;
+
+class AccumulatorSuite extends BaseTiSparkTest {
+  test("cacheInvalidateCallback does not work") {
+    val listLogAppender = new ListLogAppender
+    val logger = Logger.getRootLogger
+    logger.addAppender(listLogAppender)
+    try {
+      tidbStmt.execute("DROP TABLE IF EXISTS `t1`")
+      tidbStmt.execute("""
+                         |CREATE TABLE `t1` (
+                         |`a` BIGINT(20)  NOT NULL,
+                         |`b` varchar(255) NOT NULL,
+                         |`c` varchar(255) DEFAULT NULL
+                         |) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin
+              """.stripMargin)
+      spark.sql("SELECT * FROM t1").show()
+      tidbStmt.execute(
+        "SPLIT TABLE t1 BETWEEN (-9223372036854775808) AND (-8223372036854775808) REGIONS 300")
+      spark.sql("SELECT * FROM t1").show()
+    } finally {
+      logger.removeAppender(listLogAppender)
+    }
+    val cacheInvalidateListenerLog = listLogAppender.getLog
+      .stream()
+      .filter(e =>
+        e.getMessage.toString.contains(
+          "Failed to send notification back to driver since CacheInvalidateCallBack is null in executor node"))
+      .collect(Collectors.toList[LoggingEvent])
+    assert(cacheInvalidateListenerLog.size() == 0)
+  }
+
+  class ListLogAppender extends AppenderSkeleton {
+
+    final private val log = new util.ArrayList[LoggingEvent]()
+
+    override def requiresLayout = false
+
+    override protected def append(loggingEvent: LoggingEvent): Unit = {
+      log.add(loggingEvent)
+    }
+
+    override def close(): Unit = {}
+
+    def getLog = new util.ArrayList[LoggingEvent](log)
+  }
+}

core/src/test/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlanTestSuite.scala

@@ -81,8 +81,7 @@ class LogicalPlanTestSuite extends BasePlanTest {
       "insert into test3 values(1, 2, 3), (2, 1, 3), (2, 1, 4), (3, 2, 3), (4, 2, 1)")
     refreshConnections()
     val df =
-      spark.sql(
-        """
+      spark.sql("""
           |select t1.*, (
           |	select count(*)
           |	from test2

tikv-client/src/main/java/com/pingcap/tikv/TiSession.java

@@ -341,6 +341,8 @@ public ExecutorService getThreadPoolForDeleteRange() {
    */
   public void injectCallBackFunc(Function<CacheInvalidateEvent, Void> callBackFunc) {
     this.cacheInvalidateCallback = callBackFunc;
+    RegionManager manager = this.getRegionManager();
+    manager.setCacheInvalidateCallback(callBackFunc);
   }
 
   /**

tikv-client/src/main/java/com/pingcap/tikv/region/RegionManager.java

@@ -50,7 +50,7 @@ public class RegionManager {
   // https://github.com/pingcap/tispark/issues/1170
   private final RegionCache cache;
 
-  private final Function<CacheInvalidateEvent, Void> cacheInvalidateCallback;
+  private Function<CacheInvalidateEvent, Void> cacheInvalidateCallback;
 
   private AtomicInteger tiflashStoreIndex = new AtomicInteger(0);
 
@@ -67,6 +67,11 @@ public RegionManager(ReadOnlyPDClient pdClient) {
     this.cacheInvalidateCallback = null;
   }
 
+  public synchronized void setCacheInvalidateCallback(
+          Function<CacheInvalidateEvent, Void> cacheInvalidateCallback) {
+    this.cacheInvalidateCallback = cacheInvalidateCallback;
+  }
+
   public Function<CacheInvalidateEvent, Void> getCacheInvalidateCallback() {
     return cacheInvalidateCallback;
   }