Skip to main content

photon_protocol/codec/protobuf/
codec.rs

1use std::collections::HashMap;
2use std::time::{Duration, SystemTime, UNIX_EPOCH};
3
4use bytes::BytesMut;
5use prost::Message;
6
7use photon_core::types::id::RunId;
8use photon_core::types::metric::{Metric, MetricBatch, MetricPoint};
9use photon_core::types::query::{MetricQuery, MetricSeries, QueryRequest, QueryResponse};
10
11use crate::codec::protobuf::types::{
12    MetricBatchContent, MetricPointCompact, ProtoMetricQuery, ProtoMetricSeries,
13    ProtoQueryRequest, ProtoQueryResponse,
14};
15use crate::ports::codec::{Codec, CodecError};
16
17#[derive(Clone)]
18pub struct ProtobufCodec;
19
20impl Codec<MetricBatch> for ProtobufCodec {
21    fn encode(&self, batch: &MetricBatch, output: &mut BytesMut) -> Result<(), CodecError> {
22        let mut key_to_index: HashMap<&str, u32> = HashMap::new();
23        let mut keys: Vec<String> = Vec::new();
24
25        for p in &batch.points {
26            let key_str = p.key.as_str();
27            if !key_to_index.contains_key(key_str) {
28                let idx = keys.len() as u32;
29                key_to_index.insert(key_str, idx);
30                keys.push(key_str.to_owned());
31            }
32        }
33
34        let proto = MetricBatchContent {
35            run_id: batch.run_id.to_string(),
36            keys,
37            points: batch
38                .points
39                .iter()
40                .map(|p| MetricPointCompact {
41                    key_index: key_to_index[p.key.as_str()],
42                    value: p.value,
43                    step: p.step,
44                    timestamp_epoch_ms: system_time_to_epoch_ms(p.timestamp),
45                })
46                .collect(),
47        };
48
49        let len = proto.encoded_len();
50        output.reserve(len);
51
52        proto.encode(output).map_err(|e| CodecError::EncodeFailed {
53            reason: e.to_string(),
54        })?;
55
56        Ok(())
57    }
58
59    fn decode(&self, input: &[u8]) -> Result<MetricBatch, CodecError> {
60        let proto = MetricBatchContent::decode(input).map_err(|e| CodecError::DecodeFailed {
61            reason: e.to_string(),
62        })?;
63
64        let run_id: uuid::Uuid = proto.run_id.parse().map_err(|_| CodecError::DecodeFailed {
65            reason: format!("invalid run_id: {}", proto.run_id),
66        })?;
67
68        let metrics: Vec<Metric> = proto
69            .keys
70            .iter()
71            .map(|k| {
72                Metric::new(k).map_err(|e| CodecError::DecodeFailed {
73                    reason: format!("invalid metric key: {e}"),
74                })
75            })
76            .collect::<Result<Vec<_>, CodecError>>()?;
77
78        let points = proto
79            .points
80            .into_iter()
81            .map(|p| {
82                let key = metrics
83                    .get(p.key_index as usize)
84                    .ok_or_else(|| CodecError::DecodeFailed {
85                        reason: format!("key_index {} out of range (have {} keys)", p.key_index, metrics.len()),
86                    })?
87                    .clone();
88
89                Ok(MetricPoint {
90                    key,
91                    value: p.value,
92                    step: p.step,
93                    timestamp: epoch_ms_to_system_time(p.timestamp_epoch_ms),
94                })
95            })
96            .collect::<Result<Vec<_>, CodecError>>()?;
97
98        Ok(MetricBatch {
99            run_id: RunId::from(run_id),
100            points,
101        })
102    }
103}
104
105impl Codec<QueryRequest> for ProtobufCodec {
106    fn encode(&self, value: &QueryRequest, output: &mut BytesMut) -> Result<(), CodecError> {
107        let proto = ProtoQueryRequest::from(value);
108        let len = proto.encoded_len();
109        output.reserve(len);
110
111        proto.encode(output).map_err(|e| CodecError::EncodeFailed {
112            reason: e.to_string(),
113        })?;
114
115        Ok(())
116    }
117
118    fn decode(&self, input: &[u8]) -> Result<QueryRequest, CodecError> {
119        let proto = ProtoQueryRequest::decode(input).map_err(|e| CodecError::DecodeFailed {
120            reason: e.to_string(),
121        })?;
122
123        QueryRequest::try_from(proto).map_err(|e| CodecError::DecodeFailed {
124            reason: e.to_string(),
125        })
126    }
127}
128
129impl Codec<QueryResponse> for ProtobufCodec {
130    fn encode(&self, value: &QueryResponse, output: &mut BytesMut) -> Result<(), CodecError> {
131        let proto = ProtoQueryResponse::from(value);
132        let len = proto.encoded_len();
133        output.reserve(len);
134
135        proto.encode(output).map_err(|e| CodecError::EncodeFailed {
136            reason: e.to_string(),
137        })?;
138
139        Ok(())
140    }
141
142    fn decode(&self, input: &[u8]) -> Result<QueryResponse, CodecError> {
143        let proto = ProtoQueryResponse::decode(input).map_err(|e| CodecError::DecodeFailed {
144            reason: e.to_string(),
145        })?;
146
147        QueryResponse::try_from(proto).map_err(|e| CodecError::DecodeFailed {
148            reason: e.to_string(),
149        })
150    }
151}
152
153impl Codec<MetricQuery> for ProtobufCodec {
154    fn encode(&self, value: &MetricQuery, output: &mut BytesMut) -> Result<(), CodecError> {
155        let proto = ProtoMetricQuery::from(value);
156        let len = proto.encoded_len();
157        output.reserve(len);
158
159        proto.encode(output).map_err(|e| CodecError::EncodeFailed {
160            reason: e.to_string(),
161        })?;
162
163        Ok(())
164    }
165
166    fn decode(&self, input: &[u8]) -> Result<MetricQuery, CodecError> {
167        let proto = ProtoMetricQuery::decode(input).map_err(|e| CodecError::DecodeFailed {
168            reason: e.to_string(),
169        })?;
170
171        MetricQuery::try_from(proto).map_err(|e| CodecError::DecodeFailed {
172            reason: e.to_string(),
173        })
174    }
175}
176
177impl Codec<MetricSeries> for ProtobufCodec {
178    fn encode(&self, value: &MetricSeries, output: &mut BytesMut) -> Result<(), CodecError> {
179        let proto = ProtoMetricSeries::from(value);
180        let len = proto.encoded_len();
181        output.reserve(len);
182
183        proto.encode(output).map_err(|e| CodecError::EncodeFailed {
184            reason: e.to_string(),
185        })?;
186
187        Ok(())
188    }
189
190    fn decode(&self, input: &[u8]) -> Result<MetricSeries, CodecError> {
191        let proto = ProtoMetricSeries::decode(input).map_err(|e| CodecError::DecodeFailed {
192            reason: e.to_string(),
193        })?;
194
195        MetricSeries::try_from(proto).map_err(|e| CodecError::DecodeFailed {
196            reason: e.to_string(),
197        })
198    }
199}
200
201fn system_time_to_epoch_ms(time: SystemTime) -> u64 {
202    time.duration_since(UNIX_EPOCH)
203        .unwrap_or(Duration::ZERO)
204        .as_millis() as u64
205}
206
207fn epoch_ms_to_system_time(ms: u64) -> SystemTime {
208    UNIX_EPOCH + Duration::from_millis(ms)
209}