Akka AWS Nieuws Scala testing2 mei 2016

Write a Custom Protocol for Gatling

Gatling is a great open source load testing framework. Out of the box it has great support for HTTP. There is also support for JMS, JDBC and Redis, but I’ve never tried these. A lot of services which are written use HTTP as their external interface. Therefore HTTP support will be sufficient in a lot of cases. However, during our test with AWS Lambda we found ourself in a situation where we needed a different protocol. Gatling is written in Scala and Akka, which is a great fit for us, because at Trivento (we are a Lightbend partner) we are great Scala and Akka enthusiasts. In this blog I describe what it takes to create your own protocol. As an example I’ll use the Amazon SDK to call Lambda functions. The code can be found at https://github.com/jgordijn/lambda_gatling_test. At this moment the code is in the branch static: https://github.com/jgordijn/lambda_gatling_test/tree/static.

The search

There is not much information on the internet about writing a custom protocol. We struggled quite a bit to get it working and figured it would be useful to do a write-up about creating your own custom protocol. I will not cover the complete space of writing a protocol, but hopefully enough to get you started.

Protocol, Action & ActionBuilder

There are basically three components you need to create to be able to perform a call through your custom communication protocol:

  • The Protocol
    You can look at the protocol as the communication channel. This is the place where you can create a connection
  • The ActionBuilder
    The actionbuilder is responsible for creating the action at runtime. This is what is passed to the exec function in your Gatling DSL.
  • The Action
    The action is the class that executes the action. In our situation this is the class that really does the call to the AWS Lambda function.

If you implement these 3 files you can write something like:

class LambdaGatlingTest extends Simulation {
  val lambdaProtocol = LambdaProtocol("AWSKEY", "AWS PASS")
  val lambdaScn = scenario("Lambda call").exec(AWSLambdaBuilder("FunctionName"))

  setUp(
    lambdaScn.inject(constantUsersPerSec(15) during (5 minutes))
  ).protocols(lambdaProtocol)
}

Protocol

To create a protocol with specific configuration for the communication you need to implement io.gatling.core.config.Protocol. In this situation we use AWS Lambda and the protocol needs the AWS key and password to connect to AWS. When you create a protocol you can implement/override two functions: def warmUp(): Unit = {} and def userEnd(session: Session): Unit = {}. There is no API documentation for these functions that describe how to use them and when they are called. After looking at the code I found that the warmUp is called before the scenarios are started and can be used to create a pool or something like that. The userEnd will be called after each user request is done, but I see no use in our situation. The code becomes something like:

case class LambdaProtocol(awsAccessKeyId: String, awsSecretAccessKey: String) 
    extends Protocol {
  val credentials = new BasicAWSCredentials(awsAccessKeyId, awsSecretAccessKey)
  val lambdaClient = new AWSLambdaClient(credentials)
  lambdaClient.setRegion(Region.getRegion(Regions.EU_WEST_1))

  def call(functionName: String): Int = {
    val request = new InvokeRequest
    request.setFunctionName(functionName)
    val result = lambdaClient.invoke(request)
    result.getStatusCode
  }
}

ActionBuilder

The io.gatling.core.action.builder.ActionBuilder is responsible for creating the action. This is the place where you connect the Protocol and the Action. In your ActionBuilder you need to implement def build(next: ActorRef, protocols: Protocols): ActorRef.

  • next: ActorRef is a reference to the next Action actor in the chain.
  • protocols: Protocols is the pool of all protocols that are registered at this moment. Registration happens during setup of your test (see first code sample above).

The return type of the build is an ActorRef which is the actor that implements the Action. With the Actor DSL we create an instance of the actor and get the ActorRef.

case class AWSLambdaBuilder(functionName: String) extends ActionBuilder {
  def lambdaProtocol(protocols: Protocols) = 
    protocols.getProtocol[LambdaProtocol]
      .getOrElse(throw new UnsupportedOperationException("LambdaProtocol Protocol wasn't registered"))

  override def build(next: ActorRef, protocols: Protocols): ActorRef = {
    actor(actorName("Functioncall")) {
      new FunctionCall(functionName, lambdaProtocol(protocols), next)
    }
  }
}

Action

In Gatling you can describe a scenario by composing actions. The action is an actor that is responsible for performing the actual logic and this is the place where you can register how long it took to actually call the Lambda function. Your action should extend io.gatling.core.action.Chainable, at least when it is possible that your action can be followed by another action. You can implement io.gatling.core.action.Chainable if your action will always be the last one in a scenario.

To implement Chainable you need to implement 2 functions:

  • def execute(session: Session): Unit
    This is the actual place where the logic of the action will be executed.
  • def next: ActorRef
    This will be a reference to the next action and will be passed in through constructor parameters. In your code you will pass the session to the next action when the action is done.

In the code below you see that we perform the call to the lambda function by using the LambdaProtocol, which we got from the ActionBuilder via a Constructor parameter. We’ll log how long the call took and register it using the DataWriterClient (both the OK case and the KO case).

class FunctionCall(functionName: String, protocol: LambdaProtocol, val next: ActorRef) extends Chainable with DataWriterClient {

  override def execute(session: Session): Unit = {
    val start = TimeHelper.nowMillis
    val result = protocol.call(functionName)
    val end = TimeHelper.nowMillis
    if(result >= 200 && result <= 299)
      writeRequestData(session, "Call function", start, start, end, end, OK )
    else
      writeRequestData(session, "Call function", start, start, end, end, KO )
    next ! session
  }
}

Conclusion

Gatling is a great tool, but I missed documentation on how to call services by other protocol than HTTP. Hopefully this blog will help you to get a quickstart. To get more information I would advise to read the source code of Gatling.