From a284480f7f604230c5a1f576c7d73ef95d53a1d1 Mon Sep 17 00:00:00 2001
From: Benoit Tellier <btellier@linagora.com>
Date: Wed, 12 Apr 2023 09:51:37 +0700
Subject: [PATCH] [FIX] Connect was synchronous

 - Blocking threads in Akka actor system is a baaad practice
 - False measurement of connect timings as "connect" was a noop
---
 .../gatling/imap/protocol/ImapSessions.scala  | 43 ++++++++++++-------
 1 file changed, 27 insertions(+), 16 deletions(-)

diff --git a/src/main/scala-2.12/com/linagora/gatling/imap/protocol/ImapSessions.scala b/src/main/scala-2.12/com/linagora/gatling/imap/protocol/ImapSessions.scala
index 2ed2f12..b2025d9 100644
--- a/src/main/scala-2.12/com/linagora/gatling/imap/protocol/ImapSessions.scala
+++ b/src/main/scala-2.12/com/linagora/gatling/imap/protocol/ImapSessions.scala
@@ -7,7 +7,8 @@ import java.util.function.Consumer
 import akka.actor.{Props, Stash}
 import com.linagora.gatling.imap.protocol.command._
 import com.yahoo.imapnio.async.client.ImapAsyncSession.DebugMode
-import com.yahoo.imapnio.async.client.{ImapAsyncClient, ImapAsyncSession, ImapAsyncSessionConfig}
+import com.yahoo.imapnio.async.client.{ImapAsyncClient, ImapAsyncSession, ImapAsyncSessionConfig, ImapFuture}
+import com.yahoo.imapnio.async.internal.ImapAsyncSessionImpl
 import io.gatling.core.akka.BaseActor
 import io.gatling.core.util.NameGen
 import javax.net.ssl.SSLContext
@@ -57,27 +58,37 @@ private object ImapSession {
 private class ImapSession(client: => ImapAsyncClient, protocol: ImapProtocol) extends BaseActor with Stash with NameGen {
   val uri = new URI(s"${protocol.protocol}://${protocol.host}:${protocol.port}")
   val config: Properties = protocol.config
-  logger.debug(s"connecting to $uri with $config")
-  val session: ImapAsyncSession = {
-    val config = new ImapAsyncSessionConfig
-    config.setConnectionTimeoutMillis(50000)
-    config.setReadTimeoutMillis(60000)
-    val sniNames = null
-
-    val localAddress = null
-    client
-      .createSession(uri, config, localAddress, sniNames, DebugMode.DEBUG_OFF, "ImapSession", ImapSession.sslContext)
-      .get()
-      .getSession
-  }
+  var session: ImapAsyncSession = null
 
   override def receive: Receive = disconnected
 
   def disconnected: Receive = {
     case Command.Connect(userId) =>
       logger.debug(s"got connect request, $userId connecting to $uri")
-      context.become(connected)
-      sender() ! Response.Connected(ImapResponses.empty)
+
+      val config = new ImapAsyncSessionConfig
+      config.setConnectionTimeoutMillis(50000)
+      config.setReadTimeoutMillis(60000)
+      val sniNames = null
+      val localAddress = null
+
+      val future: ImapFuture[ImapAsyncSessionImpl] = client
+        .createSession(uri, config, localAddress, sniNames, DebugMode.DEBUG_OFF, "ImapSession", ImapSession.sslContext)
+        .asInstanceOf[ImapFuture[ImapAsyncSessionImpl]]
+
+      future.setDoneCallback(s => {
+        session = s
+        context.become(connected)
+        sender() ! Response.Connected(ImapResponses.empty)
+      })
+      val errorCallback: Consumer[Exception] = e => {
+        logger.error(s"${getClass.getSimpleName} connection failed")
+        sender ! e
+        context.stop(self)
+      }
+
+      future.setExceptionCallback(errorCallback)
+
     case Response.Disconnected(_) => ()
     case Command.Disconnect(_) => ()
     case msg =>