[GSOC2022]REFACTOR---Normalize pdclient to client-java

core/src/main/java/com/pingcap/tikv/hostmap/UriHostMapping.java

@@ -16,7 +16,6 @@
 
 package com.pingcap.tikv.hostmap;
 
-import com.pingcap.tikv.util.HostMapping;
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.util.Arrays;
@@ -25,6 +24,7 @@
 import java.util.stream.Collectors;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.tikv.common.HostMapping;
 
 public class UriHostMapping implements HostMapping {
 

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

@@ -16,7 +16,7 @@
 
 package com.pingcap.tispark
 
-import com.pingcap.tikv.util.RangeSplitter.RegionTask
+import org.tikv.common.util.RangeSplitter.RegionTask
 import org.apache.spark.Partition
 
 class TiPartition(val idx: Int, val tasks: Seq[RegionTask], val appId: String) extends Partition {

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

@@ -16,8 +16,8 @@
 
 package com.pingcap.tispark
 
-import com.pingcap.tikv.exception.TiInternalException
 import org.slf4j.LoggerFactory
+import org.tikv.common.exception.TiInternalException
 
 object TiSparkInfo {
   private final val logger = LoggerFactory.getLogger(getClass.getName)

core/src/main/scala/com/pingcap/tispark/accumulator/CacheInvalidateAccumulator.scala

@@ -17,9 +17,8 @@
 package com.pingcap.tispark.accumulator
 
 import java.util
-
-import com.pingcap.tikv.event.CacheInvalidateEvent
 import org.apache.spark.util.AccumulatorV2
+import org.tikv.common.event.CacheInvalidateEvent
 
 import scala.collection.JavaConversions._
 

core/src/main/scala/com/pingcap/tispark/handler/CacheInvalidateEventHandler.scala

@@ -18,11 +18,11 @@
 
 package com.pingcap.tispark.handler
 
-import com.pingcap.tikv.event.CacheInvalidateEvent
-import com.pingcap.tikv.event.CacheInvalidateEvent.CacheType
-import com.pingcap.tikv.region.RegionManager
+import org.tikv.common.event.CacheInvalidateEvent.CacheType
 import com.pingcap.tispark.listener.CacheInvalidateListener
 import org.slf4j.LoggerFactory
+import org.tikv.common.event.CacheInvalidateEvent
+import org.tikv.common.region.RegionManager
 
 /**
  * A CacheInvalidateEventHandler as it's name indicates what this class will do.
@@ -44,8 +44,8 @@ class CacheInvalidateEventHandler(regionManager: RegionManager) {
         case CacheType.REGION_STORE =>
           // Used for updating region/store cache in the given regionManager
           if (event.shouldUpdateRegion()) {
-            logger.info(s"Invalidating region ${event.getRegion.getId} cache at driver.")
-            val region = regionManager.getRegionByKey(event.getRegion.getStartKey)
+            logger.info(s"Invalidating region ${event.getRegionId} cache at driver.")
+            val region = regionManager.getRegionById(event.getRegionId)
             if (region != null) {
               regionManager.invalidateRegion(region)
             }
@@ -58,16 +58,16 @@ class CacheInvalidateEventHandler(regionManager: RegionManager) {
         case CacheType.LEADER =>
           // Used for updating leader information cached in the given regionManager
           logger.info(
-            s"Invalidating leader of region:${event.getRegion.getId} store:${event.getStoreId} cache at driver.")
-          val region = regionManager.getRegionByKey(event.getRegion.getStartKey)
+            s"Invalidating leader of region:${event.getRegionId} store:${event.getStoreId} cache at driver.")
+          val region = regionManager.getRegionById(event.getRegionId)
           if (region != null) {
             regionManager.updateLeader(region, event.getStoreId)
             regionManager.invalidateRegion(region)
           }
 
         case CacheType.REQ_FAILED =>
-          logger.info(s"Request failed cache invalidation for region ${event.getRegion.getId}")
-          val region = regionManager.getRegionByKey(event.getRegion.getStartKey)
+          logger.info(s"Request failed cache invalidation for region ${event.getRegionId}")
+          val region = regionManager.getRegionById(event.getRegionId)
           if (region != null) {
             regionManager.onRequestFail(region)
           }

core/src/main/scala/com/pingcap/tispark/listener/CacheInvalidateListener.scala

@@ -18,12 +18,12 @@
 
 package com.pingcap.tispark.listener
 
-import com.pingcap.tikv.event.CacheInvalidateEvent
-import com.pingcap.tikv.region.RegionManager
 import com.pingcap.tispark.accumulator.CacheInvalidateAccumulator
 import com.pingcap.tispark.handler.CacheInvalidateEventHandler
 import org.apache.spark.SparkContext
 import org.slf4j.LoggerFactory
+import org.tikv.common.event.CacheInvalidateEvent
+import org.tikv.common.region.RegionManager
 
 class CacheInvalidateListener()
     extends Serializable

core/src/main/scala/com/pingcap/tispark/statistics/StatisticsHelper.scala

@@ -18,7 +18,7 @@
 
 package com.pingcap.tispark.statistics
 
-import com.google.common.primitives.UnsignedLong
+import org.tikv.shade.com.google.common.primitives.UnsignedLong
 import com.pingcap.tikv.expression.{ByItem, ColumnRef, ComparisonBinaryExpression, Constant}
 import com.pingcap.tikv.key.Key
 import com.pingcap.tikv.meta.TiDAGRequest.PushDownType
@@ -27,6 +27,7 @@ import com.pingcap.tikv.row.Row
 import com.pingcap.tikv.statistics._
 import com.pingcap.tikv.types.{BytesType, IntegerType}
 import org.slf4j.LoggerFactory
+import org.tikv.common.meta.TiTimestamp
 
 import scala.collection.JavaConversions._
 import scala.collection.mutable

core/src/main/scala/com/pingcap/tispark/statistics/StatisticsManager.scala

@@ -18,16 +18,16 @@
 
 package com.pingcap.tispark.statistics
 
-import com.google.common.cache.CacheBuilder
 import com.pingcap.tikv.catalog.Catalog
 import com.pingcap.tikv.meta.{TiColumnInfo, TiDAGRequest, TiIndexInfo, TiTableInfo}
 import com.pingcap.tikv.row.Row
 import com.pingcap.tikv.statistics._
 import com.pingcap.tikv.types.DataType
-import com.pingcap.tikv.{Snapshot, TiSession}
+import com.pingcap.tikv.{ClientSession, Snapshot}
 import com.pingcap.tispark.statistics.StatisticsHelper.shouldUpdateHistogram
 import com.pingcap.tispark.statistics.estimate.{DefaultTableSizeEstimator, TableSizeEstimator}
 import org.slf4j.LoggerFactory
+import org.tikv.shade.com.google.common.cache.CacheBuilder
 
 import scala.collection.JavaConversions._
 import scala.collection.mutable
@@ -70,7 +70,7 @@ object StatisticsManager {
     .newBuilder()
     .build[java.lang.Long, TableStatistics]
   protected var initialized: Boolean = false
-  private var session: TiSession = _
+  private var clientSession: ClientSession = _
   private var snapshot: Snapshot = _
   private var catalog: Catalog = _
   private var dbPrefix: String = _
@@ -157,7 +157,7 @@ object StatisticsManager {
     // load count, modify_count, version info
     loadMetaToTblStats(tblId, tblStatistic)
     val req = StatisticsHelper
-      .buildHistogramsRequest(histTable, tblId, session.getTimestamp)
+      .buildHistogramsRequest(histTable, tblId, clientSession.getTiKVSession.getTimestamp)
 
     val rows = readDAGRequest(req, histTable.getId)
     if (rows.isEmpty) return
@@ -201,7 +201,10 @@ object StatisticsManager {
 
   private def loadMetaToTblStats(tableId: Long, tableStatistics: TableStatistics): Unit = {
     val req =
-      StatisticsHelper.buildMetaRequest(metaTable, tableId, session.getTimestamp)
+      StatisticsHelper.buildMetaRequest(
+        metaTable,
+        tableId,
+        clientSession.getTiKVSession.getTimestamp)
 
     val rows = readDAGRequest(req, metaTable.getId)
     if (rows.isEmpty) return
@@ -219,7 +222,10 @@ object StatisticsManager {
       tableId: Long,
       requests: Seq[StatisticsDTO]): Seq[StatisticsResult] = {
     val req =
-      StatisticsHelper.buildBucketRequest(bucketTable, tableId, session.getTimestamp)
+      StatisticsHelper.buildBucketRequest(
+        bucketTable,
+        tableId,
+        clientSession.getTiKVSession.getTimestamp)
 
     val rows = readDAGRequest(req, bucketTable.getId)
     if (rows.isEmpty) return Nil
@@ -260,28 +266,28 @@ object StatisticsManager {
 
   def setEstimator(estimator: TableSizeEstimator): Unit = tableSizeEstimator = estimator
 
-  def initStatisticsManager(tiSession: TiSession): Unit =
+  def initStatisticsManager(clientSession: ClientSession): Unit =
     if (!initialized) {
       synchronized {
         if (!initialized) {
-          initialize(tiSession)
+          initialize(clientSession)
           initialized = true
         }
       }
     }
 
-  protected def initialize(tiSession: TiSession): Unit = {
-    session = tiSession
-    snapshot = tiSession.createSnapshot()
-    catalog = tiSession.getCatalog
-    dbPrefix = tiSession.getConf.getDBPrefix
+  protected def initialize(clientSession: ClientSession): Unit = {
+    this.clientSession = clientSession
+    this.snapshot = clientSession.createSnapshot()
+    this.catalog = clientSession.getCatalog
+    this.dbPrefix = clientSession.getConf.getDBPrefix
     // An estimator used to calculate table size.
-    tableSizeEstimator = DefaultTableSizeEstimator
+    this.tableSizeEstimator = DefaultTableSizeEstimator
     val mysqlDB = catalog.getDatabaseFromCache(s"${dbPrefix}mysql")
-    metaTable = catalog.getTableFromCache(mysqlDB, "stats_meta")
-    histTable = catalog.getTableFromCache(mysqlDB, "stats_histograms")
-    bucketTable = catalog.getTableFromCache(mysqlDB, "stats_buckets")
-    statisticsMap.invalidateAll()
+    this.metaTable = catalog.getTableFromCache(mysqlDB, "stats_meta")
+    this.histTable = catalog.getTableFromCache(mysqlDB, "stats_histograms")
+    this.bucketTable = catalog.getTableFromCache(mysqlDB, "stats_buckets")
+    this.statisticsMap.invalidateAll()
   }
 
   def reset(): Unit = initialized = false

core/src/main/scala/com/pingcap/tispark/telemetry/TeleMsg.scala

@@ -16,11 +16,12 @@
 
 package com.pingcap.tispark.telemetry
 
-import com.pingcap.tikv.util.ConcreteBackOffer
-import com.pingcap.tikv.{TiConfiguration, TiSession, TwoPhaseCommitter}
+import com.pingcap.tikv.{ClientSession, TiConfiguration}
 import com.pingcap.tispark.utils.{SystemInfoUtil, TiUtil}
 import org.apache.spark.sql.SparkSession
 import org.slf4j.LoggerFactory
+import org.tikv.common.util.ConcreteBackOffer
+import org.tikv.txn.TwoPhaseCommitter
 import java.util.UUID
 
 /**
@@ -45,15 +46,15 @@ class TeleMsg(sparkSession: SparkSession) {
     try {
       val conf = TiConfiguration.createDefault(pdAddr.get)
       TiUtil.sparkConfToTiConfWithoutPD(SparkSession.active.sparkContext.getConf, conf)
-      val tiSession = TiSession.getInstance(conf)
-      val snapShot = tiSession.createSnapshot()
+      val clientSession = ClientSession.getInstance(conf)
+      val snapShot = clientSession.createSnapshot()
       val value = snapShot.get(TRACK_ID.getBytes("UTF-8"))
 
       if (value.nonEmpty)
         return new String(value, "UTF-8")
 
       val uuid = TRACK_ID_PREFIX + UUID.randomUUID().toString
-      putKeyValue(TRACK_ID, uuid, conf, tiSession)
+      putKeyValue(TRACK_ID, uuid, clientSession)
       uuid
     } catch {
       case e: Throwable =>
@@ -62,14 +63,10 @@ class TeleMsg(sparkSession: SparkSession) {
     }
   }
 
-  private def putKeyValue(
-      key: String,
-      value: String,
-      conf: TiConfiguration,
-      tiSession: TiSession): Unit = {
-    val startTS = tiSession.getTimestamp.getVersion
+  private def putKeyValue(key: String, value: String, clientSession: ClientSession): Unit = {
+    val startTS = clientSession.getTiKVSession.getTimestamp.getVersion
     try {
-      val twoPhaseCommitter = new TwoPhaseCommitter(conf, startTS)
+      val twoPhaseCommitter = new TwoPhaseCommitter(clientSession.getTiKVSession, startTS)
       val backOffer = ConcreteBackOffer.newCustomBackOff(1000)
       twoPhaseCommitter.prewritePrimaryKey(
         backOffer,
@@ -78,7 +75,7 @@ class TeleMsg(sparkSession: SparkSession) {
       twoPhaseCommitter.commitPrimaryKey(
         backOffer,
         key.getBytes("UTF-8"),
-        tiSession.getTimestamp.getVersion)
+        clientSession.getTiKVSession.getTimestamp.getVersion)
     } catch {
       case e: Throwable =>
         logger.warn("Failed to set telemetry ID to TiKV.", e.getMessage)

core/src/main/scala/com/pingcap/tispark/telemetry/TiSparkTeleInfo.scala

@@ -16,12 +16,12 @@
 
 package com.pingcap.tispark.telemetry
 
-import com.pingcap.tispark.utils.{HttpClientUtil, TiUtil}
-import com.pingcap.tispark.TiSparkVersion
-import org.apache.spark.sql.SparkSession
 import com.fasterxml.jackson.databind.ObjectMapper
 import com.pingcap.tikv.TiConfiguration
+import com.pingcap.tispark.TiSparkVersion
 import com.pingcap.tispark.auth.TiAuthorization
+import com.pingcap.tispark.utils.{HttpClientUtil, TiUtil}
+import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.internal.SQLConf
 import org.slf4j.LoggerFactory
 import scalaj.http.HttpResponse

core/src/main/scala/com/pingcap/tispark/utils/TiUtil.scala

@@ -19,8 +19,7 @@ package com.pingcap.tispark.utils
 import com.pingcap.tikv.TiConfiguration
 import com.pingcap.tikv.datatype.TypeMapping
 import com.pingcap.tikv.hostmap.UriHostMapping
-import com.pingcap.tikv.meta.{TiDAGRequest, TiTableInfo, TiTimestamp}
-import com.pingcap.tikv.region.TiStoreType
+import com.pingcap.tikv.meta.{TiDAGRequest, TiTableInfo}
 import com.pingcap.tikv.types._
 import com.pingcap.tispark._
 import org.apache.commons.lang3.StringUtils
@@ -30,6 +29,9 @@ import org.apache.spark.sql.types.{MetadataBuilder, StructField, StructType}
 import org.apache.spark.sql.{DataFrame, Row, SQLContext, SparkSession}
 import org.apache.spark.{SparkConf, sql}
 import org.slf4j.LoggerFactory
+import org.tikv.common.meta
+import org.tikv.common.meta.TiTimestamp
+import org.tikv.common.region.TiStoreType
 import org.tikv.kvproto.Kvrpcpb.{CommandPri, IsolationLevel}
 
 import java.time.{Instant, LocalDate, ZoneId}