eval: refactor flow to create context earlier (#1679)

Create context earlier to allow for performing validation
or logging related to the data sources in the firts parts
of the stream.

atlas-eval/src/main/scala/com/netflix/atlas/eval/stream/DataSourceLogger.scala

@@ -0,0 +1,49 @@
+/*
+ * Copyright 2014-2024 Netflix, Inc.
+ *
+ * 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
+ *
+ *     http://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 com.netflix.atlas.eval.stream
+
+import com.netflix.atlas.json.JsonSupport
+import com.netflix.atlas.pekko.StreamOps
+
+private[stream] trait DataSourceLogger {
+
+  def apply(ds: Evaluator.DataSource, msg: JsonSupport): Unit
+
+  def close(): Unit
+}
+
+private[stream] object DataSourceLogger {
+
+  case object Noop extends DataSourceLogger {
+
+    override def apply(ds: Evaluator.DataSource, msg: JsonSupport): Unit = {}
+
+    override def close(): Unit = {}
+  }
+
+  case class Queue(queue: StreamOps.SourceQueue[Evaluator.MessageEnvelope])
+      extends DataSourceLogger {
+
+    override def apply(ds: Evaluator.DataSource, msg: JsonSupport): Unit = {
+      val env = new Evaluator.MessageEnvelope(ds.id, msg)
+      queue.offer(env)
+    }
+
+    override def close(): Unit = {
+      queue.complete()
+    }
+  }
+}

atlas-eval/src/main/scala/com/netflix/atlas/eval/stream/EvaluatorImpl.scala

@@ -104,7 +104,9 @@ private[stream] abstract class EvaluatorImpl(
   // Execution context to use for parsing payloads coming back from lwcapi service
   private val parsingEC = ThreadPools.fixedSize(registry, "AtlasEvalParsing", parsingNumThreads)
 
-  private def newStreamContext(dsLogger: DataSourceLogger = (_, _) => ()): StreamContext = {
+  private def newStreamContext(
+    dsLogger: DataSourceLogger = DataSourceLogger.Noop
+  ): StreamContext = {
     new StreamContext(
       config,
       materializer,
@@ -113,6 +115,19 @@ private[stream] abstract class EvaluatorImpl(
     )
   }
 
+  private def createStreamContextSource: (Source[MessageEnvelope, NotUsed], StreamContext) = {
+    // Flow used for logging diagnostic messages
+    val (queue, logSrc) = StreamOps
+      .blockingQueue[MessageEnvelope](registry, "DataSourceLogger", 10)
+      .toMat(BroadcastHub.sink(1))(Keep.both)
+      .run()
+    val dsLogger = DataSourceLogger.Queue(queue)
+
+    // One manager per processor to ensure distinct stream ids on LWC
+    val context = newStreamContext(dsLogger)
+    logSrc -> context
+  }
+
   protected def validateImpl(ds: DataSource): Unit = {
     validationStreamContext.validateDataSource(ds).get
   }
@@ -176,8 +191,11 @@ private[stream] abstract class EvaluatorImpl(
         Source.single(ds)
     }
 
+    val (logSrc, context) = createStreamContextSource
     source
-      .via(createProcessorFlow)
+      .via(createProcessorFlow(context))
+      .via(new OnUpstreamFinish[MessageEnvelope](context.dsLogger.close()))
+      .merge(logSrc, eagerComplete = false)
       .map(_.message)
       .toMat(Sink.asPublisher(true))(Keep.right)
       .run()
@@ -192,15 +210,18 @@ private[stream] abstract class EvaluatorImpl(
     * to maximize throughput under heavy load by enabling operator fusion optimization.
     */
   def createStreamsFlow: Flow[DataSources, MessageEnvelope, NotUsed] = {
+    val (logSrc, context) = createStreamContextSource
     Flow[DataSources]
       .map(dss => groupByHost(dss))
       // Emit empty DataSource if no more DataSource for a host, so that the sub-stream get the info
       .via(new FillRemovedKeysWith[String, DataSources](_ => DataSources.empty()))
       .flatMapMerge(Int.MaxValue, dssMap => Source(dssMap.toList))
       .groupBy(Int.MaxValue, _._1, true) // groupBy host
       .map(_._2) // keep only DataSources
-      .via(createProcessorFlow)
+      .via(createProcessorFlow(context))
       .mergeSubstreams
+      .via(new OnUpstreamFinish[MessageEnvelope](context.dsLogger.close()))
+      .merge(logSrc, eagerComplete = false)
   }
 
   protected def groupByHost(dataSources: DataSources): scala.collection.Map[String, DataSources] = {
@@ -216,20 +237,9 @@ private[stream] abstract class EvaluatorImpl(
       Uri(dataSource.uri).authority.host.address
   }
 
-  private[stream] def createProcessorFlow: Flow[DataSources, MessageEnvelope, NotUsed] = {
-
-    // Flow used for logging diagnostic messages
-    val (queue, logSrc) = StreamOps
-      .blockingQueue[MessageEnvelope](registry, "DataSourceLogger", 10)
-      .toMat(BroadcastHub.sink(1))(Keep.both)
-      .run()
-    val dsLogger: DataSourceLogger = { (ds, msg) =>
-      val env = new MessageEnvelope(ds.id, msg)
-      queue.offer(env)
-    }
-
-    // One manager per processor to ensure distinct stream ids on LWC
-    val context = newStreamContext(dsLogger)
+  private[stream] def createProcessorFlow(
+    context: StreamContext
+  ): Flow[DataSources, MessageEnvelope, NotUsed] = {
 
     val g = GraphDSL.create() { implicit builder =>
       import GraphDSL.Implicits.*
@@ -281,8 +291,6 @@ private[stream] abstract class EvaluatorImpl(
       .via(context.monitorFlow("12_OutputSources"))
       .flatMapConcat(s => s)
       .via(context.monitorFlow("13_OutputMessages"))
-      .via(new OnUpstreamFinish[MessageEnvelope](queue.complete()))
-      .merge(logSrc, eagerComplete = false)
   }
 
   protected def createDatapointProcessorImpl(
@@ -293,17 +301,9 @@ private[stream] abstract class EvaluatorImpl(
     val stepSize = sources.stepSize()
 
     // Flow used for logging diagnostic messages
-    val (queue, logSrc) = StreamOps
-      .blockingQueue[MessageEnvelope](registry, "DataSourceLogger", 10)
-      .toMat(BroadcastHub.sink(1))(Keep.both)
-      .run()
-    val dsLogger: DataSourceLogger = { (ds, msg) =>
-      val env = new MessageEnvelope(ds.id, msg)
-      queue.offer(env)
-    }
+    val (logSrc, context) = createStreamContextSource
 
     // Initialize context with fixed data sources
-    val context = newStreamContext(dsLogger)
     context.validate(sources)
     context.setDataSources(sources)
     val interpreter = context.interpreter
@@ -320,7 +320,7 @@ private[stream] abstract class EvaluatorImpl(
       .merge(Source.single(sources), eagerComplete = false)
       .via(new FinalExprEval(interpreter))
       .flatMapConcat(s => s)
-      .via(new OnUpstreamFinish[MessageEnvelope](queue.complete()))
+      .via(new OnUpstreamFinish[MessageEnvelope](context.dsLogger.close()))
       .merge(logSrc, eagerComplete = false)
       .toProcessor
       .run()

atlas-eval/src/main/scala/com/netflix/atlas/eval/stream/StreamContext.scala

@@ -53,7 +53,7 @@ private[stream] class StreamContext(
   rootConfig: Config,
   val materializer: Materializer,
   val registry: Registry = new NoopRegistry,
-  val dsLogger: DataSourceLogger = (_, _) => ()
+  val dsLogger: DataSourceLogger = DataSourceLogger.Noop
 ) {
 
   import StreamContext.*

atlas-eval/src/main/scala/com/netflix/atlas/eval/stream/package.scala

@@ -31,6 +31,4 @@ package object stream {
   type SimpleClient = Flow[HttpRequest, Try[HttpResponse], NotUsed]
 
   type SourcesAndGroups = (DataSources, EddaSource.Groups)
-
-  type DataSourceLogger = (DataSource, JsonSupport) => Unit
 }

atlas-eval/src/test/scala/com/netflix/atlas/eval/stream/LwcToAggrDatapointSuite.scala

@@ -71,7 +71,7 @@ class LwcToAggrDatapointSuite extends FunSuite {
   private val context = new StreamContext(
     ConfigFactory.load(),
     materializer,
-    dsLogger = (_, _) => ()
+    dsLogger = DataSourceLogger.Noop
   )
 
   context.setDataSources(