Merge pull request #173 from ExpediaDotCom/traceId-krastogi

Fetching trace details from cassandra in one go

Author: vsen

reader/src/main/scala/com/expedia/www/haystack/trace/reader/stores/CassandraEsTraceStore.scala

@@ -21,7 +21,7 @@ import com.expedia.www.haystack.trace.commons.clients.cassandra.{CassandraCluste
 import com.expedia.www.haystack.trace.commons.clients.es.document.TraceIndexDoc
 import com.expedia.www.haystack.trace.commons.config.entities.{CassandraConfiguration, WhitelistIndexFieldConfiguration}
 import com.expedia.www.haystack.trace.reader.config.entities.{ElasticSearchConfiguration, ServiceMetadataReadConfiguration}
-import com.expedia.www.haystack.trace.reader.metrics.{AppMetricNames, MetricsSupport}
+import com.expedia.www.haystack.trace.reader.metrics.MetricsSupport
 import com.expedia.www.haystack.trace.reader.stores.readers.ServiceMetadataReader
 import com.expedia.www.haystack.trace.reader.stores.readers.cassandra.CassandraTraceReader
 import com.expedia.www.haystack.trace.reader.stores.readers.es.ElasticSearchReader
@@ -32,7 +32,6 @@ import org.slf4j.LoggerFactory
 
 import scala.collection.JavaConverters._
 import scala.concurrent.{ExecutionContextExecutor, Future}
-import scala.util.{Failure, Success, Try}
 
 class CassandraEsTraceStore(cassandraConfig: CassandraConfiguration,
                             serviceMetadataConfig: ServiceMetadataReadConfiguration,
@@ -41,7 +40,6 @@ class CassandraEsTraceStore(cassandraConfig: CassandraConfiguration,
   extends TraceStore with MetricsSupport with ResponseParser {
 
   private val LOGGER = LoggerFactory.getLogger(classOf[ElasticSearchReader])
-  private val traceRejected = metricRegistry.meter(AppMetricNames.SEARCH_TRACE_REJECTED)
 
   private val cassandraSession = new CassandraSession(cassandraConfig, new CassandraClusterFactory)
   private val cassandraReader: CassandraTraceReader = new CassandraTraceReader(cassandraSession, cassandraConfig)
@@ -79,36 +77,21 @@ class CassandraEsTraceStore(cassandraConfig: CassandraConfiguration,
     // go through each hit and fetch trace for parsed traceId
     val sourceList = result.getSourceAsStringList
     if (sourceList != null && sourceList.size() > 0) {
-      val traceFutures = sourceList
+      val traceIds = sourceList
         .asScala
         .map(source => extractTraceIdFromSource(source))
         .filter(!_.isEmpty)
-        .toSet[String]
-        .toSeq
-        .map(id => getTrace(id))
-
-      // wait for all Futures to complete and then map them to Traces
-      Future
-        .sequence(liftToTry(traceFutures))
-        .map(_.flatMap(retrieveTriedTrace))
+        .toSet[String] // de-dup traceIds
+        .toList
+
+      cassandraReader.readRawTraces(traceIds)
     } else {
       Future.successful(Nil)
     }
   }
 
   override def getTrace(traceId: String): Future[Trace] = cassandraReader.readTrace(traceId)
 
-  private def retrieveTriedTrace(mayBeTrace: Try[Trace]): Option[Trace] = {
-    mayBeTrace match {
-      case Success(trace) =>
-        Some(trace)
-      case Failure(ex) =>
-        LOGGER.warn("traceId not found in cassandra, rejected searched trace", ex)
-        traceRejected.mark()
-        None
-    }
-  }
-
   override def getFieldNames(): Future[Seq[String]] = {
     Future.successful(indexConfig.whitelistIndexFields.map(_.name).distinct.sorted)
   }
@@ -148,11 +131,6 @@ class CassandraEsTraceStore(cassandraConfig: CassandraConfiguration,
     cassandraReader.readRawTraces(request.getTraceIdList.asScala.toList)
   }
 
-  // convert all Futures to Try to make sure they all complete
-  private def liftToTry[T](futures: Seq[Future[T]]): Seq[Future[Try[T]]] = futures.map { f =>
-    f.map(Try(_)).recover { case t: Throwable => Failure(t) }
-  }
-
   override def close(): Unit = {
     cassandraReader.close()
     esReader.close()