Skip to main content

tfserver/codec/
length_delimited.rs

1use crate::codec::codec_trait::TfCodec;
2use crate::structures::transport::Transport;
3use async_trait::async_trait;
4use std::io;
5use tokio_util::bytes::{Bytes, BytesMut};
6use tokio_util::codec::{Decoder, Encoder};
7#[derive(Clone, Default)]
8///The wrapper aroung existing codec, to be compatible with server
9pub struct LengthDelimitedCodec {
10    codec: tokio_util::codec::LengthDelimitedCodec,
11}
12
13impl LengthDelimitedCodec {
14    pub fn new(max_message_length: usize) -> Self {
15        Self {
16            codec: tokio_util::codec::LengthDelimitedCodec::builder()
17                .max_frame_length(max_message_length)
18                .new_codec(),
19        }
20    }
21    
22    pub fn from(codec: tokio_util::codec::LengthDelimitedCodec) -> Self {
23        Self{
24            codec,
25        }
26    }
27}
28
29impl Decoder for LengthDelimitedCodec {
30    type Item = BytesMut;
31    type Error = io::Error;
32
33    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
34        self.codec.decode(src)
35    }
36}
37
38impl Encoder<Bytes> for LengthDelimitedCodec {
39    type Error = io::Error;
40    fn encode(&mut self, item: Bytes, dst: &mut BytesMut) -> Result<(), Self::Error> {
41        self.codec.encode(item, dst)
42    }
43}
44#[async_trait]
45impl TfCodec for LengthDelimitedCodec {
46    async fn initial_setup(&mut self, _: &mut Transport) -> bool {
47        true
48    }
49}