System.scala

/**
 * Base Node class, as well as system (non-data) nodes.
 */

package curiodb

import akka.actor.{ActorLogging, ActorRef, Cancellable, Props, Stash}
import akka.dispatch.ControlMessage
import akka.event.LoggingReceive
import akka.persistence._
import java.util.UUID.randomUUID
import scala.collection.JavaConversions._
import scala.collection.mutable
import scala.concurrent.duration.DurationInt
import scala.concurrent.ExecutionContext.Implicits.global
import scala.util.{Success, Failure, Try}

/**
 * Node is the base actor class that all concrete node types subclass
 * Specifically, there is a concrete Node class for each data type
 * namely: StringNode for strings, HashNode for hashes, etc. There are
 * also some special Node types, such as KeyNode which manages the key
 * space for regular Node actors, and ClientNode which manages a single
 * client connection.
 *
 * The base class here defines the core features of a Node actor,
 * namely:
 *
 *  - Its data value, typed by the type parameter that concrete
 *    subclasses must define.
 *  - Execution of the node's CommandRunner each time a Command is
 *    received.
 *  - Optionally persisting the node's value to disk (snapshotting)
 *    after a command has been handled (second point above), cleaning
 *    up older snapshots, and restoring from a snapshot on startup.
 *
 * Persistence warrants some discussion: we use akka-persistence, but
 * not completely, as event sourcing is not used, and we rely entirely
 * on its snapshotting feature, only ever keeping a single snapshot.
 * This was basically the easiest way to get persistence working.
 * We always store a reference to the last snapshot's meta-data (the
 * lastSnapshot var) so that we can delete old snapshots whenever a
 * new one is saved. As for saving, this is controlled via the config
 * var curiodb.persist-after which is the duration after which a
 * command runs that writes the node's value (described as writable in
 * the commands.conf file). When one of these commands runs, we call
 * save, which will schedule a Persist message back to the node itself.
 * This is based on the assumption that there's no guarantee an
 * actor's recieve and scheduler won't both execute at the exact same
 * time, so we have everything run through receive. The persisting
 * var stores whether persisting has been scheduled, to allow extra
 * save calls to do nothing when persisting has already been scheduled.
 */
abstract class Node[T] extends PersistentActor with CommandProcessing with ActorLogging {

  /**
   * MVCC map of transaction values for the Node mapped to client IDs,
   * including the committed "main" value, which is identified with
   * an empty string as its key. LinkedHashMap is used so that we can
   * access the most recently written value when transaction isolation
   * is configued as read-uncommitted.
   */
  var values = mutable.LinkedHashMap[String, T]("" -> emptyValue)

  /**
   * Returns the default value for the type of Node, which subclasses
   * must override. Used to set the initial value when a Node is
   * first created, and also in the edge case of resetting a
   * transaction's value if within a transaction, the Node is marked
   * deleted, and then recreated.
   */
  def emptyValue: T

  /**
   * Value getter according to the configured transaction isolation
   * level:
   *
   * - repeatable:  should always return the value of the current transaction
   * - uncommitted: always returns the newest written value, even if uncommitted
   * - committed:   always return the commit value, even inside a transaction
   *
   * Note that there is no "serializable" level since there is no notion of
   * a range query. You can read more info on isolation levels here:
   * https://en.wikipedia.org/wiki/Isolation_(database_systems)
   */
  def value: T = values(isolationLevel match {
    case "repeatable" if inTransaction => command.clientId // transaction value
    case "uncommitted"                 => values.keys.last // most recent value
    case _                             => ""               // committed value
  })

  /**
   * Value setter - writes the current value, either to the main
   * key, or a client ID if in a transaction.
   */
  def value_=(value: T): Unit =
    values(if (inTransaction) command.clientId else "") = value

  /**
   * Clones the value for use in a new transaction. Node subclasses
   * that work on mutable values must override, typically by calling
   * value.clone, otherwise we just default to the current value.
   */
  def cloneValue: T = value

  /**
   * Checks if the current Command is in a transaction.
   */
  def inTransaction: Boolean = values.contains(command.clientId)

  /**
   * Transaction isolation level: "repeatable" reads, read "committed",
   * or read "uncommitted".
   */
  val isolationLevel = context.system.settings.config.getString("curiodb.transactions.isolation")

  /**
   * The most recently saved snapshot. We store it on save and recover,
   * so that we can delete it (and any earlier snapshots) each time we
   * successfully save a new snapshot.
   */
  var lastSnapshot: Option[SnapshotMetadata] = None

  /**
   * Boolean representing whether we've scheduled an internal
   * Persist message, so that it only occurs once at a time
   * according to the duration configured by the curiodb.persist-after
   * setting.
   */
  var persisting: Boolean = false

  /**
   * Stores the duration which controls the minimum delay between
   * receiving a command that writes, and persisting the actual value
   * written.
   */
  val persistAfter = durationSetting("curiodb.persist-after")

  /**
   * Stores the duration which controls the amount of time we allow
   * a transaction to exist for. This ensures that if for some reason
   * the transaction does not complete, it can start handling write
   * commands again.
   */
  val transactionTimeout = durationSetting("curiodb.transactions.timeout")

  /**
   * Abstract definition of each Node actor's CommandRunner that must
   * be implemented. It's the partial function that matches a Command
   * name to the code that runs it.
   */
  def run: CommandRunner

  def persistenceId: String = self.path.name

  /**
   * As discussed above, saves a snapshot. If curiodb.persist-after is
   * zero or less (which means every change must cause a snapshot to be
   * saved), then just save a snapshot. Otherwise schedule the internal
   * Persist message to save a snapshot, if none has already been
   * scheduled, according to the value of the persisting var.
   */
  def save(): Unit = {
    if (isolationLevel == "uncommitted") {
      // Update insertion order.
      val key = if (inTransaction) command.clientId else ""
      values(key) = values.remove(key).get
    }
    if (!inTransaction) {
      if (persistAfter == 0) {
        saveSnapshot(value)
      } else if (persistAfter > 0 && !persisting) {
        persisting = true
        context.system.scheduler.scheduleOnce(persistAfter milliseconds) {self ! Persist}
      }
    }
  }

  /**
   * Deletes old snapshots after a new one is saved.
   */
  def deleteOldSnapshots(stopping: Boolean = false): Unit =
    if (persistAfter >= 0) {
      lastSnapshot.foreach {meta =>
        val criteria = if (stopping) SnapshotSelectionCriteria()
          else SnapshotSelectionCriteria(meta.sequenceNr, meta.timestamp - 1)
        deleteSnapshots(criteria)
      }
    }

  /**
   * Restore the Node actor's value on startup.
   */
  override def receiveRecover: Receive = {
    case SnapshotOffer(meta, snapshot) =>
      lastSnapshot = Some(meta)
      value = snapshot.asInstanceOf[T]
  }

  /**
   * Main Receive handler - deals with starting/stopping/persisting,
   * and receiving Command instances, then running the Node actor's
   * CommandRunner.
   */
  def receiveCommand: Receive = {
    case SaveSnapshotSuccess(meta) => lastSnapshot = Some(meta); deleteOldSnapshots()
    case SaveSnapshotFailure(_, e) => log.error(e, "Snapshot write failed")
    case Persist    => persisting = false; saveSnapshot(values(""))
    case Delete     => deleteOldSnapshots(stopping = true); stop()
    case Sleep      => stop()
    case c: Command =>
      command = c
      // If the KeyNode for this Node marked the command as creating
      // the Node inside a transaction, set up the transaction.
      if (command.createdInTransaction) {
        beginTransaction()
        // Remove the main value, as it conceptually shouldn't exist
        // since this Node doesn't have a committed value yet. More
        // specifically though, we use the absence of a committed value
        // as a flag for shutting down the actor if the transaction is
        // aborted and no committed value exists.
        values.remove("")
      }
      respond(Try((run orElse runForAnyNode)(command.name)) match {
        case Success(response) => if (command.writes) save(); response
        case Failure(e) => log.error(e, s"Error running: $command"); ErrorReply("Unknown error")
      })
  }

  /**
   * Moves the current value into a transaction for the current
   * command's client ID. This can occur at the start of a transaction
   * (eg _MULTI command), or if the Node actor is created inside a
   * transaction, in which case we use emptyValue, which handles the
   * edge case of a Node existing, going into a transaction, being
   * marked deleted, then being recreated.
   */
  def beginTransaction(): Any = {
    val clientId = command.clientId
    if (!inTransaction) context.system.scheduler.scheduleOnce(transactionTimeout milliseconds) {closeTransaction(clientId)}
    values(clientId) = if (inTransaction) emptyValue else cloneValue
    TransactionAck()
  }

  /**
   * Removes the transaction value and writes it back into the main
   * value. This could be called after the transaction has timed out,
   * so we only move the value if we're still in a transaction.
   */
  def commitTransaction(): Any =
    closeTransaction(committing = true) match {
      case Some(v) => value = v; TransactionAck()
      case None    => ()
    }

  /**
   * Removes and returns the transaction value. Expected to run twice
   * per transaction - once when commitTransaction is called when
   * _EXEC is received, and once when the transaction times out.
   * We also check for the lack of a committed value, which indicates
   * that the actor was created in a transaction
   */
  def closeTransaction(clientId: String = command.clientId, committing: Boolean = false): Option[T] = {
    val removed = values.remove(clientId)
    // Transaction was aborted when the Node was created in a
    // transaction, so shut down. We also check that we're actually
    // removing the transaction value, since this will run twice and
    // calling stop twice raises NullPointerExecption in Akka.
    if (!committing && values.isEmpty) stop()
    removed
  }

  /**
   * Alternate Receive handler for all Node actors - these are internal
   * commands can run for any Node type, and are not validated in any
   * way, nor are they configured in commands.conf. Currently used for
   * transaction begin (_MULTI) and commit (_EXEC).
   *
   * When _MULTI is received and a transaction is beginning, we store
   * the current value in the transaction map against the transaction's
   * client ID. Subsequent commands within the transaction will then
   * read from and write to the transaction map, until it's committed
   * via the _EXEC command, at which point we copy the value from the
   * transaction map back to the main value, and remove it from the
   * transaction map.
   *
   * The transaction may also timeout which is scheduled when _MULTI is
   * received. If this occurs and then _EXEC is received, we don't send
   * acknowledgement back to the ClientNode, so it can then deal with
   * its own timeout.
   */
  def runForAnyNode: CommandRunner = {
    case "_MULTI"   => beginTransaction()
    case "_EXEC"    => commitTransaction()
    case "_DISCARD" => closeTransaction(); ()
  }

  override def receive: Receive = LoggingReceive(super.receive)

  /**
   * Handles the RENAME/RENAMENX commands for all node types. Each
   * Node subclass is responsible for converting its value to something
   * that can be used by the relevant command given - this command
   * is typically implemented (and named) differently per node, and
   * knows how to accept fromValue. Renaming is basically a delete,
   * then create a new node, so we first delete ourselves (via the
   * internal _DEL command), then pass the fromValue arg with the given
   * command to what will be a newly created Node.
   */
  def rename(fromValue: Any, toCommand: String): Unit =
    if (command.key != args(0)) {
      route(Seq("_DEL", command.key))
      route(Seq(toCommand, args(0)) ++ (fromValue match {
        case x: Iterable[Any] => x
        case x => Seq(x)
      }))
    }

  /**
   * Handles the SORT command. Defined here since it can be run against
   * multiple Node types, namely ListNode, SetNode and SortedSetNode.
   */
  def sort(values: Iterable[String]): Any = {
    // TODO: BY/GET support.
    var sorted = if (command.indexOf("ALPHA") > -1) values.toSeq.sorted else values.toSeq.sortBy(_.toFloat)
    if (command.indexOf("DESC") > -1) sorted = sorted.reverse
    val limit = command.indexOf("LIMIT")
    if (limit > -1) sorted = sorted.slice(args(limit + 1).toInt, args(limit + 2).toInt)
    val store = command.indexOf("STORE")
    if (store > -1) {
      route(Seq("_LSTORE", args(store + 1)) ++ sorted)
      sorted.size
    } else sorted
  }

}

// Following are some messages we send to nodes.

/**
 * Message that a Node actor sends to itself when it's ready to save
 * its value to disk. See Node.save() for more detail.
 */
case object Persist extends ControlMessage

/**
 * Message that shuts a Node actor down and delete its value from disk.
 */
case object Delete extends ControlMessage

/**
 * Message that tells an actor to shut down, but not delete its value
 * from disk - the KeyNode holding its key still contains a reference
 * to it, and a new actor will be started for it the next time a
 * command is received for its key. See KeyNode for more detail.
 */
case object Sleep extends ControlMessage

/**
 * Message that a Node or KeyNode actor sends back to a ClientNode
 * actor to acknowledge the number of keys it has dealt with after
 * receiving _MULTI or _EXEC commands.
 * See ClientNode.awaitTransactionAcks for more detail.
 */
case class TransactionAck(keyCount: Int = 1)

/**
 * A KeyNode manages a subset of keys, and stores these by mapping
 * DB names to keys to nodes, where nodes are represented by a
 * NodeEntry.
 */
@SerialVersionUID(1L)
case class NodeEntry(
    val kind: String,
    @transient var node: Option[ActorRef] = None,
    @transient var expiry: Option[(Long, Cancellable)] = None,
    @transient var sleep: Option[Cancellable] = None,
    @transient var deletedInTransaction: Boolean = false)
  extends Serializable

/**
 * A KeyNode is a special type of Node in the system. It does not
 * represent any key/value provided by a client, but instead is
 * responsible for managing the keyspace for a subset of keys in its
 * value field, which contains a map of DB names, mapped to keys,
 * mapped to nodes, where each node is represented by a NodeEntry.
 * KeyNode takes on the role of a Node, as there are a variety of
 * commands that logically belong to it.
 *
 * When a ClientNode (the other special type of node, responsble
 * for handling a single client) receives a command for a particular
 * key, the command is first routed to the KeyNode responsible for that
 * key. The KeyNode maps each key to an ActorRef for the key's actual
 * node (StringNode, ListNode, etc), having first created the Node
 * actor at some point (this actually occurs when the key/node doesn't
 * exist and, a command is run that doesn't have a default defined, and
 * writes, as per in commands.conf), and forwards each command onto the
 * Node when received. The routing strategy used is configurable via
 * Akka (actor.deployment./keys.router in application.conf),
 * defaulting to the consistent-hashing-pool router.
 *
 * A KeyNode also handles many commands itself that deal with key
 * management, such as deleting, persisting, and expiring keys.
 * Each NodeEntry instance stores a Option(Long/Cancellable) pair
 * which if expiring, contains the timestamp of when the expiry will
 * occur, and the Cancellable task that will actually run, deleting
 * the node. Deletion of a node (either via expiry or the DEL command)
 * simply sends a message to the Node actor which when received,
 * shuts down the actor, and then removes the key and NodeEntry from
 * the keyspace map.
 *
 * Lastly worthy of discussion is a feature that Redis does not
 * provide, virtual memory, which simply allows a Node to persist
 * its value to disk, and shut down after a period of time (defined by
 * the curiodb.sleep-after duration value in reference.conf). The
 * difference between this occurring and a Node being deleted, is that
 * the key and NodeEntry is kept in the keyspace map. This is also why
 * the ActorRef for each Node in a NodeEntry is an Option - a value of
 * None indicates a sleeping Node. When a command is run against a key
 * mapped to a sleeping Node, a new Node actor is created, which will
 * read its previous value from disk. The idea here is to allow more
 * data to be stored in the system than can fit in memory.
 */
class KeyNode extends Node[mutable.Map[String, mutable.Map[String, NodeEntry]]] with PubSubServer with ScriptingServer {

  /**
   * Alias for keys mapped to each NodeEntry.
   */
  type DB = mutable.Map[String, NodeEntry]

  /**
   * The actual value persisted for a KeyNode is a map of db names
   * to DB instances.
   */
  override def emptyValue = mutable.Map[String, DB]()

  /**
   * Deep clone which drills down into each DB, and copies all its
   * entries.
   */
  override def cloneValue = value.map {case (dbName, entries) =>
    dbName -> entries.map {case (key, entry) => key -> entry.copy()}
  }

  /**
   * Error message sent to ClientNode when a command is issued against
   * an existing key that contains a different type of node than the
   * type that the command belongs to.
   */
  val wrongType = ErrorReply("Operation against a key holding the wrong kind of value", prefix = "WRONGTYPE")

  /**
   * Milliseconds after which a Node should persist its value to disk,
   * and shut down, (aka sleep).
   */
  val sleepAfter = durationSetting("curiodb.sleep-after")

  /**
   * Milliseconds after which a Node automatically expires itself.
   */
  val expireAfter = durationSetting("curiodb.expire-after")

  /**
   * Current keys in transaction mapped to client IDs.
   * See transactionAcknowledge/transactionClose for more detail.
   */
  var transactionKeys = mutable.Map[String, String]()

  /**
   * Shortcut for grabbing the String/NodeEntry map (aka DB) for the
   * given DB name, or the current command.
   */
  def db(name: String = command.db): DB =
    value.getOrElseUpdate(name, mutable.Map[String, NodeEntry]())

  /**
   * Cancels expiry on a node when the PERSIST command is run.
   */
  def persist(): Int = db()(command.key).expiry match {
    case Some((_, cancellable)) => cancellable.cancel(); db()(command.key).expiry = None; 1
    case None => 0
  }

  /**
   * Initiates expiry on a node when any of the relevant commands are
   * run, namely EXPIRE/PEXPIRE/EXPIREAT/PEXPIREAT, or when the
   * "expire-after" setting is configured.
   */
  def expire(when: Long): Int = {
    persist
    val expires = (when - System.currentTimeMillis).toInt milliseconds
    val cancellable = context.system.scheduler.scheduleOnce(expires) {
      self ! command.copy(Seq("_DEL", command.key), client = None)
    }
    db()(command.key).expiry = Some((when, cancellable))
    1
  }

  /**
   * Called every time a command is run, only if sleeping is enabled
   * (via the curiodb.sleep-after config value), and initiates the
   * actor's sleeping process, whereby its value is persisted to disk
   * and then shut down. The key's NodeEntry has its ActorRef then
   * set to None to signify a sleeping Node. When a command is next
   * run against the key, the actor for the Node is then recreated,
   * with its value restored from disk.
   */
  def sleep(): Unit = {
    val when = sleepAfter milliseconds
    val key = command.key
    val entry = db()(key)
    entry.sleep.foreach(_.cancel())
    entry.sleep = Some(context.system.scheduler.scheduleOnce(when) {
      db().get(key).foreach {entry =>
        entry.node.foreach(_ ! Sleep)
        entry.node = None
      }
    })
  }

  /**
   * Retrieves the milliseconds remaining until expiry occurs for a key
   * when the TTL/PTTL commands are run.
   */
  def ttl: Long = db()(command.key).expiry match {
    case Some((when, _)) => when - System.currentTimeMillis
    case None => -1
  }

  /**
   * Handles the bridge between the external SORT and internal _SORT
   * commands. This is essentially a work around given that we only
   * allow external (client facing) commands to map to a single type
   * of Node.
   */
  def sort(): Any = {
    val entry = db()(command.key)
    entry.kind match {
      case "list" | "set" | "sortedset" =>
        entry.node.get ! command.copy(Seq("_SORT", command.key) ++ command.args)
      case _ => wrongType
    }
  }

  /**
   * Validates that the key and command for the current Command can be
   * run. This boils down to ensuring the command belongs to the type
   * of Node mapped to the key, and that the Node must or musn't exist,
   * given the particular command. Optionally returns an error, or a
   * default value if the Node doesn't exist and a default is defined
   * (as per commands.conf). Otherwise if validation passes, None is
   * returned, indicating the command should be sent to the key's
   * Node. Also validates that for write commands, the keys for the
   * command aren't part of an existing transaction - if this occurs,
   * the commands gets stashed until the transaction has completed.
   */
  def validate: Option[Any] = {
    val _exists     = exists(command.key)
    val kind        = if (_exists) db()(command.key).kind else ""
    val invalidKind = (kind != "" && command.kind != kind &&
      command.kind != "keys" && !command.overwrites)
    val mustExist   = command.name == "LPUSHX" || command.name == "RPUSHX"
    val cantExist   = command.name == "SETNX"
    val wait        = (command.writes &&
      !command.keys.filter(transactionKeys.getOrElse(_, command.clientId) != command.clientId).isEmpty)
    if (invalidKind) {
      Some(wrongType)
    } else if ((_exists && cantExist) || (!_exists && mustExist) || (!_exists && !cantExist && command.default != () && command.keyed)) {
      Some(if (command.default != ()) command.default else 0)
    } else if (wait) {
      stash(); Some(())
    } else {
      None
    }
  }

  /**
   * Resets the current command to a copy of itself with the
   * createdInTransaction flag set to true. The Node actor
   * that receives this command will then use it to set up
   * its own transaction handling.
   */
  def commandCreatedInTransaction(): Unit =
    command = command.copy(createdInTransaction = true)

  /**
   * Shortcut for retrieving the ActorRef for the current Command.
   * There are actually 5 states handled here:
   *
   * 1) No key, Command belongs to a KeyNode.
   * 2) Key doesn't exist.
   * 3) Key exists and Node actor is not running (it's asleep).
   * 4) Key exists, Node is running.
   * 5) Key exists, Node is running, but marked deleted in transaction.
   */
  def node(): ActorRef = {
    if (command.kind == "keys")
      self  // 1) Command belongs to a KeyNode, just return self.
    else
      db().get(command.key).flatMap({entry =>
        // 5) The Node actor is running, but it's makred deleted in
        // a transaction, and we're conceptually recreating it even
        // though it has a running Node actor. In this case we reset
        // the deletedInTransaction flag, and mark the command as
        // creating the Node in a transaction, so that the Node can
        // use its emptyValue when setting the transaction up again.
        // This code reads a bit weirdly, since we're hijacking the
        // flatMap to modify some state.
        if (entry.deletedInTransaction) {
          entry.deletedInTransaction = false
          commandCreatedInTransaction()
        }
        entry.node
      }) match {
        // 4) and 5) Just return the running Node.
        case Some(node) => node
        // 2) and 3) The flatMap above ensures None is matched even
        // if the key exists, but contains no Node running, eg it's
        // asleep. In either case we create a new Node actor.
        case None => create(command.db, command.key, command.kind).get
      }
  }

  /**
   * Creates a new Node actor for the given DB name, key and node type.
   * The recovery arg is set to true when the KeyNode first starts
   * and restores its keyspace from disk, otherwise (the default)
   * we persists the keyspace to disk.
   */
  def create(dbName: String, key: String, kind: String, recovery: Boolean = false): Option[ActorRef] = {
    val node = if (recovery && sleepAfter > 0) None else Some(context.actorOf(kind match {
      case "string"      => Props[StringNode]
      case "bitmap"      => Props[BitmapNode]
      case "hyperloglog" => Props[HyperLogLogNode]
      case "hash"        => Props[HashNode]
      case "list"        => Props[ListNode]
      case "set"         => Props[SetNode]
      case "sortedset"   => Props[SortedSetNode]
    }, s"$dbName-$kind-$key"))
    db(dbName)(key) = new NodeEntry(kind, node)
    if (!recovery) save()
    if (transactionKeys.getOrElse(key, "") == command.clientId) commandCreatedInTransaction()
    node
  }

  /**
   * Deletes a Node.
   */
  def delete(key: String, dbName: Option[String] = None): Boolean = {
    val _db = db(dbName.getOrElse(command.db))
    _db.get(key) match {
      case Some(entry) if inTransaction =>
        val deleted = !entry.deletedInTransaction
        entry.deletedInTransaction = true
        deleted
      case Some(entry) => _db.remove(key); entry.node.foreach(_ ! Delete); true
      case None        => false
    }
  }

  /**
   * Does the key exist - rather than using db.contains, we also
   * need to account for keys marked as deleted in a transaction.
   */
  def exists(key: String): Boolean = db().get(key) match {
    case Some(entry) => !entry.deletedInTransaction
    case None        => false
  }

  /**
   * Called by beginTransaction and commitTransaction when _MULTI and
   * _EXEC are received respectively, and handled both forwarding on
   * these commands, and returning TransactionAck responses indicating
   * the number of *invalid* keys received, so that the ClientNode
   * managing the transaction can account for all keys sent. See
   * ClientNode.awaitTransactionAcks for more detail about this flow.
   */
  def forwardTransactionToNodes(): Any = {
    val keys = command.keys
    val nodes = keys.flatMap(db().get(_).flatMap(_.node))
    val keyCount = keys.size - nodes.size
    if (!nodes.isEmpty) {
      val c = command.copy(Seq(command.name))
      nodes.foreach(_ ! c)
    }
    if (keyCount > 0) TransactionAck(keyCount) else ()
  }

  /**
   * Overridden so that we can add the keys for this transaction to
   * transactionKeys, and run forwardTransactionToNodes.
   */
  override def beginTransaction(): Any = {
    if (!inTransaction) command.keys.foreach(key => transactionKeys(key) = command.clientId)
    super.beginTransaction()
    forwardTransactionToNodes()
  }

  /**
   * Overridden so that we can migrate the DB in the transaction map,
   * specifically handling the case of keys marked as deleted, as this
   * is the stage where they're permananely deleted as they would be
   * outside of a transaction.
   */
  override def commitTransaction(): Any =
    closeTransaction(committing = true) match {
      case Some(dbs) =>
        dbs.foreach {case (dbName, entries) =>
          entries.foreach {case (key, entry) =>
            if (entry.deletedInTransaction) {
              db(dbName).remove(key)
              entry.node.foreach(_ ! Delete)
            } else {
              db(dbName)(key) = entry
            }
          }
        }
        forwardTransactionToNodes()
      case None => ()
    }

  /**
   * Overridden so that we can clean up transactionKeys, and unstash
   * any write commands that were previously stashed due to conflicting
   * with keys in an existing transaction.
   */
  override def closeTransaction(clientId: String = command.clientId, committing: Boolean = false) = {
    if (inTransaction) {
      unstashAll()
      transactionKeys = transactionKeys.filter {case (_, cId) => cId != clientId}
    }
    super.closeTransaction(clientId, committing)
  }

  override def run: CommandRunner = ({
    case "_DISCARD"      => forwardTransactionToNodes(); closeTransaction(); ()
    case "_DEL"          => args.map(delete(_))
    case "_KEYS"         => pattern(db().keys, args(0))
    case "_RANDOMKEY"    => randomItem(db().keys)
    case "_FLUSHDB"      => db().keys.map(key => delete(key)); SimpleReply()
    case "_FLUSHALL"     => value.foreach {case (dbName, entries) => entries.keys.map({key => delete(key, Some(dbName))})}; SimpleReply()
    case "_EXISTS"       => args.map(exists)
    case "TTL"           => val x = ttl; if (x == -1) x else x / 1000
    case "PTTL"          => ttl
    case "EXPIRE"        => expire(System.currentTimeMillis + (args(0).toInt * 1000))
    case "PEXPIRE"       => expire(System.currentTimeMillis + args(0).toInt)
    case "EXPIREAT"      => expire(args(0).toLong / 1000)
    case "PEXPIREAT"     => expire(args(0).toLong)
    case "TYPE"          => if (exists(command.key)) db()(command.key).kind else null  // TODO: flatMap+getOrElse
    case "RENAMENX"      => val x = exists(args(0)); if (x) {run("RENAME")}; x
    case "RENAME"        => db()(command.key).node.foreach(_ ! command.copy(Seq("_RENAME" +: command.keys), client = None)); SimpleReply()
    case "PERSIST"       => persist()
    case "SORT"          => sort()
  }: CommandRunner) orElse runPubSub orElse runScripting

  /**
   * We override the KeyNode actor's Receive so as to perform validation,
   * prior to the Node parent class Receive running, which wil call
   * CommandRunner for the KeyNode.
   */
  override def receiveCommand: Receive = ({
    case Routable(c) => command = c; validate match {
      case Some(errorOrDefault) => respond(errorOrDefault)
      case None =>
        if (command.overwrites && !db().get(command.key).filter(_.kind != command.kind).isEmpty) delete(command.key)
        node() ! command
        if (command.keyed) {
          if (sleepAfter > 0) sleep()
          if (expireAfter > 0) expire(System.currentTimeMillis + expireAfter)
        }
    }
  }: Receive) orElse super.receiveCommand

  /**
   * Restores the keyspace from disk on startup, creating each Node
   * actor.
   */
  override def receiveRecover: Receive = {
    case SnapshotOffer(_, snapshot) =>
      snapshot.asInstanceOf[mutable.Map[String, DB]].foreach {case (dbName, entries) =>
        entries.foreach {case (key, entry) => create(dbName, key, entry.kind, recovery = true)}
      }
  }

}

/**
 * A ClientNode is a special type of Node in the system. It does not
 * represent any key/value provided by a client, but instead is
 * responsible for managing the life-cycle of a single client
 * connection.
 *
 * ClientNode is subclassed according to each external protocol
 * supported, such as the Redis protocol over TCP, JSON over HTTP, and
 * also the Lua pcall/call scripting API. Each subclass is responsible
 * for converting its input into a Command payload, and converting
 * Response payloads back to the relevant protocol deal with.
 *
 * ClientNode is a Node subclass, as it also handles certain commands
 * itself, such as utiilities that don't need to be routed via KeyNode
 * actors.
 */
abstract class ClientNode extends Node[Null] with PubSubClient with AggregateCommands with ScriptingClient {

  lazy val id = randomUUID.toString

  override def emptyValue = null

  /**
   * ActorRef we send final responses back to.
   */
  var client: Option[ActorRef] = None

  /**
   * The current DB name - default, or one provided by "SELECT".
   */
  var db = command.db

  /**
   * List of disabled commands.
   */
  val disabled = context.system.settings.config.getStringList("curiodb.commands.disabled").map(_.toUpperCase).toSet

  def run: CommandRunner = ({
    case "SELECT"         => db = args(0); SimpleReply()
    case "ECHO"           => args(0)
    case "PING"           => SimpleReply("PONG")
    case "TIME"           => val x = System.nanoTime; Seq(x / 1000000000, x % 1000000)
    case "SHUTDOWN"       => context.system.terminate(); SimpleReply()
    case "QUIT"           => respondFinal(SimpleReply()); self ! Delete
    case "MULTI" | "EXEC" => ()
  }: CommandRunner) orElse runPubSub orElse runAggregate orElse runScripting

  /**
   * Performs initial Command validation before sending it anywhere,
   * as per the command's definition in commands.conf, namely that the
   * command belongs to a type of Node, it contains a key if required,
   * and the number of arguments fall within the configured range.
   */
  def validate: Option[ErrorReply] = {
    if (command.kind == "") {
      Some(ErrorReply(s"unknown command '${command.name}'"))
    } else if (disabled.contains(command.name)) {
      Some(ErrorReply(s"command '${command.name}' is disabled"))
    } else if ((command.key == "" && command.keyed) || !command.argsInRange) {
      Some(ErrorReply(s"wrong number of arguments for '${command.name}' command"))
    } else {
      None
    }
  }

  /**
   * Ordered map of commands that have been received - used to
   * order responses received when multiple commands are in play,
   * eg MULTI/EXEC.
   */
  var commands = mutable.LinkedHashMap[String, Command]()

  /**
   * Map of responses that have been received inside a transaction.
   */
  var responses = mutable.Map[String, Response]()

  /**
   * Flags a validation error occurring for a command inside MULTI/EXEC
   * which is then used to abort when EXEC is called.
   */
  var transactionError = false

  /**
   * Should transactions abort if any of their commands receives
   * an error response.
   */
  val transactionAbortOnError = context.system.settings.config.getString("curiodb.transactions.on-error") == "rollback"

  /**
   * Flag used to mark the ClientNode as being in "streaming" mode,
   * namely when a PubSub channel is subscribed to, and command
   * timeouts should be bypassed.
   */
  var streaming = false

  /**
   * Cleanup performed whenever MULTI/EXEC is aborted, either
   * explicitly via DISCARD or implicitly via an error.
   */
  def abortTransaction(response: Any): Unit = {
    transactionError = false
    responses.clear
    commands.clear
    respondFinal(response)
  }

  /**
   * Determines if a transaction is triggered via MULTI.
   */
  def multi: Boolean = commands.size > 0 && commands.head._2.name == "MULTI"

  /**
   * Routes the given Command instance, initializing its timeout event.
   */
  def routeWithTimeout(command: Command) = {
    streaming = command.streaming
    if (!streaming && commandTimeout > 0) {
      val id = command.id
      context.system.scheduler.scheduleOnce(commandTimeout milliseconds) {
        commands.remove(id) match {
          case Some(c) => respondFinal(ErrorReply(s"Timeout on $c"))
          case None    =>
        }
      }
    }
    route(command)
  }

  /**
   * Internal API for subclasses to send a command once they've
   * received its full input. Here we constructs a Command payload for
   * the given input, validate it, and typically route it to the
   * relevant actor, other than in the case of transactionKeys. In this
   * case, we either buffer the commands if MULTI/EXEC is the reason
   * for the transaction (Lua scripts also run a transaction), and then
   * when committing (EXEC/EVAL/EVALSHA), we begin the transaction
   * phase - see awaitTransactionAcks for more detail.
   */
  def sendCommand(input: Seq[String], originalClientId: Option[String] = None): Unit = {
    command = Command(input, Some(self), db, originalClientId.getOrElse(id))
    client = Some(sender())
    validate match {
      case Some(error) =>
        respondFinal(error)
        // If we're in a transaction, note the error so we can abort
        // when EXEC is run.
        if (multi) transactionError = true
      case None =>
        commands(command.id) = command
        command.name match {
          case "DISCARD" =>
            // Aborting transactions.
            abortTransaction(SimpleReply())
          case "EXEC" if transactionError =>
            // Error occurred on initial validation of a command,
            // so abort transactions.
            abortTransaction(ErrorReply("Transaction discarded because of previous errors.", "EXECABORT"))
          case "EXEC" if !multi =>
            // EXEC without MULTI.
            abortTransaction(ErrorReply("EXEC without MULTI"))
          case "EXEC" if commands.size == 2 =>
            // MULTI and EXEC with no commands, return empty list.
            abortTransaction(Seq())
          case "EXEC" | "EVAL" | "EVALSHA" =>
            // Committing a transaction, start the first step.
            awaitTransactionAcks("_MULTI", {
              context.become(awaitTransactionResponses)
              commands.values.foreach(routeWithTimeout)
            })
          case _ if multi =>
            // Since SELECT modifies subsequent commands as they're
            // constructed, we need to run it immediately. Running it
            // again within a transaction (which happens) does nothing.
            if (command.name == "SELECT") run(command.name)
            // Adding commands to a transaction, send SimpleReply.
            respondFinal(SimpleReply(if (commands.size == 1) "OK" else "QUEUED"))
          case _ =>
            // No transaction, send the command.
            routeWithTimeout(commands.head._2)
        }
    }
  }

  /**
   * Handles transaction management, specifically broadcasting either
   * _MULTI or _EXEC to KeyNode (and indirectly, Node) actors, and then
   * running the next stage on completion of receiving all expected
   * replies. Following are the three stages of a transaction:
   *
   * 1) Send _MULTI and await all replies confirming we can begin
   * 2) Send actual commands and await responses from Node actors
   * 3) Send _EXEC and await all replies confirming we have committed
   *
   * The awaitTransactionAcks method therefore handles steps 1 and 3
   * (step 2 is handled by awaitTransactionResponses). In each case,
   * we perform the same brodcast optimization found throughout
   * Aggregation.scala, where we send all keys to all KeyNode actors,
   * to reduce the number of messages sent, at the cost of redundant
   * keys. When we broadcast all keys for a awaitTransactionAcks step,
   * we need all KeyNode and relevant Node actors to confirm they are
   * ready. This is managed by initially calculating the number of keys
   * being transmitted (total keys * total KeyNode actors), and
   * awaiting TransactionAck replies which each contain a keyCount.
   * KeyNode actors will return a keyCount for the number of *invalid*
   * keys it recevies, while passing on _MULTI and _EXEC to the Node
   * actors it correctly received keys for - those Node actors should
   * then return a keyCount of 1 in their TransactionAck replies.
   *
   * Once the ClientNode has received TransactionAck.keyCount values
   * matching the total number of keys, it proceeds with the next
   * step. For _MULTI, this means sending the actual queued commands
   * via awaitTransactionResponses. For _EXEC, this means sending the
   * collected Response messages back to the external client.
   */
  def awaitTransactionAcks(commandName: String, onComplete: => Unit): Unit = {
    val keys = commands.values.flatMap(_.keys).toSet
    if (keys.isEmpty) {
      // If there are no keys involved, there won't be any
      // TransactionAck messages to receive, so go straight to the
      // next step.
      onComplete
    } else {
      val timeout = context.system.scheduler.scheduleOnce(commandTimeout milliseconds) {
        abortTransaction(ErrorReply(s"Transaction timeout"))
        context.unbecome
      }
      var keyCount = keys.size * context.system.settings.config.getInt("curiodb.keynodes")
      context.become({
        case Response(TransactionAck(kc), _) =>
          keyCount -= kc
          if (keyCount == 0) {context.unbecome; timeout.cancel(); onComplete}
      })
      route(commandName +: keys.toSeq, client = Some(self))
    }
  }

  /**
   * Sends the final response once a transaction has ended, which in
   * the case of MULTI/EXEC, is a sequence of all the collected
   * responses, ordered by their corresponding commands.
   */
  def endTransaction() = {
    respondFinal(if (multi) {
      commands.keys.slice(1, commands.size - 1).toSeq.map(responses(_).value)
    } else responses.values.head.value)
    commands.clear
    responses.clear
    context.become(receiveCommand)
  }

  /**
   * Middle phase of a transaction where we receive and collect
   * Response messages for queued commands - once they're all received,
   * we construct the final response, switch state back to the default
   * receiveCommand, and send the final response back to self to be
   * received as a regular Response. Here we also check for error
   * replies if we're configured to abort transactions on errors, and
   * if any are received, we abort the transaction with _DISCARD, and
   * set all responses other than the error to null, so that clients
   * can determine which command caused the error.
   */
  def awaitTransactionResponses(): Receive = ({
    case response @ Response(_: ErrorReply, _) if transactionAbortOnError =>
      val keys = commands.values.flatMap(_.keys).toSet
      route("_DISCARD" +: keys.toSeq)
      // Set all responses to null, and then put the error response
      // against its ID - this way when the sequence of ordered
      // responses is sent back as one response, clients can determine
      // which command caused the error.
      commands.keys.foreach(id => responses(id) = Response(null, id))
      responses(response.id) = response
      endTransaction()
    case response: Response =>
      responses(response.id) = response
      // We've received all responses when there's the same number as
      // commands, omitting MULTI and EXEC.
      if (!multi || responses.size == commands.size - 2)
      awaitTransactionAcks("_EXEC", endTransaction)
  }: Receive) orElse receiveCommand  // ClientNode may still receive commands.

  /**
   * Sends a final response back to the original actor that provided the
   * input for the Command.
   */
  def respondFinal(response: Any): Unit = client.get ! formatResponse(response)

  /**
   * Hook for subclasses to override and convert a response value
   * into an appropriate format.
   */
  def formatResponse(response: Any): Any = response

  /**
   * Receives final Response, which is then formatted per ClientNode
   * subclass, and sent back to the original actor that provided the
   * input for the Command.
   */
  override def receiveCommand: Receive = ({
    case response: Response =>
      commands.remove(response.id) match {
        case None if !streaming => log.error(s"Client received response after timeout $response")
        case _                  => respondFinal(response.value)
      }
  }: Receive) orElse receivePubSub orElse super.receiveCommand

}