1use bytes::{buf::Buf, BytesMut};
7use tokio_util::codec::{Decoder, Encoder};
8
9use crate::{Error, Result};
10
11pub struct JsonLinesCodec<D, S> {
13 deser: std::marker::PhantomData<D>,
14 ser: std::marker::PhantomData<S>,
15}
16
17impl<D, S> Default for JsonLinesCodec<D, S> {
18 fn default() -> Self {
19 Self {
20 deser: std::marker::PhantomData,
21 ser: std::marker::PhantomData,
22 }
23 }
24}
25
26impl<D, S> Decoder for JsonLinesCodec<D, S>
27where
28 D: serde::de::DeserializeOwned,
29{
30 type Item = D;
31 type Error = Error;
32
33 fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>> {
34 if buf.is_empty() {
35 Ok(None)
36 } else {
37 loop {
40 if let Some(newline_offset) = memchr::memchr(b'\n', &buf[..]) {
42 let to_parse = &buf[..newline_offset];
44 if newline_offset == 0 {
45 buf.advance(newline_offset + 1);
46 continue;
48 }
49 debug_assert!(buf[newline_offset] == b'\n');
50 match serde_json::from_slice(&to_parse[..newline_offset]) {
51 Ok(msg) => {
52 buf.advance(newline_offset + 1);
53 return Ok(Some(msg));
54 }
55 Err(_e) => {
56 buf.advance(newline_offset + 1);
57 return Err(Error::DeserializeJson);
62 }
63 }
64 } else {
65 return Ok(None);
67 }
68 }
69 }
70 }
71}
72
73impl<D, S> Encoder<S> for JsonLinesCodec<D, S>
76where
77 S: serde::Serialize,
78{
79 type Error = Error;
80 fn encode(&mut self, msg: S, final_buf: &mut BytesMut) -> Result<()> {
81 let mut v = serde_json::to_vec(&msg).map_err(|_| Error::SerializeJson)?;
82 if memchr::memchr(b'\n', &v).is_some() {
83 return Err(Error::NewlineInData);
84 }
85 v.push(b'\n');
86 final_buf.extend_from_slice(&v);
87 Ok(())
88 }
89}
90
91#[cfg(test)]
92mod tests {
93 use super::*;
94 use serde::{Deserialize, Serialize};
95
96 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
97 struct MyStruct {
98 val1: u8,
99 val2: u8,
100 }
101
102 #[test]
103 fn roundtrip() -> Result<()> {
104 let msg1 = MyStruct { val1: 12, val2: 34 };
105 let msg2 = MyStruct { val1: 56, val2: 78 };
106 let mut bytes = BytesMut::new();
107 let mut codec = JsonLinesCodec::default();
108 codec.encode(msg1.clone(), &mut bytes)?;
109 codec.encode(msg2.clone(), &mut bytes)?;
110 let found1: Option<MyStruct> = codec.decode(&mut bytes)?;
111 let found2: Option<MyStruct> = codec.decode(&mut bytes)?;
112 assert_eq!(found1, Some(msg1));
113 assert_eq!(found2, Some(msg2));
114 Ok(())
115 }
116}