use crate::{ConsumerGroupId, ConsumerMemberId, GroupId, NodeId, data::CodecError, layout::Reader};
use std::num::NonZeroU64;
mod broker;
mod counters;
mod logs;
mod partitions;
use broker::{decode_broker, encode_broker};
use counters::{decode_counters, encode_counters};
use logs::{decode_logs, encode_logs};
use partitions::{decode_status, encode_status};
const DELETE_GROUP: u8 = 5;
const LOGS: u8 = 6;
const ERROR: u8 = 0xff;
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Request {
Broker,
Counters,
Partitions,
Connections,
Logs,
DeleteGroup {
topic: String,
name: String,
},
}
impl Request {
const fn kind(&self) -> u8 {
match self {
Self::Broker => 1,
Self::Counters => 2,
Self::Partitions => 3,
Self::Connections => 4,
Self::DeleteGroup { .. } => DELETE_GROUP,
Self::Logs => LOGS,
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Reply {
Broker(Broker),
Counters(Counters),
Partitions(Partitions),
Connections(Connections),
Logs(Vec<LogEntry>),
GroupDeleted(Option<ConsumerGroupId>),
Error(String),
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Broker {
pub name: String,
pub id: NodeId,
pub pid: u32,
pub uptime_ms: u64,
pub max_clients: u32,
pub shards: u32,
pub brokers: Vec<(String, NodeId)>,
pub partitions: Vec<PartitionName>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct PartitionName {
pub group: GroupId,
pub topic: String,
pub partition: u32,
pub shard: u32,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[repr(u8)]
pub enum EventKind {
Count = 0,
Peak = 1,
Gauge = 2,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Counters {
pub uptime_ms: u64,
pub events: Vec<(String, EventKind)>,
pub totals: Vec<u64>,
pub threads: Vec<(String, Vec<u64>)>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Partitions {
pub uptime_ms: u64,
pub missing_shards: u32,
pub rows: Vec<PartitionStatus>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct PartitionStatus {
pub group: GroupId,
pub state: PartitionState,
pub ready: bool,
pub disk_pending: bool,
pub view: u64,
pub leader: Option<NodeId>,
pub accepted: u64,
pub committed: u64,
pub applied: u64,
pub pending_operations: u64,
pub pending_body_bytes: u64,
pub validation_cycles: u64,
pub reclaimed_bytes: u64,
pub written: u64,
pub durable: u64,
pub leader_commit: u64,
pub followers: Vec<(NodeId, u64)>,
pub log: Option<PartitionLog>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct PartitionLog {
pub first_offset: u64,
pub next_offset: u64,
pub producers: u32,
pub retention_age_millis: Option<NonZeroU64>,
pub retention_bytes: Option<NonZeroU64>,
pub policy_revision: u64,
pub segments: u32,
pub segment_bytes: u64,
pub dedup_entries: u64,
pub dedup_limit: Option<NonZeroU64>,
pub groups: Vec<GroupStatus>,
pub cursors: Vec<GroupCursor>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct GroupStatus {
pub id: ConsumerGroupId,
pub name: String,
pub revision: u64,
pub members: Vec<(ConsumerMemberId, NodeId)>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct GroupCursor {
pub group: ConsumerGroupId,
pub member: Option<ConsumerMemberId>,
pub offset: Option<u64>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Connections {
pub brokers: Vec<(NodeId, bool)>,
pub clients: u32,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct LogEntry {
pub uptime_ms: u64,
pub level: LogLevel,
pub thread: String,
pub message: String,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[repr(u8)]
pub enum LogLevel {
Info = 0,
Warn = 1,
Error = 2,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[repr(u8)]
pub enum PartitionState {
Local = 0,
Normal = 1,
Electing = 2,
}
pub fn encode_request(id: u64, request: &Request) -> Result<Vec<u8>, CodecError> {
let mut out = Vec::with_capacity(9);
out.push(request.kind());
out.extend_from_slice(&id.to_be_bytes());
if let Request::DeleteGroup { topic, name } = request {
text(&mut out, topic)?;
text(&mut out, name)?;
}
Ok(out)
}
pub fn decode_request(input: &[u8]) -> Result<(u64, Request), CodecError> {
let mut reader = Reader::new(input);
let kind = reader.u8()?;
let id = reader.u64()?;
let request = match kind {
1 => Request::Broker,
2 => Request::Counters,
3 => Request::Partitions,
4 => Request::Connections,
LOGS => Request::Logs,
DELETE_GROUP => Request::DeleteGroup {
topic: read_text(&mut reader)?,
name: read_text(&mut reader)?,
},
_ => return Err(CodecError::Command),
};
reader.finish()?;
Ok((id, request))
}
pub fn request_id(input: &[u8]) -> Option<u64> {
let mut reader = Reader::new(input);
reader.u8().ok()?;
reader.u64().ok()
}
pub fn encode_reply(id: u64, reply: &Reply) -> Result<Vec<u8>, CodecError> {
let mut out = Vec::new();
let kind = match reply {
Reply::Broker(_) => Request::Broker.kind(),
Reply::Counters(_) => Request::Counters.kind(),
Reply::Partitions(_) => Request::Partitions.kind(),
Reply::Connections(_) => Request::Connections.kind(),
Reply::Logs(_) => LOGS,
Reply::GroupDeleted(_) => DELETE_GROUP,
Reply::Error(_) => ERROR,
};
out.push(kind);
out.extend_from_slice(&id.to_be_bytes());
match reply {
Reply::Broker(broker) => encode_broker(&mut out, broker)?,
Reply::Counters(counters) => encode_counters(&mut out, counters)?,
Reply::Partitions(partitions) => {
out.extend_from_slice(&partitions.uptime_ms.to_be_bytes());
out.extend_from_slice(&partitions.missing_shards.to_be_bytes());
count(&mut out, partitions.rows.len())?;
for row in &partitions.rows {
encode_status(&mut out, row)?;
}
}
Reply::Connections(connections) => {
count(&mut out, connections.brokers.len())?;
for (id, up) in &connections.brokers {
out.extend_from_slice(id.as_bytes());
out.push(u8::from(*up));
}
out.extend_from_slice(&connections.clients.to_be_bytes());
}
Reply::Logs(entries) => encode_logs(&mut out, entries)?,
Reply::GroupDeleted(group) => {
out.push(u8::from(group.is_some()));
if let Some(group) = group {
out.extend_from_slice(group.as_bytes());
}
}
Reply::Error(message) => text(&mut out, message)?,
}
Ok(out)
}
pub fn decode_reply(input: &[u8]) -> Result<(u64, Reply), CodecError> {
let mut reader = Reader::new(input);
let kind = reader.u8()?;
let id = reader.u64()?;
let reply = match kind {
1 => Reply::Broker(decode_broker(&mut reader)?),
2 => Reply::Counters(decode_counters(&mut reader)?),
3 => {
let uptime_ms = reader.u64()?;
let missing_shards = reader.u32()?;
let rows = list(&mut reader, decode_status)?;
Reply::Partitions(Partitions {
uptime_ms,
missing_shards,
rows,
})
}
4 => Reply::Connections(Connections {
brokers: list(&mut reader, |reader| {
let id = NodeId::from_bytes(reader.array()?);
let up = match reader.u8()? {
0 => false,
1 => true,
_ => return Err(CodecError::Profile),
};
Ok((id, up))
})?,
clients: reader.u32()?,
}),
LOGS => Reply::Logs(decode_logs(&mut reader)?),
DELETE_GROUP => Reply::GroupDeleted(match reader.u8()? {
0 => None,
1 => Some(ConsumerGroupId::from_bytes(reader.array()?)),
_ => return Err(CodecError::Profile),
}),
ERROR => Reply::Error(read_text(&mut reader)?),
_ => return Err(CodecError::Command),
};
reader.finish()?;
Ok((id, reply))
}
fn count(out: &mut Vec<u8>, length: usize) -> Result<(), CodecError> {
let length = u32::try_from(length).map_err(|_| CodecError::Length)?;
out.extend_from_slice(&length.to_be_bytes());
Ok(())
}
fn list<T>(
reader: &mut Reader<'_>,
mut item: impl FnMut(&mut Reader<'_>) -> Result<T, CodecError>,
) -> Result<Vec<T>, CodecError> {
let length = reader.u32()? as usize;
if length > reader.rest().len() {
return Err(CodecError::Length);
}
(0..length).map(|_| item(reader)).collect()
}
fn values(out: &mut Vec<u8>, values: &[u64]) {
for value in values {
out.extend_from_slice(&value.to_be_bytes());
}
}
fn text(out: &mut Vec<u8>, value: &str) -> Result<(), CodecError> {
let length = u16::try_from(value.len()).map_err(|_| CodecError::Length)?;
out.extend_from_slice(&length.to_be_bytes());
out.extend_from_slice(value.as_bytes());
Ok(())
}
fn read_text(reader: &mut Reader<'_>) -> Result<String, CodecError> {
let length = reader.u16()? as usize;
std::str::from_utf8(reader.take(length)?)
.map(str::to_owned)
.map_err(|_| CodecError::Profile)
}
#[cfg(test)]
mod tests;