Improve use of Futures to avoid suspected deadlocks (close #238)

Author: istreeter

core/src/main/scala/com.snowplowanalytics.stream/loader/clients/BulkSender.scala

@@ -56,7 +56,6 @@ trait BulkSender[A] {
 
   def send(records: List[A]): List[A]
   def close(): Unit
-  def logHealth(): Unit
   def chunkConfig(): ESChunk
 
   /**

core/src/main/scala/com.snowplowanalytics.stream/loader/sinks/KinesisSink.scala

@@ -27,18 +27,14 @@ import java.nio.charset.StandardCharsets.UTF_8
 import org.slf4j.LoggerFactory
 
 // Scala
-import scala.util.{Failure, Random, Success}
+import scala.util.{Failure, Random, Success, Try}
 
 // Amazon
 import com.amazonaws.services.kinesis.model._
 import com.amazonaws.client.builder.AwsClientBuilder.EndpointConfiguration
 import com.amazonaws.services.kinesis.AmazonKinesisClientBuilder
 import com.amazonaws.auth.DefaultAWSCredentialsProviderChain
 
-// Concurrent libraries
-import scala.concurrent.Future
-import scala.concurrent.ExecutionContext.Implicits.global
-
 import com.snowplowanalytics.stream.loader.Config.Sink.BadSink.{Kinesis => KinesisSinkConfig}
 
 /**
@@ -82,7 +78,7 @@ class KinesisSink(conf: KinesisSinkConfig) extends ISink {
       case rnfe: ResourceNotFoundException => false
     }
 
-  private def put(name: String, data: ByteBuffer, key: String): Future[PutRecordResult] = Future {
+  private def put(name: String, data: ByteBuffer, key: String): PutRecordResult = {
     val putRecordRequest = {
       val p = new PutRecordRequest()
       p.setStreamName(name)
@@ -102,11 +98,13 @@ class KinesisSink(conf: KinesisSinkConfig) extends ISink {
    * @param good Unused parameter which exists to extend ISink
    */
   def store(output: String, key: Option[String], good: Boolean): Unit =
-    put(
-      conf.streamName,
-      ByteBuffer.wrap(output.getBytes(UTF_8)),
-      key.getOrElse(Random.nextInt.toString)
-    ) onComplete {
+    Try {
+      put(
+        conf.streamName,
+        ByteBuffer.wrap(output.getBytes(UTF_8)),
+        key.getOrElse(Random.nextInt.toString)
+      )
+    } match {
       case Success(result) =>
         log.info("Writing successful")
         log.info(s"  + ShardId: ${result.getShardId}")

core/src/test/scala/com.snowplowanalytics.stream.loader/EmitterSpec.scala

@@ -51,7 +51,6 @@ case class MockElasticsearchSender(chunkConf: ESChunk) extends BulkSender[Emitte
     List.empty
   }
   override def close() = {}
-  override def logHealth(): Unit             = ()
   override def chunkConfig(): ESChunk        = chunkConf
   override val tracker                       = None
   override val maxConnectionWaitTimeMs: Long = 1000L
@@ -67,7 +66,6 @@ class EmitterSpec extends Specification {
       val fakeSender: BulkSender[EmitterJsonInput] = new BulkSender[EmitterJsonInput] {
         override def send(records: List[EmitterJsonInput]): List[EmitterJsonInput] = List.empty
         override def close(): Unit                                                 = ()
-        override def logHealth(): Unit                                             = ()
         override def chunkConfig(): ESChunk                                        = ESChunk(1L, 1L)
         override val tracker                                                       = None
         override val log: Logger                                                   = null

elasticsearch/src/main/scala/com/snowplowanalytics/stream/loader/clients/ElasticsearchBulkSender.scala

@@ -21,7 +21,6 @@ import org.slf4j.LoggerFactory
 // Scala
 import scala.concurrent.ExecutionContext.Implicits.global
 import scala.concurrent.Future
-import scala.util.{Failure => SFailure, Success => SSuccess}
 
 import org.elasticsearch.client.RestClient
 
@@ -40,7 +39,7 @@ import org.apache.http.message.BasicHeader
 import cats.Id
 import cats.effect.{IO, Timer}
 import cats.data.Validated
-import cats.syntax.validated._
+import cats.implicits._
 
 import retry.implicits._
 import retry.{RetryDetails, RetryPolicy}
@@ -142,6 +141,15 @@ class ElasticsearchBulkSender(
           onErrorHandler
         )
         .map(extractResult(records))
+        .flatTap { failures =>
+          IO.delay(log.info(s"Emitted ${esObjects.size - failures.size} records to Elasticsearch"))
+        }
+        .flatTap { failures =>
+          if (failures.nonEmpty)
+            BulkSender.futureToTask(logHealth())
+          else
+            IO.unit
+        }
         .attempt
         .unsafeRunSync() match {
         case Right(s) => s
@@ -156,9 +164,6 @@ class ElasticsearchBulkSender(
       }
     } else Nil
 
-    log.info(s"Emitted ${esObjects.size - newFailures.size} records to Elasticsearch")
-    if (newFailures.nonEmpty) logHealth()
-
     val allFailures = oldFailures ++ newFailures
 
     if (allFailures.nonEmpty) log.warn(s"Returning ${allFailures.size} records as failed")
@@ -195,20 +200,21 @@ class ElasticsearchBulkSender(
   }
 
   /** Logs the cluster health */
-  override def logHealth(): Unit =
+  def logHealth(): Future[Unit] =
     client
       .execute(clusterHealth)
       .flatMap { health =>
         health.fold(failure => Future.failed(failure.error.asException), Future.successful(_))
       }
-      .onComplete {
-        case SSuccess(result) =>
-          result.status match {
-            case "green"  => log.info("Cluster health is green")
-            case "yellow" => log.warn("Cluster health is yellow")
-            case "red"    => log.error("Cluster health is red")
-          }
-        case SFailure(e) => log.error("Couldn't retrieve cluster health", e)
+      .map { result =>
+        result.status match {
+          case "green"  => log.info("Cluster health is green")
+          case "yellow" => log.warn("Cluster health is yellow")
+          case "red"    => log.error("Cluster health is red")
+        }
+      }
+      .recover { case t: Throwable =>
+        log.error("Couldn't retrieve cluster health", t)
       }
 
   /**