node: Improve framing code

Thanks to @fintohaps for these suggestions.
This commit is contained in:
Alexis Sellier 2023-04-23 15:10:42 +02:00
parent 5a4d0fff3f
commit fd3a6583d0
No known key found for this signature in database
2 changed files with 51 additions and 22 deletions

View File

@ -85,36 +85,27 @@ impl StreamId {
/// Get the kind of stream this is.
pub fn kind(&self) -> Result<StreamKind, u8> {
let id = *self.0;
match (id >> 1) & 0b11 {
0 => Ok(StreamKind::Control),
1 => Ok(StreamKind::Gossip),
2 => Ok(StreamKind::Git),
n => Err(n as u8),
}
let kind = ((id >> 1) & 0b11) as u8;
StreamKind::try_from(kind)
}
/// Create a control identifier.
pub fn control(link: Link) -> Self {
match link {
Link::Outbound => Self(VarInt::from(0b000u8)),
Link::Inbound => Self(VarInt::from(0b001u8)),
}
let link = if link.is_outbound() { 0 } else { 1 };
Self(VarInt::from((StreamKind::Control as u8) << 1 | link))
}
/// Create a gossip identifier.
pub fn gossip(link: Link) -> Self {
match link {
Link::Outbound => Self(VarInt::from(0b010u8)),
Link::Inbound => Self(VarInt::from(0b011u8)),
}
let link = if link.is_outbound() { 0 } else { 1 };
Self(VarInt::from((StreamKind::Gossip as u8) << 1 | link))
}
/// Create a git identifier.
pub fn git(link: Link) -> Self {
match link {
Link::Outbound => Self(VarInt::from(0b100u8)),
Link::Inbound => Self(VarInt::from(0b101u8)),
}
let link = if link.is_outbound() { 0 } else { 1 };
Self(VarInt::from((StreamKind::Git as u8) << 1 | link))
}
/// Get the nth identifier while preserving the stream type and initiator.
@ -157,13 +148,27 @@ impl wire::Encode for StreamId {
/// Type of stream.
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
#[repr(u8)]
pub enum StreamKind {
/// Control stream, used to open and close streams.
Control,
Control = 0b00,
/// Gossip stream, used to exchange messages.
Gossip,
Gossip = 0b01,
/// Git stream, used for replication.
Git,
Git = 0b10,
}
impl TryFrom<u8> for StreamKind {
type Error = u8;
fn try_from(value: u8) -> Result<Self, Self::Error> {
match value {
0b00 => Ok(StreamKind::Control),
0b01 => Ok(StreamKind::Gossip),
0b10 => Ok(StreamKind::Git),
n => Err(n),
}
}
}
/// Protocol frame.
@ -357,3 +362,27 @@ impl wire::Encode for Frame {
Ok(n)
}
}
#[cfg(test)]
mod test {
use super::*;
#[test]
fn test_stream_id() {
assert_eq!(StreamId(VarInt(0b000)).kind().unwrap(), StreamKind::Control);
assert_eq!(StreamId(VarInt(0b010)).kind().unwrap(), StreamKind::Gossip);
assert_eq!(StreamId(VarInt(0b100)).kind().unwrap(), StreamKind::Git);
assert_eq!(StreamId(VarInt(0b001)).link(), Link::Inbound);
assert_eq!(StreamId(VarInt(0b000)).link(), Link::Outbound);
assert_eq!(StreamId(VarInt(0b101)).link(), Link::Inbound);
assert_eq!(StreamId(VarInt(0b100)).link(), Link::Outbound);
assert_eq!(StreamId::git(Link::Outbound), StreamId(VarInt(0b100)));
assert_eq!(StreamId::control(Link::Outbound), StreamId(VarInt(0b000)));
assert_eq!(StreamId::gossip(Link::Outbound), StreamId(VarInt(0b010)));
assert_eq!(StreamId::git(Link::Inbound), StreamId(VarInt(0b101)));
assert_eq!(StreamId::control(Link::Inbound), StreamId(VarInt(0b001)));
assert_eq!(StreamId::gossip(Link::Inbound), StreamId(VarInt(0b011)));
}
}

View File

@ -151,7 +151,7 @@ impl<T: AsRef<[u8]>> ChannelWriter<T> {
/// the protocol via the usual means, which is to close the connection. Git also doesn't have
/// any special message we can send to signal the end of the protocol.
///
/// Hence, we there's no other way for the server to know that we're done sending requests
/// Hence, there's no other way for the server to know that we're done sending requests
/// than to send a special message outside the git protocol. This message can then be processed
/// by the remote worker to end the protocol. We use the special "eof" control message for this.
pub fn eof(&self) -> Result<(), chan::SendError<ChannelEvent<T>>> {