Update metrics for read and writes via DSV2

Thanks @ymuzammil for fixing the issue with the read metrics.

Fixes SPARKC-712

Author: siadat

connector/src/main/scala/com/datastax/spark/connector/datasource/CassandraInJoinReaderFactory.scala

@@ -30,6 +30,8 @@ import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader, PartitionReaderFactory}
 import org.apache.spark.sql.sources.In
 import org.apache.spark.sql.types.{LongType, StructField, StructType}
+import org.apache.spark.metrics.InputMetricsUpdater
+import org.apache.spark.TaskContext
 
 import scala.util.{Failure, Success}
 
@@ -80,16 +82,18 @@ abstract class CassandraBaseInJoinReader(
   protected val maybeRateLimit = JoinHelper.maybeRateLimit(readConf)
   protected val requestsPerSecondRateLimiter = JoinHelper.requestsPerSecondRateLimiter(readConf)
 
+  protected val metricsUpdater = InputMetricsUpdater(TaskContext.get(), readConf)
   protected def pairWithRight(left: CassandraRow): SettableFuture[Iterator[(CassandraRow, InternalRow)]] = {
     val resultFuture = SettableFuture.create[Iterator[(CassandraRow, InternalRow)]]
     val leftSide = Iterator.continually(left)
 
     queryExecutor.executeAsync(bsb.bind(left).executeAs(readConf.executeAs)).onComplete {
       case Success(rs) =>
         val resultSet = new PrefetchingResultSetIterator(rs)
+        val iteratorWithMetrics = resultSet.map(metricsUpdater.updateMetrics)
         /* This is a much less than ideal place to actually rate limit, we are buffering
         these futures this means we will most likely exceed our threshold*/
-        val throttledIterator = resultSet.map(maybeRateLimit)
+        val throttledIterator = iteratorWithMetrics.map(maybeRateLimit)
         val rightSide = throttledIterator.map(rowReader.read(_, rowMetadata))
         resultFuture.set(leftSide.zip(rightSide))
       case Failure(throwable) =>
@@ -121,6 +125,7 @@ abstract class CassandraBaseInJoinReader(
   override def get(): InternalRow = currentRow
 
   override def close(): Unit = {
+    metricsUpdater.finish()
     session.close()
   }
 }

connector/src/main/scala/com/datastax/spark/connector/datasource/CassandraScanPartitionReaderFactory.scala

@@ -30,6 +30,8 @@ import com.datastax.spark.connector.util.Logging
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.connector.read._
 import org.apache.spark.sql.types.{LongType, StructField, StructType}
+import org.apache.spark.metrics.InputMetricsUpdater
+import org.apache.spark.TaskContext
 
 case class CassandraScanPartitionReaderFactory(
   connector: CassandraConnector,
@@ -79,6 +81,8 @@ abstract class CassandraPartitionReaderBase
   protected val rowIterator = getIterator()
   protected var lastRow: InternalRow = InternalRow()
 
+  protected val metricsUpdater = InputMetricsUpdater(TaskContext.get(), readConf)
+
   override def next(): Boolean = {
     if (rowIterator.hasNext) {
       lastRow = rowIterator.next()
@@ -91,6 +95,7 @@ abstract class CassandraPartitionReaderBase
   override def get(): InternalRow = lastRow
 
   override def close(): Unit = {
+    metricsUpdater.finish()
     scanner.close()
   }
 
@@ -125,7 +130,8 @@ abstract class CassandraPartitionReaderBase
     tokenRanges.iterator.flatMap { range =>
       val scanResult = ScanHelper.fetchTokenRange(scanner, tableDef, queryParts, range, readConf.consistencyLevel, readConf.fetchSizeInRows)
       val meta = scanResult.metadata
-      scanResult.rows.map(rowReader.read(_, meta))
+      val iteratorWithMetrics = scanResult.rows.map(metricsUpdater.updateMetrics)
+      iteratorWithMetrics.map(rowReader.read(_, meta))
     }
   }
 

connector/src/main/scala/com/datastax/spark/connector/datasource/CasssandraDriverDataWriterFactory.scala

@@ -25,6 +25,8 @@ import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.connector.write.streaming.StreamingDataWriterFactory
 import org.apache.spark.sql.connector.write.{DataWriter, DataWriterFactory, WriterCommitMessage}
 import org.apache.spark.sql.types.StructType
+import org.apache.spark.metrics.OutputMetricsUpdater
+import org.apache.spark.TaskContext
 
 case class CassandraDriverDataWriterFactory(
   connector: CassandraConnector,
@@ -54,22 +56,31 @@ case class CassandraDriverDataWriter(
 
   private val columns = SomeColumns(inputSchema.fieldNames.map(name => ColumnName(name)): _*)
 
-  private val writer =
+  private val metricsUpdater = OutputMetricsUpdater(TaskContext.get(), writeConf)
+
+  private val asycWriter =
     TableWriter(connector, tableDef, columns, writeConf, false)(unsafeRowWriterFactory)
       .getAsyncWriter()
 
+  private val writer = asycWriter.copy(
+      successHandler = Some(metricsUpdater.batchFinished(success = true, _, _, _)),
+      failureHandler = Some(metricsUpdater.batchFinished(success = false, _, _, _)))
+
   override def write(record: InternalRow): Unit = writer.write(record)
 
   override def commit(): WriterCommitMessage = {
+    metricsUpdater.finish()
     writer.close()
     CassandraCommitMessage()
   }
 
   override def abort(): Unit = {
+    metricsUpdater.finish()
     writer.close()
   }
 
   override def close(): Unit = {
+    metricsUpdater.finish()
     //Our proxy Session Handler handles double closes by ignoring them so this is fine
     writer.close()
   }