io.gatling.core.session.Session Scala Examples

package com.ringcentral.gatling.mongo.action

import com.ringcentral.gatling.mongo.command.MongoInsertCommand
import com.ringcentral.gatling.mongo.response.MongoCountResponse
import io.gatling.commons.stats.KO
import io.gatling.commons.util.TimeHelper.nowMillis
import io.gatling.commons.validation.Validation
import io.gatling.core.action.Action
import io.gatling.core.config.GatlingConfiguration
import io.gatling.core.session.{Expression, Session}
import io.gatling.core.stats.StatsEngine
import reactivemongo.api.DefaultDB
import reactivemongo.play.json.ImplicitBSONHandlers._
import reactivemongo.play.json.collection.JSONCollection

import scala.concurrent.ExecutionContext.Implicits.global
import scala.util.{Failure, Success}

class MongoInsertAction(command: MongoInsertCommand, database: DefaultDB, val statsEngine: StatsEngine, configuration: GatlingConfiguration, val next: Action) extends MongoAction(database) {

  override def name: String = genName("Mongo insert command")

  override def commandName: Expression[String] = command.commandName

  override def executeCommand(commandName: String, session: Session): Validation[Unit] = for {
    collectionName <- command.collection(session)
    resolvedDocument <- command.document(session)
    document <- string2JsObject(resolvedDocument)
  } yield {
    val sent = nowMillis
    database.collection[JSONCollection](collectionName).insert(document).onComplete {
      case Success(result) =>
        if (result.ok) {
          processResult(session, sent, nowMillis, command.checks, MongoCountResponse(result.n), next, commandName)
        } else {
          executeNext(session, sent, nowMillis, KO, next, commandName, Some(result.writeErrors.map(we => s"[${we.code}] ${we.errmsg}").mkString(", ")))
        }
      case Failure(err) =>
        executeNext(session, sent, nowMillis, KO, next, commandName, Some(err.getMessage))
    }

  }
} 

package ru.tinkoff.gatling.amqp.action

import io.gatling.commons.validation.{Validation, _}
import io.gatling.core.action.RequestAction
import io.gatling.core.controller.throttle.Throttler
import io.gatling.core.session.{Expression, Session}
import io.gatling.core.util.NameGen
import ru.tinkoff.gatling.amqp.client.AmqpPublisher
import ru.tinkoff.gatling.amqp.protocol.AmqpComponents
import ru.tinkoff.gatling.amqp.request.{AmqpAttributes, AmqpProtocolMessage}

abstract class AmqpAction(
    attributes: AmqpAttributes,
    components: AmqpComponents,
    throttler: Throttler,
    throttled: Boolean
) extends RequestAction with AmqpLogging with NameGen {
  override val requestName: Expression[String] = attributes.requestName

  private val publisher = new AmqpPublisher(attributes.destination, components)

  override def sendRequest(requestName: String, session: Session): Validation[Unit] = {
    for {
      props <- resolveProperties(attributes.messageProperties, session)
      message <- attributes.message
                  .amqpProtocolMessage(session)
                  .map(_.mergeProperties(props + ("deliveryMode" -> components.protocol.deliveryMode)))
      around <- aroundPublish(requestName, session, message)
    } yield {

      if (throttled) {
        throttler.throttle(
          session.scenario,
          () => publisher.publish(message, around, session)
        )
      } else publisher.publish(message, around, session)

    }

  }

  private def resolveProperties(
      properties: Map[Expression[String], Expression[Any]],
      session: Session
  ): Validation[Map[String, Any]] =
    properties.foldLeft(Map.empty[String, Any].success) {
      case (resolvedProperties, (key, value)) =>
        for {
          key                <- key(session)
          value              <- value(session)
          resolvedProperties <- resolvedProperties
        } yield resolvedProperties + (key -> value)
    }

  protected def aroundPublish(requestName: String, session: Session, message: AmqpProtocolMessage): Validation[Around]

} 

package ru.tinkoff.gatling.amqp.action

import io.gatling.commons.util.Clock
import io.gatling.commons.validation.Validation
import io.gatling.core.action.Action
import io.gatling.core.controller.throttle.Throttler
import io.gatling.core.session.Session
import io.gatling.core.stats.StatsEngine
import ru.tinkoff.gatling.amqp.protocol.AmqpComponents
import ru.tinkoff.gatling.amqp.request.{AmqpAttributes, AmqpProtocolMessage, _}

class RequestReply(
    attributes: AmqpAttributes,
    replyDestination: AmqpExchange,
    setAmqpReplyTo: Boolean,
    trackerDestination: Option[AmqpExchange],
    components: AmqpComponents,
    val statsEngine: StatsEngine,
    val clock: Clock,
    val next: Action,
    throttler: Throttler,
    throttled: Boolean
) extends AmqpAction(attributes, components, throttler, throttled) {

  private val replyTimeout = components.protocol.replyTimeout.getOrElse(0L)

  override def name: String = genName("amqpRequestReply")

  override protected def aroundPublish(requestName: String,
                                       session: Session,
                                       message: AmqpProtocolMessage): Validation[Around] = {
    resolveDestination(replyDestination, session).map { qName =>
      val tracker = components.trackerPool.tracker(qName,
                                                   components.protocol.consumersThreadCount,
                                                   components.protocol.messageMatcher,
                                                   components.protocol.responseTransformer)
      val id = components.protocol.messageMatcher.requestMatchId(message)
      Around(
        before = {

          if (logger.underlying.isDebugEnabled) {
            logMessage(s"Message sent matchId=$id", message)
          }
          tracker.track(id, clock.nowMillis, replyTimeout, attributes.checks, session, next, requestName)
        },
        after = {}
      )
    }

  }

  def resolveDestination(dest: AmqpExchange, session: Session): Validation[String] =
    dest match {
      case AmqpDirectExchange(name, _, _) => name(session)
      case AmqpQueueExchange(name, _)     => name(session)
    }

} 

package ru.tinkoff.gatling.amqp.action

import io.gatling.commons.stats.OK
import io.gatling.commons.util.Clock
import io.gatling.commons.validation.{Validation, _}
import io.gatling.core.action.Action
import io.gatling.core.config.GatlingConfiguration
import io.gatling.core.controller.throttle.Throttler
import io.gatling.core.session.Session
import io.gatling.core.stats.StatsEngine
import ru.tinkoff.gatling.amqp.protocol.AmqpComponents
import ru.tinkoff.gatling.amqp.request.{AmqpAttributes, AmqpProtocolMessage}

class Publish(attributes: AmqpAttributes,
              components: AmqpComponents,
              val statsEngine: StatsEngine,
              val clock: Clock,
              configuration: GatlingConfiguration,
              val next: Action,
              throttler: Throttler,
              throttled: Boolean)
    extends AmqpAction(attributes, components, throttler, throttled) {
  override val name: String = genName("amqpPublish")

  override protected def aroundPublish(requestName: String,
                                       session: Session,
                                       message: AmqpProtocolMessage): Validation[Around] =
    Around(
      before = {
        if (logger.underlying.isDebugEnabled) {
          logMessage(s"Message sent user=${session.userId} AMQPMessageID=${message.messageId}", message)
        }

        val now = clock.nowMillis

        statsEngine.logResponse(session, requestName, now, now, OK, None, None)

        next ! session
      },
      after = ()
    ).success
} 

package io.github.gatling.sql.protocol

import java.sql.Connection

import akka.actor.ActorSystem
import io.gatling.core
import io.gatling.core.CoreComponents
import io.gatling.core.config.GatlingConfiguration
import io.gatling.core.protocol.{Protocol, ProtocolComponents, ProtocolKey}
import io.gatling.core.session.Session


case class SqlProtocol(connection: Connection) extends Protocol {
  type Components = SqlComponents
}

object SqlProtocol {

  val SqlProtocolKey = new ProtocolKey {

    type Protocol = SqlProtocol
    type Components = SqlComponents

    override def protocolClass: Class[core.protocol.Protocol] = classOf[SqlProtocol].asInstanceOf[Class[io.gatling.core.protocol.Protocol]]

    override def defaultProtocolValue(configuration: GatlingConfiguration): SqlProtocol = throw new IllegalStateException("Can't provide a default value for SqlProtocol")

    override def newComponents(system: ActorSystem, coreComponents: CoreComponents): SqlProtocol => SqlComponents = {
      sqlProtocol => SqlComponents(sqlProtocol)
    }
  }
}

case class SqlComponents(sqlProtocol: SqlProtocol) extends ProtocolComponents {
  def onStart: Option[Session => Session] = None
  def onExit: Option[Session => Unit] = None
}

case class SqlProtocolBuilder(connection: Connection) {
  def build() = SqlProtocol(connection)
}

object SqlProtocolBuilder {
  def connection(connection: Connection) = SqlProtocolBuilder(connection)
} 

package io.github.gatling.sql.action

import io.github.gatling.sql.protocol.SqlProtocol
import io.github.gatling.sql.request.SqlAttributes
import io.gatling.commons.stats.{KO, OK}
import io.gatling.commons.util.ClockSingleton.nowMillis
import io.gatling.commons.validation.Validation
import io.gatling.core.action.builder.ActionBuilder
import io.gatling.core.action.{Action, ExitableAction}
import io.gatling.core.protocol.ProtocolComponentsRegistry
import io.gatling.core.session.Session
import io.gatling.core.stats.StatsEngine
import io.gatling.core.stats.message.ResponseTimings
import io.gatling.core.structure.ScenarioContext
import io.gatling.core.util.NameGen
import java.sql.PreparedStatement

import scala.concurrent.ExecutionContext.Implicits.global
import scala.concurrent.Future

class SqlActionBuilder(attr: SqlAttributes) extends ActionBuilder with NameGen {

  private def components(protocolComponentsRegistry: ProtocolComponentsRegistry) =
    protocolComponentsRegistry.components(SqlProtocol.SqlProtocolKey)

  override def build(ctx: ScenarioContext, next: Action): Action = {
    import ctx._
    val statsEngine = coreComponents.statsEngine
    val sqlComponents = components(protocolComponentsRegistry)
    new SqlAction(genName(s"SQL: ${attr.tag}"), sqlComponents.sqlProtocol, statsEngine, next, attr)
  }

}

class SqlAction(val name: String, protocol: SqlProtocol, val statsEngine: StatsEngine, val next: Action,
                val attr: SqlAttributes) extends ExitableAction {

  def execute(session: Session): Unit = {
    val stmt: Validation[PreparedStatement] = attr.statement(session)

    stmt.onFailure(err => {
      statsEngine.logResponse(session, name, ResponseTimings(nowMillis, nowMillis), KO, None, Some("Error setting up statement: " + err), Nil)
      next ! session.markAsFailed
    })

    stmt.onSuccess({ stmt =>
      val start = nowMillis

      val result = Future {
        stmt.execute()
      }

      result.onFailure { case t =>
        statsEngine.reportUnbuildableRequest(session, name, t.getMessage)
      }

      result.onSuccess { case result =>

        val requestEndDate = nowMillis

        statsEngine.logResponse(
          session,
          name,
          ResponseTimings(startTimestamp = start, endTimestamp = requestEndDate),
          if (result) OK else KO,
          None,
          if (result) None else Some("Failed... TBD")
        )

        next ! session.markAsSucceeded
      }
    })
  }
} 

package io.github.gatling.sql

import java.sql.{Connection, PreparedStatement}

import com.typesafe.scalalogging.StrictLogging
import io.github.gatling.sql.db.ConnectionPool
import io.gatling.commons.validation.Validation
import io.gatling.core.session.{Expression, Session}
import io.gatling.commons.validation._

trait SqlStatement extends StrictLogging {

  def apply(session:Session): Validation[PreparedStatement]

  def connection = ConnectionPool.connection
}

case class SimpleSqlStatement(statement: Expression[String]) extends SqlStatement {
  def apply(session: Session): Validation[PreparedStatement] = statement(session).flatMap { stmt =>
      logger.debug(s"STMT: ${stmt}")
      connection.prepareStatement(stmt).success
    }
} 

package com.ringcentral.gatling.mongo.action

import com.ringcentral.gatling.mongo.command.MongoRawCommand
import com.ringcentral.gatling.mongo.response.MongoStringResponse
import io.gatling.commons.stats.KO
import io.gatling.commons.util.TimeHelper.nowMillis
import io.gatling.commons.validation.Validation
import io.gatling.core.action.Action
import io.gatling.core.config.GatlingConfiguration
import io.gatling.core.session.{Expression, Session}
import io.gatling.core.stats.StatsEngine
import play.api.libs.json.JsObject
import reactivemongo.api.commands.Command
import reactivemongo.api.{DefaultDB, FailoverStrategy, ReadPreference}
import reactivemongo.play.json.ImplicitBSONHandlers._
import reactivemongo.play.json.JSONSerializationPack

import scala.concurrent.ExecutionContext.Implicits.global
import scala.util.{Failure, Success}

class MongoRawCommandAction(command: MongoRawCommand, database: DefaultDB, val statsEngine: StatsEngine, configuration: GatlingConfiguration, val next: Action) extends MongoAction(database) {

  override def name: String = genName("Mongo raw command")

  override def commandName: Expression[String] = command.commandName

  override def executeCommand(commandName: String, session: Session): Validation[Unit] = for {
    commandText <- command.command(session)
    commandDocument <- string2JsObject(commandText)
  } yield {
    val sent = nowMillis
    val runner = Command.run(JSONSerializationPack, FailoverStrategy.default)
    runner.apply(database, runner.rawCommand(commandDocument)).one[JsObject](ReadPreference.primaryPreferred).onComplete {
      case Success(result) => processResult(session, sent, nowMillis, command.checks, MongoStringResponse(result.toString()), next, commandName)
      case Failure(err) => executeNext(session, sent, nowMillis, KO, next, commandName, Some(err.getMessage))
    }
  }
} 

package com.ringcentral.gatling.mongo.action

import com.ringcentral.gatling.mongo.command.MongoUpdateCommand
import com.ringcentral.gatling.mongo.response.MongoCountResponse
import io.gatling.commons.stats.KO
import io.gatling.commons.util.TimeHelper.nowMillis
import io.gatling.commons.validation.Validation
import io.gatling.core.action.Action
import io.gatling.core.config.GatlingConfiguration
import io.gatling.core.session.{Expression, Session}
import io.gatling.core.stats.StatsEngine
import reactivemongo.api.DefaultDB
import reactivemongo.play.json.ImplicitBSONHandlers._
import reactivemongo.play.json.collection.JSONCollection

import scala.concurrent.ExecutionContext.Implicits.global
import scala.util.{Failure, Success}

class MongoUpdateAction(command: MongoUpdateCommand, database: DefaultDB, val statsEngine: StatsEngine, configuration: GatlingConfiguration, val next: Action) extends MongoAction(database) {

  override def name: String = genName("Mongo update command")

  override def commandName: Expression[String] = command.commandName

  override def executeCommand(commandName: String, session: Session): Validation[Unit] = for {
    collectionName <- command.collection(session)
    resolvedSelector <- command.selector(session)
    resolvedModifier <- command.modifier(session)
    upsert <- command.upsert(session)
    multi <- command.multi(session)
    selector <- string2JsObject(resolvedSelector)
    modifier <- string2JsObject(resolvedModifier)
  } yield {
    val sent = nowMillis
    database.collection[JSONCollection](collectionName).update(selector, modifier, upsert = upsert, multi = multi).onComplete {
      case Success(result) =>
        if (result.ok) {
          processResult(session, sent, nowMillis, command.checks, MongoCountResponse(result.n), next, commandName)
        } else {
          executeNext(session, sent, nowMillis, KO, next, commandName, Some(result.writeErrors.map(we => s"[${we.code}] ${we.errmsg}").mkString(", ")))
        }
      case Failure(err) =>
        executeNext(session, sent, nowMillis, KO, next, commandName, Some(err.getMessage))
    }

  }
} 

package com.ringcentral.gatling.mongo.action

import com.ringcentral.gatling.mongo.check.MongoCheck
import com.ringcentral.gatling.mongo.response.MongoResponse
import io.gatling.commons.stats.{KO, OK, Status}
import io.gatling.commons.validation
import io.gatling.commons.validation.{NoneSuccess, Validation}
import io.gatling.core.action.{Action, ExitableAction}
import io.gatling.core.check.Check
import io.gatling.core.session.{Expression, Session}
import io.gatling.core.stats.message.ResponseTimings
import io.gatling.core.util.NameGen
import play.api.libs.json._
import reactivemongo.api.DefaultDB
import reactivemongo.api.collections.GenericQueryBuilder
import reactivemongo.play.json.JSONSerializationPack

import scala.util.{Failure, Success, Try}

abstract class MongoAction(database: DefaultDB) extends ExitableAction with NameGen {

  def commandName: Expression[String]

  def executeCommand(commandName: String, session: Session): Validation[Unit]

  override def execute(session: Session): Unit = recover(session) {
    commandName(session).flatMap { resolvedCommandName =>
      val outcome = executeCommand(resolvedCommandName, session)
      outcome.onFailure(errorMessage => statsEngine.reportUnbuildableRequest(session, resolvedCommandName, errorMessage))
      outcome
    }
  }

  def string2JsObject(string: String): Validation[JsObject] = {
    Try[JsObject](Json.parse(string).as[JsObject]) match {
      case Success(json) => validation.SuccessWrapper(json).success
      case Failure(err) =>
        validation.FailureWrapper(s"Error parse JSON string: $string. ${err.getMessage}").failure
    }
  }

  def string2JsObject(optionString: Option[String]): Validation[Option[JsObject]] =
    optionString match {
      case Some(string) => string2JsObject(string).map(Some.apply)
      case None => NoneSuccess
    }

  protected def executeNext(session: Session,
                            sent: Long,
                            received: Long,
                            status: Status,
                            next: Action,
                            requestName: String,
                            message: Option[String]): Unit = {
    val timings = ResponseTimings(sent, received)
    statsEngine.logResponse(session, requestName, timings, status, None, message)
    next ! session
  }

  protected def processResult(session: Session,
                              sent: Long,
                              received: Long,
                              checks: List[MongoCheck],
                              response: MongoResponse,
                              next: Action,
                              requestName: String): Unit = {
    // run all the checks, advise the Gatling API that it is complete and move to next
    val (checkSaveUpdate, error) = Check.check(response, session, checks)
    val newSession = checkSaveUpdate(session)
    error match {
      case Some(validation.Failure(errorMessage)) => executeNext(newSession.markAsFailed, sent, received, KO, next, requestName, Some(errorMessage))
      case _ => executeNext(newSession, sent, received, OK, next, requestName, None)
    }
  }

  implicit class GenericQueryBuilderExt(b: GenericQueryBuilder[JSONSerializationPack.type]) {

    def sort(sort: Option[JsObject]): GenericQueryBuilder[JSONSerializationPack.type] = {
      sort.map(b.sort).getOrElse(b)
    }

    def hint(sort: Option[JsObject]): GenericQueryBuilder[JSONSerializationPack.type] = {
      sort.map(b.hint).getOrElse(b)
    }
  }

} 

package ru.tinkoff.gatling.amqp.client

import akka.actor.ActorRef
import io.gatling.core.action.Action
import io.gatling.core.session.Session
import ru.tinkoff.gatling.amqp.AmqpCheck
import ru.tinkoff.gatling.amqp.client.AmqpMessageTrackerActor.MessagePublished

class AmqpMessageTracker(actor: ActorRef) {

  def track(
      matchId: String,
      sent: Long,
      replyTimeout: Long,
      checks: List[AmqpCheck],
      session: Session,
      next: Action,
      requestName: String
  ): Unit =
    actor ! MessagePublished(
      matchId,
      sent,
      replyTimeout,
      checks,
      session,
      next,
      requestName
    )
} 

package com.ringcentral.gatling.mongo.action

import com.ringcentral.gatling.mongo.command.MongoCountCommand
import com.ringcentral.gatling.mongo.response.MongoCountResponse
import io.gatling.commons.stats.KO
import io.gatling.commons.util.TimeHelper.nowMillis
import io.gatling.commons.validation._
import io.gatling.core.action.Action
import io.gatling.core.config.GatlingConfiguration
import io.gatling.core.session.{Expression, Session, _}
import io.gatling.core.stats.StatsEngine
import reactivemongo.api.DefaultDB
import reactivemongo.play.json.collection.JSONCollection

//TODO remove global context everywhere
import scala.concurrent.ExecutionContext.Implicits.global
import scala.util.{Failure, Success}

class MongoCountAction(command: MongoCountCommand, database: DefaultDB, val statsEngine: StatsEngine, configuration: GatlingConfiguration, val next: Action) extends MongoAction(database) {

  override def name: String = genName("Mongo count command")

  override def commandName: Expression[String] = command.commandName

  override def executeCommand(commandName: String, session: Session): Validation[Unit] =
    for {
      collectionName <- command.collection(session)
      selectorDocument <- resolveOptionalExpression(command.selector, session)
      hint <- resolveOptionalExpression(command.hint, session)
      selector <- selectorDocument match {
        case Some(d) => string2JsObject(d).map(Some.apply)
        case None => NoneSuccess
      }
    } yield {
      val sent = nowMillis
      val collection: JSONCollection = database.collection[JSONCollection](collectionName)
      collection.count(selector, command.limit, command.skip, hint).onComplete {
        case Success(result) => processResult(session, sent, nowMillis, command.checks, MongoCountResponse(result), next, commandName)
        case Failure(err) => executeNext(session, sent, nowMillis, KO, next, commandName, Some(err.getMessage))
      }

    }
} 

package com.ringcentral.gatling.mongo.action

import com.ringcentral.gatling.mongo.command.MongoRemoveCommand
import com.ringcentral.gatling.mongo.response.MongoCountResponse
import io.gatling.commons.stats.KO
import io.gatling.commons.util.TimeHelper.nowMillis
import io.gatling.commons.validation.Validation
import io.gatling.core.action.Action
import io.gatling.core.config.GatlingConfiguration
import io.gatling.core.session.{Expression, Session}
import io.gatling.core.stats.StatsEngine
import reactivemongo.api.DefaultDB
import reactivemongo.play.json.ImplicitBSONHandlers._
import reactivemongo.play.json.collection.JSONCollection

import scala.concurrent.ExecutionContext.Implicits.global
import scala.util.{Failure, Success}

class MongoRemoveAction(command: MongoRemoveCommand, database: DefaultDB, val statsEngine: StatsEngine, configuration: GatlingConfiguration, val next: Action) extends MongoAction(database) {

  override def name: String = genName("Mongo find command")

  override def commandName: Expression[String] = command.commandName

  override def executeCommand(commandName: String, session: Session): Validation[Unit] = for {
    collectionName <- command.collection(session)
    resolvedSelector <- command.selector(session)
    selector <- string2JsObject(resolvedSelector)
  } yield {
    val sent = nowMillis
    database.collection[JSONCollection](collectionName).remove(selector).onComplete {
      case Success(result) =>
        if (result.ok) {
          processResult(session, sent, nowMillis, command.checks, MongoCountResponse(result.n), next, commandName)
        } else {
          executeNext(session, sent, nowMillis, KO, next, commandName, Some(result.writeErrors.map(we => s"[${we.code}] ${we.errmsg}").mkString(", ")))
        }

      case Failure(err) => executeNext(session, sent, nowMillis, KO, next, commandName, Some(err.getMessage))

    }

  }
} 

package com.ringcentral.gatling.mongo.action

import com.ringcentral.gatling.mongo.command.MongoFindCommand
import com.ringcentral.gatling.mongo.response.MongoStringResponse
import io.gatling.commons.stats.KO
import io.gatling.commons.util.TimeHelper.nowMillis
import io.gatling.commons.validation.Validation
import io.gatling.core.action.Action
import io.gatling.core.config.GatlingConfiguration
import io.gatling.core.session.{Expression, Session, resolveOptionalExpression}
import io.gatling.core.stats.StatsEngine
import play.api.libs.json.JsObject
import reactivemongo.api.{DefaultDB, QueryOpts, ReadPreference}
import reactivemongo.play.json.ImplicitBSONHandlers._
import reactivemongo.play.json.collection.JSONCollection
import reactivemongo.play.json.collection.JsCursor._

import scala.concurrent.ExecutionContext.Implicits.global
import scala.util.{Failure, Success}

class MongoFindAction(command: MongoFindCommand, database: DefaultDB, val statsEngine: StatsEngine, configuration: GatlingConfiguration, val next: Action) extends MongoAction(database) {

  override def name: String = genName("Mongo find command")

  override def commandName: Expression[String] = command.commandName

  override def executeCommand(commandName: String, session: Session): Validation[Unit] =
    for {
      collectionName <- command.collection(session)
      resolvedFilter <- command.query(session)
      filter <- string2JsObject(resolvedFilter)
      resolvedHint <- resolveOptionalExpression(command.hint, session)
      hint <- string2JsObject(resolvedHint)
      resolvedSort <- resolveOptionalExpression(command.sort, session)
      sort <- string2JsObject(resolvedSort)
    } yield {
      val sent = nowMillis
      database.collection[JSONCollection](collectionName).find(filter).options(QueryOpts().batchSize(command.limit)).sort(sort).hint(hint)
        .cursor[JsObject](ReadPreference.primary).jsArray(command.limit).onComplete {
        case Success(result) => processResult(session, sent, nowMillis, command.checks, MongoStringResponse(result.toString()), next, commandName)
        case Failure(err) => executeNext(session, sent, nowMillis, KO, next, commandName, Some(err.getMessage))
      }
    }
} 

package com.github.phisgr.gatling

import com.google.protobuf.Message
import io.gatling.commons.validation.{Failure, Success, Validation}
import io.gatling.core.Predef.value2Expression
import io.gatling.core.session.{Expression, Session}

package object javapb {
  private type BuilderMutation[B] = (B, Session) => Failure // nullable

  private case class MutationWithExpression[B <: Message.Builder, F](setter: B => F => Message.Builder, e: Expression[F])
    extends BuilderMutation[B] {
    override def apply(builder: B, session: Session): Failure = {
      e(session) match {
        case Success(value) =>
          setter(builder)(value)
          null
        case f@Failure(_) =>
          f
      }
    }
  }

  private[javapb] class MessageExpressionUpdater[M <: Message, B <: Message.Builder](
    private[javapb] val original: Expression[M],
    private[javapb] val reversedMutations: List[BuilderMutation[B]]
  ) {
    def update[Field](setter: B => Field => Message.Builder)(value: Expression[Field]) =
      new MessageExpressionUpdater[M, B](original, MutationWithExpression(setter, value) :: reversedMutations)
  }

  implicit def toExpression[M <: Message, B <: Message.Builder](u: MessageExpressionUpdater[M, B]): Expression[M] = {
    val mutations = u.reversedMutations.reverse.toArray
    val size = mutations.length
    val original = u.original

    s =>
      original(s) match {
        case Success(message) =>
          val builder = message.toBuilder.asInstanceOf[B]
          var ret: Validation[M] = null
          var i = 0
          do {
            if (i < size) {
              ret = mutations(i)(builder, s)
              i += 1
            } else {
              ret = Success(builder.build().asInstanceOf[M])
            }
          } while (ret eq null)
          ret
        case f@Failure(_) => f
      }
  }

  type BuilderOf[M <: Message] = {def build(): M}

  implicit def message2ExprUpdater[M <: Message, B <: Message.Builder with BuilderOf[M]](e: M)(implicit evidence: BuilderEvidence[M, B]): MessageExpressionUpdater[M, B] =
    expr2exprUpdater(value2Expression(e))

  implicit def expr2exprUpdater[M <: Message, B <: Message.Builder with BuilderOf[M]](e: Expression[M])(implicit evidence: BuilderEvidence[M, B]): MessageExpressionUpdater[M, B] =
    new MessageExpressionUpdater[M, B](e, Nil)

  implicit def builderEvidence[M <: Message, B]: BuilderEvidence[M, B] = macro BuilderEvidence.impl[M, B]

} 

package com.github.phisgr.gatling.pb.bench

import java.util.concurrent.TimeUnit

import com.github.phisgr.gatling.grpc.Predef._
import com.github.phisgr.gatling.javapb._
import com.github.phisgr.gatling.pb.Test
import com.github.phisgr.gatling.pb.test._
import io.gatling.commons.validation.Validation
import io.gatling.core.session.{Expression, Session}
import org.openjdk.jmh.annotations.{Benchmark, OutputTimeUnit}

@OutputTimeUnit(TimeUnit.MILLISECONDS)
class TestUpdateExpr {

  import TestUpdateExpr._

  @Benchmark
  def updateSimpleExpr(): Validation[SimpleMessage] = {
    SimpleExpr(Session1)
  }

  @Benchmark
  def updateComplexExpr(): Validation[ComplexMessage] = {
    ComplexExpr(Session1)
  }

  @Benchmark
  def updateSimpleExprJava(): Validation[Test.SimpleMessage] = {
    SimpleExprJava(Session1)
  }

  @Benchmark
  def updateComplexExprJava(): Validation[Test.ComplexMessage] = {
    ComplexExprJava(Session1)
  }

  @Benchmark
  def lambdaSimpleExprJava(): Validation[Test.SimpleMessage] = {
    SimpleExprJavaLambda(Session1)
  }

  @Benchmark
  def lambdaComplexExprJava(): Validation[Test.ComplexMessage] = {
    ComplexExprJavaLambda(Session1)
  }
}

object TestUpdateExpr {
  private val SimpleExpr = SimpleMessage.defaultInstance.updateExpr(
    _.s :~ $("name")
  )

  private val SimpleExprJava: Expression[Test.SimpleMessage] =
    Test.SimpleMessage.getDefaultInstance
      .update(_.setS)($("name"))

  private val SimpleExprJavaLambda: Expression[Test.SimpleMessage] = { s: Session =>
    for {
      name <- s("name").validate[String]
    } yield {
      val builder = Test.SimpleMessage.newBuilder()
      builder.setS(name)
      builder.build()
    }
  }

  private val ComplexExpr = ComplexMessage.defaultInstance.updateExpr(
    _.m.s :~ $("name"),
    _.i :~ $("count"),
  )

  private val ComplexExprJava: Expression[Test.ComplexMessage] =
    Test.ComplexMessage.getDefaultInstance
      .update(_.getMBuilder.setS)($("name"))
      .update(_.setI)($("count"))

  private val ComplexExprJavaLambda: Expression[Test.ComplexMessage] = { s: Session =>
    for {
      name <- s("name").validate[String]
      count <- s("count").validate[Int]
    } yield {
      val builder = Test.ComplexMessage.newBuilder()
      builder.getMBuilder.setS(name)
      builder.setI(count)
      builder.build()
    }
  }

  private val Session1 = Session(
    scenario = "Scenario",
    userId = 1L,
    attributes = Map(
      "name" -> "Asdf Qwer",
      "count" -> 123
    ),
    startDate = System.currentTimeMillis()
  )
} 

package com.github.phisgr.gatling.grpc.action

import com.github.phisgr.gatling.grpc.HeaderPair
import com.github.phisgr.gatling.grpc.check.{GrpcCheck, ResponseExtract}
import com.github.phisgr.gatling.grpc.protocol.GrpcProtocol
import io.gatling.commons.validation.Success
import io.gatling.core.action.Action
import io.gatling.core.action.builder.ActionBuilder
import io.gatling.core.check.{MultipleFindCheckBuilder, ValidatorCheckBuilder}
import io.gatling.core.session.{Expression, ExpressionSuccessWrapper, Session}
import io.gatling.core.structure.ScenarioContext
import io.grpc.{CallOptions, Metadata, MethodDescriptor}

import scala.collection.breakOut

case class GrpcCallActionBuilder[Req, Res](
  requestName: Expression[String],
  method: MethodDescriptor[Req, Res],
  payload: Expression[Req],
  callOptions: Expression[CallOptions] = CallOptions.DEFAULT.expressionSuccess,
  reversedHeaders: List[HeaderPair[_]] = Nil,
  checks: List[GrpcCheck[Res]] = Nil,
  protocolOverride: Option[GrpcProtocol] = None,
  isSilent: Boolean = false
) extends ActionBuilder {
  override def build(ctx: ScenarioContext, next: Action): Action = GrpcCallAction(this, ctx, next)

  def callOptions(callOptions: Expression[CallOptions]) = copy(
    callOptions = callOptions
  )

  
  def callOptions(callOptions: => CallOptions) = copy(
    callOptions = { _: Session => Success(callOptions) }
  )

  def header[T](key: Metadata.Key[T])(value: Expression[T]) = copy(
    reversedHeaders = HeaderPair(key, value) :: reversedHeaders
  )

  private def mapToList[T, U](s: Seq[T])(f: T => U) = s.map[U, List[U]](f)(breakOut)

  def check(checks: GrpcCheck[Res]*) = copy(
    checks = this.checks ::: checks.toList
  )

  // In fact they can be added to checks using .check
  // but the type Res cannot be inferred there
  def extract[X](
    f: Res => Option[X])(
    ts: (ValidatorCheckBuilder[ResponseExtract, Res, X] => GrpcCheck[Res])*
  ) = {
    val e = ResponseExtract.extract(f)
    copy(
      checks = checks ::: mapToList(ts)(_.apply(e))
    )
  }

  def exists[X](f: Res => Option[X]) = extract(f)(_.exists.build(ResponseExtract.materializer))

  def extractMultiple[X](
    f: Res => Option[Seq[X]])(
    ts: (MultipleFindCheckBuilder[ResponseExtract, Res, X] => GrpcCheck[Res])*
  ) = {
    val e = ResponseExtract.extractMultiple[Res, X](f)
    copy(
      checks = checks ::: mapToList(ts)(_.apply(e))
    )
  }

  def target(protocol: GrpcProtocol) = copy(protocolOverride = Some(protocol))

  def silent = copy(isSilent = true)

} 

package com.github.jeanadrien.gatling.mqtt.protocol

import com.github.jeanadrien.gatling.mqtt.client.MqttQoS.MqttQoS
import com.github.jeanadrien.gatling.mqtt.client.{MqttClientConfiguration, Will}
import io.gatling.commons.validation.{Failure, Success, Validation}
import io.gatling.core.session.{Session, _}
import com.github.jeanadrien.gatling.mqtt.client.ConfigurationUtils._


case class ConnectionSettings(
    clientId     : Option[Expression[String]],
    cleanSession : Option[Boolean],
    userName     : Option[Expression[String]],
    password     : Option[Expression[String]],
    willTopic    : Option[Expression[String]],
    willMessage  : Option[Expression[String]],
    willQos      : Option[MqttQoS],
    willRetain   : Option[Boolean]
) {

    private def configureWill(session : Session)
        (mqtt1 : MqttClientConfiguration) : Validation[MqttClientConfiguration] = (willTopic, willMessage) match {
        case (Some(wte), Some(wme)) =>
            for {
                wt <- wte(session)
                wm <- wme(session)
            } yield {
                var will = Will(
                    topic = wt,
                    message = wm
                )
                will = willQos.map(qos => will.copy(qos = qos)).getOrElse(will)
                will = willRetain.map(r => will.copy(willRetain = r)).getOrElse(will)
                mqtt1.copy(will = Some(will))
            }
        case (None, None) =>
            Success(mqtt1)
        case _ =>
            Failure("Both will topic and message must be defined")
    }

    private[protocol] def configureMqtt(session : Session)
        (mqtt1 : MqttClientConfiguration) : Validation[MqttClientConfiguration] = {
        var mqtt = mqtt1
        mqtt = cleanSession.map(cs => mqtt.copy(cleanSession = cs)).getOrElse(mqtt)

        Success(mqtt).flatMap { mqtt =>
            realize[String, MqttClientConfiguration](clientId) { (config, value) => config.copy(clientId = value)
            }(mqtt, session)
        } flatMap { mqtt =>
            realize[String, MqttClientConfiguration](userName) { (config, value) => config.copy(username = value)
            }(mqtt, session)
        } flatMap { mqtt =>
            realize[String, MqttClientConfiguration](password) { (config, value) => config.copy(password = value)
            }(mqtt, session)
        } flatMap {
            configureWill(session)
        }
    }

} 

package com.github.jeanadrien.gatling.mqtt.actions

import akka.actor.ActorRef
import akka.pattern.AskTimeoutException
import akka.util.Timeout
import com.github.jeanadrien.gatling.mqtt.client.MqttCommands
import com.github.jeanadrien.gatling.mqtt.protocol.MqttComponents
import io.gatling.core.CoreComponents
import io.gatling.core.action.Action
import io.gatling.core.session.Session

import scala.concurrent.duration.FiniteDuration
import scala.util.{Failure, Success}


class WaitForMessagesAction(
    mqttComponents : MqttComponents,
    coreComponents : CoreComponents,
    timeout        : FiniteDuration,
    val next       : Action
) extends MqttAction(mqttComponents, coreComponents) {

    import akka.pattern.ask
    import mqttComponents.system.dispatcher

    override val name = genName("mqttWaitForMessage")

    override def execute(session : Session) : Unit = recover(session)(for {
        connection <- session("engine").validate[ActorRef]
        connectionId <- session("connectionId").validate[String]
    } yield {
        implicit val messageTimeout = Timeout(timeout)
        logger.debug(s"Wait for message action started... Timeout: ${timeout}")
        (connection ? MqttCommands.WaitForMessages).mapTo[MqttCommands] onComplete {
            case Success(MqttCommands.WaitForMessagesDone) =>
                logger.info(s"${connectionId} : Done waitForMessage.")
                next ! session
            case Failure(t) if t.isInstanceOf[AskTimeoutException] =>
                logger.warn("Wait for remaining messages timed out")
                next ! session.markAsFailed
            case Failure(t) =>
                logger.warn("Wait for remaining messages error:", t)
                next ! session.markAsFailed
        }
    })
} 

package io.gatling.amqp.event

import io.gatling.amqp.data._
import io.gatling.core.session.Session

sealed trait AmqpAction
object AmqpPublishAction extends AmqpAction
object AmqpConsumeAction extends AmqpAction

sealed trait AmqpEvent {
  def action: AmqpAction
}


abstract class AmqpConsumevent extends AmqpEvent {
  def action: AmqpAction = AmqpConsumeAction
}
case class AmqpConsumeRequest(req: ConsumeRequest, session: Session) extends AmqpConsumevent

case class AmqpConsuming(consumerName: String, no: Int, startedAt: Long, req: ConsumeRequest, session: Session) extends AmqpConsumevent {
  def eventId: String = s"$consumerName-$no"
}
case class AmqpConsumed(consumerName: String, no: Int, stoppedAt: Long, event: AmqpConsuming) extends AmqpConsumevent {
  def eventId: String = s"$consumerName-$no"
}
case class AmqpConsumeFailed(consumerName: String, no: Int, stoppedAt: Long, e: Throwable) extends AmqpConsumevent {
  def eventId: String = s"$consumerName-$no"
}
case class AmqpConsumeAcked(consumerName: String, no: Int, multiple: Boolean, stoppedAt: Long) extends AmqpConsumevent {
  def eventId: String = s"$consumerName-$no"
}
case class AmqpConsumeNacked(consumerName: String, no: Int, multiple: Boolean, stoppedAt: Long) extends AmqpConsumevent {
  def eventId: String = s"$consumerName-$no"
} 

package io.gatling.keycloak

import akka.actor.ActorRef
import akka.actor.ActorDSL._
import io.gatling.core.action.builder.ActionBuilder
import io.gatling.core.action.Interruptable
import io.gatling.core.config.Protocols
import io.gatling.core.result.writer.DataWriterClient
import io.gatling.core.session.{Expression, Session}
import io.gatling.core.validation.Validation
import org.keycloak.admin.client.resource.RealmResource
import org.keycloak.representations.idm.UserRepresentation

import scala.collection.JavaConversions._

case class FindUserAttributes(
  requestName: Expression[String],
  realm: Expression[RealmResource],
  username: Option[Expression[String]] = None,
  firstName: Option[Expression[String]] = None,
  lastName: Option[Expression[String]] = None,
  email: Option[Expression[String]] = None,
  firstResult: Option[Expression[Integer]] = None,
  maxResults: Option[Expression[Integer]] = None,
  use: Option[(Session, List[UserRepresentation]) => Session] = None
) {}

object FindUserBuilder {
  implicit def toActionBuilder(builder: FindUserBuilder) = new FindUserActionBuilder(builder.attributes)
}

class FindUserBuilder(private val attributes: FindUserAttributes) {
  def username(username: Expression[String]) = new FindUserBuilder(attributes.copy(username = Some(username)))
  def firstName(firstName: Expression[String]) = new FindUserBuilder(attributes.copy(firstName = Some(firstName)))
  def lastName(lastName: Expression[String]) = new FindUserBuilder(attributes.copy(lastName = Some(lastName)))
  def email(email: Expression[String]) = new FindUserBuilder(attributes.copy(email = Some(email)))
  def firstResult(firstResult: Expression[Integer]) = new FindUserBuilder(attributes.copy(firstResult = Some(firstResult)))
  def maxResults(maxResults: Expression[Integer]) = new FindUserBuilder(attributes.copy(maxResults = Some(maxResults)))
  def use(use: (Session, List[UserRepresentation]) => Session) = new FindUserActionBuilder(attributes.copy(use = Some(use)))
}

class FindUserActionBuilder(attributes: FindUserAttributes) extends ActionBuilder {
  override def build(next: ActorRef, protocols: Protocols): ActorRef =
    actor(actorName("find-user"))(new FindUserAction(attributes, next))
}

class FindUserAction(
                      attributes: FindUserAttributes,
                      val next: ActorRef
                    ) extends Interruptable with ExitOnFailure with DataWriterClient {
  override def executeOrFail(session: Session): Validation[_] =
    attributes.realm(session).flatMap(realm =>
      Blocking(() =>
        Stopwatch(() => {
          realm.users().search(
            attributes.username.map(a => a(session).get).orNull,
            attributes.firstName.map(a => a(session).get).orNull,
            attributes.lastName.map(a => a(session).get).orNull,
            attributes.email.map(a => a(session).get).orNull,
            attributes.firstResult.map(a => a(session).get).orNull,
            attributes.maxResults.map(a => a(session).get).orNull)
        })
          .recordAndContinue(this, session, attributes.requestName(session).get, users => {
            attributes.use.map(use => use(session, users.toList)).getOrElse(session)
          })
      )
    )
} 

package io.gatling.keycloak

import akka.actor.ActorRef
import io.gatling.core.action.Interruptable
import io.gatling.core.result.writer.DataWriterClient
import io.gatling.core.session.{Expression, Session}
import io.gatling.core.validation.Validation
import org.keycloak.admin.client.resource.RealmResource


class RemoveUserAction(
                        requestName: Expression[String],
                        realm: Expression[RealmResource],
                        userId: Expression[String],
                        val next: ActorRef
                      ) extends Interruptable with ExitOnFailure with DataWriterClient {
  override def executeOrFail(session: Session): Validation[_] = {
    realm(session).map(realm =>
      Blocking(() =>
        Stopwatch(() => realm.users.get(userId(session).get).remove())
          .recordAndContinue(this, session, requestName(session).get)
      )
    )
  }
} 

package io.gatling.keycloak

import javax.ws.rs.core.{HttpHeaders, Response}

import akka.actor.ActorDSL._
import akka.actor.ActorRef
import io.gatling.core.action.Interruptable
import io.gatling.core.action.builder.ActionBuilder
import io.gatling.core.config.Protocols
import io.gatling.core.result.writer.DataWriterClient
import io.gatling.core.session.{Expression, Session}
import io.gatling.core.validation.{Success, Validation}
import org.keycloak.admin.client.resource.RealmResource
import org.keycloak.representations.idm.{CredentialRepresentation, UserRepresentation}

case class AddUserAttributes(
  requestName: Expression[String],
  realm: Expression[RealmResource],
  username: Expression[String],
  enabled: Expression[Boolean] = _ => Success(true),
  firstName: Option[Expression[String]] = None,
  lastName: Option[Expression[String]] = None,
  email: Option[Expression[String]] = None,
  password: Option[Expression[String]] = None,
  passwordTemporary: Option[Expression[Boolean]] = None,
  save: Option[(Session, String) => Session] = None
) {}


case class AddUserActionBuilder(
                                 attributes: AddUserAttributes
                               ) extends ActionBuilder {
  def newInstance(attributes: AddUserAttributes) = new AddUserActionBuilder(attributes)

  def enabled(enabled: Expression[Boolean]) = newInstance(attributes.copy(enabled = enabled))
  def password(password: Expression[String], temporary: Expression[Boolean]) =
    newInstance(attributes.copy(password = Some(password), passwordTemporary = Some(temporary)))
  def firstName(firstName: Expression[String]) = newInstance(attributes.copy(firstName = Some(firstName)))
  def lastName(lastName: Expression[String]) = newInstance(attributes.copy(lastName = Some(lastName)))
  def email(email: Expression[String]) = newInstance(attributes.copy(email = Some(email)))
  def saveWith[T](save: (Session, String) => Session) = newInstance(attributes.copy(save = Some(save)))

  def saveAs(name: String) = saveWith((session, id) => session.set(name, id))

  override def build(next: ActorRef, protocols: Protocols): ActorRef = {
    actor(actorName("add-user"))(new AddUserAction(attributes, next))
  }
}

class AddUserAction(
                     attributes: AddUserAttributes,
                     val next: ActorRef
                   ) extends Interruptable with ExitOnFailure with DataWriterClient {
  override def executeOrFail(session: Session): Validation[_] = {
    val user = new UserRepresentation
    user.setUsername(attributes.username(session).get)
    user.setEnabled(attributes.enabled(session).get)
    attributes.firstName.map(fn => user.setFirstName(fn(session).get))
    attributes.lastName.map(ln => user.setLastName(ln(session).get))
    attributes.email.map(e => user.setEmail(e(session).get))
    attributes.realm(session).map(realm =>
      Blocking(() =>
        Stopwatch(() => realm.users.create(user))
          .check(response => response.getStatus == 201, response => {
            val status = response.getStatusInfo.toString
            response.close()
            status
          })
          .recordAndStopOnFailure(this, session, attributes.requestName(session).get + ".create-user")
          .onSuccess(response => {
            val id = getUserId(response)
            response.close()
            val newSession = attributes.save.map(s => s(session, id)).getOrElse(session)

            attributes.password.map(password => {
              val credentials = new CredentialRepresentation
              credentials.setType(CredentialRepresentation.PASSWORD)
              credentials.setValue(password(newSession).get)
              attributes.passwordTemporary.map(a => credentials.setTemporary(a(newSession).get))
              Stopwatch(() => realm.users.get(id).resetPassword(credentials))
                .recordAndContinue(this, newSession, attributes.requestName(session).get + ".reset-password")
            }).getOrElse(next ! newSession)
          })
      )
    )
  }

  def getUserId(u: Response): String = {
    val location = u.getHeaderString(HttpHeaders.LOCATION)
    val lastSlash = location.lastIndexOf('/');
    if (lastSlash < 0) null else location.substring(lastSlash + 1)
  }
} 

package io.gatling.keycloak

import akka.actor.ActorDSL._
import akka.actor.ActorRef
import io.gatling.core.action.Interruptable
import io.gatling.core.action.builder.ActionBuilder
import io.gatling.core.config.Protocols
import io.gatling.core.result.writer.DataWriterClient
import io.gatling.core.session.{Session, Expression}
import io.gatling.core.validation.Validation


class RefreshTokenActionBuilder(requestName: Expression[String]) extends ActionBuilder{
  override def build(next: ActorRef, protocols: Protocols): ActorRef = {
    actor(actorName("refresh-token"))(new RefreshTokenAction(requestName, next))
  }
}

class RefreshTokenAction(
                          requestName: Expression[String],
                          val next: ActorRef
                        ) extends Interruptable with ExitOnFailure with DataWriterClient {
  override def executeOrFail(session: Session): Validation[_] = {
    val requestAuth: MockRequestAuthenticator = session(MockRequestAuthenticator.KEY).as[MockRequestAuthenticator]
    Blocking(() =>
      Stopwatch(() => requestAuth.getKeycloakSecurityContext.refreshExpiredToken(false))
        .check(identity, _ => "Could not refresh token")
        .recordAndContinue(this, session, requestName(session).get)
    )
  }
} 

package io.gatling.amqp.infra

import io.gatling.amqp.config._
import io.gatling.core.result.message.{KO, OK}
import io.gatling.core.session.Session

trait Stats { this: AmqpActor =>
  implicit val amqp: AmqpProtocol
  private lazy val statsEngine = amqp.statsEngine

  protected def statsOk(session: Session, startedAt: Long, stoppedAt: Long, title: String, code: Option[String] = None): Unit = {
    amqp.tracer ! WriteStat(session, startedAt, stoppedAt, title, OK, code, None)
  }

  protected def statsNg(session: Session, startedAt: Long, stoppedAt: Long, title: String, code: Option[String], mes: String): Unit = {
    amqp.tracer ! WriteStat(session, startedAt, stoppedAt, title, KO, code, Some(mes))
  }
} 

package io.gatling.amqp.infra

import akka.actor._
import io.gatling.amqp.config._
import io.gatling.amqp.event._
import io.gatling.core.result.message.{KO, OK, ResponseTimings, Status}
import io.gatling.core.result.writer.StatsEngine
import io.gatling.core.session.Session

case class MessageOk(event: AmqpPublishing, stoppedAt: Long, title: String)
case class MessageNg(event: AmqpPublishing, stoppedAt: Long, title: String, message: Option[String])
case class WriteStat(session: Session, startedAt: Long, stoppedAt: Long, title: String, status: Status, code: Option[String], mes: Option[String])


class AmqpTracer(statsEngine: StatsEngine)(implicit amqp: AmqpProtocol) extends Actor with Logging {
  def receive = {
    case WriteStat(session, startedAt, stoppedAt, title, status, code, mes) =>
      val timings = ResponseTimings(startedAt, stoppedAt, stoppedAt, stoppedAt)
      statsEngine.logResponse(session, title, timings, OK, code, mes)

    case MessageOk(event, stoppedAt, title) =>
      import event._
      val timings = ResponseTimings(startedAt, stoppedAt, stoppedAt, stoppedAt)
      statsEngine.logResponse(session, title, timings, OK, None, None)

    case MessageNg(event, stoppedAt, title, message) =>
      import event._
      val timings = ResponseTimings(startedAt, stoppedAt, stoppedAt, stoppedAt)
      statsEngine.logResponse(session, title, timings, KO, None, message)
  }
}

object AmqpTracer {
  def props(statsEngine : StatsEngine, amqp: AmqpProtocol) = Props(classOf[AmqpTracer], statsEngine, amqp)
} 

package io.gatling.amqp.infra

import akka.actor._
import akka.routing._
import io.gatling.amqp.config._
import io.gatling.amqp.data._
import io.gatling.amqp.event._
import io.gatling.core.result.writer.StatsEngine
import io.gatling.core.session.Session

import scala.collection.mutable
import scala.util._

class AmqpRouter(statsEngine: StatsEngine)(implicit amqp: AmqpProtocol) extends Actor with Logging {
  private var publishers = Router(RoundRobinRoutingLogic(), Vector[Routee]())

  // create one consumer for one session
  private val consumerActors = mutable.HashMap[String, ActorRef]()  // UserId -> ref(AmqpConsumer)
  private def consumerActorFor(session: Session): ActorRef = {
    val name = s"AmqpConsumer-user-${session.userId}"
    consumerActors.getOrElseUpdate(session.userId, context.actorOf(AmqpConsumer.props(name, session, amqp), name))
  }

  override def preStart(): Unit = {
    super.preStart()
  }

  private def initializePublishersOnce(): Unit = {
    if (publishers.routees.isEmpty) {
      for(i <- 1 to amqp.connection.poolSize) { addPublisher(i) }
    }
  }

  def receive: Receive = {
    case m: AmqpPublishRequest =>
      initializePublishersOnce()
      publishers.route(m, sender())

    case m: AmqpConsumeRequest =>
      consumerActorFor(m.session).forward(m)

    case m: WaitTermination if consumerActors.isEmpty =>
      sender() ! Success("no consumers")

    case m: WaitTermination =>
      consumerActorFor(m.session).forward(m)

    case Terminated(ref) =>
      publishers = publishers.removeRoutee(ref)
  }

  private def addPublisher(i: Int): Unit = {
    val name = s"AmqpPublisher-$i"
    val ref = context.actorOf(AmqpPublisher.props(name, amqp), name)
    context watch ref
    publishers = publishers.addRoutee(ref)
  }
}

object AmqpRouter {
  def props(statsEngine : StatsEngine, amqp: AmqpProtocol) = Props(classOf[AmqpRouter], statsEngine, amqp)
} 

package io.gatling.amqp.infra

import java.util.concurrent.atomic._

import akka.actor.Props
import com.rabbitmq.client._
import io.gatling.amqp.config._
import io.gatling.amqp.data._
import io.gatling.amqp.event._
import io.gatling.core.result.writer.StatsEngine
import io.gatling.core.session.Session
import io.gatling.core.util.TimeHelper.nowMillis

import scala.util._

class AmqpPublisher(actorName: String)(implicit amqp: AmqpProtocol) extends AmqpActor {
  private val nacker = amqp.nacker
  private val isConfirmMode = amqp.isConfirmMode
  private def sendEvent(event: AmqpEvent): Unit = nacker ! event
//  private def sendEvent(event: AmqpEvent): Unit = amqp.event.publish(event)

  override def preStart(): Unit = {
    super.preStart()
    if (isConfirmMode) {
      channel.confirmSelect()

      channel.addConfirmListener(new ConfirmListener() {
        def handleAck (no: Long, multi: Boolean): Unit = {
          sendEvent(AmqpPublishAcked (actorName, no.toInt, multi, nowMillis))
        }

        def handleNack(no: Long, multi: Boolean): Unit =
          sendEvent(AmqpPublishNacked(actorName, no.toInt, multi, nowMillis))
      })
    }
  }

  private val localPublishSeqNoCounter = new AtomicInteger(1)
  private def getNextPublishSeqNo: Int = {
    if (isConfirmMode)
      channel.getNextPublishSeqNo.toInt
    else
      localPublishSeqNoCounter.getAndIncrement
  }

  override def receive = {
    case AmqpPublishRequest(req, session) if isConfirmMode =>
      publishAsync(req, session)

    case AmqpPublishRequest(req, session) =>
      publishSync(req, session)
  }

  protected def publishSync(req: PublishRequest, session: Session): Unit = {
    import req._
    val startedAt = nowMillis
    val no: Int = getNextPublishSeqNo
    val event = AmqpPublishing(actorName, no, nowMillis, req, session)
    Try {
      channel.basicPublish(exchange.name, routingKey, props, bytes)
    } match {
      case Success(_) =>
        sendEvent(AmqpPublished(actorName, no, nowMillis, event))
      case Failure(e) =>
        sendEvent(AmqpPublishFailed(actorName, no, nowMillis, e))
        log.error(s"basicPublish($exchange) failed", e)
    }
  }

  protected def publishAsync(req: PublishRequest, session: Session): Unit = {
    import req._
    val no: Int = getNextPublishSeqNo
    sendEvent(AmqpPublishing(actorName, no, nowMillis, req, session))
    try {
      channel.basicPublish(exchange.name, routingKey, props, bytes)
    } catch {
      case e: Exception =>
        sendEvent(AmqpPublishFailed(actorName, no, nowMillis, e))
        log.error(s"basicPublish($exchange) failed", e)
    }
  }
}

object AmqpPublisher {
  def props(name: String, amqp: AmqpProtocol) = Props(classOf[AmqpPublisher], name, amqp)
} 

package io.gatling.keycloak

import com.typesafe.scalalogging.StrictLogging
import io.gatling.core.action.{UserEnd, Chainable}
import io.gatling.core.akka.GatlingActorSystem
import io.gatling.core.result.message.{OK, KO, Status}
import io.gatling.core.result.writer.{DataWriter, DataWriterClient}
import io.gatling.core.session.Session
import io.gatling.core.util.TimeHelper
import io.gatling.core.validation.{Validation, Success, Failure}


object Stopwatch extends StrictLogging {
  @volatile var recording: Boolean = true;
  GatlingActorSystem.instance.registerOnTermination(() => recording = true)

  def apply[T](f: () => T): Result[T] = {
    val start = TimeHelper.nowMillis
    try {
      val result = f()
      Result(Success(result), OK, start, TimeHelper.nowMillis, false)
    } catch {
      case ie: InterruptedException => {
        Result(Failure("Interrupted"), KO, start, start, true)
      }
      case e: Throwable => {
        Stopwatch.log.error("Operation failed with exception", e)
        Result(Failure(e.toString), KO, start, TimeHelper.nowMillis, false)
      }
    }
  }

  def log = logger;
}

case class Result[T](
                      val value: Validation[T],
                      val status: Status,
                      val startTime: Long,
                      val endTime: Long,
                      val interrupted: Boolean
) {
  def check(check: T => Boolean, fail: T => String): Result[T] = {
     value match {
       case Success(v) =>
         if (!check(v)) {
           Result(Failure(fail(v)), KO, startTime, endTime, interrupted);
         } else {
           this
         }
       case _ => this
     }
  }

  def isSuccess =
    value match {
      case Success(_) => true
      case _ => false
    }

  private def record(client: DataWriterClient, session: Session, name: String): Validation[T] = {
    if (!interrupted && Stopwatch.recording) {
      client.writeRequestData(session, name, startTime, startTime, endTime, endTime, status)
    }
    value
  }

  def recordAndStopOnFailure(client: DataWriterClient with Chainable, session: Session, name: String): Validation[T] = {
    val validation = record(client, session, name)
    validation.onFailure(message => {
        Stopwatch.log.error(s"'${client.self.path.name}', ${session.userId} failed to execute: $message")
        UserEnd.instance ! session.markAsFailed
    })
    validation
  }

  def recordAndContinue(client: DataWriterClient with Chainable, session: Session, name: String): Unit = {
    // can't specify follow function as default arg since it uses another parameter
    recordAndContinue(client, session, name, _ => session);
  }

  def recordAndContinue(client: DataWriterClient with Chainable, session: Session, name: String, follow: T => Session): Unit = {
    // 'follow' intentionally does not get session as arg, since caller site already has the reference
    record(client, session, name) match {
      case Success(value) => try {
        client.next ! follow(value)
      } catch {
        case t: Throwable => {
          Stopwatch.log.error(s"'${client.self.path.name}' failed processing", t)
          UserEnd.instance ! session.markAsFailed
      }
    }
      case Failure(message) => {
        Stopwatch.log.error(s"'${client.self.path.name}' failed to execute: $message")
        UserEnd.instance ! session.markAsFailed
      }
    }
  }
} 

package io.gatling.amqp.config

import akka.pattern.ask
import akka.util.Timeout
import io.gatling.amqp.data._
import io.gatling.core.session.Session
import pl.project13.scala.rainbow._

import scala.concurrent.Await
import scala.concurrent.duration._
import scala.util._


trait AmqpTermination { this: AmqpProtocol =>
  private val publishTimeout: Timeout = Timeout(1 hour)
  private val consumeTimeout: Timeout = Timeout(1 hour)

  protected def awaitTerminationFor(session: Session): Unit = {
    // wait nacker to ensure all confirms has been fired
    Await.result((nacker ask WaitTermination(session))(publishTimeout), Duration.Inf) match {
      case Success(m) => logger.debug(s"amqp: $m".green)
      case Failure(e) => throw e
    }

    // wait consumers
    Await.result((router ask WaitTermination(session))(consumeTimeout), Duration.Inf) match {
      case Success(m) => logger.debug(s"amqp: $m".green)
      case Failure(e) => throw e
    }
  }
} 

package ch.tamedia.gatling.actions

import akka.actor.Props
import ch.tamedia.gatling.GrpcProtocol
import ch.tamedia.gatling.grpc.GrpcCheck
import com.trueaccord.scalapb.GeneratedMessage
import io.gatling.commons.stats.{KO, OK}
import io.gatling.commons.util.TimeHelper
import io.gatling.commons.validation.Failure
import io.gatling.core.action.{Action, ActionActor}
import io.gatling.core.check.Check
import io.gatling.core.session.Session
import io.gatling.core.stats.StatsEngine
import io.gatling.core.stats.message.ResponseTimings

import scala.concurrent.Future


    def logResult(maybeResult: Option[GeneratedMessage], error: Option[Throwable] = None) = {
      val endTime = TimeHelper.nowMillis
      val timings = ResponseTimings(startTime, endTime)

      if (error.isEmpty) {
        val result = maybeResult.get
        if (Option(result).nonEmpty) {
          val (newSession, error) = Check.check(result, session, checks)
          error match {
            case None => {
              statsEngine.logResponse(session, action.name, timings, OK, None, None)
              next ! newSession(session)
            }
            case Some(Failure(errorMessage)) => {
              statsEngine.logResponse(session, action.name, timings, KO, None, Some(errorMessage))
              next ! newSession(session).markAsFailed
            }
          }
        }
        else {
          statsEngine.logResponse(session, action.name, timings, KO, None, Some(s"Error during the call!"))
          next ! session.markAsFailed
        }
      }
      else {
        val throwable = error.get
        statsEngine.logResponse(session, action.name, timings, KO, None, Some(throwable.getMessage))
        next ! session.markAsFailed
      }
    }
  }
} 

package com.linagora.gatling.imap.check

import java.util

import com.linagora.gatling.imap.protocol.ImapResponses
import io.gatling.commons.validation.{Failure, Validation}
import io.gatling.core.check.CheckResult
import io.gatling.core.session.Session

object ImapSimpleCheck {
  val DefaultMessage = "Imap check failed"
}

case class ImapSimpleCheck(validate: ImapResponses => Boolean, errorMessage: ImapResponses => String = _ => ImapSimpleCheck.DefaultMessage) extends ImapCheck {
  override def check(responses: ImapResponses, session: Session)(implicit preparedCache: util.Map[Any, Any]): Validation[CheckResult] = {
    if (validate(responses))
      CheckResult.NoopCheckResultSuccess
    else
      Failure(errorMessage(responses))
  }
} 

package com.linagora.gatling.imap.protocol

import java.util.Properties
import java.util.UUID.randomUUID

import akka.actor.ActorRef
import com.linagora.gatling.imap.protocol.Command.Disconnect
import io.gatling.core.CoreComponents
import io.gatling.core.config.GatlingConfiguration
import io.gatling.core.protocol.{Protocol, ProtocolComponents, ProtocolKey}
import io.gatling.core.session.Session

object ImapProtocol {
  val ImapProtocolKey = new ProtocolKey[ImapProtocol, ImapComponents] {

    override def protocolClass: Class[io.gatling.core.protocol.Protocol] = classOf[ImapProtocol].asInstanceOf[Class[io.gatling.core.protocol.Protocol]]

    override def defaultProtocolValue(configuration: GatlingConfiguration): ImapProtocol = throw new IllegalStateException("Can't provide a default value for ImapProtocol")

    override def newComponents(coreComponents: CoreComponents): ImapProtocol => ImapComponents = { protocol =>
      val sessions: ActorRef = coreComponents.actorSystem.actorOf(ImapSessions.props(protocol), "imapsessions_" + randomUUID().toString)
      ImapComponents(protocol, sessions)
    }
  }
}

case class ImapComponents(protocol: ImapProtocol, sessions: ActorRef) extends ProtocolComponents {
  override def onStart: Session => Session = s => s

  override def onExit: Session => Unit = session => sessions ! Disconnect(UserId(session.userId))
}

case class ImapProtocol(host: String,
                        port: Int = 143,
                        config: Properties = new Properties()
                       ) extends Protocol 

package com.linagora.gatling.imap.action

import akka.actor.Props
import com.linagora.gatling.imap.protocol.Response.Disconnected
import com.linagora.gatling.imap.protocol.{ImapResponses, Response}
import io.gatling.commons.stats.{KO, OK}
import io.gatling.commons.validation.Failure
import io.gatling.core.akka.BaseActor
import io.gatling.core.check.Check
import io.gatling.core.session.Session

import scala.collection.immutable.Seq

object ImapResponseHandler {
  def props(imapActionContext: ImapActionContext, requestName: String, session: Session, start: Long, checks: Seq[Check[ImapResponses]]): Props = {
    Props(new ImapResponseHandler(imapActionContext, requestName, session, start, checks))
  }
}

class ImapResponseHandler(imapActionContext: ImapActionContext, requestName: String, session: Session, start: Long, checks: Seq[Check[ImapResponses]]) extends BaseActor {
  private val statsEngine = imapActionContext.statsEngine
  private val next = imapActionContext.next

  override def receive: Receive = {
    case Response(responses) =>
      val (newSession, error) = Check.check(responses, session, checks.toList)
      error.fold(ok(newSession, start))(ko(session, start))
    case e: Exception =>
      ko(session, start)(Failure(e.getMessage))
    case Disconnected(e) =>
      ko(session, start)(Failure(e.getMessage))
    case msg =>
      logger.error(s"received unexpected message $msg")
  }

  def ko(session: Session, start: Long)(failure: Failure) = {
    statsEngine.logResponse(session, requestName, start, imapActionContext.clock.nowMillis, KO, None, Some(failure.message))
    next ! session.markAsFailed
    context.stop(self)
  }

  def ok(session: Session, start: Long) = {
    statsEngine.logResponse(session, requestName, start, imapActionContext.clock.nowMillis, OK, None, None)
    next ! session
    context.stop(self)
  }
} 

package ru.tinkoff.gatling.amqp.checks

import java.util.{Map => JMap}

import io.gatling.commons.validation._
import io.gatling.core.check.CheckResult
import io.gatling.core.session.Session
import ru.tinkoff.gatling.amqp.AmqpCheck
import ru.tinkoff.gatling.amqp.request.AmqpProtocolMessage

object AmqpMessageCheck {

  private val AmqpMessageCheckFailure = "AMQP check failed".failure
}

case class AmqpMessageCheck(func: AmqpProtocolMessage => Boolean) extends AmqpCheck {
  override def check(response: AmqpProtocolMessage, session: Session, preparedCache: JMap[Any, Any]): Validation[CheckResult] =
    if (func(response)) {
      CheckResult.NoopCheckResultSuccess
    } else {
      AmqpMessageCheck.AmqpMessageCheckFailure
    }
} 

package ru.tinkoff.gatling.amqp

import com.rabbitmq.client.MessageProperties
import io.gatling.commons.validation.Validation
import io.gatling.core.session.{Expression, Session}

package object request {

  sealed trait AmqpExchange

  case class AmqpDirectExchange(name: Expression[String], routingKey: Expression[String], durable: Boolean = false)
      extends AmqpExchange
  case class AmqpQueueExchange(name: Expression[String], durable: Boolean = false) extends AmqpExchange

  sealed trait AmqpMessage {
    private[amqp] def amqpProtocolMessage(session: Session): Validation[AmqpProtocolMessage]
  }

  case class TextAmqpMessage(text: Expression[String]) extends AmqpMessage {
    override private[amqp] def amqpProtocolMessage(session: Session) =
      text(session).map(str => AmqpProtocolMessage(MessageProperties.MINIMAL_BASIC, str.getBytes()))
  }
  case class BytesAmqpMessage(bytes: Expression[Array[Byte]]) extends AmqpMessage {
    override private[amqp] def amqpProtocolMessage(session: Session) =
      bytes(session).map(AmqpProtocolMessage(MessageProperties.MINIMAL_BASIC, _))
  }

} 

package ru.tinkoff.gatling.amqp.client

import io.gatling.core.session.Session
import javax.jms.DeliveryMode
import ru.tinkoff.gatling.amqp.action.Around
import ru.tinkoff.gatling.amqp.protocol.AmqpComponents
import ru.tinkoff.gatling.amqp.request.{AmqpDirectExchange, AmqpExchange, AmqpProtocolMessage, AmqpQueueExchange}

import scala.collection.JavaConverters._

class AmqpPublisher(destination: AmqpExchange, components: AmqpComponents) extends WithAmqpChannel {
  def publish(message: AmqpProtocolMessage, around: Around, session: Session): Unit = {

    val protocolDurable = components.protocol.deliveryMode == DeliveryMode.PERSISTENT

    destination match {
      case AmqpDirectExchange(name, routingKey, durable) =>
        for {
          exName <- name(session)
          exKey  <- routingKey(session)
        } withChannel { channel =>
          channel.exchangeDeclare(exName, "direct", durable || protocolDurable)
          around(channel.basicPublish(exName, exKey, message.amqpProperties, message.payload))
        }

      case AmqpQueueExchange(name, durable) =>
        name(session).foreach(qName =>
          withChannel { channel =>
            channel.queueDeclare(qName, durable || protocolDurable, false, false, Map.empty[String, Object].asJava)
            around(channel.basicPublish("", qName, message.amqpProperties, message.payload))
        })
    }
  }

  override protected val pool: AmqpConnectionPool = components.connectionPool
}