<commit_before>package controllers

import akka.actor._
import akka.pattern.ask
import akka.util.Timeout
import nl.surfnet.safnari._
import scala.concurrent._
import scala.concurrent.duration.{ Duration, DurationInt }
import scala.concurrent.stm._
import play.Logger

class ConnectionManager(connectionFactory: (ConnectionId, Reserve) => (ActorRef, ConnectionEntity)) {
  implicit val timeout = Timeout(2.seconds)

  private val connections = TMap.empty[ConnectionId, ActorRef]
  private val messageStore = new MessageStore[Message]()
  @volatile private var replaying = false

  def add(connectionId: ConnectionId, connection: ActorRef) = connections.single(connectionId) = connection

  def get(connectionId: ConnectionId): Option[ActorRef] = connections.single.get(connectionId)

  def find(connectionIds: Seq[String]): Seq[ActorRef] = connections.single.filterKeys(connectionIds.contains).values.toSeq

  def findBySegment(connectionId: ConnectionId)(implicit executionContext: ExecutionContext): Future[Option[ActorRef]] = {
    val knowsAboutSegment = connections.single.values.map(c => c ? SegmentKnown(connectionId) map (c -> _.asInstanceOf[Boolean]))
    Future.find(knowsAboutSegment)(_._2) map (_.map(_._1))
  }

  def all: Seq[ActorRef] = connections.single.values.toSeq

  def restore(implicit actorSystem: ActorSystem, executionContext: ExecutionContext) {
    replaying = true
    Await.ready(Future.sequence(for {
      (connectionId, messages @ (FromRequester(initialReserve: Reserve) +: _)) <- messageStore.loadEverything()
      connection = createConnection(connectionId, initialReserve)
      message <- messages
    } yield {
      Logger.info(s"Replaying message: ${message.getClass.getName}")
      connection ? message
    }), Duration.Inf)

    replaying = false
    Logger.info("Replay completed")
  }

  def findOrCreateConnection(request: NsiProviderOperation)(implicit actorSystem: ActorSystem): (ConnectionId, Option[ActorRef]) = (request, request.optionalConnectionId) match {
    case (_, Some(connectionId)) =>
      (connectionId, get(connectionId))
    case (initialReserve: Reserve, None) =>
      val connectionId = newConnectionId
      val connection = createConnection(connectionId, initialReserve)
      (connectionId, Some(connection))
    case _ =>
      sys.error("illegal initial message")
  }

  private def createConnection(connectionId: ConnectionId, initialReserve: Reserve)(implicit actorSystem: ActorSystem): ActorRef = {
    val (outbound, connection) = connectionFactory(connectionId, initialReserve)

    val storingActor = actorSystem.actorOf(Props(new Actor {
      override def receive = {
        case 'query                     => sender ! connection.query
        case 'querySegments             => sender ! connection.segments

        case SegmentKnown(connectionId) => sender ! connection.rsm.segmentKnown(connectionId)

        case message: InboundMessage =>
          if (!replaying) {
            val store = message match {
              case message: FromRequester => Some(message)
              case message: FromProvider  => Some(message)
              case message: FromPce       => Some(message)
              case message =>
                message.pp("Not persisted")
                None
            }
            store foreach { messageStore.store(connectionId, _) }
          }

          val output = connection.process(message)

          output match {
            case None =>
              sender ! messageNotApplicable(message)
            case Some(messages) =>
              if (!replaying) messages.foreach(outbound ! _)
              message match {<commit_msg>Simplified case, all inbound messages are stored<commit_after>package controllers

import akka.actor._
import akka.pattern.ask
import akka.util.Timeout
import nl.surfnet.safnari._
import scala.concurrent._
import scala.concurrent.duration.{ Duration, DurationInt }
import scala.concurrent.stm._
import play.Logger

class ConnectionManager(connectionFactory: (ConnectionId, Reserve) => (ActorRef, ConnectionEntity)) {
  implicit val timeout = Timeout(2.seconds)

  private val connections = TMap.empty[ConnectionId, ActorRef]
  private val messageStore = new MessageStore[Message]()
  @volatile private var replaying = false

  def add(connectionId: ConnectionId, connection: ActorRef) = connections.single(connectionId) = connection

  def get(connectionId: ConnectionId): Option[ActorRef] = connections.single.get(connectionId)

  def find(connectionIds: Seq[String]): Seq[ActorRef] = connections.single.filterKeys(connectionIds.contains).values.toSeq

  def findBySegment(connectionId: ConnectionId)(implicit executionContext: ExecutionContext): Future[Option[ActorRef]] = {
    val knowsAboutSegment = connections.single.values.map(c => c ? SegmentKnown(connectionId) map (c -> _.asInstanceOf[Boolean]))
    Future.find(knowsAboutSegment)(_._2) map (_.map(_._1))
  }

  def all: Seq[ActorRef] = connections.single.values.toSeq

  def restore(implicit actorSystem: ActorSystem, executionContext: ExecutionContext) {
    replaying = true
    Await.ready(Future.sequence(for {
      (connectionId, messages @ (FromRequester(initialReserve: Reserve) +: _)) <- messageStore.loadEverything()
      connection = createConnection(connectionId, initialReserve)
      message <- messages
    } yield {
      Logger.info(s"Replaying message: ${message.getClass.getName}")
      connection ? message
    }), Duration.Inf)

    replaying = false
    Logger.info("Replay completed")
  }

  def findOrCreateConnection(request: NsiProviderOperation)(implicit actorSystem: ActorSystem): (ConnectionId, Option[ActorRef]) = (request, request.optionalConnectionId) match {
    case (_, Some(connectionId)) =>
      (connectionId, get(connectionId))
    case (initialReserve: Reserve, None) =>
      val connectionId = newConnectionId
      val connection = createConnection(connectionId, initialReserve)
      (connectionId, Some(connection))
    case _ =>
      sys.error("illegal initial message")
  }

  private def createConnection(connectionId: ConnectionId, initialReserve: Reserve)(implicit actorSystem: ActorSystem): ActorRef = {
    val (outbound, connection) = connectionFactory(connectionId, initialReserve)

    val storingActor = actorSystem.actorOf(Props(new Actor {
      override def receive = {
        case 'query                     => sender ! connection.query
        case 'querySegments             => sender ! connection.segments

        case SegmentKnown(connectionId) => sender ! connection.rsm.segmentKnown(connectionId)

        case message: InboundMessage =>
          if (!replaying) {
            messageStore.store(connectionId, message)
          }

          val output = connection.process(message)

          output match {
            case None =>
              sender ! messageNotApplicable(message)
            case Some(messages) =>
              if (!replaying) messages.foreach(outbound ! _)
              message match {
{'fraction_non_alphanumeric': 0.09889952836930113, 'fraction_numerical': 0.0008575103615835358, 'mean_word_length': 3.6436629064366293, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 31, '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': '3939091', 'n_tokens_mistral': 1890, 'n_tokens_neox': 1813, 'n_words': 428}