<commit_before>package drt.chroma

import akka.stream._
import akka.stream.stage._
import drt.shared.CrunchApi.MillisSinceEpoch
import drt.shared.{Arrival, UniqueArrival}
import org.slf4j.{Logger, LoggerFactory}
import server.feeds.{ArrivalsFeedFailure, ArrivalsFeedResponse, ArrivalsFeedSuccess}
import services.metrics.{Metrics, StageTimer}

import scala.collection.mutable


final class ArrivalsDiffingStage(initialKnownArrivals: mutable.SortedMap[UniqueArrival, Arrival], forecastMaxMillis: () => MillisSinceEpoch) extends GraphStage[FlowShape[ArrivalsFeedResponse, ArrivalsFeedResponse]] {
  val in: Inlet[ArrivalsFeedResponse] = Inlet[ArrivalsFeedResponse]("DiffingStage.in")
  val out: Outlet[ArrivalsFeedResponse] = Outlet[ArrivalsFeedResponse]("DiffingStage.out")
  val stageName = "arrivals-diffing"

  val log: Logger = LoggerFactory.getLogger(getClass)

  override val shape: FlowShape[ArrivalsFeedResponse, ArrivalsFeedResponse] = FlowShape.of(in, out)

  override def createLogic(inheritedAttributes: Attributes): GraphStageLogic = new GraphStageLogic(shape) {
    val knownArrivals: mutable.SortedMap[UniqueArrival, Arrival] = mutable.SortedMap[UniqueArrival, Arrival]()
    var maybeResponseToPush: Option[ArrivalsFeedResponse] = None

    knownArrivals ++= initialKnownArrivals

    override def preStart(): Unit = {
      log.info(s"Started with ${knownArrivals.size} known arrivals")
      super.preStart()
    }

    setHandlers(in, out, new InHandler with OutHandler {
      override def onPush(): Unit = {
        val timer = StageTimer(stageName, in)
        if (maybeResponseToPush.isEmpty) {
          log.info(s"Incoming ArrivalsFeedResponse")
          maybeResponseToPush = processFeedResponse(grab(in))
        } else log.info(s"Not ready to grab until we've pushed")

        if (isAvailable(out)) pushAndClear()
        timer.stopAndReport()
      }

      override def onPull(): Unit = {
        val timer = StageTimer(stageName, in)
        pushAndClear()

        if (!hasBeenPulled(in)) pull(in)
        timer.stopAndReport()
      }
    })

    def pushAndClear(): Unit = {
      maybeResponseToPush.collect {
        case afs: ArrivalsFeedSuccess if afs.arrivals.flights.nonEmpty =>
          Metrics.counter(s"$stageName.arrival-updates", afs.arrivals.flights.length)
          push(out, afs)
      }
      maybeResponseToPush = None
    }

    def processFeedResponse(arrivalsFeedResponse: ArrivalsFeedResponse): Option[ArrivalsFeedResponse] = arrivalsFeedResponse match {
      case afs@ArrivalsFeedSuccess(latestArrivals, _) =>
        val maxScheduledMillis = forecastMaxMillis()
        val incomingArrivals: Seq[(UniqueArrival, Arrival)] = latestArrivals.flights.filter(_.Scheduled <= maxScheduledMillis).map(a => (UniqueArrival(a), a))
        val newUpdates: Seq[(UniqueArrival, Arrival)] = filterArrivalsWithUpdates(knownArrivals, incomingArrivals)
        log.info(s"Got ${newUpdates.size} new arrival updates")
        knownArrivals.clear<commit_msg>DRT-5881 Fix diffing stage bug<commit_after>package drt.chroma

import akka.stream._
import akka.stream.stage._
import drt.shared.CrunchApi.MillisSinceEpoch
import drt.shared.{Arrival, UniqueArrival}
import org.slf4j.{Logger, LoggerFactory}
import server.feeds.{ArrivalsFeedFailure, ArrivalsFeedResponse, ArrivalsFeedSuccess}
import services.metrics.{Metrics, StageTimer}

import scala.collection.mutable


final class ArrivalsDiffingStage(initialKnownArrivals: mutable.SortedMap[UniqueArrival, Arrival], forecastMaxMillis: () => MillisSinceEpoch) extends GraphStage[FlowShape[ArrivalsFeedResponse, ArrivalsFeedResponse]] {
  val in: Inlet[ArrivalsFeedResponse] = Inlet[ArrivalsFeedResponse]("DiffingStage.in")
  val out: Outlet[ArrivalsFeedResponse] = Outlet[ArrivalsFeedResponse]("DiffingStage.out")
  val stageName = "arrivals-diffing"

  val log: Logger = LoggerFactory.getLogger(getClass)

  override val shape: FlowShape[ArrivalsFeedResponse, ArrivalsFeedResponse] = FlowShape.of(in, out)

  override def createLogic(inheritedAttributes: Attributes): GraphStageLogic = new GraphStageLogic(shape) {
    val knownArrivals: mutable.SortedMap[UniqueArrival, Arrival] = mutable.SortedMap[UniqueArrival, Arrival]()
    var maybeResponseToPush: Option[ArrivalsFeedResponse] = None

    knownArrivals ++= initialKnownArrivals

    override def preStart(): Unit = {
      log.info(s"Started with ${knownArrivals.size} known arrivals")
      super.preStart()
    }

    setHandlers(in, out, new InHandler with OutHandler {
      override def onPush(): Unit = {
        val timer = StageTimer(stageName, in)
        maybeResponseToPush = processFeedResponse(grab(in))

        if (isAvailable(out)) pushAndClear()
        timer.stopAndReport()
      }

      override def onPull(): Unit = {
        val timer = StageTimer(stageName, in)
        pushAndClear()

        if (!hasBeenPulled(in)) pull(in)
        timer.stopAndReport()
      }
    })

    def pushAndClear(): Unit = {
      maybeResponseToPush.collect {
        case afs: ArrivalsFeedSuccess if !afs.isEmpty =>
          Metrics.counter(s"$stageName.arrival-updates", afs.arrivals.flights.length)
          push(out, afs)
          maybeResponseToPush = None
        case empty =>
          push(out, empty)
      }
    }

    def processFeedResponse(arrivalsFeedResponse: ArrivalsFeedResponse): Option[ArrivalsFeedResponse] = arrivalsFeedResponse match {
      case afs@ArrivalsFeedSuccess(latestArrivals, _) =>
        val maxScheduledMillis = forecastMaxMillis()
        val incomingArrivals: Seq[(UniqueArrival, Arrival)] = latestArrivals.flights.filter(_.Scheduled <= maxScheduledMillis).map(a => (UniqueArrival(a), a))
        val newUpdates: Seq[(UniqueArrival, Arrival)] = filterArrivalsWithUpdates(knownArrivals, incomingArrivals)
        log.info(s"Got ${newUpdates.size} new arrival updates")
        knownArrivals.clear
{'fraction_non_alphanumeric': 0.09969429347826086, 'fraction_numerical': 0.0010190217391304348, 'mean_word_length': 4.3439201451905625, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 12, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5525110', 'n_tokens_mistral': 1864, 'n_tokens_neox': 1806, 'n_words': 345}