Add retry to store activations. (#5277)

Author: style95

common/scala/src/main/resources/application.conf

@@ -348,7 +348,7 @@ whisk {
 
     # ActivationStore related configuration
     # For example:
-    # activationStore {
+    # activation-store {
     #   elasticsearch {
     #      protocol         =       # "http" or "https"
     #      hosts            =       # the hosts address of ES, can be multi hosts combined with commas, like "172.17.0.1:9200,172.17.0.2:9200,172.17.0.3:9200"
@@ -361,6 +361,15 @@ whisk {
     #   }
     # }
 
+    activation-store {
+        retry-config {
+            max-tries = 3
+        }
+        elasticsearch {
+            keep-alive = 13 minutes
+        }
+    }
+
     azure-blob {
         # Config property when using AzureBlobAttachmentStore
         # whisk {

common/scala/src/main/scala/org/apache/openwhisk/core/database/Batcher.scala

@@ -49,9 +49,10 @@ import akka.stream.scaladsl.{Sink, Source}
  * @tparam T the type to be batched
  * @tparam R return type of a single element after operation
  */
-class Batcher[T, R](batchSize: Int, concurrency: Int)(operation: Seq[T] => Future[Seq[R]])(implicit
-                                                                                           system: ActorSystem,
-                                                                                           ec: ExecutionContext) {
+class Batcher[T, R](batchSize: Int, concurrency: Int, retry: Int)(operation: (Seq[T], Int) => Future[Seq[R]])(
+  implicit
+  system: ActorSystem,
+  ec: ExecutionContext) {
 
   val cm: PartialFunction[Any, CompletionStrategy] = {
     case Done =>
@@ -69,7 +70,7 @@ class Batcher[T, R](batchSize: Int, concurrency: Int)(operation: Seq[T] => Futur
       val elements = els.map(_._1)
       val promises = els.map(_._2)
 
-      val f = operation(elements)
+      val f = operation(elements, retry)
       f.onComplete {
         case Success(results) => results.zip(promises).foreach { case (result, p) => p.success(result) }
         case Failure(e)       => promises.foreach(_.failure(e))

common/scala/src/main/scala/org/apache/openwhisk/core/database/CouchDbRestStore.scala

@@ -22,6 +22,7 @@ import akka.event.Logging.ErrorLevel
 import akka.http.scaladsl.model._
 import akka.stream.scaladsl._
 import akka.util.ByteString
+
 import scala.concurrent.Await
 import scala.concurrent.duration._
 import spray.json._
@@ -30,6 +31,7 @@ import org.apache.openwhisk.core.database.StoreUtils._
 import org.apache.openwhisk.core.entity.Attachments.Attached
 import org.apache.openwhisk.core.entity.{BulkEntityResult, DocInfo, DocumentReader, UUID}
 import org.apache.openwhisk.http.Messages
+import pureconfig.loadConfigOrThrow
 
 import scala.concurrent.Future
 import scala.util.Try
@@ -74,8 +76,9 @@ class CouchDbRestStore[DocumentAbstraction <: DocumentSerializer](dbProtocol: St
   // and more documents need to be stored, then all arriving documents will be put into batches (if enabled) to avoid a long queue.
   private val maxOpenDbRequests = system.settings.config.getInt("akka.http.host-connection-pool.max-connections") / 2
 
+  private val maxRetry = loadConfigOrThrow[Int]("whisk.activation-store.retry-config.max-tries")
   private val batcher: Batcher[JsObject, Either[ArtifactStoreException, DocInfo]] =
-    new Batcher(500, maxOpenDbRequests)(put(_)(TransactionId.dbBatcher))
+    new Batcher(500, maxOpenDbRequests, maxRetry)(put(_, _)(TransactionId.dbBatcher))
 
   override protected[database] def put(d: DocumentAbstraction)(implicit transid: TransactionId): Future[DocInfo] = {
     val asJson = d.toDocumentRecord
@@ -137,7 +140,7 @@ class CouchDbRestStore[DocumentAbstraction <: DocumentSerializer](dbProtocol: St
         transid.failed(this, start, s"[PUT] '$dbName' internal error, failure: '${failure.getMessage}'", ErrorLevel))
   }
 
-  private def put(ds: Seq[JsObject])(
+  private def put(ds: Seq[JsObject], retry: Int)(
     implicit transid: TransactionId): Future[Seq[Either[ArtifactStoreException, DocInfo]]] = {
     val count = ds.size
     val start = transid.started(this, LoggingMarkers.DATABASE_BULK_SAVE, s"'$dbName' saving $count documents")
@@ -166,10 +169,16 @@ class CouchDbRestStore[DocumentAbstraction <: DocumentSerializer](dbProtocol: St
       }
     }
 
-    reportFailure(
-      f,
-      failure =>
-        transid.failed(this, start, s"[PUT] '$dbName' internal error, failure: '${failure.getMessage}'", ErrorLevel))
+    f.recoverWith {
+      case t: ArtifactStoreException => Future.failed(t)
+      case _ if retry > 0 =>
+        transid.failed(this, start, s"failed to store an activation to CouchDB")
+        put(ds, retry - 1)
+      case t =>
+        transid
+          .failed(this, start, s"[PUT] '$dbName' internal error, failure: '${t.getMessage}'", ErrorLevel)
+        Future.failed(t)
+    }
   }
 
   override protected[database] def del(doc: DocInfo)(implicit transid: TransactionId): Future[Boolean] = {

common/scala/src/main/scala/org/apache/openwhisk/core/database/elasticsearch/ElasticSearchActivationStore.scala

@@ -17,35 +17,39 @@
 
 package org.apache.openwhisk.core.database.elasticsearch
 
-import java.time.Instant
-import java.util.concurrent.TimeUnit
-
-import scala.language.postfixOps
 import akka.actor.ActorSystem
 import akka.event.Logging.ErrorLevel
 import akka.http.scaladsl.model._
 import akka.stream.scaladsl.Flow
+import com.google.common.base.Throwables
 import com.sksamuel.elastic4s.http.search.SearchHit
 import com.sksamuel.elastic4s.http.{ElasticClient, ElasticProperties, NoOpRequestConfigCallback}
 import com.sksamuel.elastic4s.indexes.IndexRequest
 import com.sksamuel.elastic4s.searches.queries.RangeQuery
 import com.sksamuel.elastic4s.searches.queries.matches.MatchPhrase
+import org.apache.http
 import org.apache.http.auth.{AuthScope, UsernamePasswordCredentials}
+import org.apache.http.conn.ConnectionKeepAliveStrategy
 import org.apache.http.impl.client.BasicCredentialsProvider
 import org.apache.http.impl.nio.client.HttpAsyncClientBuilder
-import pureconfig.loadConfigOrThrow
-import pureconfig.generic.auto._
-import spray.json._
+import org.apache.http.protocol.HttpContext
 import org.apache.openwhisk.common.{Logging, LoggingMarkers, TransactionId}
 import org.apache.openwhisk.core.ConfigKeys
 import org.apache.openwhisk.core.containerpool.logging.ElasticSearchJsonProtocol._
-import org.apache.openwhisk.core.database._
 import org.apache.openwhisk.core.database.StoreUtils._
+import org.apache.openwhisk.core.database._
 import org.apache.openwhisk.core.entity._
 import org.apache.openwhisk.http.Messages
 import org.elasticsearch.client.RestClientBuilder.HttpClientConfigCallback
+import pureconfig.loadConfigOrThrow
+import pureconfig.generic.auto._
+import spray.json._
 
+import java.time.Instant
+import java.util.concurrent.TimeUnit
+import scala.concurrent.duration.FiniteDuration
 import scala.concurrent.{ExecutionContextExecutor, Future, Promise}
+import scala.language.postfixOps
 import scala.util.Try
 
 case class ElasticSearchActivationStoreConfig(protocol: String,
@@ -60,8 +64,8 @@ class ElasticSearchActivationStore(
   useBatching: Boolean = false)(implicit actorSystem: ActorSystem, override val logging: Logging)
     extends ActivationStore {
 
-  import com.sksamuel.elastic4s.http.ElasticDsl._
   import ElasticSearchActivationStore.{generateIndex, httpClientCallback}
+  import com.sksamuel.elastic4s.http.ElasticDsl._
 
   private implicit val executionContextExecutor: ExecutionContextExecutor = actorSystem.dispatcher
 
@@ -74,8 +78,9 @@ class ElasticSearchActivationStore(
   private val esType = "_doc"
   private val maxOpenDbRequests = actorSystem.settings.config
     .getInt("akka.http.host-connection-pool.max-connections") / 2
+  private val maxRetry = loadConfigOrThrow[Int]("whisk.activation-store.retry-config.max-tries")
   private val batcher: Batcher[IndexRequest, Either[ArtifactStoreException, DocInfo]] =
-    new Batcher(500, maxOpenDbRequests)(doStore(_)(TransactionId.dbBatcher))
+    new Batcher(500, maxOpenDbRequests, maxRetry)(doStore(_, _)(TransactionId.dbBatcher))
 
   private val minStart = 0L
   private val maxStart = Instant.now.toEpochMilli + TimeUnit.DAYS.toMillis(365 * 100) //100 years from now
@@ -128,10 +133,10 @@ class ElasticSearchActivationStore(
         throw PutException("error on 'put'")
     }
 
-    reportFailure(res, start, failure => s"[PUT] 'activations' internal error, failure: '${failure.getMessage}'")
+    res
   }
 
-  private def doStore(ops: Seq[IndexRequest])(
+  private def doStore(ops: Seq[IndexRequest], retry: Int)(
     implicit transid: TransactionId): Future[Seq[Either[ArtifactStoreException, DocInfo]]] = {
     val count = ops.size
     val start = transid.started(this, LoggingMarkers.DATABASE_BULK_SAVE, s"'activations' saving $count documents")
@@ -169,7 +174,20 @@ class ElasticSearchActivationStore(
         }
       }
 
-    reportFailure(res, start, failure => s"[PUT] 'activations' internal error, failure: '${failure.getMessage}'")
+    res.recoverWith {
+      case t: ArtifactStoreException => Future.failed(t)
+      case _ if retry > 0 =>
+        transid.failed(this, start, s"store activation to ElasticSearch failed")
+        doStore(ops, retry - 1)
+      case t =>
+        transid.failed(
+          this,
+          start,
+          s"[PUT] 'activations' internal error, failure: '${t.getMessage}' [${t.getClass.getSimpleName}]\n" + Throwables
+            .getStackTraceAsString(t),
+          ErrorLevel)
+        Future.failed(t)
+    }
   }
 
   override def get(activationId: ActivationId, context: UserContext)(
@@ -426,6 +444,7 @@ object ElasticSearchActivationStore {
         AuthScope.ANY,
         new UsernamePasswordCredentials(elasticSearchConfig.username, elasticSearchConfig.password))
       httpClientBuilder.setDefaultCredentialsProvider(provider)
+      httpClientBuilder.setKeepAliveStrategy(new CustomKeepAliveStrategy())
     }
   }
 
@@ -442,3 +461,9 @@ object ElasticSearchActivationStoreProvider extends ActivationStoreProvider {
       actorSystem,
       logging)
 }
+
+class CustomKeepAliveStrategy extends ConnectionKeepAliveStrategy {
+  override def getKeepAliveDuration(response: http.HttpResponse, context: HttpContext): Long = {
+    loadConfigOrThrow[FiniteDuration]("whisk.activation-store.elasticsearch.keep-alive").toMillis
+  }
+}

tests/src/test/scala/org/apache/openwhisk/core/database/test/BatcherTests.scala

@@ -54,32 +54,32 @@ class BatcherTests extends FlatSpec with Matchers with WskActorSystem {
 
     val transform = (i: Int) => i + 1
 
-    val batchOperation = LoggedFunction((els: Seq[Int]) => {
+    val batchOperation = LoggedFunction((els: Seq[Int], retry: Int) => {
       batchPromises.dequeue().future.map(_ => els.map(transform))
     })
 
-    val batcher = new Batcher[Int, Int](2, 1)(batchOperation)
+    val batcher = new Batcher[Int, Int](2, 1, 1)(batchOperation)
 
     val values = 1 to 5
     val results = values.map(batcher.put)
 
     // First "batch"
     retry(batchOperation.calls should have size 1, (promiseDelay.toMillis * 2).toInt)
-    batchOperation.calls(0) should have size 1
+    batchOperation.calls(0)._1 should have size 1
 
     // Allow batch to build up
     resolveDelayed(ps(0))
 
     // Second batch
     retry(batchOperation.calls should have size 2, (promiseDelay.toMillis * 2).toInt)
-    batchOperation.calls(1) should have size 2
+    batchOperation.calls(1)._1 should have size 2
 
     // Allow batch to build up
     resolveDelayed(ps(1))
 
     // Third batch
     retry(batchOperation.calls should have size 3, (promiseDelay.toMillis * 2).toInt)
-    batchOperation.calls(2) should have size 2
+    batchOperation.calls(2)._1 should have size 2
     ps(2).success(())
 
     await(Future.sequence(results)) shouldBe values.map(transform)
@@ -90,7 +90,7 @@ class BatcherTests extends FlatSpec with Matchers with WskActorSystem {
     val parallel = new AtomicInteger(0)
     val concurrency = 2
 
-    val batcher = new Batcher[Int, Int](1, concurrency)(els => {
+    val batcher = new Batcher[Int, Int](1, concurrency, 1)((els, _) => {
       parallel.incrementAndGet()
       p.future.map(_ => els)
     })
@@ -108,7 +108,7 @@ class BatcherTests extends FlatSpec with Matchers with WskActorSystem {
   }
 
   it should "complete batched values with the thrown exception" in {
-    val batcher = new Batcher[Int, Int](2, 1)(_ => Future.failed(new Exception))
+    val batcher = new Batcher[Int, Int](2, 1, 1)((_, _) => Future.failed(new Exception))
 
     val r1 = batcher.put(1)
     val r2 = batcher.put(2)
@@ -123,4 +123,71 @@ class BatcherTests extends FlatSpec with Matchers with WskActorSystem {
     an[Exception] should be thrownBy await(r3)
     an[Exception] should be thrownBy await(r4)
   }
+
+  it should "complete batched values with max retry limit" in {
+    val p = Promise[Unit]()
+
+    val maxRetry = 3
+    val batchSize = 1
+    val concurrency = 1
+
+    var retryCount = new AtomicInteger(0)
+
+    def doStore(els: Seq[Int], retry: Int): Future[Seq[Int]] = {
+      val result = if (retry > 0) {
+        Future.failed(new Exception)
+      } else {
+        p.future.map(_ => els)
+      }
+
+      result.recoverWith {
+        case _ if retry > 0 =>
+          retryCount.incrementAndGet()
+          doStore(els, retry - 1)
+        case e =>
+          Future.failed(e)
+      }
+
+    }
+    val batcher = new Batcher[Int, Int](batchSize, concurrency, maxRetry)(doStore)
+
+    val values = List(1)
+    val results = values.map(batcher.put)
+
+    p.success(())
+
+    await(Future.sequence(results)) shouldBe values
+
+    retryCount.get() shouldBe maxRetry
+  }
+
+  it should "complete batched values with the thrown exception with max retry limit" in {
+    val p = Promise[Unit]()
+
+    val maxRetry = 3
+    val batchSize = 1
+    val concurrency = 1
+
+    val retryCount = new AtomicInteger(0)
+
+    def doStore(els: Seq[Int], retry: Int): Future[Seq[Int]] = {
+      val result = Future.failed(new Exception)
+
+      result.recoverWith {
+        case _ if retry > 0 =>
+          retryCount.incrementAndGet()
+          doStore(els, retry - 1)
+        case e =>
+          Future.failed(e)
+      }
+
+    }
+    val batcher = new Batcher[Int, Int](batchSize, concurrency, maxRetry)(doStore)
+
+    val r1 = batcher.put(1)
+
+    an[Exception] should be thrownBy await(r1)
+
+    retryCount.get() shouldBe maxRetry
+  }
 }