Atlas Cloud Watch: Add the Firehose endpoint that will receive messag… (

Netflix-Skunkworks#415)

* Atlas Cloud Watch: Add the Firehose endpoint that will receive messages from
Firehose, parse and process them. Includes the RequestId class that stores the
timestamp and request ID that firehose requires in the response to accept the
data as successfully received.

atlas-cloudwatch/src/main/scala/com/netflix/atlas/webapi/CloudWatchFirehoseEndpoint.scala

@@ -0,0 +1,261 @@
+/*
+ * Copyright 2014-2023 Netflix, 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.netflix.atlas.webapi
+
+import akka.actor.ActorSystem
+import akka.http.scaladsl.model.HttpRequest
+import akka.http.scaladsl.model.StatusCodes
+import akka.http.scaladsl.server.Directives._
+import akka.http.scaladsl.server.Route
+import com.fasterxml.jackson.core.JsonParser
+import com.netflix.atlas.akka.CustomDirectives.customJson
+import com.netflix.atlas.akka.CustomDirectives.endpointPath
+import com.netflix.atlas.akka.CustomDirectives.parseEntity
+import com.netflix.atlas.akka.WebApi
+import com.netflix.atlas.cloudwatch.CloudWatchMetricsProcessor
+import com.netflix.atlas.cloudwatch.FirehoseMetric
+import com.netflix.atlas.json.Json
+import com.netflix.atlas.json.JsonParserHelper.foreachField
+import com.netflix.atlas.json.JsonParserHelper.foreachItem
+import com.netflix.atlas.webapi.CloudWatchFirehoseEndpoint.decodeCommonTags
+import com.netflix.atlas.webapi.CloudWatchFirehoseEndpoint.decodeMetricJson
+import com.netflix.spectator.api.Registry
+import com.typesafe.scalalogging.StrictLogging
+import software.amazon.awssdk.services.cloudwatch.model.Datapoint
+import software.amazon.awssdk.services.cloudwatch.model.Dimension
+
+import java.time.Instant
+import java.util.Base64
+import scala.collection.mutable
+import scala.util.Using
+
+class CloudWatchFirehoseEndpoint(
+  registry: Registry,
+  processor: CloudWatchMetricsProcessor
+)(implicit val system: ActorSystem)
+    extends WebApi
+    with StrictLogging {
+
+  private val dataReceived = registry.counter("atlas.cloudwatch.firehose.parse.dps")
+  private val unknownField = registry.counter("atlas.cloudwatch.firehose.parse.unknown")
+  private val parseException = registry.createId("atlas.cloudwatch.firehose.parse.exception")
+
+  override def routes: Route = {
+    (post | put) {
+      endpointPath("api" / "v1" / "firehose") {
+        handleReq
+      } ~
+      complete(StatusCodes.NotFound)
+    }
+  }
+
+  private def handleReq: Route = {
+    extractRequestContext { ctx =>
+      val commonTags = decodeCommonTags(ctx.request)
+      parseEntity(customJson(p => parseFirehoseMessage(p, commonTags))) { response =>
+        if (response.exception.isDefined) {
+          complete(StatusCodes.BadRequest, response.getEntity)
+        } else {
+          complete(StatusCodes.OK, response.getEntity)
+        }
+      }
+    }
+  }
+
+  private def parseFirehoseMessage(parser: JsonParser, commonTags: List[Dimension]): RequestId = {
+    var requestId: String = null
+    var timestamp: Long = 0
+    val receivedTimestamp = System.currentTimeMillis()
+    try {
+      foreachField(parser) {
+        case "requestId" => requestId = parser.nextTextValue()
+        case "timestamp" => timestamp = parser.nextLongValue(0L)
+        case "records" =>
+          foreachItem(parser) {
+            foreachField(parser) {
+              case "data" =>
+                try {
+                  // catch the exception here in the case one data set is errant but the others are alright.
+                  // At least we should get _some_ data that way.
+                  Using.resource(
+                    Json.newJsonParser(Base64.getDecoder.decode(parser.nextTextValue()))
+                  ) { p =>
+                    val datapoints = decodeMetricJson(p, commonTags)
+                    processor.processDatapoints(datapoints, receivedTimestamp)
+                    dataReceived.increment(datapoints.size)
+                  }
+                } catch {
+                  case ex: Exception =>
+                    logger.warn("Unexpected data parse exception", ex)
+                    incrementException(ex, true)
+                }
+              case unknown => // Ignore unknown fields
+                logger.debug("Skipping unknown firehose record field: {}", unknown)
+                parser.nextToken()
+                parser.skipChildren()
+                unknownField.increment()
+            }
+          }
+        case unknown => // Ignore unknown fields
+          logger.debug("Skipping unknown firehose field: {}", unknown)
+          parser.nextToken()
+          parser.skipChildren()
+          unknownField.increment()
+      }
+      RequestId(requestId, timestamp)
+    } catch {
+      case ex: Exception =>
+        logger.error("Failed to parse request", ex)
+        incrementException(ex)
+        RequestId(requestId, timestamp, Some(ex))
+    }
+  }
+
+  private def incrementException(ex: Exception, data: Boolean = false): Unit = {
+    val id = if (data) "data" else "outer"
+    registry.counter(parseException.withTags("id", id, "ex", ex.getClass.getSimpleName)).increment()
+  }
+}
+
+object CloudWatchFirehoseEndpoint extends StrictLogging {
+
+  private[webapi] def decodeMetricJson(
+    parser: JsonParser,
+    commonTags: List[Dimension]
+  ): List[FirehoseMetric] = {
+    // JSON stream format, i.e. new line separated list of objects.
+    var t = parser.nextToken()
+    val results = List.newBuilder[FirehoseMetric]
+
+    var dp = Datapoint.builder()
+    var metricStreamName: String = null
+    var accountId: String = null
+    var region: String = null
+    var namespace: String = null
+    var metricName: String = null
+    var dimensions = List.newBuilder[Dimension]
+    while (t != null) {
+      foreachField(parser) {
+        case "metric_stream_name" => metricStreamName = parser.nextTextValue()
+        case "account_id"         => accountId = parser.nextTextValue
+        case "region"             => region = parser.nextTextValue
+        case "namespace"          => namespace = parser.nextTextValue
+        case "metric_name"        => metricName = parser.nextTextValue
+        case "unit"               => dp.unit(parser.nextTextValue())
+        case "dimensions"         => decodeDimensions(parser, dimensions)
+        case "timestamp" =>
+          dp.timestamp(Instant.ofEpochMilli(parser.nextLongValue(0L)))
+        case "value" => decodeValue(parser, dp)
+        case unknown => // Ignore unknown fields
+          logger.debug("Skipping unknown firehose metric field: {}", unknown)
+          parser.nextToken()
+          parser.skipChildren()
+      }
+
+      if (accountId != null) {
+        dimensions += Dimension
+          .builder()
+          .name("nf.account")
+          .value(accountId)
+          .build()
+      }
+      if (region != null) {
+        dimensions += Dimension
+          .builder()
+          .name("nf.region")
+          .value(region)
+          .build()
+      }
+      dimensions ++= commonTags
+
+      results += FirehoseMetric(
+        metricStreamName,
+        namespace,
+        metricName,
+        dimensions.result(),
+        dp.build()
+      )
+
+      // reset
+      dp = Datapoint.builder()
+      metricStreamName = null
+      accountId = null
+      region = null
+      namespace = null
+      metricName = null
+      dimensions = List.newBuilder[Dimension]
+
+      // advance
+      t = parser.nextToken()
+    }
+    results.result()
+  }
+
+  private[webapi] def decodeDimensions(
+    parser: JsonParser,
+    dimensions: mutable.Builder[Dimension, List[Dimension]]
+  ): Unit = {
+    foreachField(parser) {
+      case key =>
+        val value = parser.nextTextValue()
+        if (value != null) {
+          dimensions += Dimension
+            .builder()
+            .name(key)
+            .value(value)
+            .build()
+        }
+    }
+  }
+
+  private[webapi] def decodeValue(parser: JsonParser, dp: Datapoint.Builder): Unit = {
+    foreachField(parser) {
+      case "sum" =>
+        parser.nextToken()
+        dp.sum(parser.getDoubleValue)
+      case "count" =>
+        parser.nextToken()
+        dp.sampleCount(parser.getDoubleValue)
+      case "max" =>
+        parser.nextToken()
+        dp.maximum(parser.getDoubleValue)
+      case "min" =>
+        parser.nextToken()
+        dp.minimum(parser.getDoubleValue)
+      case unknown => // Ignore unknown fields
+        // TODO - Extend this if we need to record any of the other fields offered by AWS.
+        // Note that each field is an additional financial cost though.
+        logger.debug("Skipping unknown value field: {}", unknown)
+        parser.nextToken()
+        parser.skipChildren()
+    }
+  }
+
+  private[webapi] def decodeCommonTags(request: HttpRequest): List[Dimension] = {
+    val hdr = request.getHeader("X-Amz-Firehose-Common-Attributes")
+    if (hdr.isPresent) {
+      Using.resource(Json.newJsonParser(hdr.get().value())) { parser =>
+        val dimensions = List.newBuilder[Dimension]
+        foreachField(parser) {
+          case "commonAttributes" => decodeDimensions(parser, dimensions)
+        }
+        dimensions.result()
+      }
+    } else {
+      List.empty
+    }
+  }
+}

atlas-cloudwatch/src/main/scala/com/netflix/atlas/webapi/RequestId.scala

@@ -0,0 +1,69 @@
+/*
+ * Copyright 2014-2023 Netflix, 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.netflix.atlas.webapi
+
+import akka.http.scaladsl.model.ContentTypes
+import akka.http.scaladsl.model.HttpEntity
+import com.netflix.atlas.json.Json
+import com.netflix.atlas.webapi.RequestId.getByteArrayOutputStream
+
+import java.io.ByteArrayOutputStream
+import scala.util.Using
+
+/**
+  * This is the class we have to use to respond to the Firehose publisher. The request ID and timestamp must match
+  * those of the message received in order for Firehose to consider the data successfully processed.
+  *
+  * @param requestId
+  *     The request ID given by AWS Firehose.
+  * @param timestamp
+  *     The timestamp in unix epoch milliseconds given by AWS Firehose.
+  * @param exception
+  *     An optional exception to encode if processing failed. This will be stored in Cloud Watch logs.
+  */
+case class RequestId(requestId: String, timestamp: Long, exception: Option[Exception] = None) {
+
+  def getEntity: HttpEntity.Strict = {
+    val stream = getByteArrayOutputStream
+    Using.resource(Json.newJsonGenerator(stream)) { json =>
+      json.writeStartObject()
+      json.writeStringField("requestId", requestId)
+      json.writeNumberField("timestamp", timestamp)
+      exception.map { ex =>
+        json.writeStringField("errorMessage", ex.getMessage)
+      }
+      json.writeEndObject()
+    }
+    HttpEntity(ContentTypes.`application/json`, stream.toByteArray)
+  }
+
+}
+
+object RequestId {
+
+  private val byteArrayStreams = new ThreadLocal[ByteArrayOutputStream]
+
+  private[atlas] def getByteArrayOutputStream: ByteArrayOutputStream = {
+    var baos = byteArrayStreams.get()
+    if (baos == null) {
+      baos = new ByteArrayOutputStream()
+      byteArrayStreams.set(baos)
+    } else {
+      baos.reset()
+    }
+    baos
+  }
+}