Merge pull request #250 from ExpediaDotCom/multi-backend-reads

doctorXWrites · web-flow · commit 253b4e56a9cc · 2019-03-29T15:01:05.000+05:30
adding support to read from multiple backend stores
diff --git a/commons/src/main/scala/com/expedia/www/haystack/trace/commons/config/entities/TraceStoreBackends.scala b/commons/src/main/scala/com/expedia/www/haystack/trace/commons/config/entities/TraceStoreBackends.scala
@@ -17,9 +17,16 @@
 
 package com.expedia.www.haystack.trace.commons.config.entities
 
+
 /**
   * defines the configuration parameters for trace-backend  *
   * @param host              : trace backend grpc hostname
   * @param port              : trace backend grpc port
   */
-case class TraceBackendClientConfiguration(host: String, port:Int)
+case class GrpcClientConfig(host: String, port: Int)
+
+/**
+  * multiple store backends
+  * @param backends configuration of all trace store backends
+  */
+case class TraceStoreBackends(backends: Seq[GrpcClientConfig])
diff --git a/indexer/src/main/scala/com/expedia/www/haystack/trace/indexer/config/ProjectConfiguration.scala b/indexer/src/main/scala/com/expedia/www/haystack/trace/indexer/config/ProjectConfiguration.scala
@@ -33,6 +33,7 @@ import org.apache.kafka.clients.producer.ProducerConfig
 import org.apache.kafka.common.serialization.{ByteArraySerializer, StringDeserializer, StringSerializer}
 
 import scala.collection.JavaConverters._
+import scala.collection.mutable
 import scala.util.Try
 
 class ProjectConfiguration extends AutoCloseable {
@@ -126,13 +127,19 @@ class ProjectConfiguration extends AutoCloseable {
     * trace backend configuration object
     */
   val backendConfig: TraceBackendConfiguration = {
+    val traceBackendConfig = config.getConfig("backend")
 
+    val grpcClients = traceBackendConfig.entrySet().asScala
+      .map(k => StringUtils.split(k.getKey, '.')(0)).toSeq
+      .map(cl => traceBackendConfig.getConfig(cl))
+      .filter(cl => cl.hasPath("host") && cl.hasPath("port"))
+      .map(cl => GrpcClientConfig(cl.getString("host"), cl.getInt("port")))
 
-    val traceBackendConfig = config.getConfig("backend")
-    val clientConfig = traceBackendConfig.getConfig("client")
+    // we dont support multiple backends for write operations
+    require(grpcClients.size == 1)
 
     TraceBackendConfiguration(
-      TraceBackendClientConfiguration(clientConfig.getString("host"), clientConfig.getInt("port")),
+      TraceStoreBackends(grpcClients),
       maxInFlightRequests = traceBackendConfig.getInt("max.inflight.requests"))
 
   }
diff --git a/indexer/src/main/scala/com/expedia/www/haystack/trace/indexer/config/entities/TraceBackendConfiguration.scala b/indexer/src/main/scala/com/expedia/www/haystack/trace/indexer/config/entities/TraceBackendConfiguration.scala
@@ -17,11 +17,11 @@
 
 package com.expedia.www.haystack.trace.indexer.config.entities
 
-import com.expedia.www.haystack.trace.commons.config.entities.TraceBackendClientConfiguration
+import com.expedia.www.haystack.trace.commons.config.entities.TraceStoreBackends
 
 /**
   * @param clientConfig defines the grpc client configuration for connecting to the trace backend
   * @param maxInFlightRequests defines the max parallel writes to trace-backend
   */
-case class TraceBackendConfiguration(clientConfig: TraceBackendClientConfiguration,
+case class TraceBackendConfiguration(clientConfig: TraceStoreBackends,
                                      maxInFlightRequests: Int)
diff --git a/indexer/src/main/scala/com/expedia/www/haystack/trace/indexer/writers/grpc/GrpcTraceWriter.scala b/indexer/src/main/scala/com/expedia/www/haystack/trace/indexer/writers/grpc/GrpcTraceWriter.scala
@@ -41,9 +41,12 @@ class GrpcTraceWriter(config: TraceBackendConfiguration)(implicit val dispatcher
   private val writeTimer = metricRegistry.timer(AppMetricNames.BACKEND_WRITE_TIME)
   private val writeFailures = metricRegistry.meter(AppMetricNames.BACKEND_WRITE_FAILURE)
 
-  private val channel =  ManagedChannelBuilder.forAddress(config.clientConfig.host,config.clientConfig.port)
-    .usePlaintext(true)
-    .build()
+  private val channel =  {
+    val grpcConfig = config.clientConfig.backends.head
+    ManagedChannelBuilder.forAddress(grpcConfig.host, grpcConfig.port)
+      .usePlaintext(true)
+      .build()
+  }
   private val client = StorageBackendGrpc.newStub(channel)
 
   // this semaphore controls the parallel writes to trace-backend
diff --git a/indexer/src/test/scala/com/expedia/www/haystack/trace/indexer/integration/clients/GrpcTestClient.scala b/indexer/src/test/scala/com/expedia/www/haystack/trace/indexer/integration/clients/GrpcTestClient.scala
@@ -20,7 +20,7 @@ package com.expedia.www.haystack.trace.indexer.integration.clients
 import java.util.concurrent.Executors
 
 import com.expedia.open.tracing.backend.{ReadSpansRequest, StorageBackendGrpc, TraceRecord}
-import com.expedia.www.haystack.trace.commons.config.entities.TraceBackendClientConfiguration
+import com.expedia.www.haystack.trace.commons.config.entities.{GrpcClientConfig, TraceStoreBackends}
 import com.expedia.www.haystack.trace.indexer.config.entities.TraceBackendConfiguration
 import com.expedia.www.haystack.trace.indexer.integration.TraceDescription
 import com.expedia.www.haystack.trace.storage.backends.memory.Service
@@ -44,7 +44,7 @@ class GrpcTestClient {
 
 
   def buildConfig = TraceBackendConfiguration(
-    TraceBackendClientConfiguration("localhost", port), 10)
+    TraceStoreBackends(Seq(GrpcClientConfig("localhost", port))), 10)
 
   def queryTraces(traceDescriptions: Seq[TraceDescription]): Seq[TraceRecord] = {
     val traceIds = traceDescriptions.map(traceDescription => traceDescription.traceId).toList
diff --git a/reader/src/main/resources/config/base.conf b/reader/src/main/resources/config/base.conf
@@ -15,6 +15,14 @@ backend {
     host = "localhost"
     port = 8090
   }
+
+  # we support multiple grpc based backends, to provide another one use something like following.
+  # you are required to provide host and port
+
+  # another_client {
+  #   host = "localhost"
+  #  port = 8092
+  # }
 }
 
 elasticsearch {
diff --git a/reader/src/main/scala/com/expedia/www/haystack/trace/reader/config/ProviderConfiguration.scala b/reader/src/main/scala/com/expedia/www/haystack/trace/reader/config/ProviderConfiguration.scala
@@ -25,6 +25,7 @@ import com.expedia.www.haystack.trace.reader.config.entities._
 import com.expedia.www.haystack.trace.reader.readers.transformers.{PartialSpanTransformer, SpanTreeTransformer, TraceTransformer}
 import com.expedia.www.haystack.trace.reader.readers.validators.TraceValidator
 import com.typesafe.config.Config
+import org.apache.commons.lang3.StringUtils
 
 import scala.collection.JavaConverters._
 import scala.reflect.ClassTag
@@ -44,9 +45,18 @@ class ProviderConfiguration {
   /**
     * trace backend configuration object
     */
-  val traceBackendConfiguration: TraceBackendClientConfiguration = {
-    val clientConfig = config.getConfig("backend.client")
-    TraceBackendClientConfiguration(clientConfig.getString("host"), clientConfig.getInt("port"))
+  val traceBackendConfiguration: TraceStoreBackends = {
+    val traceBackendConfig = config.getConfig("backend")
+
+    val grpcClients = traceBackendConfig.entrySet().asScala
+      .map(k => StringUtils.split(k.getKey, '.')(0)).toSeq
+      .map(cl => traceBackendConfig.getConfig(cl))
+      .filter(cl => cl.hasPath("host") && cl.hasPath("port"))
+      .map(cl => GrpcClientConfig(cl.getString("host"), cl.getInt("port")))
+
+    require(grpcClients.nonEmpty)
+
+    TraceStoreBackends(grpcClients)
   }
 
 
diff --git a/reader/src/main/scala/com/expedia/www/haystack/trace/reader/readers/utils/TraceMerger.scala b/reader/src/main/scala/com/expedia/www/haystack/trace/reader/readers/utils/TraceMerger.scala
@@ -0,0 +1,35 @@
+/*
+ *  Copyright 2019 Expedia, Inc.
+ *
+ *       Licensed 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.expedia.www.haystack.trace.reader.readers.utils
+
+import com.expedia.open.tracing.api.Trace
+
+object TraceMerger {
+  def merge(traces: Seq[Trace]): Seq[Trace] = {
+    traces.groupBy(_.getTraceId).mapValues {
+      seq => {
+        if (seq.size == 1) {
+          seq.head
+        } else {
+          val head = seq.head.toBuilder
+          seq.tail.foreach(t => head.addAllChildSpans(t.getChildSpansList))
+          head.build()
+        }
+      }
+    }.values.toSeq
+  }
+}
diff --git a/reader/src/main/scala/com/expedia/www/haystack/trace/reader/stores/EsIndexedTraceStore.scala b/reader/src/main/scala/com/expedia/www/haystack/trace/reader/stores/EsIndexedTraceStore.scala
@@ -19,25 +19,25 @@ package com.expedia.www.haystack.trace.reader.stores
 import com.expedia.open.tracing.api._
 import com.expedia.www.haystack.commons.metrics.MetricsSupport
 import com.expedia.www.haystack.trace.commons.clients.es.document.TraceIndexDoc
-import com.expedia.www.haystack.trace.commons.config.entities.{TraceBackendClientConfiguration, WhitelistIndexField, WhitelistIndexFieldConfiguration}
+import com.expedia.www.haystack.trace.commons.config.entities.{TraceStoreBackends, WhitelistIndexFieldConfiguration}
 import com.expedia.www.haystack.trace.reader.config.entities.ElasticSearchConfiguration
 import com.expedia.www.haystack.trace.reader.stores.readers.es.ElasticSearchReader
 import com.expedia.www.haystack.trace.reader.stores.readers.es.query.{FieldValuesQueryGenerator, ServiceMetadataQueryGenerator, TraceCountsQueryGenerator, TraceSearchQueryGenerator}
-import com.expedia.www.haystack.trace.reader.stores.readers.grpc.GrpcTraceReader
+import com.expedia.www.haystack.trace.reader.stores.readers.grpc.GrpcTraceReaders
 import io.searchbox.core.SearchResult
 import org.elasticsearch.index.IndexNotFoundException
 import org.slf4j.LoggerFactory
 
 import scala.collection.JavaConverters._
 import scala.concurrent.{ExecutionContextExecutor, Future}
 
-class EsIndexedTraceStore(traceBackendConfig: TraceBackendClientConfiguration,
+class EsIndexedTraceStore(traceStoreBackendConfig: TraceStoreBackends,
                           elasticSearchConfiguration: ElasticSearchConfiguration,
                           whitelistedFieldsConfiguration: WhitelistIndexFieldConfiguration)(implicit val executor: ExecutionContextExecutor)
   extends TraceStore with MetricsSupport with ResponseParser {
   private val LOGGER = LoggerFactory.getLogger(classOf[ElasticSearchReader])
 
-  private val traceReader: GrpcTraceReader = new GrpcTraceReader(traceBackendConfig)
+  private val traceReader: GrpcTraceReaders = new GrpcTraceReaders(traceStoreBackendConfig)
   private val esReader: ElasticSearchReader = new ElasticSearchReader(elasticSearchConfiguration.clientConfiguration)
   private val traceSearchQueryGenerator = new TraceSearchQueryGenerator(elasticSearchConfiguration.spansIndexConfiguration, ES_NESTED_DOC_NAME, whitelistedFieldsConfiguration)
   private val traceCountsQueryGenerator = new TraceCountsQueryGenerator(elasticSearchConfiguration.spansIndexConfiguration, ES_NESTED_DOC_NAME, whitelistedFieldsConfiguration)
diff --git a/reader/src/main/scala/com/expedia/www/haystack/trace/reader/stores/readers/grpc/GrpcTraceReaders.scala b/reader/src/main/scala/com/expedia/www/haystack/trace/reader/stores/readers/grpc/GrpcTraceReaders.scala
@@ -19,37 +19,68 @@ package com.expedia.www.haystack.trace.reader.stores.readers.grpc
 import com.expedia.open.tracing.api.Trace
 import com.expedia.open.tracing.backend.{ReadSpansRequest, StorageBackendGrpc}
 import com.expedia.www.haystack.commons.metrics.MetricsSupport
-import com.expedia.www.haystack.trace.commons.config.entities.TraceBackendClientConfiguration
+import com.expedia.www.haystack.trace.commons.config.entities.TraceStoreBackends
+import com.expedia.www.haystack.trace.reader.exceptions.TraceNotFoundException
 import com.expedia.www.haystack.trace.reader.metrics.AppMetricNames
-import io.grpc.ManagedChannelBuilder
+import com.expedia.www.haystack.trace.reader.readers.utils.TraceMerger
+import io.grpc.{ManagedChannel, ManagedChannelBuilder}
 import org.slf4j.LoggerFactory
 
 import scala.collection.JavaConverters._
 import scala.concurrent.{ExecutionContextExecutor, Future, Promise}
 
-class GrpcTraceReader(config: TraceBackendClientConfiguration)
-                     (implicit val dispatcher: ExecutionContextExecutor) extends MetricsSupport with AutoCloseable {
-  private val LOGGER = LoggerFactory.getLogger(classOf[GrpcTraceReader])
+class GrpcTraceReaders(config: TraceStoreBackends)
+                      (implicit val dispatcher: ExecutionContextExecutor) extends MetricsSupport with AutoCloseable {
+  private val LOGGER = LoggerFactory.getLogger(classOf[GrpcTraceReaders])
 
   private val readTimer = metricRegistry.timer(AppMetricNames.BACKEND_READ_TIME)
   private val readFailures = metricRegistry.meter(AppMetricNames.BACKEND_READ_FAILURES)
   private val tracesFailures = metricRegistry.meter(AppMetricNames.BACKEND_TRACES_FAILURE)
-  private val channel = ManagedChannelBuilder
-    .forAddress(config.host, config.port)
-    .usePlaintext(true)
-    .build()
 
-  val client: StorageBackendGrpc.StorageBackendFutureStub = StorageBackendGrpc.newFutureStub(channel)
+  private val clients: Seq[GrpcChannelClient] =  config.backends.map {
+    backend => {
+      val channel = ManagedChannelBuilder
+        .forAddress(backend.host, backend.port)
+        .usePlaintext(true)
+        .build()
 
+      val client = StorageBackendGrpc.newFutureStub(channel)
+      GrpcChannelClient(channel, client)
+    }
+  }
 
   def readTraces(traceIds: List[String]): Future[Seq[Trace]] = {
+    val allFutures = clients.map {
+      client =>
+        readTraces(traceIds, client.stub) recoverWith  {
+          case _: Exception => Future.successful(Seq.empty[Trace])
+        }
+    }
+
+    Future.sequence(allFutures)
+      .map(traceSeq => traceSeq.flatten)
+      .map {
+        traces =>
+          if (traces.isEmpty) throw new TraceNotFoundException() else TraceMerger.merge(traces)
+      }
+  }
+
+  private def readTraces(traceIds: List[String], client: StorageBackendGrpc.StorageBackendFutureStub): Future[Seq[Trace]] = {
     val timer = readTimer.time()
     val promise = Promise[Seq[Trace]]
 
     try {
-      val readSpansRequest = ReadSpansRequest.newBuilder().addAllTraceIds(traceIds.asJava).build()
+      val readSpansRequest = ReadSpansRequest.newBuilder().addAllTraceIds(traceIds.asJavaCollection).build()
       val futureResponse = client.readSpans(readSpansRequest)
-      futureResponse.addListener(new ReadSpansResponseListener(futureResponse, promise, timer, readFailures, tracesFailures, traceIds.size), dispatcher)
+      futureResponse.addListener(new ReadSpansResponseListener(
+        futureResponse,
+        promise,
+        timer,
+        readFailures,
+        tracesFailures,
+        traceIds.size), dispatcher)
+
+      // return the future with the results for the given client
       promise.future
     } catch {
       case ex: Exception =>
@@ -61,6 +92,8 @@ class GrpcTraceReader(config: TraceBackendClientConfiguration)
   }
 
   override def close(): Unit = {
-    channel.shutdown()
+    clients.foreach(_.channel.shutdown())
   }
+
+  case class GrpcChannelClient(channel: ManagedChannel, stub: StorageBackendGrpc.StorageBackendFutureStub)
 }
diff --git a/reader/src/test/scala/com/expedia/www/haystack/trace/reader/unit/config/ConfigurationLoaderSpec.scala b/reader/src/test/scala/com/expedia/www/haystack/trace/reader/unit/config/ConfigurationLoaderSpec.scala
@@ -75,8 +75,8 @@ class ConfigurationLoaderSpec extends BaseUnitTestSpec {
 
     it("should load trace backend configuration") {
       val traceBackendConfig = new ProviderConfiguration().traceBackendConfiguration
-      traceBackendConfig.host shouldEqual "localhost"
-      traceBackendConfig.port shouldEqual 8090
+      traceBackendConfig.backends.head.host shouldEqual "localhost"
+      traceBackendConfig.backends.head.port shouldEqual 8090
     }
   }
 }
diff --git a/reader/src/test/scala/com/expedia/www/haystack/trace/reader/unit/readers/TraceMergerSpec.scala b/reader/src/test/scala/com/expedia/www/haystack/trace/reader/unit/readers/TraceMergerSpec.scala
@@ -0,0 +1,61 @@
+/*
+ *  Copyright 2019 Expedia, Inc.
+ *
+ *       Licensed 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.expedia.www.haystack.trace.reader.unit.readers
+
+import com.expedia.open.tracing.Span
+import com.expedia.open.tracing.api.Trace
+import com.expedia.www.haystack.trace.reader.readers.utils.TraceMerger
+import com.expedia.www.haystack.trace.reader.unit.BaseUnitTestSpec
+
+class TraceMergerSpec extends BaseUnitTestSpec {
+  describe("Trace Merger") {
+    it("should merge the traces by traceId") {
+      val trace_1 = buildTrace("t1", "s1", "svc1", "op1")
+      val trace_2 = buildTrace("t2", "s2", "svc2", "op1")
+      val trace_3 = buildTrace("t1", "s3", "svc3", "op1")
+      val trace_4 = buildTrace("t2", "s4", "svc4", "op1")
+      val trace_5 = buildTrace("t3", "s5", "svc5", "op1")
+
+      val mergedTraces = TraceMerger.merge(Seq(trace_1, trace_2, trace_3, trace_4, trace_5))
+      mergedTraces.size shouldBe 3
+      val t1 = mergedTraces.find(_.getTraceId == "t1").head
+      t1.getChildSpansCount shouldBe 2
+      t1.getChildSpans(0).getTraceId shouldBe "t1"
+      t1.getChildSpans(0).getSpanId shouldBe "s1"
+      t1.getChildSpans(1).getTraceId shouldBe "t1"
+      t1.getChildSpans(1).getSpanId shouldBe "s3"
+      t1.getChildSpans(0).getServiceName shouldBe "svc1"
+      t1.getChildSpans(1).getServiceName shouldBe "svc3"
+
+      val t2 = mergedTraces.find(_.getTraceId == "t2").head
+      t2.getChildSpansCount shouldBe 2
+      t2.getChildSpans(0).getSpanId shouldBe "s2"
+      t2.getChildSpans(1).getSpanId shouldBe "s4"
+
+      mergedTraces.find(_.getTraceId == "t3").head.getChildSpansCount shouldBe 1
+    }
+  }
+
+  private def buildTrace(traceId: String, spanId: String, serviceName: String, operationName: String): Trace = {
+    Trace.newBuilder().setTraceId(traceId).addChildSpans(
+      Span.newBuilder()
+        .setSpanId(spanId)
+        .setTraceId(traceId)
+        .setServiceName(serviceName)
+        .setOperationName(operationName)).build()
+  }
+}
diff --git a/reader/src/test/scala/com/expedia/www/haystack/trace/reader/unit/readers/TraceProcessorSpec.scala b/reader/src/test/scala/com/expedia/www/haystack/trace/reader/unit/readers/TraceProcessorSpec.scala
@@ -1,3 +1,19 @@
+/*
+ *  Copyright 2017 Expedia, Inc.
+ *
+ *       Licensed 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.expedia.www.haystack.trace.reader.unit.readers
 
 import com.expedia.open.tracing.api.Trace

Original file line number	Diff line number	Diff line change
`@@ -15,6 +15,14 @@ backend {`
`15`	`15`	`host = "localhost"`
`16`	`16`	`port = 8090`
`17`	`17`	`}`
	`18`	`+`
	`19`	`+ # we support multiple grpc based backends, to provide another one use something like following.`
	`20`	`+ # you are required to provide host and port`
	`21`	`+`
	`22`	`+ # another_client {`
	`23`	`+ # host = "localhost"`
	`24`	`+ # port = 8092`
	`25`	`+ # }`
`18`	`26`	`}`
`19`	`27`
`20`	`28`	`elasticsearch {`