Support service safe point (#2634) (#2645)

Author: ti-chi-bot

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

@@ -75,4 +75,12 @@ object TiConfigConst {
   val DEFAULT_STALE_READ: String = ""
 
   val TELEMETRY_ENABEL: String = "spark.tispark.telemetry.enable"
+
+  val NEW_COLLATION_ENABLE: String = "spark.tispark.new_collation_enable"
+
+  // 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,97 @@
+package com.pingcap.tispark.safepoint
+
+import com.pingcap.tikv.TiSession
+import com.pingcap.tikv.exception.TiInternalException
+import com.pingcap.tikv.meta.TiTimestamp
+import com.pingcap.tikv.util.{BackOffer, ConcreteBackOffer}
+import org.slf4j.LoggerFactory
+
+import java.util.concurrent.{Executors, ScheduledExecutorService, TimeUnit}
+
+case class ServiceSafePoint(
+    serviceId: String,
+    ttl: Long,
+    GCMaxWaitTime: Long,
+    tiSession: TiSession) {
+
+  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 = tiSession.getPDClient.updateServiceGCSafePoint(
+          serviceId,
+          ttl,
+          minStartTs,
+          ConcreteBackOffer.newCustomBackOff(BackOffer.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 = tiSession.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 = tiSession.getPDClient.updateServiceGCSafePoint(
+      serviceId,
+      ttl,
+      minStartTs,
+      ConcreteBackOffer.newCustomBackOff(BackOffer.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 = tiSession.getPDClient.updateServiceGCSafePoint(
+      serviceId,
+      ttl,
+      startTs,
+      ConcreteBackOffer.newCustomBackOff(BackOffer.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 = {
+    minStartTs = Long.MaxValue
+    tiSession.getPDClient.updateServiceGCSafePoint(
+      serviceId,
+      ttl,
+      Long.MaxValue,
+      ConcreteBackOffer.newCustomBackOff(BackOffer.PD_UPDATE_SAFE_POINT_BACKOFF))
+    service.shutdownNow()
+  }
+}

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

@@ -196,6 +196,7 @@ class TiBatchWrite(
     val startTimeStamp = tiSession.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,19 +21,20 @@ import com.pingcap.tikv.{TiConfiguration, TiSession}
 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
 import org.apache.spark.internal.Logging
 import org.apache.spark.scheduler.{SparkListener, SparkListenerApplicationEnd}
-import org.apache.spark.sql.catalyst.catalog._
 import org.json4s.DefaultFormats
 import org.json4s.JsonAST._
 import org.json4s.JsonDSL._
 import org.json4s.jackson.JsonMethods._
 import scalaj.http.Http
 
 import java.lang
+import java.util.UUID
 import scala.collection.JavaConverters._
 import scala.collection.mutable
 
@@ -49,11 +50,28 @@ class TiContext(val sparkSession: SparkSession) extends Serializable with Loggin
     } else Option.empty)
   final val tiSession: TiSession = TiSession.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,
+      tiSession)
 
   sparkSession.sparkContext.addSparkListener(new SparkListener() {
     override def onApplicationEnd(applicationEnd: SparkListenerApplicationEnd): Unit = {
       if (tiSession != null) {
         try {
+          serviceSafePoint.stopRegisterSafePoint()
           tiSession.close()
         } catch {
           case e: Throwable => logWarning("fail to close TiSession!", e)

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

@@ -107,6 +107,7 @@ case class TiStrategy(getOrCreateTiContext: SparkSession => TiContext)(sparkSess
     } else {
       tiContext.tiSession.getSnapshotTimestamp
     }
+    tiContext.serviceSafePoint.updateStartTs(ts)
 
     if (plan.isStreaming) {
       // We should use a new timestamp for next batch execution.

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

@@ -110,6 +110,8 @@ case class TiStrategy(getOrCreateTiContext: SparkSession => TiContext)(sparkSess
     } else {
       tiContext.tiSession.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

@@ -110,6 +110,7 @@ case class TiStrategy(getOrCreateTiContext: SparkSession => TiContext)(sparkSess
     } else {
       tiContext.tiSession.getSnapshotTimestamp
     }
+    tiContext.serviceSafePoint.updateStartTs(ts)
 
     if (plan.isStreaming) {
       // We should use a new timestamp for next batch execution.