tfserver/codec/
length_delimited.rs1use 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)]
8pub 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}