fix(Metric.scala):close_instrumentation_client (#356)

YotpoLtd · Aug 10, 2020 · 7467fa6 · 7467fa6
1 parent e6daaed
commit 7467fa6
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 3 deletions.
diff --git a/src/main/scala/com/yotpo/metorikku/Metorikku.scala b/src/main/scala/com/yotpo/metorikku/Metorikku.scala
@@ -15,7 +15,17 @@ object Metorikku extends App {
     case Some(periodic) => {
       executePeriodicTask(periodic)
     }
-    case _ => runMetrics(session)
+    case _ => {
+      runMetrics(session)
+      try
+        {
+          session.instrumentationClient.close()
+        }
+      catch
+        {
+          case e: Throwable => log.error(s"Got exception while closing connection to instrumentationClient", e)
+        }
+    }
   }
 
   private def executePeriodicTask(periodic: Periodic) = {

diff --git a/src/main/scala/com/yotpo/metorikku/metric/Metric.scala b/src/main/scala/com/yotpo/metorikku/metric/Metric.scala
@@ -136,7 +136,6 @@ case class Metric(configuration: Configuration, metricDir: Option[File], metricN
           val dataFrameName = outputConfig.dataFrameName
           val dataFrame = repartition(outputConfig, job.sparkSession.table(dataFrameName))
 
-
           val outputOptions = Option(outputConfig.outputOptions).getOrElse(Map())
           outputOptions.get("protectFromEmptyOutput").asInstanceOf[Option[Boolean]] match {
             case Some(true) => {
@@ -158,7 +157,6 @@ case class Metric(configuration: Configuration, metricDir: Option[File], metricN
               job.config.cacheCountOnOutput)
           }
         })
-
         for ((dataFrameName, streamingConfig) <- streamingWriterList) writeStream(dataFrameName,
           streamingConfig.streamingWritingConfiguration, job.config.streaming, job.instrumentationClient)
 
@@ -167,3 +165,4 @@ case class Metric(configuration: Configuration, metricDir: Option[File], metricN
     }
   }
 }
+