<commit_before>package akka.persistence.inmemory.snapshot

import akka.actor.{Actor, ActorLogging, Props}
import akka.pattern.ask
import akka.persistence.snapshot.SnapshotStore
import akka.persistence.{SelectedSnapshot, SnapshotMetadata, SnapshotSelectionCriteria}
import akka.util.Timeout

import scala.concurrent.Future
import scala.concurrent.duration._

trait SnapshotEvent

case class SaveSnapshot(metadata: SnapshotMetadata, snapshot: Any) extends SnapshotEvent

case class DeleteSnapshotByMetadata(metadata: SnapshotMetadata) extends SnapshotEvent

case class DeleteSnapshotByCriteria(persistenceId: String, criteria: SnapshotSelectionCriteria) extends SnapshotEvent

// API
case class LoadSnapshot(persistenceId: String, criteria: SnapshotSelectionCriteria)

case class LoadSnapshotResult(selectedSnapshot: Option[SelectedSnapshot])

// general ack
case object SnapshotAck

case class SnapshotCache(cache: Map[SnapshotMetadata, Any]) {
  def update(event: SnapshotEvent): SnapshotCache = event match {
    case SaveSnapshot(metadata, snapshot) =>
      copy(cache = cache + (metadata -> snapshot))

    case DeleteSnapshotByMetadata(metadata) =>
      copy(cache = cache - metadata)

    case DeleteSnapshotByCriteria(persistenceId, criteria) =>
      copy(cache = cache.filterNot {
        case (meta, _) =>
          meta.persistenceId == persistenceId && meta.sequenceNr <= criteria.maxSequenceNr
      })
  }
}

class SnapshotActor extends Actor {
  var snapshots = SnapshotCache(Map.empty[SnapshotMetadata, Any])

  override def receive: Receive = {
    case event: SnapshotEvent =>
      snapshots = snapshots.update(event)
      sender() ! SnapshotAck

    case LoadSnapshot(persistenceId, criteria) =>
      val ss = snapshots.cache.filter {
        case (meta, _) =>
          meta.persistenceId == persistenceId &&
          meta.sequenceNr <= criteria.maxSequenceNr &&
          meta.timestamp <= criteria.maxTimestamp
      }.map {
        case (meta, snap) => SelectedSnapshot(meta, snap)
      }.toSeq
        .sortBy(_.metadata.sequenceNr)
        .reverse
        .headOption

      sender() ! LoadSnapshotResult(ss)
  }
}

class InMemorySnapshotStore extends SnapshotStore with ActorLogging {
  implicit val timeout = Timeout(100.seconds)
  implicit val ec = context.system.dispatcher
  val snapshots = context.actorOf(Props(new SnapshotActor))

  override def loadAsync(persistenceId: String, criteria: SnapshotSelectionCriteria): Future[Option[SelectedSnapshot]] = {
    log.debug("loading for persistenceId: {}, criteria: {}", persistenceId, criteria)
    (snapshots ? LoadSnapshot(persistenceId, criteria)).mapTo[LoadSnapshotResult].map(_.selectedSnapshot)
  }

  override def saveAsync(metadata: SnapshotMetadata, snapshot: Any): Future[Unit] = {
    log.debug("Saving metadata: {}, snapshot: {}", metadata, snapshot)
    (snapshots ? SaveSnapshot(metadata, snapshot)).map(_ => ())
  }

  override def deleteAsync(metadata: SnapshotMetadata): Future[Unit] = {
    log.debug("Deleting: {}", metadata)
    (snapshots ? DeleteSnapshotByMetadata(metadata)).map(_ => ())
  }<commit_msg>InMemorySnapshotStore fix failing tests<commit_after>package akka.persistence.inmemory.snapshot

import akka.actor.{Actor, ActorLogging, Props}
import akka.pattern.ask
import akka.persistence.snapshot.SnapshotStore
import akka.persistence.{SelectedSnapshot, SnapshotMetadata, SnapshotSelectionCriteria}
import akka.util.Timeout

import scala.concurrent.Future
import scala.concurrent.duration._

trait SnapshotEvent

case class SaveSnapshot(metadata: SnapshotMetadata, snapshot: Any) extends SnapshotEvent

case class DeleteSnapshotByMetadata(metadata: SnapshotMetadata) extends SnapshotEvent

case class DeleteSnapshotByCriteria(persistenceId: String, criteria: SnapshotSelectionCriteria) extends SnapshotEvent

// API
case class LoadSnapshot(persistenceId: String, criteria: SnapshotSelectionCriteria)

case class LoadSnapshotResult(selectedSnapshot: Option[SelectedSnapshot])

// general ack
case object SnapshotAck

case class SnapshotCache(cache: Map[SnapshotCacheKey, SnapshotCacheValue]) {

  def find(persistenceId: String, criteria: SnapshotSelectionCriteria): Option[SelectedSnapshot] = {
    cache.filter {
      case (k, v) =>
        k.persistenceId == persistenceId &&
          k.sequenceNr <= criteria.maxSequenceNr &&
          v.metadata.timestamp <= criteria.maxTimestamp
    }.map {
      case (k, v) => SelectedSnapshot(v.metadata, v.snapshot)
    }.toSeq
      .sortBy(_.metadata.sequenceNr)
      .reverse
      .headOption
  }

  def update(event: SnapshotEvent): SnapshotCache = event match {
    case SaveSnapshot(metadata, snapshot) =>
      copy(cache = cache + (new SnapshotCacheKey(metadata) -> new SnapshotCacheValue(metadata, snapshot)))

    case DeleteSnapshotByMetadata(metadata) =>
      copy(cache = cache - new SnapshotCacheKey(metadata))

    case DeleteSnapshotByCriteria(persistenceId, criteria) =>
      copy(cache = cache.filterNot {
        case (k, v) =>
          k.persistenceId == persistenceId &&
            k.sequenceNr <= criteria.maxSequenceNr &&
            v.metadata.timestamp <= criteria.maxTimestamp
      })
  }
}

case class SnapshotCacheKey(persistenceId: String, sequenceNr: Long) {
  def this(metadata: SnapshotMetadata) = this(metadata.persistenceId, metadata.sequenceNr)
}

case class SnapshotCacheValue(metadata: SnapshotMetadata, snapshot: Any)


class SnapshotActor extends Actor {
  var snapshots = SnapshotCache(Map.empty[SnapshotCacheKey, SnapshotCacheValue])

  override def receive: Receive = {
    case event: SnapshotEvent =>
      snapshots = snapshots.update(event)
      sender() ! SnapshotAck

    case LoadSnapshot(persistenceId, criteria) =>
      val ss = snapshots.find(persistenceId, criteria)
      sender() ! LoadSnapshotResult(ss)
  }
}

class InMemorySnapshotStore extends SnapshotStore with ActorLogging {
  implicit val timeout = Timeout(100.seconds)
  implicit val ec = context.system.dispatcher
  val snapshots = context.actorOf(Props(new SnapshotActor))

  override def loadAsync(persistenceId: String, criteria: SnapshotSelectionCriteria): Future[Option[SelectedSnapshot]] = {
    log.debug("loading for persistenceId: {}, criteria: {}", persistenceId, criteria)
    (snapshots ? LoadSnapshot(persistenceId, criteria)).mapTo[LoadSnapshotResult].map(_.selectedSnapshot)
  }

  override def saveAsync(metadata: SnapshotMetadata, snapshot: Any): Future[Unit] = {
    log.debug("Saving metadata: {}, snapshot: {}", metadata, snapshot)
    (snapshots ? SaveSnapshot(metadata, snapshot)).map(_ => ())
  }

  override def deleteAsync(metadata: SnapshotMetadata): Future[Unit] = {
    log.debug("Deleting: {}", metadata)
    (snapshots ? DeleteSnapshotByMetadata(metadata)).map(_ => ())
  }
{'fraction_non_alphanumeric': 0.10107558567850303, 'fraction_numerical': 0.0008840430234271402, 'mean_word_length': 4.66138448707256, 'pattern_counts': {'":': 0, '<': 10, '<?xml version=': 0, '>': 25, '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': '536365', 'n_tokens_mistral': 1842, 'n_tokens_neox': 1841, 'n_words': 430}