Support for async non blocking predict method in Prediction IO and async non blocking ES access

build.sbt

@@ -14,7 +14,7 @@ scalaVersion in ThisBuild := "2.11.11"
 
 val mahoutVersion = "0.13.0"
 
-val pioVersion = "0.12.0-incubating"
+val pioVersion = "0.14.0-SNAPSHOT"
 
 val elasticsearchVersion = "5.5.2"
 

src/main/scala/EsClient.scala

@@ -22,13 +22,13 @@ import java.util
 
 import grizzled.slf4j.Logger
 import org.apache.http.util.EntityUtils
-import org.apache.predictionio.data.storage.{ DataMap, Storage, StorageClientConfig }
+import org.apache.predictionio.data.storage.{DataMap, Storage, StorageClientConfig}
 import org.apache.predictionio.workflow.CleanupFunctions
 import org.apache.spark.SparkContext
 import org.apache.spark.rdd.RDD
 import org.elasticsearch.client.RestClient
 import org.apache.http.HttpHost
-import org.apache.http.auth.{ AuthScope, UsernamePasswordCredentials }
+import org.apache.http.auth.{AuthScope, UsernamePasswordCredentials}
 import org.apache.http.entity.ContentType
 import org.apache.http.entity.StringEntity
 import org.apache.http.impl.client.BasicCredentialsProvider
@@ -42,6 +42,9 @@ import org.elasticsearch.spark._
 import org.json4s.JValue
 import org.json4s.DefaultFormats
 import org.json4s.JsonAST.JString
+import ScalaRestClient.ExtendedScalaRestClient
+
+import scala.concurrent.{ExecutionContext, Future}
 // import org.json4s.native.Serialization.writePretty
 import com.actionml.helpers.{ ItemID, ItemProps }
 
@@ -367,20 +370,23 @@ object EsClient {
    *  @param indexName the index to search
    *  @return a [PredictedResults] collection
    */
-  def search(query: String, indexName: String): Option[JValue] = {
+  def search(query: String, indexName: String)(implicit ec: ExecutionContext): Future[Option[JValue]] = {
     logger.info(s"Query:\n${query}")
-    val response = client.performRequest(
+    val responseFuture = client.performRequestFuture(
       "POST",
       s"/$indexName/_search",
-      Map.empty[String, String].asJava,
+      Map.empty[String, String],
       new StringEntity(query, ContentType.APPLICATION_JSON))
-    response.getStatusLine.getStatusCode match {
-      case 200 =>
-        logger.info(s"Got source from query: ${query}")
-        Some(parse(EntityUtils.toString(response.getEntity)))
-      case _ =>
-        logger.info(s"Query: ${query}\nproduced status code: ${response.getStatusLine.getStatusCode}")
-        None
+    responseFuture.map {
+      response =>
+        response.getStatusLine.getStatusCode match {
+          case 200 =>
+            logger.info(s"Got source from query: ${query}")
+            Some(parse(EntityUtils.toString(response.getEntity)))
+          case _ =>
+            logger.info(s"Query: ${query}\nproduced status code: ${response.getStatusLine.getStatusCode}")
+            None
+        }
     }
   }
 

src/main/scala/FutureTimeout.scala

@@ -0,0 +1,70 @@
+/*
+ * Copyright ActionML, LLC under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * ActionML licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.actionml
+
+import java.util.concurrent.TimeoutException
+import java.util.{Timer, TimerTask}
+
+import scala.concurrent.duration.FiniteDuration
+import scala.concurrent.{ExecutionContext, Future, Promise}
+import scala.language.postfixOps
+import scala.util.{Failure, Success}
+
+object FutureUtil {
+
+  // All Future's that use futureWithTimeout will use the same Timer object
+  // it is thread safe and scales to thousands of active timers
+  // The true parameter ensures that timeout timers are daemon threads and do not stop
+  // the program from shutting down
+
+  val timer: Timer = new Timer(true)
+
+  /**
+    * Returns the result of the provided future within the given time or a timeout exception, whichever is first
+    * This uses Java Timer which runs a single thread to handle all futureWithTimeouts and does not block like a
+    * Thread.sleep would
+    * @param future Caller passes a future to execute
+    * @param timeout Time before we return a Timeout exception instead of future's outcome
+    * @return Future[T]
+    */
+  def futureWithTimeout[T](future : Future[T], timeout : FiniteDuration)(implicit ec: ExecutionContext): Future[T] = {
+
+    // Promise will be fulfilled with either the callers Future or the timer task if it times out
+    var p = Promise[T]
+
+    // and a Timer task to handle timing out
+
+    val timerTask = new TimerTask() {
+      def run() : Unit = p.tryFailure(new TimeoutException())
+    }
+
+    // Set the timeout to check in the future
+    timer.schedule(timerTask, timeout.toMillis)
+
+    future.onComplete {
+      case Success(value) =>
+        p.trySuccess(value)
+        timerTask.cancel()
+      case Failure(e) =>
+        p.tryFailure(e)
+        timerTask.cancel()
+    }
+
+    p.future
+  }
+}

src/main/scala/ScalaRestClient.scala

@@ -0,0 +1,23 @@
+package com.actionml
+
+import org.apache.http.{Header, HttpEntity}
+import org.elasticsearch.client.{Response, ResponseListener, RestClient}
+import scala.collection.JavaConverters._
+import scala.concurrent.{Future, Promise}
+
+object ScalaRestClient {
+
+  implicit class ExtendedScalaRestClient(restClient: RestClient) {
+
+    def performRequestFuture(method: String, endpoint: String, params: Map[String, String],
+                             entity: HttpEntity, headers: Header*): Future[Response] = {
+      val promise: Promise[Response] = Promise()
+      val responseListener = new ResponseListener {
+        override def onSuccess(response: Response): Unit = promise.success(response)
+        override def onFailure(exception: Exception): Unit = promise.failure(exception)
+      }
+      restClient.performRequestAsync(method, endpoint, params.asJava, entity, responseListener, headers: _*)
+      promise.future
+    }
+  }
+}