+per #18192 leveldb impl of EventsByTag query

* also refactoring of EventsByPersistenceIdPublisher

* increase test timeouts

akka-persistence/src/main/scala/akka/persistence/journal/leveldb/LeveldbJournal.scala

@@ -5,12 +5,16 @@ package akka.persistence.journal.leveldb
 
 import scala.concurrent.duration._
 import scala.language.postfixOps
-
 import akka.actor._
 import akka.persistence.Persistence
 import akka.persistence.journal._
 import akka.util.Timeout
 import akka.util.Helpers.ConfigOps
+import akka.persistence.PersistentRepr
+import scala.concurrent.Future
+import akka.persistence.JournalProtocol.RecoverySuccess
+import akka.persistence.JournalProtocol.ReplayMessagesFailure
+import akka.pattern.pipe
 
 /**
  * INTERNAL API.
@@ -21,12 +25,36 @@ private[persistence] class LeveldbJournal extends { val configPath = "akka.persi
   import LeveldbJournal._
 
   override def receivePluginInternal: Receive = {
+    case r @ ReplayTaggedMessages(fromSequenceNr, toSequenceNr, max, tag, replyTo) ⇒
+      import context.dispatcher
+      asyncReadHighestSequenceNr(tagAsPersistenceId(tag), fromSequenceNr)
+        .flatMap { highSeqNr ⇒
+          val toSeqNr = math.min(toSequenceNr, highSeqNr)
+          if (highSeqNr == 0L || fromSequenceNr > toSeqNr)
+            Future.successful(highSeqNr)
+          else {
+            asyncReplayTaggedMessages(tag, fromSequenceNr, toSeqNr, max) {
+              case ReplayedTaggedMessage(p, tag, offset) ⇒
+                adaptFromJournal(p).foreach { adaptedPersistentRepr ⇒
+                  replyTo.tell(ReplayedTaggedMessage(adaptedPersistentRepr, tag, offset), Actor.noSender)
+                }
+            }.map(_ ⇒ highSeqNr)
+          }
+        }.map {
+          highSeqNr ⇒ RecoverySuccess(highSeqNr)
+        }.recover {
+          case e ⇒ ReplayMessagesFailure(e)
+        }.pipeTo(replyTo)
+
     case SubscribePersistenceId(persistenceId: String) ⇒
       addPersistenceIdSubscriber(sender(), persistenceId)
       context.watch(sender())
     case SubscribeAllPersistenceIds ⇒
       addAllPersistenceIdsSubscriber(sender())
       context.watch(sender())
+    case SubscribeTag(tag: String) ⇒
+      addTagSubscriber(sender(), tag)
+      context.watch(sender())
     case Terminated(ref) ⇒
       removeSubscriber(ref)
   }
@@ -43,18 +71,31 @@ private[persistence] object LeveldbJournal {
    * Used by query-side. The journal will send [[EventAppended]] messages to
    * the subscriber when `asyncWriteMessages` has been called.
    */
-  case class SubscribePersistenceId(persistenceId: String) extends SubscriptionCommand
-  case class EventAppended(persistenceId: String) extends DeadLetterSuppression
+  final case class SubscribePersistenceId(persistenceId: String) extends SubscriptionCommand
+  final case class EventAppended(persistenceId: String) extends DeadLetterSuppression
 
   /**
-   * Subscribe the `sender` to changes (appended events) for a specific `persistenceId`.
+   * Subscribe the `sender` to current and new persistenceIds.
    * Used by query-side. The journal will send one [[CurrentPersistenceIds]] to the
    * subscriber followed by [[PersistenceIdAdded]] messages when new persistenceIds
    * are created.
    */
-  case object SubscribeAllPersistenceIds extends SubscriptionCommand
-  case class CurrentPersistenceIds(allPersistenceIds: Set[String]) extends DeadLetterSuppression
-  case class PersistenceIdAdded(persistenceId: String) extends DeadLetterSuppression
+  final case object SubscribeAllPersistenceIds extends SubscriptionCommand
+  final case class CurrentPersistenceIds(allPersistenceIds: Set[String]) extends DeadLetterSuppression
+  final case class PersistenceIdAdded(persistenceId: String) extends DeadLetterSuppression
+
+  /**
+   * Subscribe the `sender` to changes (appended events) for a specific `tag`.
+   * Used by query-side. The journal will send [[TaggedEventAppended]] messages to
+   * the subscriber when `asyncWriteMessages` has been called.
+   */
+  final case class SubscribeTag(tag: String) extends SubscriptionCommand
+  final case class TaggedEventAppended(tag: String) extends DeadLetterSuppression
+
+  final case class ReplayTaggedMessages(fromSequenceNr: Long, toSequenceNr: Long, max: Long,
+                                        tag: String, replyTo: ActorRef) extends SubscriptionCommand
+  final case class ReplayedTaggedMessage(persistent: PersistentRepr, tag: String, offset: Long)
+    extends DeadLetterSuppression with NoSerializationVerificationNeeded
 }
 
 /**

akka-persistence/src/main/scala/akka/persistence/journal/leveldb/LeveldbRecovery.scala

@@ -6,10 +6,10 @@
 package akka.persistence.journal.leveldb
 
 import scala.concurrent.Future
-
 import akka.persistence._
 import akka.persistence.journal.AsyncRecovery
 import org.iq80.leveldb.DBIterator
+import akka.persistence.journal.leveldb.LeveldbJournal.ReplayedTaggedMessage
 
 /**
  * INTERNAL API.
@@ -74,6 +74,38 @@ private[persistence] trait LeveldbRecovery extends AsyncRecovery { this: Leveldb
     }
   }
 
+  def asyncReplayTaggedMessages(tag: String, fromSequenceNr: Long, toSequenceNr: Long, max: Long)(replayCallback: ReplayedTaggedMessage ⇒ Unit): Future[Unit] = {
+    val tagNid = tagNumericId(tag)
+    Future(replayTaggedMessages(tag, tagNid, fromSequenceNr: Long, toSequenceNr, max: Long)(replayCallback))(replayDispatcher)
+  }
+
+  def replayTaggedMessages(tag: String, tagNid: Int, fromSequenceNr: Long, toSequenceNr: Long, max: Long)(
+    replayCallback: ReplayedTaggedMessage ⇒ Unit): Unit = {
+
+    @scala.annotation.tailrec
+    def go(iter: DBIterator, key: Key, ctr: Long, replayCallback: ReplayedTaggedMessage ⇒ Unit) {
+      if (iter.hasNext) {
+        val nextEntry = iter.next()
+        val nextKey = keyFromBytes(nextEntry.getKey)
+        if (nextKey.sequenceNr > toSequenceNr) {
+          // end iteration here
+        } else if (key.persistenceId == nextKey.persistenceId) {
+          val msg = persistentFromBytes(nextEntry.getValue)
+          if (ctr < max) {
+            replayCallback(ReplayedTaggedMessage(msg, tag, nextKey.sequenceNr))
+            go(iter, nextKey, ctr + 1L, replayCallback)
+          }
+        }
+      }
+    }
+
+    withIterator { iter ⇒
+      val startKey = Key(tagNid, if (fromSequenceNr < 1L) 1L else fromSequenceNr, 0)
+      iter.seek(keyToBytes(startKey))
+      go(iter, startKey, 0L, replayCallback)
+    }
+  }
+
   def readHighestSequenceNr(persistenceId: Int) = {
     val ro = leveldbSnapshot()
     try {

akka-persistence/src/main/scala/akka/persistence/journal/leveldb/LeveldbStore.scala

@@ -34,8 +34,12 @@ private[persistence] trait LeveldbStore extends Actor with WriteJournalBase with
   var leveldb: DB = _
 
   private val persistenceIdSubscribers = new mutable.HashMap[String, mutable.Set[ActorRef]] with mutable.MultiMap[String, ActorRef]
+  private val tagSubscribers = new mutable.HashMap[String, mutable.Set[ActorRef]] with mutable.MultiMap[String, ActorRef]
   private var allPersistenceIdsSubscribers = Set.empty[ActorRef]
 
+  private var tagSequenceNr = Map.empty[String, Long]
+  private val tagPersistenceIdPrefix = "$$$"
+
   def leveldbFactory =
     if (nativeLeveldb) org.fusesource.leveldbjni.JniDBFactory.factory
     else org.iq80.leveldb.impl.Iq80DBFactory.factory
@@ -46,22 +50,34 @@ private[persistence] trait LeveldbStore extends Actor with WriteJournalBase with
 
   def asyncWriteMessages(messages: immutable.Seq[AtomicWrite]): Future[immutable.Seq[Try[Unit]]] = {
     var persistenceIds = Set.empty[String]
-    val hasSubscribers = hasPersistenceIdSubscribers
+    var allTags = Set.empty[String]
+
     val result = Future.fromTry(Try {
       withBatch(batch ⇒ messages.map { a ⇒
         Try {
-          a.payload.foreach(message ⇒ addToMessageBatch(message, batch))
-          if (hasSubscribers)
+          a.payload.foreach { p ⇒
+            val (p2, tags) = p.payload match {
+              case Tagged(payload, tags) ⇒
+                (p.withPayload(payload), tags)
+              case _ ⇒ (p, Set.empty[String])
+            }
+            if (tags.nonEmpty && hasTagSubscribers)
+              allTags ++= tags
+            addToMessageBatch(p2, tags, batch)
+          }
+          if (hasPersistenceIdSubscribers)
             persistenceIds += a.persistenceId
         }
       })
     })
 
-    if (hasSubscribers) {
+    if (hasPersistenceIdSubscribers) {
       persistenceIds.foreach { pid ⇒
         notifyPersistenceIdChange(pid)
       }
     }
+    if (hasTagSubscribers && allTags.nonEmpty)
+      allTags.foreach(notifyTagChange)
     result
   }
 
@@ -112,12 +128,35 @@ private[persistence] trait LeveldbStore extends Actor with WriteJournalBase with
   def persistentToBytes(p: PersistentRepr): Array[Byte] = serialization.serialize(p).get
   def persistentFromBytes(a: Array[Byte]): PersistentRepr = serialization.deserialize(a, classOf[PersistentRepr]).get
 
-  private def addToMessageBatch(persistent: PersistentRepr, batch: WriteBatch): Unit = {
+  private def addToMessageBatch(persistent: PersistentRepr, tags: Set[String], batch: WriteBatch): Unit = {
+    val persistentBytes = persistentToBytes(persistent)
     val nid = numericId(persistent.persistenceId)
     batch.put(keyToBytes(counterKey(nid)), counterToBytes(persistent.sequenceNr))
-    batch.put(keyToBytes(Key(nid, persistent.sequenceNr, 0)), persistentToBytes(persistent))
+    batch.put(keyToBytes(Key(nid, persistent.sequenceNr, 0)), persistentBytes)
+
+    tags.foreach { tag ⇒
+      val tagNid = tagNumericId(tag)
+      val tagSeqNr = nextTagSequenceNr(tag)
+      batch.put(keyToBytes(counterKey(tagNid)), counterToBytes(tagSeqNr))
+      batch.put(keyToBytes(Key(tagNid, tagSeqNr, 0)), persistentBytes)
+    }
   }
 
+  private def nextTagSequenceNr(tag: String): Long = {
+    val n = tagSequenceNr.get(tag) match {
+      case Some(n) ⇒ n
+      case None    ⇒ readHighestSequenceNr(tagNumericId(tag))
+    }
+    tagSequenceNr = tagSequenceNr.updated(tag, n + 1)
+    n + 1
+  }
+
+  def tagNumericId(tag: String): Int =
+    numericId(tagAsPersistenceId(tag))
+
+  def tagAsPersistenceId(tag: String): String =
+    tagPersistenceIdPrefix + tag
+
   override def preStart() {
     leveldb = leveldbFactory.open(leveldbDir, if (nativeLeveldb) leveldbOptions else leveldbOptions.compressionType(CompressionType.NONE))
     super.preStart()
@@ -137,9 +176,17 @@ private[persistence] trait LeveldbStore extends Actor with WriteJournalBase with
     val keys = persistenceIdSubscribers.collect { case (k, s) if s.contains(subscriber) ⇒ k }
     keys.foreach { key ⇒ persistenceIdSubscribers.removeBinding(key, subscriber) }
 
+    val tagKeys = tagSubscribers.collect { case (k, s) if s.contains(subscriber) ⇒ k }
+    tagKeys.foreach { key ⇒ tagSubscribers.removeBinding(key, subscriber) }
+
     allPersistenceIdsSubscribers -= subscriber
   }
 
+  protected def hasTagSubscribers: Boolean = tagSubscribers.nonEmpty
+
+  protected def addTagSubscriber(subscriber: ActorRef, tag: String): Unit =
+    tagSubscribers.addBinding(tag, subscriber)
+
   protected def hasAllPersistenceIdsSubscribers: Boolean = allPersistenceIdsSubscribers.nonEmpty
 
   protected def addAllPersistenceIdsSubscriber(subscriber: ActorRef): Unit = {
@@ -153,8 +200,14 @@ private[persistence] trait LeveldbStore extends Actor with WriteJournalBase with
       persistenceIdSubscribers(persistenceId).foreach(_ ! changed)
     }
 
+  private def notifyTagChange(tag: String): Unit =
+    if (tagSubscribers.contains(tag)) {
+      val changed = LeveldbJournal.TaggedEventAppended(tag)
+      tagSubscribers(tag).foreach(_ ! changed)
+    }
+
   override protected def newPersistenceIdAdded(id: String): Unit = {
-    if (hasAllPersistenceIdsSubscribers) {
+    if (hasAllPersistenceIdsSubscribers && !id.startsWith(tagPersistenceIdPrefix)) {
       val added = LeveldbJournal.PersistenceIdAdded(id)
       allPersistenceIdsSubscribers.foreach(_ ! added)
     }

akka-persistence/src/main/scala/akka/persistence/journal/leveldb/Tagged.scala

@@ -0,0 +1,14 @@
+/**
+ * Copyright (C) 2015 Typesafe Inc. <http://www.typesafe.com>
+ */
+package akka.persistence.journal.leveldb
+
+/**
+ * The LevelDB journal supports tagging of events that are used
+ * by the `EventsByTag` query. To specify the tags you create an
+ * [[akka.persistence.journal.EventAdapter]] that wraps the events
+ * in a `Tagged` with the given `tags`.
+ *
+ * The journal will unwrap the event and store the `payload`.
+ */
+case class Tagged(payload: Any, tags: Set[String])