split kafka codec (#1052)

shotover · Feb 21, 2023 · d3b0300 · d3b0300
1 parent 161419d
commit d3b0300
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 22 deletions.
diff --git a/shotover-proxy/src/codec/kafka.rs b/shotover-proxy/src/codec/kafka.rs
@@ -5,28 +5,31 @@ use anyhow::Result;
 use bytes::{Buf, BytesMut};
 use tokio_util::codec::{Decoder, Encoder};
 
-impl CodecBuilder for KafkaCodec {
- type Decoder = KafkaCodec;
- type Encoder = KafkaCodec;
- fn build(&self) -> (KafkaCodec, KafkaCodec) {
- (KafkaCodec::new(), KafkaCodec::new())
+#[derive(Clone, Default)]
+pub struct KafkaCodecBuilder {}
+
+impl KafkaCodecBuilder {
+ pub fn new() -> Self {
+ Default::default()
  }
 }
 
-#[derive(Debug, Clone)]
-pub struct KafkaCodec {
- messages: Messages,
+impl CodecBuilder for KafkaCodecBuilder {
+ type Decoder = KafkaDecoder;
+ type Encoder = KafkaEncoder;
+ fn build(&self) -> (KafkaDecoder, KafkaEncoder) {
+ (KafkaDecoder::new(), KafkaEncoder::new())
+ }
 }
 
-impl Default for KafkaCodec {
- fn default() -> Self {
- Self::new()
- }
+#[derive(Default)]
+pub struct KafkaDecoder {
+ messages: Messages,
 }
 
-impl KafkaCodec {
- pub fn new() -> KafkaCodec {
- KafkaCodec { messages: vec![] }
+impl KafkaDecoder {
+ pub fn new() -> Self {
+ KafkaDecoder::default()
  }
 }
 
@@ -43,7 +46,7 @@ fn get_length_of_full_message(src: &mut BytesMut) -> Option<usize> {
  }
 }
 
-impl Decoder for KafkaCodec {
+impl Decoder for KafkaDecoder {
  type Item = Messages;
  type Error = CodecReadError;
 
@@ -66,13 +69,21 @@ impl Decoder for KafkaCodec {
  }
 }
 
-impl Encoder<Messages> for KafkaCodec {
+#[derive(Default)]
+pub struct KafkaEncoder {}
+
+impl KafkaEncoder {
+ pub fn new() -> Self {
+ KafkaEncoder::default()
+ }
+}
+
+impl Encoder<Messages> for KafkaEncoder {
  type Error = anyhow::Error;
 
  fn encode(&mut self, item: Messages, dst: &mut BytesMut) -> Result<()> {
  item.into_iter().try_for_each(|m| {
  let start = dst.len();
- // TODO: MessageType::Kafka
  let result = match m.into_encodable(MessageType::Kafka)? {
  Encodable::Bytes(bytes) => {
  dst.extend_from_slice(&bytes);

diff --git a/shotover-proxy/src/sources/kafka.rs b/shotover-proxy/src/sources/kafka.rs
@@ -1,4 +1,4 @@
-use crate::codec::kafka::KafkaCodec;
+use crate::codec::kafka::KafkaCodecBuilder;
 use crate::server::TcpCodecListener;
 use crate::sources::Sources;
 use crate::tls::{TlsAcceptor, TlsAcceptorConfig};
@@ -67,7 +67,7 @@ impl KafkaSource {
  name.to_string(),
  listen_addr.clone(),
  hard_connection_limit.unwrap_or(false),
- KafkaCodec::new(),
+ KafkaCodecBuilder::new(),
  Arc::new(Semaphore::new(connection_limit.unwrap_or(512))),
  trigger_shutdown_rx.clone(),
  tls.map(TlsAcceptor::new).transpose()?,

diff --git a/shotover-proxy/src/transforms/kafka/sink_single.rs b/shotover-proxy/src/transforms/kafka/sink_single.rs
@@ -1,4 +1,4 @@
-use crate::codec::kafka::KafkaCodec;
+use crate::codec::kafka::KafkaCodecBuilder;
 use crate::error::ChainResponse;
 use crate::message::Messages;
 use crate::tcp;
@@ -87,7 +87,7 @@ pub struct KafkaSinkSingle {
 impl Transform for KafkaSinkSingle {
  async fn transform<'a>(&'a mut self, message_wrapper: Wrapper<'a>) -> ChainResponse {
  if self.outbound.is_none() {
- let codec = KafkaCodec::new();
+ let codec = KafkaCodecBuilder::new();
  let tcp_stream = tcp::tcp_stream(self.connect_timeout, &self.address).await?;
  let (rx, tx) = tcp_stream.into_split();
  self.outbound = Some(spawn_read_write_tasks(&codec, rx, tx));