Add force_run_chain Notify (#1525)

custom-transforms-example/src/redis_get_rewrite.rs

@@ -3,6 +3,7 @@ use async_trait::async_trait;
 use serde::{Deserialize, Serialize};
 use shotover::frame::{Frame, RedisFrame};
 use shotover::message::{MessageIdSet, Messages};
+use shotover::transforms::TransformContextBuilder;
 use shotover::transforms::{
  Transform, TransformBuilder, TransformConfig, TransformContextConfig, Wrapper,
 };
@@ -33,7 +34,7 @@ pub struct RedisGetRewriteBuilder {
 }
 
 impl TransformBuilder for RedisGetRewriteBuilder {
- fn build(&self) -> Box<dyn Transform> {
+ fn build(&self, _transform_context: TransformContextBuilder) -> Box<dyn Transform> {
  Box::new(RedisGetRewrite {
  get_requests: MessageIdSet::default(),
  result: self.result.clone(),

shotover/benches/benches/chain.rs

@@ -19,7 +19,9 @@ use shotover::transforms::protect::{KeyManagerConfig, ProtectConfig};
 use shotover::transforms::redis::cluster_ports_rewrite::RedisClusterPortsRewrite;
 use shotover::transforms::redis::timestamp_tagging::RedisTimestampTagger;
 use shotover::transforms::throttling::RequestThrottlingConfig;
-use shotover::transforms::{TransformConfig, TransformContextConfig, Wrapper};
+use shotover::transforms::{
+ TransformConfig, TransformContextBuilder, TransformContextConfig, Wrapper,
+};
 
 fn criterion_benchmark(c: &mut Criterion) {
  crate::init();
@@ -38,7 +40,7 @@ fn criterion_benchmark(c: &mut Criterion) {
  group.bench_function("loopback", |b| {
  b.to_async(&rt).iter_batched(
  || BenchInput {
- chain: chain.build(),
+ chain: chain.build(TransformContextBuilder::new()),
  wrapper: wrapper.clone(),
  },
  BenchInput::bench,
@@ -57,7 +59,7 @@ fn criterion_benchmark(c: &mut Criterion) {
  group.bench_function("nullsink", |b| {
  b.to_async(&rt).iter_batched(
  || BenchInput {
- chain: chain.build(),
+ chain: chain.build(TransformContextBuilder::new()),
  wrapper: wrapper.clone(),
  },
  BenchInput::bench,
@@ -95,7 +97,7 @@ fn criterion_benchmark(c: &mut Criterion) {
  group.bench_function("redis_filter", |b| {
  b.to_async(&rt).iter_batched(
  || BenchInput {
- chain: chain.build(),
+ chain: chain.build(TransformContextBuilder::new()),
  wrapper: wrapper.clone(),
  },
  BenchInput::bench,
@@ -129,7 +131,7 @@ fn criterion_benchmark(c: &mut Criterion) {
  group.bench_function("redis_timestamp_tagger_untagged", |b| {
  b.to_async(&rt).iter_batched(
  || BenchInput {
- chain: chain.build(),
+ chain: chain.build(TransformContextBuilder::new()),
  wrapper: wrapper_set.clone(),
  },
  BenchInput::bench,
@@ -148,7 +150,7 @@ fn criterion_benchmark(c: &mut Criterion) {
  group.bench_function("redis_timestamp_tagger_tagged", |b| {
  b.to_async(&rt).iter_batched(
  || BenchInput {
- chain: chain.build(),
+ chain: chain.build(TransformContextBuilder::new()),
  wrapper: wrapper_get.clone(),
  },
  BenchInput::bench,
@@ -177,7 +179,7 @@ fn criterion_benchmark(c: &mut Criterion) {
  group.bench_function("redis_cluster_ports_rewrite", |b| {
  b.to_async(&rt).iter_batched(
  || BenchInput {
- chain: chain.build(),
+ chain: chain.build(TransformContextBuilder::new()),
  wrapper: wrapper.clone(),
  },
  BenchInput::bench,
@@ -224,7 +226,7 @@ fn criterion_benchmark(c: &mut Criterion) {
  group.bench_function("cassandra_request_throttling_unparsed", |b| {
  b.to_async(&rt).iter_batched(
  || BenchInput {
- chain: chain.build(),
+ chain: chain.build(TransformContextBuilder::new()),
  wrapper: wrapper.clone(),
  },
  BenchInput::bench,
@@ -278,7 +280,7 @@ fn criterion_benchmark(c: &mut Criterion) {
  group.bench_function("cassandra_rewrite_peers_passthrough", |b| {
  b.to_async(&rt).iter_batched(
  || BenchInput {
- chain: chain.build(),
+ chain: chain.build(TransformContextBuilder::new()),
  wrapper: wrapper.clone(),
  },
  BenchInput::bench,
@@ -324,7 +326,7 @@ fn criterion_benchmark(c: &mut Criterion) {
  group.bench_function("cassandra_protect_unprotected", |b| {
  b.to_async(&rt).iter_batched(
  || BenchInput {
- chain: chain.build(),
+ chain: chain.build(TransformContextBuilder::new()),
  wrapper: wrapper.clone(),
  },
  BenchInput::bench,
@@ -339,7 +341,7 @@ fn criterion_benchmark(c: &mut Criterion) {
  group.bench_function("cassandra_protect_protected", |b| {
  b.to_async(&rt).iter_batched(
  || BenchInput {
- chain: chain.build(),
+ chain: chain.build(TransformContextBuilder::new()),
  wrapper: wrapper.clone(),
  },
  BenchInput::bench,

shotover/src/server.rs

@@ -4,7 +4,7 @@ use crate::message::{Message, Messages};
 use crate::sources::Transport;
 use crate::tls::{AcceptError, TlsAcceptor};
 use crate::transforms::chain::{TransformChain, TransformChainBuilder};
-use crate::transforms::{TransformContextConfig, Wrapper};
+use crate::transforms::{TransformContextBuilder, TransformContextConfig, Wrapper};
 use anyhow::{anyhow, Context, Result};
 use bytes::BytesMut;
 use futures::future::join_all;
@@ -16,7 +16,7 @@ use std::sync::Arc;
 use tokio::io::{AsyncRead, AsyncWrite};
 use tokio::net::{TcpListener, TcpStream};
 use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};
-use tokio::sync::{mpsc, watch, OwnedSemaphorePermit, Semaphore};
+use tokio::sync::{mpsc, watch, Notify, OwnedSemaphorePermit, Semaphore};
 use tokio::task::JoinHandle;
 use tokio::time;
 use tokio::time::Duration;
@@ -190,10 +190,15 @@ impl<C: CodecBuilder + 'static> TcpCodecListener<C> {
  let (pushed_messages_tx, pushed_messages_rx) =
  tokio::sync::mpsc::unbounded_channel::<Messages>();
 
+ let force_run_chain = Arc::new(Notify::new());
+ let context = TransformContextBuilder {
+ force_run_chain: force_run_chain.clone(),
+ };
+
  let handler = Handler {
  chain: self
  .chain_builder
- .build_with_pushed_messages(pushed_messages_tx),
+ .build_with_pushed_messages(pushed_messages_tx, context),
  codec: self.codec.clone(),
  shutdown: Shutdown::new(self.trigger_shutdown_rx.clone()),
  tls: self.tls.clone(),
@@ -206,7 +211,7 @@ impl<C: CodecBuilder + 'static> TcpCodecListener<C> {
  self.connection_handles.push(tokio::spawn(
  async move {
  // Process the connection. If an error is encountered, log it.
- if let Err(err) = handler.run(stream, transport).await {
+ if let Err(err) = handler.run(stream, transport, force_run_chain).await {
  error!(
  "{:?}",
  err.context("connection was unexpectedly terminated")
@@ -576,7 +581,12 @@ impl<C: CodecBuilder + 'static> Handler<C> {
  ///
  /// When the shutdown signal is received, the connection is processed until
  /// it reaches a safe state, at which point it is terminated.
- pub async fn run(mut self, stream: TcpStream, transport: Transport) -> Result<()> {
+ pub async fn run(
+ mut self,
+ stream: TcpStream,
+ transport: Transport,
+ force_run_chain: Arc<Notify>,
+ ) -> Result<()> {
  stream.set_nodelay(true)?;
 
  let client_details = stream
@@ -658,7 +668,7 @@ impl<C: CodecBuilder + 'static> Handler<C> {
  };
 
  let result = self
- .process_messages(&client_details, local_addr, in_rx, out_tx)
+ .process_messages(&client_details, local_addr, in_rx, out_tx, force_run_chain)
  .await;
 
  // Flush messages regardless of if we are shutting down due to a failure or due to application shutdown
@@ -700,13 +710,32 @@ impl<C: CodecBuilder + 'static> Handler<C> {
  local_addr: SocketAddr,
  mut in_rx: mpsc::Receiver<Messages>,
  out_tx: mpsc::UnboundedSender<Messages>,
+ force_run_chain: Arc<Notify>,
  ) -> Result<()> {
  // As long as the shutdown signal has not been received, try to read a
  // new request frame.
  while !self.shutdown.is_shutdown() {
  // While reading a request frame, also listen for the shutdown signal
  debug!("Waiting for message {client_details}");
  let responses = tokio::select! {
+ biased;
+ _ = self.shutdown.recv() => {
+ // If a shutdown signal is received, return from `run`.
+ // This will result in the task terminating.
+ return Ok(());
+ }
+ Some(responses) = self.pushed_messages_rx.recv() => {
+ debug!("Received unrequested responses from destination {:?}", responses);
+ self.process_backward(client_details, local_addr, responses).await?
+ }
+ () = force_run_chain.notified() => {
+ let mut requests = vec!();
+ while let Ok(x) = in_rx.try_recv() {
+ requests.extend(x);
+ }
+ debug!("A transform in the chain requested that a chain run occur, requests {:?}", requests);
+ self.process_forward(client_details, local_addr, &out_tx, requests).await?
+ },
  requests = Self::receive_with_timeout(self.timeout, &mut in_rx, client_details) => {
  match requests {
  Some(mut requests) => {
@@ -722,15 +751,6 @@ impl<C: CodecBuilder + 'static> Handler<C> {
  }
  }
  },
- Some(responses) = self.pushed_messages_rx.recv() => {
- debug!("Received unrequested responses from destination {:?}", responses);
- self.process_backward(client_details, local_addr, responses).await?
- }
- _ = self.shutdown.recv() => {
- // If a shutdown signal is received, return from `run`.
- // This will result in the task terminating.
- return Ok(());
- }
  };
 
  debug!("sending response to client: {:?}", responses);

shotover/src/transforms/cassandra/peers_rewrite.rs

@@ -1,6 +1,8 @@
 use crate::message::{Message, MessageIdMap, Messages};
 use crate::transforms::cassandra::peers_rewrite::CassandraOperation::Event;
-use crate::transforms::{Transform, TransformBuilder, TransformConfig, Wrapper};
+use crate::transforms::{
+ Transform, TransformBuilder, TransformConfig, TransformContextBuilder, Wrapper,
+};
 use crate::{
  frame::{
  value::{GenericValue, IntSize},
@@ -52,7 +54,7 @@ impl CassandraPeersRewrite {
 }
 
 impl TransformBuilder for CassandraPeersRewrite {
- fn build(&self) -> Box<dyn Transform> {
+ fn build(&self, _transform_context: TransformContextBuilder) -> Box<dyn Transform> {
  Box::new(self.clone())
  }
 

shotover/src/transforms/cassandra/sink_cluster/mod.rs

@@ -6,7 +6,8 @@ use crate::message::{Message, MessageIdMap, Messages, Metadata};
 use crate::tls::{TlsConnector, TlsConnectorConfig};
 use crate::transforms::cassandra::connection::{CassandraConnection, Response, ResponseError};
 use crate::transforms::{
- Transform, TransformBuilder, TransformConfig, TransformContextConfig, Wrapper,
+ Transform, TransformBuilder, TransformConfig, TransformContextBuilder, TransformContextConfig,
+ Wrapper,
 };
 use anyhow::{anyhow, Context, Result};
 use async_trait::async_trait;
@@ -160,7 +161,7 @@ impl CassandraSinkClusterBuilder {
 }
 
 impl TransformBuilder for CassandraSinkClusterBuilder {
- fn build(&self) -> Box<dyn Transform> {
+ fn build(&self, _transform_context: TransformContextBuilder) -> Box<dyn Transform> {
  Box::new(CassandraSinkCluster {
  contact_points: self.contact_points.clone(),
  message_rewriter: self.message_rewriter.clone(),

shotover/src/transforms/cassandra/sink_single.rs

@@ -5,7 +5,8 @@ use crate::message::{Messages, Metadata};
 use crate::tls::{TlsConnector, TlsConnectorConfig};
 use crate::transforms::cassandra::connection::Response;
 use crate::transforms::{
- Transform, TransformBuilder, TransformConfig, TransformContextConfig, Wrapper,
+ Transform, TransformBuilder, TransformConfig, TransformContextBuilder, TransformContextConfig,
+ Wrapper,
 };
 use anyhow::{anyhow, Result};
 use async_trait::async_trait;
@@ -83,7 +84,7 @@ impl CassandraSinkSingleBuilder {
 }
 
 impl TransformBuilder for CassandraSinkSingleBuilder {
- fn build(&self) -> Box<dyn Transform> {
+ fn build(&self, _transform_context: TransformContextBuilder) -> Box<dyn Transform> {
  Box::new(CassandraSinkSingle {
  outbound: None,
  version: self.version,

shotover/src/transforms/chain.rs

@@ -9,6 +9,8 @@ use tokio::sync::{mpsc, oneshot};
 use tokio::time::{Duration, Instant};
 use tracing::{debug, error, info, trace, Instrument};
 
+use super::TransformContextBuilder;
+
 type InnerChain = Vec<TransformAndMetrics>;
 
 #[derive(Debug)]
@@ -233,9 +235,9 @@ pub struct TransformBuilderAndMetrics {
 }
 
 impl TransformBuilderAndMetrics {
- fn build(&self) -> TransformAndMetrics {
+ fn build(&self, context: TransformContextBuilder) -> TransformAndMetrics {
  TransformAndMetrics {
- transform: self.builder.build(),
+ transform: self.builder.build(context),
  transform_total: self.transform_total.clone(),
  transform_failures: self.transform_failures.clone(),
  transform_latency: self.transform_latency.clone(),
@@ -331,7 +333,11 @@ impl TransformChainBuilder {
  errors
  }
 
- pub fn build_buffered(&self, buffer_size: usize) -> BufferedChain {
+ pub fn build_buffered(
+ &self,
+ buffer_size: usize,
+ context: TransformContextBuilder,
+ ) -> BufferedChain {
  let (tx, mut rx) = mpsc::channel::<BufferedChainMessages>(buffer_size);
 
  #[cfg(test)]
@@ -341,7 +347,7 @@ impl TransformChainBuilder {
 
  // Even though we don't keep the join handle, this thread will wrap up once all corresponding senders have been dropped.
 
- let mut chain = self.build();
+ let mut chain = self.build(context);
  let _jh = tokio::spawn(
  async move {
  while let Some(BufferedChainMessages {
@@ -398,8 +404,12 @@ impl TransformChainBuilder {
  }
 
  /// Clone the chain while adding a producer for the pushed messages channel
- pub fn build(&self) -> TransformChain {
- let chain = self.chain.iter().map(|x| x.build()).collect();
+ pub fn build(&self, context: TransformContextBuilder) -> TransformChain {
+ let chain = self
+ .chain
+ .iter()
+ .map(|x| x.build(context.clone()))
+ .collect();
 
  TransformChain {
  name: self.name,
@@ -414,12 +424,13 @@ impl TransformChainBuilder {
  pub fn build_with_pushed_messages(
  &self,
  pushed_messages_tx: mpsc::UnboundedSender<Messages>,
+ context: TransformContextBuilder,
  ) -> TransformChain {
  let chain = self
  .chain
  .iter()
  .map(|x| {
- let mut transform = x.build();
+ let mut transform = x.build(context.clone());
  transform
  .transform
  .set_pushed_messages_tx(pushed_messages_tx.clone());

shotover/src/transforms/coalesce.rs

@@ -1,4 +1,4 @@
-use super::TransformContextConfig;
+use super::{TransformContextBuilder, TransformContextConfig};
 use crate::message::Messages;
 use crate::transforms::{Transform, TransformBuilder, TransformConfig, Wrapper};
 use anyhow::Result;
@@ -39,7 +39,7 @@ impl TransformConfig for CoalesceConfig {
 }
 
 impl TransformBuilder for Coalesce {
- fn build(&self) -> Box<dyn Transform> {
+ fn build(&self, _transform_context: TransformContextBuilder) -> Box<dyn Transform> {
  Box::new(self.clone())
  }
 

shotover/src/transforms/debug/force_parse.rs

@@ -7,6 +7,7 @@ use crate::message::Messages;
 /// It could also be used to ensure that messages round trip correctly when parsed.
 #[cfg(feature = "alpha-transforms")]
 use crate::transforms::TransformConfig;
+use crate::transforms::TransformContextBuilder;
 #[cfg(feature = "alpha-transforms")]
 use crate::transforms::TransformContextConfig;
 use crate::transforms::{Transform, TransformBuilder, Wrapper};
@@ -76,7 +77,7 @@ pub struct DebugForceParse {
 }
 
 impl TransformBuilder for DebugForceParse {
- fn build(&self) -> Box<dyn Transform> {
+ fn build(&self, _transform_context: TransformContextBuilder) -> Box<dyn Transform> {
  Box::new(self.clone())
  }