Added throttling support

mnogu · Feb 15, 2017 · 7a0a978 · 7a0a978
1 parent eea0617
commit 7a0a978
Show file tree

Hide file tree

Showing 5 changed files with 76 additions and 33 deletions.
diff --git a/build.sbt b/build.sbt
@@ -6,8 +6,15 @@ version := "0.1.2-SNAPSHOT"
 
 scalaVersion := "2.11.8"
 
+enablePlugins(GatlingPlugin)
+
+val gatlingVersion = "2.2.3"
+
 libraryDependencies ++= Seq(
   "io.gatling" % "gatling-core" % "2.2.3" % "provided",
+  "io.gatling"            % "gatling-test-framework"    % gatlingVersion % "test",
+  "io.gatling.highcharts" % "gatling-charts-highcharts" % gatlingVersion % "test",
+
   ("org.apache.kafka" % "kafka-clients" % "0.10.1.1")
     // Gatling contains slf4j-api
     .exclude("org.slf4j", "slf4j-api")

diff --git a/project/plugins.sbt b/project/plugins.sbt
@@ -1,5 +1,7 @@
-logLevel := Level.Warn
+logLevel := Level.Info
 
 addSbtPlugin("org.xerial.sbt" % "sbt-sonatype" % "1.0")
 
 addSbtPlugin("com.jsuereth" % "sbt-pgp" % "1.0.0")
+
+addSbtPlugin("io.gatling" % "gatling-sbt" % "2.2.0")
diff --git a/src/main/scala/com/github/mnogu/gatling/kafka/action/KafkaRequestAction.scala b/src/main/scala/com/github/mnogu/gatling/kafka/action/KafkaRequestAction.scala
@@ -13,44 +13,32 @@ import io.gatling.core.stats.message.ResponseTimings
 import org.apache.kafka.clients.producer._
 
 
-
 class KafkaRequestAction[K,V]( val producer: KafkaProducer[K,V],
                                val kafkaAttributes: KafkaAttributes[K,V],
                                val coreComponents: CoreComponents,
                                val kafkaProtocol: KafkaProtocol,
+                               val throttled: Boolean,
                                val next: Action )
   extends ExitableAction with NameGen {
 
   val statsEngine = coreComponents.statsEngine
-
   override val name = genName("kafkaRequest")
 
   override def execute(session: Session): Unit = recover(session) {
 
-    kafkaAttributes.requestName(session).flatMap { resolvedRequestName =>
-      val payload = kafkaAttributes.payload
-
-      val outcome = kafkaAttributes.key match {
-        case Some(k) => k(session).flatMap { resolvedKey =>
-          sendRequest(
-            resolvedRequestName,
-            producer,
-            Some(resolvedKey),
-            payload,
-            session )
-        }
-        case None =>
-          sendRequest(
-            resolvedRequestName,
-            producer,
-            None,
-            payload,
-            session )
-      }
+    kafkaAttributes requestName session flatMap { requestName =>
+
+      val outcome =
+        sendRequest(
+          requestName,
+          producer,
+          kafkaAttributes,
+          throttled,
+          session)
 
       outcome.onFailure(
         errorMessage =>
-          statsEngine.reportUnbuildableRequest(session, resolvedRequestName, errorMessage)
+          statsEngine.reportUnbuildableRequest(session, requestName, errorMessage)
       )
 
       outcome
@@ -61,14 +49,17 @@ class KafkaRequestAction[K,V]( val producer: KafkaProducer[K,V],
 
   private def sendRequest( requestName: String,
                            producer: Producer[K,V],
-                           key: Option[K],
-                           payload: Expression[V],
+                           kafkaAttributes: KafkaAttributes[K,V],
+                           throttled: Boolean,
                            session: Session ): Validation[Unit] = {
 
-    payload(session).map { resolvedPayload =>
-      val record = key match {
-        case Some(k) => new ProducerRecord[K,V](kafkaProtocol.topic, k, resolvedPayload)
-        case None => new ProducerRecord[K,V](kafkaProtocol.topic, resolvedPayload)
+      kafkaAttributes payload session map { payload =>
+
+      val record = kafkaAttributes.key match {
+        case Some(k) =>
+          new ProducerRecord[K, V](kafkaProtocol.topic, k(session).get, payload)
+        case None =>
+          new ProducerRecord[K, V](kafkaProtocol.topic, payload)
       }
 
       val requestStartDate = nowMillis
@@ -87,13 +78,17 @@ class KafkaRequestAction[K,V]( val producer: KafkaProducer[K,V],
             if (e == null) None else Some(e.getMessage)
           )
 
-          next ! session
+          if (throttled) {
+            coreComponents.throttler.throttle(session.scenario, () => next ! session)
+          } else {
+            next ! session
+          }
 
         }
-
       })
 
     }
+
   }
 
 }
diff --git a/src/main/scala/com/github/mnogu/gatling/kafka/action/KafkaRequestActionBuilder.scala b/src/main/scala/com/github/mnogu/gatling/kafka/action/KafkaRequestActionBuilder.scala
@@ -13,7 +13,7 @@ import scala.collection.JavaConverters._
 class KafkaRequestActionBuilder[K,V](kafkaAttributes: KafkaAttributes[K,V]) extends ActionBuilder {
 
   override def build( ctx: ScenarioContext, next: Action ): Action = {
-    import ctx.{protocolComponentsRegistry, system, coreComponents}
+    import ctx.{protocolComponentsRegistry, system, coreComponents, throttled}
 
     val kafkaComponents: KafkaComponents = protocolComponentsRegistry.components(KafkaProtocol.KafkaProtocolKey)
 
@@ -26,6 +26,7 @@ class KafkaRequestActionBuilder[K,V](kafkaAttributes: KafkaAttributes[K,V]) exte
       kafkaAttributes,
       coreComponents,
       kafkaComponents.kafkaProtocol,
+      throttled,
       next
     )
 

diff --git a/src/test/scala/com/github/mnogu/gatling/kafka/test/ThrottledSimulation.scala b/src/test/scala/com/github/mnogu/gatling/kafka/test/ThrottledSimulation.scala
@@ -0,0 +1,38 @@
+package com.github.mnogu.gatling.kafka.test
+
+import io.gatling.core.Predef._
+import org.apache.kafka.clients.producer.ProducerConfig
+import scala.concurrent.duration._
+
+import com.github.mnogu.gatling.kafka.Predef._
+
+class ThrottledSimulation extends Simulation {
+  val kafkaConf = kafka
+    // Kafka topic name
+    .topic("test")
+    // Kafka producer configs
+    .properties(
+    Map(
+      ProducerConfig.ACKS_CONFIG -> "1",
+      // list of Kafka broker hostname and port pairs
+      ProducerConfig.BOOTSTRAP_SERVERS_CONFIG -> "localhost:9092",
+
+      // in most cases, StringSerializer or ByteArraySerializer
+      ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG ->
+        "org.apache.kafka.common.serialization.StringSerializer",
+      ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG ->
+        "org.apache.kafka.common.serialization.StringSerializer"))
+
+  val scn = scenario("Kafka Test")
+    .forever(
+      exec(
+        kafka("request")
+          // message to send
+          .send[String]("foo"))
+    )
+
+  setUp(
+    scn.inject(atOnceUsers(10)))
+    .throttle(jumpToRps(10), holdFor(30 seconds))
+    .protocols(kafkaConf)
+}