Gatling-插件开发

Posted 2021-02-28 csunshine

tags:

篇首语：本文由小常识网(cha138.com)小编为大家整理，主要介绍了Gatling-插件开发相关的知识，希望对你有一定的参考价值。

插件开发
使用中遇到的问题

插件开发

Action

package io.gatling.ext.redis.action

import java.util.concurrent.Executors

import com.typesafe.scalalogging.StrictLogging
import io.gatling.commons.stats.{KO, OK}
import io.gatling.commons.util.Clock
import io.gatling.core.CoreComponents
import io.gatling.core.action.{Action, ExitableAction}
import io.gatling.core.session.{Expression, Session}
import io.gatling.core.stats.StatsEngine
import io.gatling.core.util.NameGen
import io.gatling.ext.redis.protocol.RedisComponents
import redis.clients.jedis.Jedis

import scala.concurrent._
import scala.util.{Failure, Success}


class RedisGetAction(
                      key: Expression[String],
                      coreComponents: CoreComponents,
                      val redisComponents: RedisComponents,
                      throttled: Boolean,
                      val next: Action
                    ) extends ExitableAction with NameGen with StrictLogging {

  // 可以自定义也可以引入默认的执行器
  // import scala.concurrent.ExecutionContext.Implicits.global
  implicit val ec: ExecutionContextExecutor = ExecutionContext.fromExecutor(Executors.newFixedThreadPool(50))

  override def statsEngine: StatsEngine = coreComponents.statsEngine

  override def clock: Clock = coreComponents.clock

  override def name: String = genName("redis-get-action")

  override def execute(session: Session): Unit = {

    key(session) map {

      keyS => {

        val requestStartDate = clock.nowMillis

        val jedis: Jedis = redisComponents.jedisMap(session.userId)

        val future: Future[String] = Future[String] {
          jedis.get(keyS)
        }

        future onComplete {
          case Success(value) => statsEngine.logResponse(
            session,
            keyS,
            requestStartDate,
            clock.nowMillis,
            OK,
            Some(value),
            Some(value)
          )
          case Failure(exception) => {
            logger.error(exception.toString)
            statsEngine.logResponse(
              session,
              keyS,
              requestStartDate,
              clock.nowMillis,
              KO,
              Some(exception.getMessage),
              Some(exception.getMessage)
            )
          }
        }

        if (throttled) {
          coreComponents.throttler.throttle(session.scenario, () => next ! session)
        } else {
          next ! session
        }
      }
    }
  }
}

ActionBuilder

package io.gatling.ext.redis.builder

import com.typesafe.scalalogging.StrictLogging
import io.gatling.core.action.Action
import io.gatling.core.action.builder.ActionBuilder
import io.gatling.core.session.Expression
import io.gatling.core.structure.ScenarioContext
import io.gatling.ext.redis.action.RedisGetAction
import io.gatling.ext.redis.protocol.RedisProtocol

class RedisActionBuilder(key: Expression[String]) extends ActionBuilder with StrictLogging {

  override def build(ctx: ScenarioContext, next: Action): Action = {

    import ctx._
    new RedisGetAction(
      key,
      coreComponents,
      protocolComponentsRegistry.components(RedisProtocol.redisProtocolKey),
      throttled,
      next
    )
  }
}

Protocol & Components & Components

package io.gatling.ext.redis.protocol

import io.gatling.core.config.GatlingConfiguration
import io.gatling.core.protocol.{Protocol, ProtocolComponents, ProtocolKey}
import io.gatling.core.session.Session
import io.gatling.core.{CoreComponents, protocol}
import redis.clients.jedis.Jedis

case class RedisProtocol(host: String, port: Int, password: Option[String]) extends Protocol {
  type Components = RedisComponents
}

object RedisProtocol {

  def apply(host: String, port: Int, password: Option[String] = None): RedisProtocol = new RedisProtocol(host, port, password)

  val redisProtocolKey: ProtocolKey[RedisProtocol, RedisComponents] = new ProtocolKey[RedisProtocol, RedisComponents] {

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

    override def defaultProtocolValue(configuration: GatlingConfiguration): RedisProtocol = RedisProtocol("127.0.0.1", 6379)

    override def newComponents(coreComponents: CoreComponents): RedisProtocol => RedisComponents = {
      redisProtocol => RedisComponents(redisProtocol)
    }
  }
}

case class RedisComponents(redisProtocol: RedisProtocol) extends ProtocolComponents {

  val jedisMap: scala.collection.mutable.Map[Long, Jedis] = scala.collection.mutable.HashMap()

  override def onStart: Session => Session = {
    session => {
      println("RedisComponents start: " + session)
      val jedis: Jedis = new Jedis(redisProtocol.host, redisProtocol.port)

      redisProtocol.password match {
        case Some(value) => jedis.auth(value)
        case None =>
      }

      jedis.set(String.valueOf(session.userId), session.toString)

      jedisMap(session.userId) = jedis

      session
    }
  }

  override def onExit: Session => Unit = {
    session => {
      println("RedisComponents exit: " + session)
    }
  }
}

case class RedisProtocolBuilder(host: String, port: Int, password: Option[String]) {
  def build: RedisProtocol = RedisProtocol(host, port, password)
}

object RedisProtocolBuilder {
  implicit def toRedisProtocol(builder: RedisProtocolBuilder): RedisProtocol = builder.build
}

Predef

package io.gatling.ext.redis

import io.gatling.core.session.Expression
import io.gatling.ext.redis.builder.RedisActionBuilder
import io.gatling.ext.redis.protocol.RedisProtocolBuilder

object Predef {

  implicit def string2Option(string: String): Option[String] = Some(string)
  
  def redis(
             host: String,
             port: Int,
             password: Option[String] = None
           ): RedisProtocolBuilder = RedisProtocolBuilder(host, port, password)

  def redis(key: Expression[String]): RedisActionBuilder = new RedisActionBuilder(key)
}