go/runtime + agent: introduce /notify/shard-failure API

The runtime invokes a new /notify/shard-failure control-plane API which
is told of shard failures that have occurred within a data-plane.

At the moment, this API verifies the data-plane token and logs the
failure, but takes no further action.

Update the taskBase.heartbeatLoop() to perform this notification if the
shard's primary loop exits with a non-cancellation error.

Issue #1666

crates/agent/src/api/authorize_task.rs

@@ -15,28 +15,7 @@ pub async fn authorize_task(
 
 #[tracing::instrument(skip(app), err(level = tracing::Level::WARN))]
 async fn do_authorize_task(app: &App, Request { token }: &Request) -> anyhow::Result<Response> {
- let jsonwebtoken::TokenData { header, mut claims }: jsonwebtoken::TokenData<
- proto_gazette::Claims,
- > = {
- // In this pass we do not validate the signature,
- // because we don't yet know which data-plane the JWT is signed by.
- let empty_key = jsonwebtoken::DecodingKey::from_secret(&[]);
- let mut validation = jsonwebtoken::Validation::default();
- validation.insecure_disable_signature_validation();
- jsonwebtoken::decode(token, &empty_key, &validation)
- }?;
- tracing::debug!(?claims, ?header, "decoded authorization request");
-
- let shard_id = claims.sub.as_str();
- if shard_id.is_empty() {
- anyhow::bail!("missing required shard ID (`sub` claim)");
- }
-
- let shard_data_plane_fqdn = claims.iss.as_str();
- if shard_data_plane_fqdn.is_empty() {
- anyhow::bail!("missing required shard data-plane FQDN (`iss` claim)");
- }
-
+ let (header, mut claims) = super::parse_untrusted_data_plane_claims(token)?;
  let journal_name_or_prefix = labels::expect_one(claims.sel.include(), "name")?.to_owned();
 
  // Require the request was signed with the AUTHORIZE capability,
@@ -62,8 +41,8 @@ async fn do_authorize_task(app: &App, Request { token }: &Request) -> anyhow::Re
  match Snapshot::evaluate(&app.snapshot, claims.iat, |snapshot: &Snapshot| {
  evaluate_authorization(
  snapshot,
- shard_id,
- shard_data_plane_fqdn,
+ &claims.sub,
+ &claims.iss,
  token,
  &journal_name_or_prefix,
  required_role,
@@ -104,47 +83,13 @@ fn evaluate_authorization(
  journal_name_or_prefix: &str,
  required_role: models::Capability,
 ) -> anyhow::Result<(jsonwebtoken::EncodingKey, String, String)> {
- // Map `claims.sub`, a Shard ID, into its task.
- let task = tasks
- .binary_search_by(|task| {
- if shard_id.starts_with(&task.shard_template_id) {
- std::cmp::Ordering::Equal
- } else {
- task.shard_template_id.as_str().cmp(shard_id)
- }
- })
- .ok()
- .map(|index| &tasks[index]);
-
- // Map `claims.iss`, a data-plane FQDN, into its task-matched data-plane.
- let task_data_plane = task.and_then(|task| {
- data_planes
- .get_by_key(&task.data_plane_id)
- .filter(|data_plane| data_plane.data_plane_fqdn == shard_data_plane_fqdn)
- });
-
- let (Some(task), Some(task_data_plane)) = (task, task_data_plane) else {
- anyhow::bail!(
- "task shard {shard_id} within data-plane {shard_data_plane_fqdn} is not known"
- )
- };
-
- // Attempt to find an HMAC key of this data-plane which validates against the request token.
- let validation = jsonwebtoken::Validation::default();
- let mut verified = false;
-
- for hmac_key in &task_data_plane.hmac_keys {
- let key = jsonwebtoken::DecodingKey::from_base64_secret(hmac_key)
- .context("invalid data-plane hmac key")?;
-
- if jsonwebtoken::decode::<proto_gazette::Claims>(token, &key, &validation).is_ok() {
- verified = true;
- break;
- }
- }
- if !verified {
- anyhow::bail!("no data-plane keys validated against the token signature");
- }
+ let (task, task_data_plane) = super::verify_data_plane_claims(
+ data_planes,
+ tasks,
+ shard_id,
+ shard_data_plane_fqdn,
+ token,
+ )?;
 
  // Map a required `name` journal label selector into its collection.
  let Some(collection) = collections

crates/agent/src/api/mod.rs

@@ -5,6 +5,7 @@ mod authorize_task;
 mod authorize_user_collection;
 mod authorize_user_task;
 mod create_data_plane;
+mod notify_shard_failure;
 mod snapshot;
 mod update_l2_reporting;
 
@@ -125,6 +126,10 @@ pub fn build_router(
  post(update_l2_reporting::update_l2_reporting)
  .route_layer(axum::middleware::from_fn_with_state(app.clone(), authorize)),
  )
+ .route(
+ "/notify/shard-failure",
+ post(notify_shard_failure::notify_shard_failure),
+ )
  .layer(tower_http::trace::TraceLayer::new_for_http())
  .layer(cors)
  .with_state(app);
@@ -239,3 +244,81 @@ fn maybe_rewrite_address(external: bool, address: &str) -> String {
  address.to_string()
  }
 }
+
+// Parse a data-plane claims token without verifying it's signature.
+fn parse_untrusted_data_plane_claims(
+ token: &str,
+) -> anyhow::Result<(jsonwebtoken::Header, proto_gazette::Claims)> {
+ let jsonwebtoken::TokenData { header, claims }: jsonwebtoken::TokenData<proto_gazette::Claims> =
+ {
+ // In this pass we do not validate the signature,
+ // because we don't yet know which data-plane the JWT is signed by.
+ let empty_key = jsonwebtoken::DecodingKey::from_secret(&[]);
+ let mut validation = jsonwebtoken::Validation::default();
+ validation.insecure_disable_signature_validation();
+ jsonwebtoken::decode(token, &empty_key, &validation)
+ }?;
+
+ if claims.sub.is_empty() {
+ anyhow::bail!("missing required shard ID (`sub` claim)");
+ }
+ if claims.iss.is_empty() {
+ anyhow::bail!("missing required shard data-plane FQDN (`iss` claim)");
+ }
+
+ tracing::debug!(?claims, ?header, "decoded authorization request");
+
+ Ok((header, claims))
+}
+
+fn verify_data_plane_claims<'s>(
+ data_planes: &'s tables::DataPlanes,
+ tasks: &'s [snapshot::SnapshotTask],
+ shard_id: &str,
+ shard_data_plane_fqdn: &str,
+ token: &str,
+) -> anyhow::Result<(&'s snapshot::SnapshotTask, &'s tables::DataPlane)> {
+ // Map `shard_id` into its task.
+ let task = tasks
+ .binary_search_by(|task| {
+ if shard_id.starts_with(&task.shard_template_id) {
+ std::cmp::Ordering::Equal
+ } else {
+ task.shard_template_id.as_str().cmp(shard_id)
+ }
+ })
+ .ok()
+ .map(|index| &tasks[index]);
+
+ // Map `shard_data_plane_fqdn` into its task-matched data-plane.
+ let task_data_plane = task.and_then(|task| {
+ data_planes
+ .get_by_key(&task.data_plane_id)
+ .filter(|data_plane| data_plane.data_plane_fqdn == shard_data_plane_fqdn)
+ });
+
+ let (Some(task), Some(task_data_plane)) = (task, task_data_plane) else {
+ anyhow::bail!(
+ "task shard {shard_id} within data-plane {shard_data_plane_fqdn} is not known"
+ )
+ };
+
+ // Attempt to find an HMAC key of this data-plane which validates against the request token.
+ let validation = jsonwebtoken::Validation::default();
+ let mut verified = false;
+
+ for hmac_key in &task_data_plane.hmac_keys {
+ let key = jsonwebtoken::DecodingKey::from_base64_secret(hmac_key)
+ .context("invalid data-plane hmac key")?;
+
+ if jsonwebtoken::decode::<proto_gazette::Claims>(token, &key, &validation).is_ok() {
+ verified = true;
+ break;
+ }
+ }
+ if !verified {
+ anyhow::bail!("no data-plane keys validated against the token signature");
+ }
+
+ Ok((task, task_data_plane))
+}

crates/agent/src/api/notify_shard_failure.rs

@@ -0,0 +1,69 @@
+use super::{App, Snapshot};
+use std::sync::Arc;
+
+/// Request sent by data-plane reactors to notify the data-plane of a shard failure.
+#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
+#[serde(rename_all = "camelCase")]
+pub struct Request {
+ /// # JWT token which identifies the shard and authorizes the request.
+ /// The token subject is the shard ID.
+ pub token: String,
+ /// # Error encountered by the shard.
+ pub error: String,
+}
+
+#[derive(Debug, Default, serde::Serialize, schemars::JsonSchema)]
+#[serde(rename_all = "camelCase")]
+pub struct Response {
+ /// # Number of milliseconds to wait before retrying the request.
+ /// Zero if the request was successful.
+ pub retry_millis: u64,
+}
+
+#[axum::debug_handler]
+pub async fn notify_shard_failure(
+ axum::extract::State(app): axum::extract::State<Arc<App>>,
+ axum::Json(request): axum::Json<Request>,
+) -> axum::response::Response {
+ super::wrap(async move { do_notify_shard_failure(&app, &request).await }).await
+}
+
+#[tracing::instrument(skip(app), err(level = tracing::Level::WARN))]
+async fn do_notify_shard_failure(
+ app: &App,
+ Request { token, error }: &Request,
+) -> anyhow::Result<Response> {
+ let (_header, claims) = super::parse_untrusted_data_plane_claims(token)?;
+
+ match Snapshot::evaluate(
+ &app.snapshot,
+ claims.iat,
+ |Snapshot {
+ data_planes, tasks, ..
+ }: &Snapshot| {
+ _ = super::verify_data_plane_claims(
+ data_planes,
+ tasks,
+ &claims.sub,
+ &claims.iss,
+ token,
+ )?;
+ Ok(())
+ },
+ ) {
+ Ok(()) => {
+ // TODO(johnny): This is a placeholder for enqueuing an automation to perform
+ // shard restart.
+ tracing::info!(%error, %claims.sub, %claims.iss, "notified of shard failure");
+
+ Ok(Response {
+ ..Default::default()
+ })
+ }
+ Err(Ok(retry_millis)) => Ok(Response {
+ retry_millis,
+ ..Default::default()
+ }),
+ Err(Err(err)) => Err(err),
+ }
+}

go/runtime/flow_consumer.go

@@ -61,6 +61,8 @@ func (c *FlowConsumerConfig) Execute(args []string) error {
 type FlowConsumer struct {
  // Configuration of this FlowConsumer.
  config *FlowConsumerConfig
+ // Control plane of this consumer, or nil in testing contexts.
+ controlPlane *controlPlane
  // Running consumer.service.
  service *consumer.Service
  // Shared catalog builds.
@@ -208,11 +210,10 @@ func (f *FlowConsumer) InitApplication(args runconsumer.InitArgs) error {
  return fmt.Errorf("catalog builds service: %w", err)
  }
 
- var controlPlane *controlPlane
  var localAuthorizer = args.Service.Authorizer
 
  if keyedAuth, ok := localAuthorizer.(*auth.KeyedAuth); ok && !config.Flow.TestAPIs {
- controlPlane = newControlPlane(
+ f.controlPlane = newControlPlane(
  keyedAuth,
  config.Flow.DataPlaneFQDN,
  config.Flow.ControlAPI,
@@ -221,7 +222,7 @@ func (f *FlowConsumer) InitApplication(args runconsumer.InitArgs) error {
  // Wrap the underlying KeyedAuth Authorizer to use the control-plane's Authorize API.
  // Next unwrap the raw JournalClient from its current AuthJournalClient,
  // and then replace it with one built using our wrapped Authorizer.
- args.Service.Authorizer = newControlPlaneAuthorizer(controlPlane)
+ args.Service.Authorizer = newControlPlaneAuthorizer(f.controlPlane)
  var rawClient = args.Service.Journals.(*pb.ComposedRoutedJournalClient).JournalClient.(*pb.AuthJournalClient).Inner
  args.Service.Journals.(*pb.ComposedRoutedJournalClient).JournalClient = pb.NewAuthJournalClient(rawClient, args.Service.Authorizer)
  }

go/runtime/task.go

@@ -22,6 +22,8 @@ import (
  "github.com/estuary/flow/go/shuffle"
  "github.com/gogo/protobuf/proto"
  "github.com/gogo/protobuf/types"
+ "github.com/golang-jwt/jwt/v5"
+ "github.com/sirupsen/logrus"
  "go.gazette.dev/core/allocator"
  "go.gazette.dev/core/broker/client"
  pb "go.gazette.dev/core/broker/protocol"
@@ -254,11 +256,12 @@ func (t *taskReader[TaskSpec]) Coordinator() *shuffle.Coordinator { return t.coo
 // and then logs the final exit status of the shard.
 func (t *taskBase[TaskSpec]) heartbeatLoop(shard consumer.Shard) {
  var (
+ id = shard.Spec().Id
  // Period between regularly-published stat intervals.
  // This period must cleanly divide into one hour!
  period = 3 * time.Minute
  // Jitters when interval stats are written cluster-wide.
- jitter = intervalJitter(period, shard.FQN())
+ jitter = intervalJitter(period, id)
  // Op notified when the shard fails.
  op = shard.PrimaryLoop()
  )
@@ -289,8 +292,33 @@ func (t *taskBase[TaskSpec]) heartbeatLoop(shard consumer.Shard) {
  "assignment", shard.Assignment().Decoded,
  )
 
- // TODO(johnny): Notify control-plane of failure.
+ // Notify control-plane of the failure.
+ var now = time.Now()
 
+ var claims = pb.Claims{
+ RegisteredClaims: jwt.RegisteredClaims{
+ Subject: id.String(),
+ IssuedAt: jwt.NewNumericDate(now),
+ ExpiresAt: jwt.NewNumericDate(now.Add(time.Hour)),
+ },
+ }
+ var token, postErr = t.host.controlPlane.signClaims(claims)
+
+ var request = struct {
+ Token string `json:"token"`
+ Error string `json:"error"`
+ }{token, err.Error()}
+
+ var response struct {
+ // No response fields.
+ }
+
+ if postErr == nil {
+ postErr = callControlAPI(shard.Context(), t.host.controlPlane, "/notify/shard-failure", &request, &response)
+ }
+ if postErr != nil {
+ logrus.WithField("err", postErr).Error("failed to notify control plane of shard failure")
+ }
  return
  }
  }
@@ -325,9 +353,9 @@ func durationToNextInterval(now time.Time, period time.Duration) time.Duration {
 
 // intervalJitter returns a globally consistent, unique jitter offset for `name`
 // so that heartbeats are uniformly distributed over time, in aggregate.
-func intervalJitter(period time.Duration, name string) time.Duration {
+func intervalJitter(period time.Duration, id pc.ShardID) time.Duration {
  var w = fnv.New32()
- w.Write([]byte(name))
+ w.Write([]byte(id))
  return time.Duration(w.Sum32()%uint32(period.Seconds())) * time.Second
 }
 

go/runtime/task_test.go

@@ -7,16 +7,17 @@ import (
  pf "github.com/estuary/flow/go/protocols/flow"
  "github.com/estuary/flow/go/protocols/ops"
  "github.com/stretchr/testify/require"
+ pc "go.gazette.dev/core/consumer/protocol"
 )
 
 func TestIntervalJitterAndDurations(t *testing.T) {
  const period = time.Minute
 
  for _, tc := range []struct {
- n string
- i time.Duration
+ id pc.ShardID
+ i  time.Duration
  }{{"foo", 35}, {"bar", 52}, {"baz", 0}, {"bing", 39}, {"quip", 56}} {
- require.Equal(t, time.Second*tc.i, intervalJitter(period, tc.n), tc.n)
+ require.Equal(t, time.Second*tc.i, intervalJitter(period, tc.id), tc.id)
 
  }