service_safe_point (#2642)

core/src/main/scala/com/pingcap/tispark/TiConfigConst.scala

@@ -90,4 +90,9 @@ object TiConfigConst {
   val REPLICA_READ_ADDRESS_WHITELIST = "spark.tispark.replica_read.address_whitelist"
   val REPLICA_READ_ADDRESS_BLACKLIST = "spark.tispark.replica_read.address_blacklist"
   val REPLICA_READ_ADDRESS_DEFAULT = ""
+
+  // GC
+  val GC_MAX_WAIT_TIME: String = "spark.tispark.gc_max_wait_time"
+  val DEFAULT_GC_MAX_WAIT_TIME: Long = 24 * 60 * 60
+  val DEFAULT_GC_SAFE_POINT_TTL: Int = 5 * 60
 }

core/src/main/scala/com/pingcap/tispark/safepoint/ServiceSafePoint.scala

@@ -0,0 +1,119 @@
+/*
+ * Copyright 2023 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.safepoint
+
+import com.pingcap.tikv.ClientSession
+import org.slf4j.LoggerFactory
+import org.tikv.common.exception.TiInternalException
+import org.tikv.common.meta.TiTimestamp
+import org.tikv.common.util.ConcreteBackOffer
+
+import java.util.concurrent.{Executors, ScheduledExecutorService, TimeUnit}
+
+case class ServiceSafePoint(
+    serviceId: String,
+    ttl: Long,
+    GCMaxWaitTime: Long,
+    clientSession: ClientSession) {
+
+  private val PD_UPDATE_SAFE_POINT_BACKOFF: Int = 20 * 1000
+  private final val logger = LoggerFactory.getLogger(getClass.getName)
+  private var minStartTs = Long.MaxValue
+  val service: ScheduledExecutorService = Executors.newSingleThreadScheduledExecutor()
+  service.scheduleAtFixedRate(
+    () => {
+      if (minStartTs != Long.MaxValue) {
+        val safePoint = clientSession.getTiKVSession.getPDClient.updateServiceGCSafePoint(
+          serviceId,
+          ttl,
+          minStartTs,
+          ConcreteBackOffer.newCustomBackOff(PD_UPDATE_SAFE_POINT_BACKOFF))
+        if (safePoint > minStartTs) {
+          // will not happen unless someone delete the TiSpark service safe point in PD compulsively
+          logger.error(
+            s"Failed to register service GC safe point because the current minimum safe point $safePoint is newer than what we assume $minStartTs. Maybe you delete the TiSpark safe point in PD.")
+        } else {
+          logger.info(s"register service GC safe point $minStartTs success.")
+        }
+      }
+    },
+    0,
+    1,
+    TimeUnit.MINUTES)
+
+  // TiSpark can only decrease minStartTs now. Because we can not known which transaction is finished, so we can not increase minStartTs.
+  def updateStartTs(startTimeStamp: TiTimestamp): Unit = {
+    this.synchronized {
+      val now = clientSession.getTiKVSession.getTimestamp
+      if (now.getPhysical - startTimeStamp.getPhysical >= GCMaxWaitTime * 1000) {
+        throw new TiInternalException(
+          s"Can not pause GC more than spark.tispark.gc_max_wait_time=$GCMaxWaitTime s. start_ts: ${startTimeStamp.getVersion}, now: ${now.getVersion}. You can adjust spark.tispark.gc_max_wait_time to increase the gc max wait time.")
+      }
+      val startTs = startTimeStamp.getVersion
+      if (startTs >= minStartTs) {
+        // minStartTs >= safe point, so startTs must >= safe point. Check it in case some one delete the TiSpark service safe point in PD compulsively.
+        checkServiceSafePoint(startTs)
+      } else {
+        // applyServiceSafePoint may throw exception. Consider startTs < safePoint < minStartTs.
+        applyServiceSafePoint(startTs)
+        // let minStartTs = startTs after applyServiceSafePoint success
+        minStartTs = startTs
+      }
+    }
+  }
+
+  private def checkServiceSafePoint(startTs: Long): Unit = {
+    val safePoint = clientSession.getTiKVSession.getPDClient.updateServiceGCSafePoint(
+      serviceId,
+      ttl,
+      minStartTs,
+      ConcreteBackOffer.newCustomBackOff(PD_UPDATE_SAFE_POINT_BACKOFF))
+    if (safePoint > startTs) {
+      throw new TiInternalException(
+        s"Failed to check service GC safe point because the current minimum safe point $safePoint is newer than start_ts $startTs.")
+    }
+    logger.info(s"check start_ts $startTs success.")
+  }
+
+  private def applyServiceSafePoint(startTs: Long): Unit = {
+    val safePoint = clientSession.getTiKVSession.getPDClient.updateServiceGCSafePoint(
+      serviceId,
+      ttl,
+      startTs,
+      ConcreteBackOffer.newCustomBackOff(PD_UPDATE_SAFE_POINT_BACKOFF))
+    if (safePoint > startTs) {
+      throw new TiInternalException(
+        s"Failed to register service GC safe point because the current minimum safe point $safePoint is newer than what we assume $startTs.")
+    }
+    logger.info(s"register service GC safe point $startTs success.")
+  }
+
+  def stopRegisterSafePoint(): Unit = {
+    try {
+      minStartTs = Long.MaxValue
+      clientSession.getTiKVSession.getPDClient.updateServiceGCSafePoint(
+        serviceId,
+        ttl,
+        Long.MaxValue,
+        ConcreteBackOffer.newCustomBackOff(PD_UPDATE_SAFE_POINT_BACKOFF))
+    } catch {
+      case e: Exception => logger.error("Failed to stop register service GC safe point", e)
+    } finally {
+      service.shutdownNow()
+    }
+  }
+}

core/src/main/scala/com/pingcap/tispark/write/TiBatchWrite.scala

@@ -192,6 +192,7 @@ class TiBatchWrite(
     val startTimeStamp = clientSession.getTiKVSession.getTimestamp
     startTs = startTimeStamp.getVersion
     logger.info(s"startTS: $startTs")
+    tiContext.serviceSafePoint.updateStartTs(startTimeStamp)
 
     // pre calculate
     val shuffledRDD: RDD[(SerializableKey, Array[Byte])] = {

core/src/main/scala/org/apache/spark/sql/TiContext.scala

@@ -21,6 +21,7 @@ import com.pingcap.tikv.{ClientSession, TiConfiguration}
 import com.pingcap.tispark._
 import com.pingcap.tispark.auth.TiAuthorization
 import com.pingcap.tispark.listener.CacheInvalidateListener
+import com.pingcap.tispark.safepoint.ServiceSafePoint
 import com.pingcap.tispark.statistics.StatisticsManager
 import com.pingcap.tispark.utils.TiUtil
 import org.apache.spark.SparkConf
@@ -33,6 +34,7 @@ import org.json4s.jackson.JsonMethods._
 import scalaj.http.Http
 
 import java.lang
+import java.util.UUID
 import scala.collection.JavaConverters._
 import scala.collection.mutable
 
@@ -48,11 +50,28 @@ class TiContext(val sparkSession: SparkSession) extends Serializable with Loggin
     } else Option.empty)
   final val clientSession = ClientSession.getInstance(tiConf)
   lazy val sqlContext: SQLContext = sparkSession.sqlContext
+  // GC
+  val GCMaxWaitTime: Long =
+    try {
+      conf
+        .get(TiConfigConst.GC_MAX_WAIT_TIME, TiConfigConst.DEFAULT_GC_MAX_WAIT_TIME.toString)
+        .toLong
+    } catch {
+      case _: Exception => TiConfigConst.DEFAULT_GC_MAX_WAIT_TIME
+    }
+
+  val serviceSafePoint: ServiceSafePoint =
+    ServiceSafePoint(
+      "tispark_" + UUID.randomUUID,
+      TiConfigConst.DEFAULT_GC_SAFE_POINT_TTL,
+      GCMaxWaitTime,
+      clientSession)
 
   sparkSession.sparkContext.addSparkListener(new SparkListener() {
     override def onApplicationEnd(applicationEnd: SparkListenerApplicationEnd): Unit = {
       if (clientSession != null) {
         try {
+          serviceSafePoint.stopRegisterSafePoint()
           clientSession.close()
         } catch {
           case e: Throwable => logWarning("fail to close ClientSession!", e)

spark-wrapper/spark-3.0/src/main/scala/org/apache/spark/sql/extensions/TiStrategy.scala

@@ -108,7 +108,7 @@ case class TiStrategy(getOrCreateTiContext: SparkSession => TiContext)(sparkSess
     } else {
       tiContext.clientSession.getSnapshotTimestamp
     }
-
+    tiContext.serviceSafePoint.updateStartTs(ts)
     if (plan.isStreaming) {
       // We should use a new timestamp for next batch execution.
       // Otherwise Spark Structure Streaming will not see new data in TiDB.

spark-wrapper/spark-3.1/src/main/scala/org/apache/spark/sql/extensions/TiStrategy.scala

@@ -111,6 +111,7 @@ case class TiStrategy(getOrCreateTiContext: SparkSession => TiContext)(sparkSess
     } else {
       tiContext.clientSession.getSnapshotTimestamp
     }
+    tiContext.serviceSafePoint.updateStartTs(ts)
     if (plan.isStreaming) {
       // We should use a new timestamp for next batch execution.
       // Otherwise Spark Structure Streaming will not see new data in TiDB.

spark-wrapper/spark-3.2/src/main/scala/org/apache/spark/sql/extensions/TiStrategy.scala

@@ -111,7 +111,7 @@ case class TiStrategy(getOrCreateTiContext: SparkSession => TiContext)(sparkSess
     } else {
       tiContext.clientSession.getSnapshotTimestamp
     }
-
+    tiContext.serviceSafePoint.updateStartTs(ts)
     if (plan.isStreaming) {
       // We should use a new timestamp for next batch execution.
       // Otherwise Spark Structure Streaming will not see new data in TiDB.

spark-wrapper/spark-3.3/src/main/scala/org/apache/spark/sql/extensions/TiStrategy.scala

@@ -111,7 +111,7 @@ case class TiStrategy(getOrCreateTiContext: SparkSession => TiContext)(sparkSess
     } else {
       tiContext.clientSession.getSnapshotTimestamp
     }
-
+    tiContext.serviceSafePoint.updateStartTs(ts)
     if (plan.isStreaming) {
       // We should use a new timestamp for next batch execution.
       // Otherwise Spark Structure Streaming will not see new data in TiDB.

tikv-client/pom.xml

@@ -183,7 +183,7 @@
         <dependency>
             <groupId>org.tikv</groupId>
             <artifactId>tikv-client-java</artifactId>
-            <version>3.3.2</version>
+            <version>3.3.3</version>
         </dependency>
       <dependency>
         <groupId>commons-codec</groupId>