Support service safe point (#2634)

Author: shiyuhang0

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

@@ -79,4 +79,10 @@ object TiConfigConst {
   val HOST_MAPPING: String = "spark.tispark.host_mapping"
 
   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

@@ -186,6 +186,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.

spark-wrapper/spark-3.3/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.

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

@@ -87,6 +87,7 @@
 import org.tikv.kvproto.Pdpb.Timestamp;
 import org.tikv.kvproto.Pdpb.TsoRequest;
 import org.tikv.kvproto.Pdpb.TsoResponse;
+import org.tikv.kvproto.Pdpb.UpdateServiceGCSafePointRequest;
 
 public class PDClient extends AbstractGRPCClient<PDBlockingStub, PDStub>
     implements ReadOnlyPDClient {
@@ -298,6 +299,17 @@ private Supplier<GetAllStoresRequest> buildGetAllStoresReq() {
     return () -> GetAllStoresRequest.newBuilder().setHeader(header).build();
   }
 
+  private Supplier<UpdateServiceGCSafePointRequest> buildUpdateServiceGCSafePointRequest(
+      ByteString serviceId, long ttl, long safePoint) {
+    return () ->
+        UpdateServiceGCSafePointRequest.newBuilder()
+            .setHeader(header)
+            .setSafePoint(safePoint)
+            .setServiceId(serviceId)
+            .setTTL(ttl)
+            .build();
+  }
+
   private <T> PDErrorHandler<GetStoreResponse> buildPDErrorHandler() {
     return new PDErrorHandler<>(
         r -> r.getHeader().hasError() ? buildFromPdpbError(r.getHeader().getError()) : null, this);
@@ -336,6 +348,20 @@ public List<Store> getAllStores(BackOffer backOffer) {
         .getStoresList();
   }
 
+  @Override
+  public Long updateServiceGCSafePoint(
+      String serviceId, long ttl, long safePoint, BackOffer backOffer) {
+    return callWithRetry(
+            backOffer,
+            PDGrpc.getUpdateServiceGCSafePointMethod(),
+            buildUpdateServiceGCSafePointRequest(
+                ByteString.copyFromUtf8(serviceId), ttl, safePoint),
+            new PDErrorHandler<>(
+                r -> r.getHeader().hasError() ? buildFromPdpbError(r.getHeader().getError()) : null,
+                this))
+        .getMinSafePoint();
+  }
+
   @Override
   public void close() throws InterruptedException {
     etcdClient.close();

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

@@ -67,4 +67,16 @@ public interface ReadOnlyPDClient {
   Future<Store> getStoreAsync(BackOffer backOffer, long storeId);
 
   List<Store> getAllStores(BackOffer backOffer);
+
+  /**
+   * Update ServiceGCSafePoint
+   *
+   * @param serviceId ServiceId
+   * @param ttl TTL in seconds
+   * @param safePoint The TiTimestamp you want to set. Set to start_ts.getPrevious() is a good
+   *     practice
+   * @return the MinSafePoint of all service. If this value is greater than safePoint, it means
+   *     update failed.
+   */
+  Long updateServiceGCSafePoint(String serviceId, long ttl, long safePoint, BackOffer backOffer);
 }