Bigtable sink metrics (#67)

Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com>

spark/ingestion/src/main/scala/feast/ingestion/stores/bigtable/BigTableSinkRelation.scala

@@ -36,6 +36,8 @@ import org.apache.spark.sql.sources.{BaseRelation, InsertableRelation}
 import org.apache.spark.sql.types.{StringType, StructType}
 import feast.ingestion.stores.serialization.Serializer
 import org.apache.hadoop.security.UserGroupInformation
+import org.apache.spark.SparkEnv
+import org.apache.spark.metrics.source.BigTableSinkMetricSource
 
 class BigTableSinkRelation(
     override val sqlContext: SQLContext,
@@ -180,6 +182,24 @@ object BigTableSinkRelation {
         column,
         schemaReference ++ r.getAs[Array[Byte]]("value")
       )
+
+      metricSource.foreach(source => {
+        val lag = System.currentTimeMillis() - r.getAs[java.sql.Timestamp]("ts").getTime
+        source.METRIC_TOTAL_ROWS_INSERTED.inc()
+        source.METRIC_ROWS_LAG.update(lag)
+      })
+
       (null, put)
     }
+
+  lazy val metricSource: Option[BigTableSinkMetricSource] = {
+    if (SparkEnv.get.metricsSystem.getSourcesByName(BigTableSinkMetricSource.sourceName).isEmpty) {
+      SparkEnv.get.metricsSystem.registerSource(new BigTableSinkMetricSource)
+    }
+
+    SparkEnv.get.metricsSystem.getSourcesByName(BigTableSinkMetricSource.sourceName) match {
+      case Seq(source: BigTableSinkMetricSource) => Some(source)
+      case _                                     => None
+    }
+  }
 }

spark/ingestion/src/main/scala/org/apache/spark/metrics/source/BigTableSinkMetricSource.scala

@@ -0,0 +1,31 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ * Copyright 2018-2020 The Feast Authors
+ *
+ * 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
+ *
+ *     https://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 org.apache.spark.metrics.source
+
+class BigTableSinkMetricSource extends BaseMetricSource {
+  override val sourceName: String = RedisSinkMetricSource.sourceName
+
+  val METRIC_TOTAL_ROWS_INSERTED =
+    metricRegistry.counter(counterWithLabels("feature_row_ingested_count"))
+
+  val METRIC_ROWS_LAG =
+    metricRegistry.histogram(metricWithLabels("feature_row_lag_ms"))
+}
+
+object BigTableSinkMetricSource {
+  val sourceName = "bigtable_sink"
+}