A simple trait to work with bson/json events in akka's persistence

PersistentJson.scala

import play.api.libs.json._
import play.modules.reactivemongo.json.ImplicitBSONHandlers._
import reactivemongo.bson.BSONDocument

trait PersistentJson extends akka.persistence.PersistentActor {

  private[this] var classMap: Map[Manifest[_], Writes[_]] = Map.empty
  private[this] var namesMap: Map[String, Reads[_]] = Map.empty

  case class ClassFormat[T](cls: Class[T], format: Format[T])

  def registerEvent[T: Format : Manifest](name: String): Unit = {
    registerEventFmt[T](name, implicitly[Format[T]])
  }

  def registerEventFmt[T: Manifest](name: String, fmt: Format[T]): Unit = {
    classMap = classMap + (manifest[T] -> (__ \ name).write[T](fmt))
    namesMap = namesMap + (name -> (__ \ name).read[T](fmt))
  }

  def persistJson[T: Manifest](e: T)(f: T => Unit) =
    classMap.get(manifest[T]).map(_.asInstanceOf[Writes[T]].writes(e)).map(JsValueWriter.write) match {
      case Some(bson) =>
        persist(bson)(_ => f(e))
      case None =>
        throw new IllegalArgumentException("Event is not registered for persistence: " + e)
    }

  def receiveRecoverJson: Receive

  def decodeBsonEvent: PartialFunction[Any, Any] = {
    case bson: BSONDocument =>
      val json = JsObjectReader.read(bson)
      val key = json.value.keys.head
      namesMap(key).reads(json) match {
        case JsSuccess(v, _) => v
        case e: JsError =>
          throw new IllegalArgumentException("Cannot read json: " + json)
      }
      
    case o =>
      play.api.Logger.warn("non-bson event: "+o)
      o
  }

  override def receiveRecover: Receive = decodeBsonEvent andThen receiveRecoverJson.orElse {
    case m => play.api.Logger.warn(s"${self.path} Unhandled recovery event: "+m)
  }
}