fix: refresh ttlcache upon failures (#964)

Author: hzding621

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

@@ -58,27 +58,29 @@ trait KVStore {
 
   // helper method to blocking read a string - used for fetching metadata & not in hotpath.
   def getString(key: String, dataset: String, timeoutMillis: Long): Try[String] = {
-    val response = getResponse(key, dataset, timeoutMillis)
-    if (response.values.isFailure) {
-      Failure(new RuntimeException(s"Request for key ${key} in dataset ${dataset} failed", response.values.failed.get))
-    } else {
-      Success(new String(response.latest.get.bytes, Constants.UTF8))
-    }
+    val bytesTry = getResponse(key, dataset, timeoutMillis)
+    bytesTry.map(bytes => new String(bytes, Constants.UTF8))
   }
 
   def getStringArray(key: String, dataset: String, timeoutMillis: Long): Try[Seq[String]] = {
-    val response = getResponse(key, dataset, timeoutMillis)
-    if (response.values.isFailure) {
-      Failure(new RuntimeException(s"Request for key ${key} in dataset ${dataset} failed", response.values.failed.get))
-    } else {
-      Success(StringArrayConverter.bytesToStrings(response.latest.get.bytes))
-    }
+    val bytesTry = getResponse(key, dataset, timeoutMillis)
+    bytesTry.map(bytes => StringArrayConverter.bytesToStrings(bytes))
   }
 
-  private def getResponse(key: String, dataset: String, timeoutMillis: Long): GetResponse = {
+  private def getResponse(key: String, dataset: String, timeoutMillis: Long): Try[Array[Byte]] = {
     val fetchRequest = KVStore.GetRequest(key.getBytes(Constants.UTF8), dataset)
     val responseFutureOpt = get(fetchRequest)
-    Await.result(responseFutureOpt, Duration(timeoutMillis, MILLISECONDS))
+    def buildException(e: Throwable) = new RuntimeException(s"Request for key ${key} in dataset ${dataset} failed", e)
+    Try(Await.result(responseFutureOpt, Duration(timeoutMillis, MILLISECONDS))) match {
+      case Failure(e) =>
+        Failure(buildException(e))
+      case Success(resp) =>
+        if (resp.values.isFailure) {
+          Failure(buildException(resp.values.failed.get))
+        } else {
+          Success(resp.latest.get.bytes)
+        }
+    }
   }
   def get(request: GetRequest): Future[GetResponse] = {
     multiGet(Seq(request))

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

@@ -203,11 +203,14 @@ class Fetcher(val kvStore: KVStore,
   lazy val getJoinCodecs = new TTLCache[String, Try[(JoinCodec, Boolean)]](
     { joinName: String =>
       val startTimeMs = System.currentTimeMillis()
-      val result: Try[(JoinCodec, Boolean)] = getJoinConf(joinName)
-        .map(_.join)
-        .map(join => buildJoinCodec(join, refreshOnFail = true))
-        .recoverWith {
+      val result: Try[(JoinCodec, Boolean)] =
+        try {
+          getJoinConf(joinName)
+            .map(_.join)
+            .map(join => buildJoinCodec(join, refreshOnFail = true))
+        } catch {
           case th: Throwable =>
+            getJoinConf.refresh(joinName)
             Failure(
               new RuntimeException(
                 s"Couldn't fetch joinName = ${joinName} or build join codec due to ${th.traceString}",
@@ -317,6 +320,7 @@ class Fetcher(val kvStore: KVStore,
                 response
               case Failure(exception) =>
                 // more validation logic will be covered in compile.py to avoid this case
+                getJoinCodecs.refresh(joinName)
                 ctx.incrementException(exception)
                 ResponseWithContext(internalResponse.request,
                                     Map("join_codec_fetch_exception" -> exception.traceString),
@@ -490,6 +494,7 @@ class Fetcher(val kvStore: KVStore,
       val joinName = request.name
       val joinConfTry: Try[JoinOps] = getJoinConf(request.name)
       if (joinConfTry.isFailure) {
+        getJoinConf.refresh(request.name)
         resultMap.update(
           request,
           Failure(
@@ -509,8 +514,11 @@ class Fetcher(val kvStore: KVStore,
     // step-2 dedup external requests across joins
     val externalToJoinRequests: Seq[ExternalToJoinRequest] = validRequests
       .flatMap { joinRequest =>
-        val parts =
-          getJoinConf(joinRequest.name).get.join.onlineExternalParts // cheap since it is cached, valid since step-1
+        val joinConf = getJoinConf(joinRequest.name)
+        if (joinConf.isFailure) {
+          getJoinConf.refresh(joinRequest.name)
+        }
+        val parts = joinConf.get.join.onlineExternalParts // cheap since it is cached, valid since step-1
         parts.iterator().asScala.map { part =>
           val externalRequest = Try(part.applyMapping(joinRequest.keys)) match {
             case Success(mappedKeys)                     => Left(Request(part.source.metadata.name, mappedKeys))

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

@@ -351,7 +351,15 @@ class FetcherBase(kvStore: KVStore,
   def fetchGroupBys(requests: scala.collection.Seq[Request]): Future[scala.collection.Seq[Response]] = {
     // split a groupBy level request into its kvStore level requests
     val groupByRequestToKvRequest: Seq[(Request, Try[GroupByRequestMeta])] = requests.iterator.map { request =>
-      val groupByRequestMetaTry: Try[GroupByRequestMeta] = getGroupByServingInfo(request.name)
+      val groupByServingInfoTry = getGroupByServingInfo(request.name)
+        .recover {
+          case ex: Throwable =>
+            getGroupByServingInfo.refresh(request.name)
+            logger.error(s"Couldn't fetch GroupByServingInfo for ${request.name}", ex)
+            request.context.foreach(_.incrementException(ex))
+            throw ex
+        }
+      val groupByRequestMetaTry: Try[GroupByRequestMeta] = groupByServingInfoTry
         .map { groupByServingInfo =>
           val context =
             request.context.getOrElse(Metrics.Context(Metrics.Environment.GroupByFetching, groupByServingInfo.groupBy))
@@ -583,7 +591,11 @@ class FetcherBase(kvStore: KVStore,
       requests.map { request =>
         // get join conf from metadata store if not passed in
         val joinTry: Try[JoinOps] = if (joinConf.isEmpty) {
-          getJoinConf(request.name)
+          val joinConfTry = getJoinConf(request.name)
+          if (joinConfTry.isFailure) {
+            getJoinConf.refresh(request.name)
+          }
+          joinConfTry
         } else {
           logger.debug(s"Using passed in join configuration: ${joinConf.get.metaData.getName}")
           Success(JoinOps(joinConf.get))
@@ -671,7 +683,7 @@ class FetcherBase(kvStore: KVStore,
                         if (debug || Math.random() < 0.001) {
                           logger.error(s"Failed to fetch $groupByRequest", ex)
                         }
-                        Map(groupByRequest.name + "_exception" -> ex.traceString)
+                        Map(prefix + "_exception" -> ex.traceString)
                     }
                     .get
                 }

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

@@ -22,6 +22,7 @@ import scala.concurrent.{ExecutionContext, ExecutionContextExecutor}
 object FlexibleExecutionContext {
   def buildExecutor: ThreadPoolExecutor = buildExecutor()
   def buildExecutionContext: ExecutionContextExecutor = buildExecutionContext()
+
   // users can also provide a custom execution context override in [MetadataStore]
   def buildExecutor(corePoolSize: Int = 20,
                     maxPoolSize: Int = 1000,

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

@@ -134,6 +134,7 @@ class MetadataStore(kvStore: KVStore,
         if (result.isSuccess) Metrics.Context(Metrics.Environment.MetaDataFetching, result.get.join)
         else Metrics.Context(Metrics.Environment.MetaDataFetching, join = name)
       // Throw exception after metrics. No join metadata is bound to be a critical failure.
+      // This will ensure that a Failure is never cached in the getJoinConf TTLCache
       if (result.isFailure) {
         context.withSuffix("join").incrementException(result.failed.get)
         throw result.failed.get
@@ -147,6 +148,7 @@ class MetadataStore(kvStore: KVStore,
   def validateJoinExist(team: String, name: String): Boolean = {
     val activeJoinList: Try[Seq[String]] = getJoinListByTeam(team)
     if (activeJoinList.isFailure) {
+      getJoinListByTeam.refresh(team)
       logger.error(s"Failed to fetch active join list for team $team")
       false
     } else {
@@ -164,6 +166,7 @@ class MetadataStore(kvStore: KVStore,
   def validateGroupByExist(team: String, name: String): Boolean = {
     val activeGroupByList: Try[Seq[String]] = getGroupByListByTeam(team)
     if (activeGroupByList.isFailure) {
+      getGroupByListByTeam.refresh(team)
       logger.error(s"Failed to fetch active group_by list for team $team")
       false
     } else {