feat(observability): API health checks + GraphQL

Cargo.toml

@@ -32,7 +32,7 @@ members = [
   "lib/tracing-limit",
   "lib/vector-wasm",
   "lib/k8s-test-framework",
-  "lib/prometheus-parser",
+  "lib/prometheus-parser"
 ]
 
 [dependencies]
@@ -91,6 +91,10 @@ prost-types = "0.6.1"
 goauth = { version = "0.7.1", optional = true }
 smpl_jwt = { version = "0.5.0", optional = true }
 
+# API
+async-graphql = { version = "1.17.11", optional = true }
+async-graphql-warp = { version = "1.17.3", optional = true }
+
 # External libs
 derivative = "1.0"
 chrono = { version = "0.4.6", features = ["serde"] }
@@ -131,7 +135,7 @@ uuid = { version = "0.7", features = ["serde", "v4"], optional = true }
 exitcode = "1.1.2"
 snafu = { version = "0.6", features = ["futures-01", "futures"] }
 url = "2.1.1"
-base64 = { version = "0.10.1", optional = true }
+base64 = { version = "0.12.3", optional = true }
 bollard = { version = "0.8.0", optional = true }
 listenfd = { version = "0.3.3", optional = true }
 inventory = "0.1"
@@ -197,13 +201,13 @@ k8s-test-framework = { version = "0.1", path = "lib/k8s-test-framework" }
 
 [features]
 # Default features for *-unknown-linux-gnu and *-apple-darwin
-default = ["sources", "transforms", "sinks", "vendor-all", "unix", "leveldb-plain", "rdkafka-plain"]
-default-musl = ["sources", "transforms", "sinks", "vendor-all", "unix", "leveldb-cmake", "rdkafka-cmake"]
+default = ["api", "sources", "transforms", "sinks", "vendor-all", "unix", "leveldb-plain", "rdkafka-plain"]
+default-musl = ["api", "sources", "transforms", "sinks", "vendor-all", "unix", "leveldb-cmake", "rdkafka-cmake"]
 # Default features for *-unknown-linux-* which make use of `cmake` for dependencies
-default-cmake = ["sources", "transforms", "sinks", "vendor-all", "unix", "leveldb-cmake", "rdkafka-cmake"]
+default-cmake = ["api", "sources", "transforms", "sinks", "vendor-all", "unix", "leveldb-cmake", "rdkafka-cmake"]
 # Default features for *-pc-windows-msvc
 # TODO: Enable SASL https://github.com/timberio/vector/pull/3081#issuecomment-659298042
-default-msvc = ["sources", "transforms", "sinks", "vendor-openssl", "vendor-libz", "leveldb-cmake", "rdkafka-cmake"]
+default-msvc = ["api", "sources", "transforms", "sinks", "vendor-openssl", "vendor-libz", "leveldb-cmake", "rdkafka-cmake"]
 
 # Enables features that work only on systems providing `cfg(unix)`
 unix = ["jemallocator"]
@@ -231,6 +235,9 @@ wasm-timings = ["wasm"]
 # transforms and sinks should depend on this feature.
 kubernetes = ["k8s-openapi", "evmap"]
 
+# Enables the API
+api = ["warp", "async-graphql", "async-graphql-warp"]
+
 # Sources
 sources = [
   "sources-docker",

src/api/handler.rs

@@ -0,0 +1,13 @@
+use serde::Serialize;
+use warp::reply::json;
+use warp::{Rejection, Reply};
+
+#[derive(Serialize)]
+struct Health {
+    ok: bool,
+}
+
+// health handler, responds with { ok: true }
+pub async fn health() -> Result<impl Reply, Rejection> {
+    Ok(json(&Health { ok: true }))
+}

src/api/macros.rs

@@ -0,0 +1,68 @@
+/// Takes a list of handler expressions and `or`s them together
+/// in a balanced tree. That is, instead of `a.or(b).or(c).or(d)`,
+/// it produces `(a.or(b)).or(c.or(d))`, thus nesting the types
+/// less deeply, which provides improvements in compile time.
+///
+/// It also applies `::warp::Filter::boxed` to each handler expression
+/// when in `debug_assertions` mode, improving compile time further.
+//
+// The basic list splitting algorithm here is based on this gist:
+// https://gist.github.com/durka/9fc479de2555225a787f
+// It uses a counter from which two items are removed each time,
+// stopping when the counter reaches 0. At each step, one item
+// is moved from the left to the right, and thus at the end,
+// there will be the same number of items in each list.
+//
+// The flow is as follows:
+// - If there is one handler expression, debug_box it and return.
+// - If there is more than one handler expression:
+//   - First, copy the list into two: the one that will go into the
+//     right side of the `or`, and one that will serve as a counter.
+//     Recurse with these separated by semicolons, plus an empty `left`
+//     list before the first semicolon.
+//   - Then, as long as there are at least two items in the counter
+//     list, remove them and move the first item on the right side of
+//     the first semicolon (`head`) to the left side of the first semicolon.
+//   - Finally, when there are one or zero items left in the counter,
+//     move one last item to the left, make the call this macro on both the
+//     left and right sides, and `or` the two sides together.
+//
+// For example, balanced_or_tree!(a, b, c, d, e) would take the following steps:
+//
+// - balanced_or_tree!(a, b, c, d, e)
+// - balanced_or_tree!(@internal ; a, b, c, d, e ; a, b, c, d, e) // initialise lists
+// - balanced_or_tree!(@internal a ; b, c, d, e ; c, d, e) // move one elem; remove two
+// - balanced_or_tree!(@internal a, b ; c, d, e ; e) // now only one elem in counter
+// - balanced_or_tree!(a, b, c).or(balanced_or_tree(d, e)) // recurse on each sublist
+#[macro_export]
+macro_rules! balanced_or_tree {
+    // Base case: just a single expression, return it wrapped in `debug_boxed`
+    ($x:expr $(,)?) => { debug_boxed!($x) };
+    // Multiple expressions: recurse with three lists: left, right and counter.
+    ($($x:expr),+ $(,)?) => {
+        balanced_or_tree!(@internal  ;     $($x),+; $($x),+)
+        //                          ^ left ^ right  ^ counter
+    };
+    // Counter 1 or 2; move one more item and recurse on each sublist, and or them together
+    (@internal $($left:expr),*; $head:expr, $($tail:expr),+; $a:expr $(,$b:expr)?) => {
+        (balanced_or_tree!($($left,)* $head)).or(balanced_or_tree!($($tail),+))
+    };
+    // Counter > 2; move one item from the right to the left and subtract two from the counter
+    (@internal $($left:expr),*; $head:expr, $($tail:expr),+; $a:expr, $b:expr, $($more:expr),+) => {
+        balanced_or_tree!(@internal $($left,)* $head; $($tail),+; $($more),+)
+    };
+}
+
+#[cfg(debug_assertions)]
+macro_rules! debug_boxed {
+    ($x:expr) => {
+        ::warp::Filter::boxed($x)
+    };
+}
+
+#[cfg(not(debug_assertions))]
+macro_rules! debug_boxed {
+    ($x:expr) => {
+        $x
+    };
+}

src/api/mod.rs

@@ -0,0 +1,7 @@
+#[macro_use]
+mod macros;
+mod handler;
+mod schema;
+mod server;
+
+pub use server::Server;

src/api/schema/health.rs

@@ -0,0 +1,40 @@
+use async_graphql::{validators::IntRange, Object, SimpleObject, Subscription};
+use chrono::{DateTime, Utc};
+use tokio::stream::{Stream, StreamExt};
+use tokio::time::Duration;
+
+#[SimpleObject]
+pub struct Heartbeat {
+    utc: DateTime<Utc>,
+}
+
+impl Heartbeat {
+    fn new() -> Self {
+        Heartbeat { utc: Utc::now() }
+    }
+}
+
+#[derive(Default)]
+pub struct HealthQuery;
+
+#[Object]
+impl HealthQuery {
+    /// Returns `true` to denote the GraphQL server is reachable
+    async fn health(&self) -> bool {
+        true
+    }
+}
+
+#[derive(Default)]
+pub struct HealthSubscription;
+
+#[Subscription]
+impl HealthSubscription {
+    /// Heartbeat, containing the UTC timestamp of the last server-sent payload
+    async fn heartbeat(
+        &self,
+        #[arg(default = 1000, validator(IntRange(min = "100", max = "60_000")))] interval: i32,
+    ) -> impl Stream<Item = Heartbeat> {
+        tokio::time::interval(Duration::from_millis(interval as u64)).map(|_| Heartbeat::new())
+    }
+}

src/api/schema/metrics.rs

@@ -0,0 +1,88 @@
+use crate::event::{Event, Metric, MetricValue};
+use crate::metrics::{capture_metrics, get_controller};
+use async_graphql::{validators::IntRange, Interface, Object, Subscription};
+use async_stream::stream;
+use chrono::{DateTime, Utc};
+use tokio::stream::{Stream, StreamExt};
+use tokio::time::Duration;
+
+pub struct EventsProcessed(Metric);
+
+#[Object]
+impl EventsProcessed {
+    /// Metric timestamp
+    async fn timestamp(&self) -> Option<DateTime<Utc>> {
+        self.0.timestamp
+    }
+
+    /// Number of events processed
+    async fn events_processed(&self) -> f64 {
+        match self.0.value {
+            MetricValue::Counter { value } => value,
+            _ => 0.00,
+        }
+    }
+}
+
+impl Into<EventsProcessed> for Metric {
+    fn into(self) -> EventsProcessed {
+        EventsProcessed(self)
+    }
+}
+
+pub struct BytesProcessed(Metric);
+
+#[Object]
+impl BytesProcessed {
+    /// Metric timestamp
+    async fn timestamp(&self) -> Option<DateTime<Utc>> {
+        self.0.timestamp
+    }
+
+    /// Number of bytes processed
+    async fn bytes_processed(&self) -> f64 {
+        match self.0.value {
+            MetricValue::Counter { value } => value,
+            _ => 0.00,
+        }
+    }
+}
+
+#[Interface(field(name = "timestamp", type = "Option<DateTime<Utc>>"))]
+pub enum MetricType {
+    EventsProcessed(EventsProcessed),
+    BytesProcessed(BytesProcessed),
+}
+
+#[derive(Default)]
+pub struct MetricsSubscription;
+
+#[Subscription]
+impl MetricsSubscription {
+    /// Returns all Vector metrics, aggregated at the provided millisecond interval
+    async fn metrics(
+        &self,
+        #[arg(validator(IntRange(min = "100", max = "60_000")))] interval: i32,
+    ) -> impl Stream<Item = MetricType> {
+        get_metrics(interval).filter_map(|m| match m.name.as_str() {
+            "events_processed" => Some(MetricType::EventsProcessed(m.into())),
+            _ => None,
+        })
+    }
+}
+
+/// Returns a stream of `Metric`s, collected at the provided millisecond interval
+fn get_metrics(interval: i32) -> impl Stream<Item = Metric> {
+    let controller = get_controller().unwrap();
+    let mut interval = tokio::time::interval(Duration::from_millis(interval as u64));
+
+    stream! {
+        while let _ = interval.tick().await {
+            for ev in capture_metrics(&controller) {
+                if let Event::Metric(m) = ev {
+                    yield m
+                }
+            }
+        }
+    }
+}

src/api/schema/mod.rs

@@ -0,0 +1,15 @@
+mod health;
+mod metrics;
+
+use async_graphql::{EmptyMutation, GQLMergedObject, GQLMergedSubscription, Schema, SchemaBuilder};
+
+#[derive(GQLMergedObject, Default)]
+pub struct Query(health::HealthQuery);
+
+#[derive(GQLMergedSubscription, Default)]
+pub struct Subscription(health::HealthSubscription, metrics::MetricsSubscription);
+
+/// Build a new GraphQL schema, comprised of Query, Mutation and Subscription types
+pub fn build_schema() -> SchemaBuilder<Query, EmptyMutation, Subscription> {
+    Schema::build(Query::default(), EmptyMutation, Subscription::default())
+}