Handling exceptions from user-supplied lambdas in GraphRequestAction

src/main/scala/com/datastax/gatling/plugin/model/DseGraphStatements.scala

@@ -9,6 +9,7 @@ package com.datastax.gatling.plugin.model
 import com.datastax.driver.dse.graph.{GraphStatement, SimpleGraphStatement}
 import com.datastax.dse.graph.api.DseGraph
 import com.datastax.gatling.plugin.exceptions.DseGraphStatementException
+import com.typesafe.scalalogging.StrictLogging
 import io.gatling.commons.validation._
 import io.gatling.core.session.{Expression, Session}
 import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.GraphTraversal
@@ -49,9 +50,17 @@ case class GraphFluentStatement(statement: GraphStatement) extends DseGraphState
   * @param lambda Scala lambda that takes a Gatling User Session (from which it can retrieve parameters)
   *               and returns a fluent Graph Statement
   */
-case class GraphFluentStatementFromScalaLambda(lambda: Session => GraphStatement) extends DseGraphStatement {
+case class GraphFluentStatementFromScalaLambda(lambda: Session => GraphStatement) extends DseGraphStatement with StrictLogging {
   def buildFromSession(gatlingSession: Session): Validation[GraphStatement] = {
-    lambda(gatlingSession).success
+    Try {
+      lambda(gatlingSession)
+    } match {
+      case TrySuccess(stmt) => stmt.success
+      case TryFailure(ex) => {
+        logger.error("Failed to generate GraphStatement", ex)
+        ex.getMessage.failure
+      }
+    }
   }
 }
 

src/main/scala/com/datastax/gatling/plugin/request/GraphRequestAction.scala

@@ -12,12 +12,14 @@ import java.util.concurrent.ExecutorService
 import java.util.concurrent.TimeUnit.MICROSECONDS
 
 import akka.actor.ActorSystem
+import com.datastax.driver.dse.graph.GraphStatement
 import com.datastax.gatling.plugin.DseProtocol
 import com.datastax.gatling.plugin.metrics.MetricsLogger
 import com.datastax.gatling.plugin.model.DseGraphAttributes
 import com.datastax.gatling.plugin.response.GraphResponseHandler
 import com.datastax.gatling.plugin.utils._
 import io.gatling.commons.stats.KO
+import io.gatling.commons.validation.Validation
 import io.gatling.core.action.{Action, ExitableAction}
 import io.gatling.core.session.Session
 import io.gatling.core.stats.StatsEngine
@@ -69,7 +71,8 @@ class GraphRequestAction(val name: String,
       ThroughputVerifier.checkForGatlingOverloading(session, gatlingTimingSource)
       GatlingResponseTime.startedByGatling(session, gatlingTimingSource)
     }
-    val stmt = dseAttributes.statement.buildFromSession(session)
+
+    val stmt: Validation[GraphStatement] = dseAttributes.statement.buildFromSession(session)
 
     stmt.onFailure(err => {
       val responseTime = responseTimeBuilder.build()

src/test/scala/com/datastax/gatling/plugin/DseGraphStatementSpec.scala

@@ -1,14 +1,18 @@
 package com.datastax.gatling.plugin
 
+import ch.qos.logback.classic.{Level, Logger}
+import ch.qos.logback.classic.spi.ILoggingEvent
+import ch.qos.logback.core.read.ListAppender
 import com.datastax.driver.dse.DseSession
 import com.datastax.driver.dse.graph.SimpleGraphStatement
 import com.datastax.dse.graph.api.DseGraph
 import com.datastax.gatling.plugin.base.BaseSpec
-import com.datastax.gatling.plugin.model.{GraphBoundStatement, GraphFluentStatement, GraphStringStatement}
+import com.datastax.gatling.plugin.model.{GraphBoundStatement, GraphFluentStatement, GraphFluentStatementFromScalaLambda, GraphStringStatement}
 import io.gatling.commons.validation.{Failure, Success}
 import io.gatling.core.session.Session
 import io.gatling.core.session.el.ElCompiler
 import org.easymock.EasyMock.reset
+import org.slf4j.LoggerFactory
 
 
 class DseGraphStatementSpec extends BaseSpec {
@@ -52,6 +56,36 @@ class DseGraphStatementSpec extends BaseSpec {
     }
   }
 
+  describe("FluentStatementFromScalaLambda") {
+
+    val fakeExceptionMessage = "fake testing exception (safe to ignore)"
+    val target = GraphFluentStatementFromScalaLambda((_:Session) => {
+        throw new RuntimeException(fakeExceptionMessage)
+    })
+
+    it("should catch and log a RuntimeException thrown by its encapsulated lambda") {
+
+      val classLogger = LoggerFactory.getLogger(classOf[GraphFluentStatementFromScalaLambda]).asInstanceOf[Logger]
+      val listAppender: ListAppender[ILoggingEvent] = new ListAppender[ILoggingEvent]
+
+      listAppender.start()
+      classLogger.addAppender(listAppender)
+
+      val result = target.buildFromSession(validGatlingSession)
+      result shouldBe a[Failure]
+
+      listAppender.list.size() shouldBe 1
+
+      val logEntry = listAppender.list.get(0)
+
+      logEntry.getLevel shouldBe Level.ERROR
+      logEntry.getFormattedMessage should include ("Failed to generate GraphStatement")
+      logEntry.getThrowableProxy should not be null
+      logEntry.getThrowableProxy.getMessage() shouldBe fakeExceptionMessage
+      logEntry.getThrowableProxy.getClassName() shouldBe classOf[RuntimeException].getCanonicalName()
+    }
+  }
+
   describe("GraphBoundStatement") {
 
     val graphStatement = new SimpleGraphStatement("g.addV(label, vertexLabel).property('type', myType)")

src/test/scala/com/datastax/gatling/plugin/request/GraphRequestActionSpec.scala

@@ -4,6 +4,9 @@ import java.util.concurrent.{Executor, TimeUnit}
 
 import akka.actor.{ActorSystem, Props}
 import akka.testkit.TestKitBase
+import ch.qos.logback.classic.{Level, Logger}
+import ch.qos.logback.classic.spi.ILoggingEvent
+import ch.qos.logback.core.read.ListAppender
 import com.datastax.driver.core._
 import com.datastax.driver.dse.DseSession
 import com.datastax.driver.dse.graph.{GraphResultSet, RegularGraphStatement, SimpleGraphStatement}
@@ -14,12 +17,13 @@ import com.datastax.gatling.plugin.DseProtocol
 import com.datastax.gatling.plugin.model.{DseGraphStatement, DseGraphAttributes}
 import com.google.common.util.concurrent.{Futures, ListenableFuture}
 import io.gatling.commons.validation.SuccessWrapper
-import io.gatling.core.action.Exit
+import io.gatling.core.action.{Action, ChainableAction, Exit}
 import io.gatling.core.config.GatlingConfiguration
 import io.gatling.core.session.Session
 import io.gatling.core.stats.StatsEngine
 import org.easymock.EasyMock
 import org.easymock.EasyMock._
+import org.slf4j.LoggerFactory
 
 class GraphRequestActionSpec extends BaseSpec with TestKitBase {
   implicit lazy val system = ActorSystem()
@@ -30,10 +34,11 @@ class GraphRequestActionSpec extends BaseSpec with TestKitBase {
   val statsEngine: StatsEngine = mock[StatsEngine]
   val gatlingSession = Session("scenario", 1)
 
-  def getTarget(dseAttributes: DseGraphAttributes): GraphRequestAction = {
+  def getTarget(dseAttributes: DseGraphAttributes,
+                next: Action = new Exit(system.actorOf(Props[DseRequestActor]), statsEngine)): GraphRequestAction = {
     new GraphRequestAction(
       "sample-dse-request",
-      new Exit(system.actorOf(Props[DseRequestActor]), statsEngine),
+      next,
       system,
       statsEngine,
       DseProtocol(dseSession),
@@ -62,24 +67,27 @@ class GraphRequestActionSpec extends BaseSpec with TestKitBase {
   override protected def afterAll(): Unit = {
     shutdown(system)
   }
-
+
+  private def getTestGraphAttributes(): DseGraphAttributes =
+    DseGraphAttributes("test", dseGraphStatement,
+      cl = Some(ConsistencyLevel.ANY),
+      userOrRole = Some("test_user"),
+      readTimeout = Some(12),
+      defaultTimestamp = Some(1498167845000L),
+      idempotent = Some(true),
+      readCL = Some(ConsistencyLevel.LOCAL_QUORUM),
+      writeCL = Some(ConsistencyLevel.LOCAL_QUORUM),
+      graphName = Some("MyGraph"),
+      graphLanguage = Some("english"),
+      graphSource = Some("mysource"),
+      graphInternalOptions = Some(Seq(("get", "this"))),
+      graphTransformResults = None
+    )
+
   describe("Graph") {
     val statementCapture = EasyMock.newCapture[RegularGraphStatement]
     it("should enable all the Graph Attributes in DseAttributes") {
-      val graphAttributes = DseGraphAttributes("test", dseGraphStatement,
-        cl = Some(ConsistencyLevel.ANY),
-        userOrRole = Some("test_user"),
-        readTimeout = Some(12),
-        defaultTimestamp = Some(1498167845000L),
-        idempotent = Some(true),
-        readCL = Some(ConsistencyLevel.LOCAL_QUORUM),
-        writeCL = Some(ConsistencyLevel.LOCAL_QUORUM),
-        graphName = Some("MyGraph"),
-        graphLanguage = Some("english"),
-        graphSource = Some("mysource"),
-        graphInternalOptions = Some(Seq(("get", "this"))),
-        graphTransformResults = None
-      )
+      val graphAttributes = getTestGraphAttributes()
 
       expecting {
         dseGraphStatement.buildFromSession(gatlingSession).andReturn(new SimpleGraphStatement("g.V()").success)