chore: pull metrics into its own target (#592)

## Summary


- Metrics is actually used in many places, so let's pull it into its own
slim target first.
- Follow up in the spark offline side to further reduce the dep on
online.

## Checklist
- [ ] Added Unit Tests
- [ ] Covered by existing CI
- [ ] Integration tested
- [ ] Documentation update
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **New Features**
- Launched a dedicated metrics library to enhance performance monitoring
and reporting.
- **Refactor**
- Consolidated and updated metrics imports and context handling across
services for improved consistency.
- **Tests**
- Aligned test configurations with the new metrics framework to ensure
robust validation.

These improvements streamline performance tracking and dependency
management internally while bolstering the overall reliability of the
application.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

<!-- av pr metadata
This information is embedded by the av CLI when creating PRs to track
the status of stacks when using Aviator. Please do not delete or edit
this section of the PR.
```
{"parent":"main","parentHead":"","trunk":"main"}
```
-->

---------

Co-authored-by: Thomas Chow <[email protected]>

Author: tchow-zlai

cloud_aws/src/main/scala/ai/chronon/integrations/aws/DynamoDBKVStoreImpl.scala

@@ -9,8 +9,8 @@ import ai.chronon.online.KVStore.ListRequest
 import ai.chronon.online.KVStore.ListResponse
 import ai.chronon.online.KVStore.ListValue
 import ai.chronon.online.KVStore.TimedValue
-import ai.chronon.online.Metrics
-import ai.chronon.online.Metrics.Context
+import ai.chronon.online.metrics.Metrics.Context
+import ai.chronon.online.metrics.Metrics
 import com.google.common.util.concurrent.RateLimiter
 import software.amazon.awssdk.core.SdkBytes
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient

cloud_gcp/src/main/scala/ai/chronon/integrations/cloud_gcp/BigTableKVStoreImpl.scala

@@ -13,7 +13,7 @@ import ai.chronon.online.KVStore
 import ai.chronon.online.KVStore.ListRequest
 import ai.chronon.online.KVStore.ListResponse
 import ai.chronon.online.KVStore.ListValue
-import ai.chronon.online.Metrics
+import ai.chronon.online.metrics.Metrics
 import com.google.cloud.RetryOption
 import com.google.cloud.bigquery.BigQuery
 import com.google.cloud.bigquery.BigQueryErrorMessages

online/BUILD.bazel

@@ -1,3 +1,21 @@
+scala_library(
+    name = "metrics_lib",
+    srcs = glob(["src/main/scala/ai/chronon/online/metrics/*.scala"]),
+    format = select({
+        "//tools/config:scala_2_13": False,  # Disable for 2.13
+        "//conditions:default": True,  # Enable for other versions
+    }),
+    visibility = ["//visibility:public"],
+    deps = [
+        "//api:lib",
+        "//api:thrift_java",
+        maven_artifact("com.datadoghq:java-dogstatsd-client"),
+        maven_artifact("org.slf4j:slf4j-api"),
+        maven_artifact("org.apache.logging.log4j:log4j-api"),
+        maven_artifact("org.apache.logging.log4j:log4j-core"),
+    ],
+)
+
 scala_library(
     name = "lib",
     srcs = glob(["src/main/**/*.scala"]) + glob(["src/main/**/*.java"]),
@@ -7,6 +25,7 @@ scala_library(
     }),
     visibility = ["//visibility:public"],
     deps = [
+        ":metrics_lib",
         "//aggregator:lib",
         "//api:lib",
         "//api:thrift_java",

online/src/main/java/ai/chronon/online/JavaFetcher.java

@@ -27,6 +27,7 @@
 import scala.concurrent.Future;
 import scala.concurrent.ExecutionContext;
 import scala.util.Try;
+import ai.chronon.online.metrics.Metrics;
 
 import java.util.ArrayList;
 import java.util.List;

online/src/main/scala/ai/chronon/online/Api.scala

@@ -60,7 +60,7 @@ object KVStore {
 // used for streaming writes, batch bulk uploads & fetching
 trait KVStore {
   @transient lazy val logger: Logger = LoggerFactory.getLogger(getClass)
-  implicit val executionContext: ExecutionContext = FlexibleExecutionContext.buildExecutionContext
+  implicit val executionContext: ExecutionContext = metrics.FlexibleExecutionContext.buildExecutionContext
   def create(dataset: String): Unit
 
   def create(dataset: String, props: Map[String, Any]): Unit = create(dataset)
@@ -205,7 +205,7 @@ trait StreamBuilder {
 }
 
 object ExternalSourceHandler {
-  private[ExternalSourceHandler] val executor = FlexibleExecutionContext.buildExecutionContext
+  private[ExternalSourceHandler] val executor = metrics.FlexibleExecutionContext.buildExecutionContext
 }
 
 // user facing class that needs to be implemented for external sources defined in a join

online/src/main/scala/ai/chronon/online/ExternalSourceRegistry.scala

@@ -18,13 +18,10 @@ package ai.chronon.online
 
 import ai.chronon.api.Constants
 import ai.chronon.online.fetcher.Fetcher.{Request, Response}
-
-import scala.collection.Seq
-import scala.collection.mutable
-import scala.concurrent.ExecutionContext
-import scala.concurrent.Future
-import scala.util.Failure
-import scala.util.Success
+import ai.chronon.online.metrics.Metrics.Context
+import scala.collection.{Seq, mutable}
+import scala.concurrent.{ExecutionContext, Future}
+import scala.util.{Failure, Success}
 
 // users can simply register external endpoints with a lambda that can return the future of a response given keys
 // keys and values need to match schema in ExternalSource - chronon will validate automatically
@@ -53,8 +50,7 @@ class ExternalSourceRegistry extends Serializable {
   // 1. keys match
   // 2. report missing & extra values
   // 3. schema integrity of returned values
-  def fetchRequests(requests: Seq[Request], context: Metrics.Context)(implicit
-      ec: ExecutionContext): Future[Seq[Response]] = {
+  def fetchRequests(requests: Seq[Request], context: Context)(implicit ec: ExecutionContext): Future[Seq[Response]] = {
     val startTime = System.currentTimeMillis()
     // we make issue one batch request per external source and flatten out it later
     val responsesByNameF: List[Future[Seq[Response]]] = requests

online/src/main/scala/ai/chronon/online/fetcher/FetchContext.scala

@@ -1,7 +1,8 @@
 package ai.chronon.online.fetcher
 import ai.chronon.api.Constants.MetadataDataset
 import ai.chronon.api.ScalaJavaConversions.JMapOps
-import ai.chronon.online.{FlagStore, FlagStoreConstants, FlexibleExecutionContext, KVStore}
+import ai.chronon.online.metrics.FlexibleExecutionContext
+import ai.chronon.online.{FlagStore, FlagStoreConstants, KVStore}
 
 import scala.concurrent.ExecutionContext
 

online/src/main/scala/ai/chronon/online/fetcher/Fetcher.scala

@@ -21,10 +21,10 @@ import ai.chronon.api
 import ai.chronon.api.Constants.UTF8
 import ai.chronon.api.Extensions.{ExternalPartOps, JoinOps, StringOps, ThrowableOps}
 import ai.chronon.api._
-import ai.chronon.online.Metrics.Environment
 import ai.chronon.online.OnlineDerivationUtil.applyDeriveFunc
+import ai.chronon.online._
 import ai.chronon.online.fetcher.Fetcher.{JoinSchemaResponse, Request, Response, ResponseWithContext}
-import ai.chronon.online.{serde, _}
+import ai.chronon.online.metrics.{Metrics, TTLCache}
 import com.google.gson.Gson
 import com.timgroup.statsd.Event
 import com.timgroup.statsd.Event.AlertType
@@ -41,10 +41,13 @@ import scala.concurrent.{ExecutionContext, Future}
 import scala.util.{Failure, Success, Try}
 
 object Fetcher {
+
+  import ai.chronon.online.metrics
+
   case class Request(name: String,
                      keys: Map[String, AnyRef],
                      atMillis: Option[Long] = None,
-                     context: Option[Metrics.Context] = None)
+                     context: Option[metrics.Metrics.Context] = None)
 
   case class PrefixedRequest(prefix: String, request: Request)
   case class Response(request: Request, values: Try[Map[String, AnyRef]])
@@ -59,17 +62,18 @@ object Fetcher {
                         prefix: Option[String],
                         keyMapping: Option[Map[String, AnyRef]])
 
-  def logResponseStats(response: Response, context: Metrics.Context): Unit = {
+  def logResponseStats(response: Response, context: metrics.Metrics.Context): Unit = {
+    import ai.chronon.online.metrics
     val responseMap = response.values.get
     var exceptions = 0
     var nulls = 0
     responseMap.foreach { case (_, v) =>
       if (v == null) nulls += 1
       else if (v.isInstanceOf[Throwable]) exceptions += 1
     }
-    context.distribution(Metrics.Name.FetchNulls, nulls)
-    context.distribution(Metrics.Name.FetchExceptions, exceptions)
-    context.distribution(Metrics.Name.FetchCount, responseMap.size)
+    context.distribution(metrics.Metrics.Name.FetchNulls, nulls)
+    context.distribution(metrics.Metrics.Name.FetchExceptions, exceptions)
+    context.distribution(metrics.Metrics.Name.FetchCount, responseMap.size)
   }
 
   /** Response for a join schema request
@@ -107,7 +111,7 @@ class Fetcher(val kvStore: KVStore,
   private def reportCallerNameFetcherVersion(): Unit = {
     val message =
       s"CallerName: ${Option(callerName).getOrElse("N/A")}, FetcherVersion: ${BuildInfo.version}"
-    val ctx = Metrics.Context(Environment.Fetcher)
+    val ctx = Metrics.Context(Metrics.Environment.Fetcher)
     val event = Event
       .builder()
       .withTitle("FetcherInitialization")
@@ -141,6 +145,7 @@ class Fetcher(val kvStore: KVStore,
     val combinedResponsesF =
       internalResponsesF.zip(externalResponsesF).map { case (internalResponses, externalResponses) =>
         internalResponses.zip(externalResponses).map { case (internalResponse, externalResponse) =>
+          import ai.chronon.online.metrics
           if (debug) {
             logger.info(internalResponse.values.get.keys.toSeq.mkString(","))
             logger.info(externalResponse.values.get.keys.toSeq.mkString(","))
@@ -161,7 +166,7 @@ class Fetcher(val kvStore: KVStore,
             Map("external_part_fetch_exception" -> externalResponse.values.failed.get.traceString))
           val derivationStartTs = System.currentTimeMillis()
           val joinName = internalResponse.request.name
-          val ctx = Metrics.Context(Environment.JoinFetching, join = joinName)
+          val ctx = Metrics.Context(Metrics.Environment.JoinFetching, join = joinName)
           val joinCodecTry = joinCodecCache(internalResponse.request.name)
           joinCodecTry match {
             case Success(joinCodec) =>
@@ -323,6 +328,7 @@ class Fetcher(val kvStore: KVStore,
 
   // Pulling external features in a batched fashion across services in-parallel
   private def fetchExternal(joinRequests: Seq[Request]): Future[Seq[Response]] = {
+    import ai.chronon.online.metrics
     val startTime = System.currentTimeMillis()
     val resultMap = new mutable.LinkedHashMap[Request, Try[mutable.HashMap[String, Any]]]
     var invalidCount = 0
@@ -374,7 +380,7 @@ class Fetcher(val kvStore: KVStore,
       .toMap
 
     val context =
-      Metrics.Context(environment = Environment.JoinFetching,
+      Metrics.Context(environment = Metrics.Environment.JoinFetching,
                       join = validRequests.iterator.map(_.name.sanitize).toSeq.distinct.mkString(","))
     context.distribution("response.external_pre_processing.latency", System.currentTimeMillis() - startTime)
     context.count("response.external_invalid_joins.count", invalidCount)
@@ -417,7 +423,7 @@ class Fetcher(val kvStore: KVStore,
       // step-4 convert the resultMap into Responses
       joinRequests.map { req =>
         Metrics
-          .Context(Environment.JoinFetching, join = req.name)
+          .Context(Metrics.Environment.JoinFetching, join = req.name)
           .distribution("external.latency.millis", System.currentTimeMillis() - startTime)
         Response(req, resultMap(req).map(_.mapValues(_.asInstanceOf[AnyRef]).toMap))
       }
@@ -426,7 +432,7 @@ class Fetcher(val kvStore: KVStore,
 
   def fetchJoinSchema(joinName: String): Try[JoinSchemaResponse] = {
     val startTime = System.currentTimeMillis()
-    val ctx = Metrics.Context(Environment.JoinSchemaFetching, join = joinName)
+    val ctx = Metrics.Context(Metrics.Environment.JoinSchemaFetching, join = joinName)
 
     val joinCodecTry = joinCodecCache(joinName)
 
@@ -470,7 +476,8 @@ class Fetcher(val kvStore: KVStore,
   private case class ExternalToJoinRequest(externalRequest: Either[Request, KeyMissingException],
                                            joinRequest: Request,
                                            part: ExternalPart) {
+
     lazy val context: Metrics.Context =
-      Metrics.Context(Environment.JoinFetching, join = joinRequest.name, groupBy = part.fullName)
+      Metrics.Context(Metrics.Environment.JoinFetching, join = joinRequest.name, groupBy = part.fullName)
   }
 }

online/src/main/scala/ai/chronon/online/fetcher/FetcherCache.scala

@@ -4,7 +4,8 @@ import ai.chronon.aggregator.windowing.FinalBatchIr
 import ai.chronon.api.GroupBy
 import ai.chronon.online.KVStore.{GetRequest, TimedValue}
 import ai.chronon.online.fetcher.FetcherCache._
-import ai.chronon.online.{GroupByServingInfoParsed, Metrics}
+import ai.chronon.online.GroupByServingInfoParsed
+import ai.chronon.online.metrics.Metrics
 import com.github.benmanes.caffeine.cache.{Cache => CaffeineCache}
 import org.slf4j.{Logger, LoggerFactory}
 

online/src/main/scala/ai/chronon/online/fetcher/GroupByFetcher.scala

@@ -6,7 +6,7 @@ import ai.chronon.api.Extensions._
 import ai.chronon.api._
 import ai.chronon.online.KVStore.{GetRequest, GetResponse, TimedValue}
 import ai.chronon.online.OnlineDerivationUtil.{applyDeriveFunc, buildRenameOnlyDerivationFunction}
-import ai.chronon.online._
+import ai.chronon.online.{metrics, _}
 import ai.chronon.online.fetcher.Fetcher.{ColumnSpec, PrefixedRequest, Request, Response}
 import ai.chronon.online.fetcher.FetcherCache.{BatchResponses, CachedBatchResponse}
 import org.slf4j.{Logger, LoggerFactory}
@@ -45,7 +45,8 @@ class GroupByFetcher(fetchContext: FetchContext, metadataStore: MetadataStore)
     .getGroupByServingInfo(request.name)
     .map { groupByServingInfo =>
       val context =
-        request.context.getOrElse(Metrics.Context(Metrics.Environment.GroupByFetching, groupByServingInfo.groupBy))
+        request.context.getOrElse(
+          metrics.Metrics.Context(metrics.Metrics.Environment.GroupByFetching, groupByServingInfo.groupBy))
       context.increment("group_by_request.count")
       var batchKeyBytes: Array[Byte] = null
       var streamingKeyBytes: Array[Byte] = null
@@ -338,4 +339,4 @@ case class LambdaKvRequest(groupByServingInfoParsed: GroupByServingInfoParsed,
                            batchRequest: GetRequest,
                            streamingRequestOpt: Option[GetRequest],
                            endTs: Option[Long],
-                           context: Metrics.Context)
+                           context: metrics.Metrics.Context)

online/src/main/scala/ai/chronon/online/fetcher/GroupByResponseHandler.scala

@@ -4,10 +4,11 @@ import ai.chronon.aggregator.windowing.{FinalBatchIr, SawtoothOnlineAggregator,
 import ai.chronon.api.Extensions.WindowOps
 import ai.chronon.api.ScalaJavaConversions.{IteratorOps, JMapOps}
 import ai.chronon.api.{DataModel, Row, Window}
-import ai.chronon.online.{AvroConversions, GroupByServingInfoParsed, Metrics}
+import ai.chronon.online.{AvroConversions, GroupByServingInfoParsed}
 import ai.chronon.online.KVStore.TimedValue
-import ai.chronon.online.Metrics.Name
+import ai.chronon.online.metrics.Metrics.Name
 import ai.chronon.online.fetcher.FetcherCache.{BatchResponses, CachedBatchResponse, KvStoreBatchResponse}
+import ai.chronon.online.metrics.Metrics
 import com.google.gson.Gson
 import org.slf4j.{Logger, LoggerFactory}
 

online/src/main/scala/ai/chronon/online/fetcher/JoinPartFetcher.scala

@@ -61,18 +61,21 @@ class JoinPartFetcher(fetchContext: FetchContext, metadataStore: MetadataStore)
     val joinDecomposed: Seq[(Request, Try[Seq[Either[PrefixedRequest, KeyMissingException]]])] =
       requests.map { request =>
         // use passed-in join or fetch one
+        import ai.chronon.online.metrics
         val joinTry: Try[JoinOps] = joinConf
           .map(conf => Success(JoinOps(conf)))
           .getOrElse(metadataStore.getJoinConf(request.name))
 
-        var joinContext: Option[Metrics.Context] = None
+        var joinContext: Option[metrics.Metrics.Context] = None
 
         val decomposedTry = joinTry.map { join =>
-          joinContext = Some(Metrics.Context(Metrics.Environment.JoinFetching, join.join))
+          import ai.chronon.online.metrics
+          joinContext = Some(metrics.Metrics.Context(metrics.Metrics.Environment.JoinFetching, join.join))
           joinContext.get.increment("join_request.count")
 
           join.joinPartOps.map { part =>
-            val joinContextInner = Metrics.Context(joinContext.get, part)
+            import ai.chronon.online.metrics
+            val joinContextInner = metrics.Metrics.Context(joinContext.get, part)
             val missingKeys = part.leftToRight.keys.filterNot(request.keys.contains)
 
             if (missingKeys.nonEmpty) {

online/src/main/scala/ai/chronon/online/fetcher/LRUCache.scala

@@ -1,6 +1,6 @@
 package ai.chronon.online.fetcher
 
-import ai.chronon.online.Metrics
+import ai.chronon.online.metrics.Metrics
 import com.github.benmanes.caffeine.cache.{Caffeine, Cache => CaffeineCache}
 import org.slf4j.{Logger, LoggerFactory}