Introduce TransformContextConfig (#1490)

custom-transforms-example/src/redis_get_rewrite.rs

@@ -3,7 +3,9 @@ use async_trait::async_trait;
 use serde::{Deserialize, Serialize};
 use shotover::frame::{Frame, RedisFrame};
 use shotover::message::{MessageIdSet, Messages};
-use shotover::transforms::{Transform, TransformBuilder, TransformConfig, Wrapper};
+use shotover::transforms::{
+    Transform, TransformBuilder, TransformConfig, TransformContextConfig, Wrapper,
+};
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
 #[serde(deny_unknown_fields)]
@@ -15,7 +17,10 @@ const NAME: &str = "RedisGetRewrite";
 #[typetag::serde(name = "RedisGetRewrite")]
 #[async_trait(?Send)]
 impl TransformConfig for RedisGetRewriteConfig {
-    async fn get_builder(&self, _chain_name: String) -> Result<Box<dyn TransformBuilder>> {
+    async fn get_builder(
+        &self,
+        _transform_context: TransformContextConfig,
+    ) -> Result<Box<dyn TransformBuilder>> {
         Ok(Box::new(RedisGetRewriteBuilder {
             result: self.result.clone(),
         }))

shotover/benches/benches/chain.rs

@@ -5,8 +5,8 @@ use criterion::{criterion_group, BatchSize, Criterion};
 use hex_literal::hex;
 use shotover::codec::CodecState;
 use shotover::frame::cassandra::{parse_statement_single, Tracing};
-use shotover::frame::RedisFrame;
 use shotover::frame::{CassandraFrame, CassandraOperation, Frame};
+use shotover::frame::{MessageType, RedisFrame};
 use shotover::message::{Message, MessageIdMap, QueryType};
 use shotover::transforms::cassandra::peers_rewrite::CassandraPeersRewrite;
 use shotover::transforms::chain::{TransformChain, TransformChainBuilder};
@@ -19,7 +19,7 @@ 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, Wrapper};
+use shotover::transforms::{TransformConfig, TransformContextConfig, Wrapper};
 
 fn criterion_benchmark(c: &mut Criterion) {
     crate::init();
@@ -194,7 +194,10 @@ fn criterion_benchmark(c: &mut Criterion) {
                         // an absurdly large value is given so that all messages will pass through
                         max_requests_per_second: std::num::NonZeroU32::new(100_000_000).unwrap(),
                     }
-                    .get_builder("".to_owned()),
+                    .get_builder(TransformContextConfig {
+                        chain_name: "".into(),
+                        protocol: MessageType::Redis,
+                    }),
                 )
                 .unwrap(),
                 Box::<NullSink>::default(),
@@ -303,7 +306,10 @@ fn criterion_benchmark(c: &mut Criterion) {
                             kek_id: "".to_string(),
                         },
                     }
-                    .get_builder("".to_owned()),
+                    .get_builder(TransformContextConfig {
+                        chain_name: "".into(),
+                        protocol: MessageType::Redis,
+                    }),
                 )
                 .unwrap(),
                 Box::<NullSink>::default(),

shotover/src/config/chain.rs

@@ -1,5 +1,5 @@
 use crate::transforms::chain::TransformChainBuilder;
-use crate::transforms::{TransformBuilder, TransformConfig};
+use crate::transforms::{TransformBuilder, TransformConfig, TransformContextConfig};
 use anyhow::Result;
 use serde::de::{DeserializeSeed, Deserializer, MapAccess, SeqAccess, Visitor};
 use serde::{Deserialize, Serialize};
@@ -14,12 +14,18 @@ pub struct TransformChainConfig(
 );
 
 impl TransformChainConfig {
-    pub async fn get_builder(&self, name: String) -> Result<TransformChainBuilder> {
+    pub async fn get_builder(
+        &self,
+        transform_context: TransformContextConfig,
+    ) -> Result<TransformChainBuilder> {
         let mut transforms: Vec<Box<dyn TransformBuilder>> = Vec::new();
         for tc in &self.0 {
-            transforms.push(tc.get_builder(name.clone()).await?)
+            transforms.push(tc.get_builder(transform_context.clone()).await?)
         }
-        Ok(TransformChainBuilder::new(transforms, name.leak()))
+        Ok(TransformChainBuilder::new(
+            transforms,
+            transform_context.chain_name.leak(),
+        ))
     }
 }
 

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::Wrapper;
+use crate::transforms::{TransformContextConfig, Wrapper};
 use anyhow::{anyhow, Context, Result};
 use bytes::BytesMut;
 use futures::future::join_all;
@@ -92,8 +92,12 @@ impl<C: CodecBuilder + 'static> TcpCodecListener<C> {
             gauge!("shotover_available_connections_count", "source" => source_name.clone());
         available_connections_gauge.set(limit_connections.available_permits() as f64);
 
+        let chain_usage_config = TransformContextConfig {
+            chain_name: source_name.clone(),
+            protocol: codec.protocol(),
+        };
         let chain_builder = chain_config
-            .get_builder(source_name.clone())
+            .get_builder(chain_usage_config)
             .await
             .map_err(|x| vec![format!("{x:?}")])?;
 

shotover/src/transforms/cassandra/peers_rewrite.rs

@@ -1,10 +1,13 @@
-use crate::frame::{
-    value::{GenericValue, IntSize},
-    CassandraOperation, CassandraResult, Frame,
-};
 use crate::message::{Message, Messages};
 use crate::transforms::cassandra::peers_rewrite::CassandraOperation::Event;
 use crate::transforms::{Transform, TransformBuilder, TransformConfig, Wrapper};
+use crate::{
+    frame::{
+        value::{GenericValue, IntSize},
+        CassandraOperation, CassandraResult, Frame,
+    },
+    transforms::TransformContextConfig,
+};
 use anyhow::Result;
 use async_trait::async_trait;
 use cassandra_protocol::frame::events::{ServerEvent, StatusChange};
@@ -23,7 +26,10 @@ const NAME: &str = "CassandraPeersRewrite";
 #[typetag::serde(name = "CassandraPeersRewrite")]
 #[async_trait(?Send)]
 impl TransformConfig for CassandraPeersRewriteConfig {
-    async fn get_builder(&self, _chain_name: String) -> Result<Box<dyn TransformBuilder>> {
+    async fn get_builder(
+        &self,
+        _transform_context: TransformContextConfig,
+    ) -> Result<Box<dyn TransformBuilder>> {
         Ok(Box::new(CassandraPeersRewrite::new(self.port)))
     }
 }

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

@@ -5,7 +5,9 @@ use crate::frame::{CassandraFrame, CassandraOperation, CassandraResult, Frame};
 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, Wrapper};
+use crate::transforms::{
+    Transform, TransformBuilder, TransformConfig, TransformContextConfig, Wrapper,
+};
 use anyhow::{anyhow, Context, Result};
 use async_trait::async_trait;
 use cassandra_protocol::events::ServerEvent;
@@ -66,7 +68,10 @@ const NAME: &str = "CassandraSinkCluster";
 #[typetag::serde(name = "CassandraSinkCluster")]
 #[async_trait(?Send)]
 impl TransformConfig for CassandraSinkClusterConfig {
-    async fn get_builder(&self, chain_name: String) -> Result<Box<dyn TransformBuilder>> {
+    async fn get_builder(
+        &self,
+        transform_context: TransformContextConfig,
+    ) -> Result<Box<dyn TransformBuilder>> {
         let tls = self.tls.clone().map(TlsConnector::new).transpose()?;
         let mut shotover_nodes = self.shotover_nodes.clone();
         let index = self
@@ -84,7 +89,7 @@ impl TransformConfig for CassandraSinkClusterConfig {
         Ok(Box::new(CassandraSinkClusterBuilder::new(
             self.first_contact_points.clone(),
             shotover_nodes,
-            chain_name,
+            transform_context.chain_name,
             local_node,
             tls,
             self.connect_timeout_ms,

shotover/src/transforms/cassandra/sink_single.rs

@@ -4,7 +4,9 @@ use crate::frame::cassandra::CassandraMetadata;
 use crate::message::{Messages, Metadata};
 use crate::tls::{TlsConnector, TlsConnectorConfig};
 use crate::transforms::cassandra::connection::Response;
-use crate::transforms::{Transform, TransformBuilder, TransformConfig, Wrapper};
+use crate::transforms::{
+    Transform, TransformBuilder, TransformConfig, TransformContextConfig, Wrapper,
+};
 use anyhow::{anyhow, Result};
 use async_trait::async_trait;
 use cassandra_protocol::frame::Version;
@@ -29,11 +31,14 @@ const NAME: &str = "CassandraSinkSingle";
 #[typetag::serde(name = "CassandraSinkSingle")]
 #[async_trait(?Send)]
 impl TransformConfig for CassandraSinkSingleConfig {
-    async fn get_builder(&self, chain_name: String) -> Result<Box<dyn TransformBuilder>> {
+    async fn get_builder(
+        &self,
+        transform_context: TransformContextConfig,
+    ) -> Result<Box<dyn TransformBuilder>> {
         let tls = self.tls.clone().map(TlsConnector::new).transpose()?;
         Ok(Box::new(CassandraSinkSingleBuilder::new(
             self.address.clone(),
-            chain_name,
+            transform_context.chain_name,
             tls,
             self.connect_timeout_ms,
             self.read_timeout,

shotover/src/transforms/coalesce.rs

@@ -1,3 +1,4 @@
+use super::TransformContextConfig;
 use crate::message::Messages;
 use crate::transforms::{Transform, TransformBuilder, TransformConfig, Wrapper};
 use anyhow::Result;
@@ -24,7 +25,10 @@ const NAME: &str = "Coalesce";
 #[typetag::serde(name = "Coalesce")]
 #[async_trait(?Send)]
 impl TransformConfig for CoalesceConfig {
-    async fn get_builder(&self, _chain_name: String) -> Result<Box<dyn TransformBuilder>> {
+    async fn get_builder(
+        &self,
+        _transform_context: TransformContextConfig,
+    ) -> Result<Box<dyn TransformBuilder>> {
         Ok(Box::new(Coalesce {
             buffer: Vec::with_capacity(self.flush_when_buffered_message_count.unwrap_or(0)),
             flush_when_buffered_message_count: self.flush_when_buffered_message_count,

shotover/src/transforms/debug/force_parse.rs

@@ -7,6 +7,8 @@ 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;
+#[cfg(feature = "alpha-transforms")]
+use crate::transforms::TransformContextConfig;
 use crate::transforms::{Transform, TransformBuilder, Wrapper};
 use anyhow::Result;
 use async_trait::async_trait;
@@ -25,7 +27,10 @@ pub struct DebugForceParseConfig {
 #[typetag::serde(name = "DebugForceParse")]
 #[async_trait(?Send)]
 impl TransformConfig for DebugForceParseConfig {
-    async fn get_builder(&self, _chain_name: String) -> Result<Box<dyn TransformBuilder>> {
+    async fn get_builder(
+        &self,
+        _transform_context: TransformContextConfig,
+    ) -> Result<Box<dyn TransformBuilder>> {
         Ok(Box::new(DebugForceParse {
             parse_requests: self.parse_requests,
             parse_responses: self.parse_responses,
@@ -49,7 +54,10 @@ const NAME: &str = "DebugForceEncode";
 #[typetag::serde(name = "DebugForceEncode")]
 #[async_trait(?Send)]
 impl TransformConfig for DebugForceEncodeConfig {
-    async fn get_builder(&self, _chain_name: String) -> Result<Box<dyn TransformBuilder>> {
+    async fn get_builder(
+        &self,
+        _transform_context: TransformContextConfig,
+    ) -> Result<Box<dyn TransformBuilder>> {
         Ok(Box::new(DebugForceParse {
             parse_requests: self.encode_requests,
             parse_responses: self.encode_responses,

shotover/src/transforms/debug/log_to_file.rs

@@ -17,7 +17,10 @@ const NAME: &str = "DebugLogToFile";
 #[typetag::serde(name = "DebugLogToFile")]
 #[async_trait(?Send)]
 impl crate::transforms::TransformConfig for DebugLogToFileConfig {
-    async fn get_builder(&self, _chain_name: String) -> Result<Box<dyn TransformBuilder>> {
+    async fn get_builder(
+        &self,
+        _transform_context: crate::transforms::TransformContextConfig,
+    ) -> Result<Box<dyn TransformBuilder>> {
         // This transform is used for debugging a specific run, so we clean out any logs left over from the previous run
         std::fs::remove_dir_all("message-log").ok();
 

shotover/src/transforms/debug/printer.rs

@@ -1,5 +1,7 @@
 use crate::message::Messages;
-use crate::transforms::{Transform, TransformBuilder, TransformConfig, Wrapper};
+use crate::transforms::{
+    Transform, TransformBuilder, TransformConfig, TransformContextConfig, Wrapper,
+};
 use anyhow::Result;
 use async_trait::async_trait;
 use serde::{Deserialize, Serialize};
@@ -13,7 +15,10 @@ const NAME: &str = "DebugPrinter";
 #[typetag::serde(name = "DebugPrinter")]
 #[async_trait(?Send)]
 impl TransformConfig for DebugPrinterConfig {
-    async fn get_builder(&self, _chain_name: String) -> Result<Box<dyn TransformBuilder>> {
+    async fn get_builder(
+        &self,
+        _transform_context: TransformContextConfig,
+    ) -> Result<Box<dyn TransformBuilder>> {
         Ok(Box::new(DebugPrinter::new()))
     }
 }

shotover/src/transforms/debug/returner.rs

@@ -1,5 +1,7 @@
 use crate::message::{Message, Messages};
-use crate::transforms::{Transform, TransformBuilder, TransformConfig, Wrapper};
+use crate::transforms::{
+    Transform, TransformBuilder, TransformConfig, TransformContextConfig, Wrapper,
+};
 use anyhow::{anyhow, Result};
 use async_trait::async_trait;
 use serde::{Deserialize, Serialize};
@@ -15,7 +17,10 @@ const NAME: &str = "DebugReturner";
 #[typetag::serde(name = "DebugReturner")]
 #[async_trait(?Send)]
 impl TransformConfig for DebugReturnerConfig {
-    async fn get_builder(&self, _chain_name: String) -> Result<Box<dyn TransformBuilder>> {
+    async fn get_builder(
+        &self,
+        _transform_context: TransformContextConfig,
+    ) -> Result<Box<dyn TransformBuilder>> {
         Ok(Box::new(DebugReturner::new(self.response.clone())))
     }
 }

shotover/src/transforms/distributed/tuneable_consistency_scatter.rs

@@ -2,7 +2,9 @@ use crate::config::chain::TransformChainConfig;
 use crate::frame::{Frame, RedisFrame};
 use crate::message::{Message, Messages, QueryType};
 use crate::transforms::chain::{BufferedChain, TransformChainBuilder};
-use crate::transforms::{Transform, TransformBuilder, TransformConfig, Wrapper};
+use crate::transforms::{
+    Transform, TransformBuilder, TransformConfig, TransformContextConfig, Wrapper,
+};
 use anyhow::Result;
 use async_trait::async_trait;
 use futures::stream::FuturesUnordered;
@@ -23,12 +25,19 @@ const NAME: &str = "TuneableConsistencyScatter";
 #[typetag::serde(name = "TuneableConsistencyScatter")]
 #[async_trait(?Send)]
 impl TransformConfig for TuneableConsistencyScatterConfig {
-    async fn get_builder(&self, _chain_name: String) -> Result<Box<dyn TransformBuilder>> {
+    async fn get_builder(
+        &self,
+        transform_context: TransformContextConfig,
+    ) -> Result<Box<dyn TransformBuilder>> {
         let mut route_map = Vec::with_capacity(self.route_map.len());
         warn!("Using this transform is considered unstable - Does not work with REDIS pipelines");
 
         for (key, value) in &self.route_map {
-            route_map.push(value.get_builder(key.clone()).await?);
+            let chain_config = TransformContextConfig {
+                chain_name: key.clone(),
+                protocol: transform_context.protocol,
+            };
+            route_map.push(value.get_builder(chain_config).await?);
         }
         route_map.sort_by_key(|x| x.name);
 

shotover/src/transforms/filter.rs

@@ -1,3 +1,4 @@
+use super::TransformContextConfig;
 use crate::message::{Message, MessageIdMap, Messages, QueryType};
 use crate::transforms::{Transform, TransformBuilder, TransformConfig, Wrapper};
 use anyhow::Result;
@@ -28,7 +29,10 @@ const NAME: &str = "QueryTypeFilter";
 #[typetag::serde(name = "QueryTypeFilter")]
 #[async_trait(?Send)]
 impl TransformConfig for QueryTypeFilterConfig {
-    async fn get_builder(&self, _chain_name: String) -> Result<Box<dyn TransformBuilder>> {
+    async fn get_builder(
+        &self,
+        _transform_context: TransformContextConfig,
+    ) -> Result<Box<dyn TransformBuilder>> {
         Ok(Box::new(QueryTypeFilter {
             filter: self.filter.clone(),
             filtered_requests: MessageIdMap::default(),

shotover/src/transforms/kafka/sink_cluster.rs

@@ -6,8 +6,8 @@ use crate::message::{Message, Messages};
 use crate::tcp;
 use crate::transforms::util::cluster_connection_pool::{spawn_read_write_tasks, Connection};
 use crate::transforms::util::{Request, Response};
-use crate::transforms::TransformConfig;
 use crate::transforms::{Transform, TransformBuilder, Wrapper};
+use crate::transforms::{TransformConfig, TransformContextConfig};
 use anyhow::{anyhow, Result};
 use async_trait::async_trait;
 use dashmap::DashMap;
@@ -45,11 +45,14 @@ const NAME: &str = "KafkaSinkCluster";
 #[typetag::serde(name = "KafkaSinkCluster")]
 #[async_trait(?Send)]
 impl TransformConfig for KafkaSinkClusterConfig {
-    async fn get_builder(&self, chain_name: String) -> Result<Box<dyn TransformBuilder>> {
+    async fn get_builder(
+        &self,
+        transform_context: TransformContextConfig,
+    ) -> Result<Box<dyn TransformBuilder>> {
         Ok(Box::new(KafkaSinkClusterBuilder::new(
             self.first_contact_points.clone(),
             self.shotover_nodes.clone(),
-            chain_name,
+            transform_context.chain_name,
             self.connect_timeout_ms,
             self.read_timeout,
         )))