<commit_before>package io.flow.event.v2

import java.util.concurrent.Executors
import javax.inject.{Inject, Singleton}

import io.flow.event.{Record, StreamNames}
import org.joda.time.DateTime
import play.api.libs.json.JsValue

import scala.concurrent.ExecutionContext
import scala.concurrent.duration.{FiniteDuration, MILLISECONDS}
import scala.reflect.runtime.universe._

@Singleton
class MockQueue @Inject()() extends Queue {

  private[this] val streams = scala.collection.mutable.Map[String, MockStream]()

  override def producer[T: TypeTag](
    numberShards: Int = 1,
    partitionKeyFieldName: String = "event_id"
  ): Producer = {
    MockProducer(stream[T])
  }

  override def consume[T: TypeTag](
    f: Record => Unit,
    pollTime: FiniteDuration = FiniteDuration(20, MILLISECONDS)
  )(
    implicit ec: ExecutionContext
  ) {
    val s = stream[T]
    Executors.newSingleThreadExecutor().execute(new Runnable() {
      override def run(): Unit = {
        while (true) {
          s.consume().foreach(f)
          Thread.sleep(pollTime.toMillis)
        }
      }
    })
  }

  override def shutdown(implicit ec: ExecutionContext): Unit = {
    streams.clear()
  }

  override def shutdownConsumers(implicit ec: ExecutionContext): Unit = {
    // no-op
  }


  def stream[T: TypeTag]: MockStream = {
    streams.get(streamName[T]).getOrElse {
      val stream = MockStream()
      streams.put(streamName[T], stream)
      stream
    }
  }

  private[this] def streamName[T: TypeTag] = {
    StreamNames.fromType[T] match {
      case Left(errors) => sys.error(errors.mkString(", "))
      case Right(name) => name
    }
  }

}

case class MockStream() {

  private[this] val pendingRecords = scala.collection.mutable.ListBuffer[Record]()
  private[this] val consumedRecords = scala.collection.mutable.ListBuffer[Record]()

  def publish(record: Record): Unit = {
    pendingRecords.append(record)
  }

  /**
    * Consumes the next event in the stream, if any
    */
  def consume(): Option[Record] = {
    synchronized {
      pendingRecords.headOption.map { r =>
        pendingRecords.remove(0)
        consumedRecords.append(r)
        r
      }
    }
  }

  /**
    * Finds the event w/ the specified id. Returns none if
    * we have not yet received this event.
    */
  def findByEventId(eventId: String): Option[Record] = {
    all.find(_.eventId == eventId)
  }

  /**
    * Returns all records seen - pending and consumed
    */
  def all: Seq[Record] = {
    pendingRecords ++ consumedRecords
  }

  def pending: Seq[Record] = pendingRecords
  def consumed: Seq[Record] = consumedRecords

}

case class MockProducer(stream: MockStream) extends Producer {

  def publish(event: JsValue)(implicit ec: ExecutionContext): Unit = {
    stream.publish(
      Record.fromJsValue(
        arrivalTimestamp = DateTime.now,
        js = event
      )
    )
  }

  def shutdown(implicit ec: ExecutionContext): Unit = {}

}
<commit_msg>Synchronize on publish and all<commit_after>package io.flow.event.v2

import java.util.concurrent.Executors
import javax.inject.{Inject, Singleton}

import io.flow.event.{Record, StreamNames}
import org.joda.time.DateTime
import play.api.libs.json.JsValue

import scala.concurrent.ExecutionContext
import scala.concurrent.duration.{FiniteDuration, MILLISECONDS}
import scala.reflect.runtime.universe._

@Singleton
class MockQueue @Inject()() extends Queue {

  private[this] val streams = scala.collection.mutable.Map[String, MockStream]()

  override def producer[T: TypeTag](
    numberShards: Int = 1,
    partitionKeyFieldName: String = "event_id"
  ): Producer = {
    MockProducer(stream[T])
  }

  override def consume[T: TypeTag](
    f: Record => Unit,
    pollTime: FiniteDuration = FiniteDuration(20, MILLISECONDS)
  )(
    implicit ec: ExecutionContext
  ) {
    val s = stream[T]
    Executors.newSingleThreadExecutor().execute(new Runnable() {
      override def run(): Unit = {
        while (true) {
          s.consume().foreach(f)
          Thread.sleep(pollTime.toMillis)
        }
      }
    })
  }

  override def shutdown(implicit ec: ExecutionContext): Unit = {
    streams.clear()
  }

  override def shutdownConsumers(implicit ec: ExecutionContext): Unit = {
    // no-op
  }


  def stream[T: TypeTag]: MockStream = {
    streams.get(streamName[T]).getOrElse {
      val stream = MockStream()
      streams.put(streamName[T], stream)
      stream
    }
  }

  private[this] def streamName[T: TypeTag] = {
    StreamNames.fromType[T] match {
      case Left(errors) => sys.error(errors.mkString(", "))
      case Right(name) => name
    }
  }

}

case class MockStream() {

  private[this] val pendingRecords = scala.collection.mutable.ListBuffer[Record]()
  private[this] val consumedRecords = scala.collection.mutable.ListBuffer[Record]()

  def publish(record: Record): Unit = {
    synchronized {
      pendingRecords.append(record)
    }
  }

  /**
    * Consumes the next event in the stream, if any
    */
  def consume(): Option[Record] = {
    pendingRecords.headOption.flatMap { r =>
      // only synchronize if we got a record, which is rare in our usage
      synchronized {
        pendingRecords.headOption.map { r =>
          pendingRecords.remove(0)
          consumedRecords.append(r)
          r
        }
      }
    }
  }

  /**
    * Finds the event w/ the specified id. Returns none if
    * we have not yet received this event.
    */
  def findByEventId(eventId: String): Option[Record] = {
    all.find(_.eventId == eventId)
  }

  /**
    * Returns all records seen - pending and consumed
    */
  def all: Seq[Record] = {
    synchronized {
      pendingRecords ++ consumedRecords
    }
  }

  def pending: Seq[Record] = pendingRecords
  def consumed: Seq[Record] = consumedRecords

}

case class MockProducer(stream: MockStream) extends Producer {

  def publish(event: JsValue)(implicit ec: ExecutionContext): Unit = {
    stream.publish(
      Record.fromJsValue(
        arrivalTimestamp = DateTime.now,
        js = event
      )
    )
  }

  def shutdown(implicit ec: ExecutionContext): Unit = {}

}

{'fraction_non_alphanumeric': 0.11565516113201374, 'fraction_numerical': 0.0016358580075249468, 'mean_word_length': 3.0065530799475755, 'pattern_counts': {'":': 0, '<': 3, '<?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': '1063910', 'n_tokens_mistral': 1933, 'n_tokens_neox': 1918, 'n_words': 437}