Merge pull request #1597 from ZettaScaleLabs/feat/fragment_marker

feat: add start/stop extension markers for fragment chain
eclipse-zenoh · Dec 5, 2024 · 64e8caa · 64e8caa
2 parents b3ccf82 + 13cf0de
commit 64e8caa
Show file tree

Hide file tree

Showing 21 changed files with 899 additions and 23 deletions.
diff --git a/commons/zenoh-codec/src/transport/fragment.rs b/commons/zenoh-codec/src/transport/fragment.rs
@@ -39,6 +39,8 @@ where
             more,
             sn,
             ext_qos,
+            ext_first,
+            ext_drop,
         } = x;
 
         // Header
@@ -49,7 +51,10 @@ where
         if *more {
             header |= flag::M;
         }
-        if ext_qos != &ext::QoSType::DEFAULT {
+        let mut n_exts = (ext_qos != &ext::QoSType::DEFAULT) as u8
+            + ext_first.is_some() as u8
+            + ext_drop.is_some() as u8;
+        if n_exts != 0 {
             header |= flag::Z;
         }
         self.write(&mut *writer, header)?;
@@ -59,7 +64,16 @@ where
 
         // Extensions
         if ext_qos != &ext::QoSType::DEFAULT {
-            self.write(&mut *writer, (*ext_qos, false))?;
+            n_exts -= 1;
+            self.write(&mut *writer, (*ext_qos, n_exts != 0))?;
+        }
+        if let Some(first) = ext_first {
+            n_exts -= 1;
+            self.write(&mut *writer, (first, n_exts != 0))?
+        }
+        if let Some(drop) = ext_drop {
+            n_exts -= 1;
+            self.write(&mut *writer, (drop, n_exts != 0))?
         }
 
         Ok(())
@@ -99,6 +113,8 @@ where
 
         // Extensions
         let mut ext_qos = ext::QoSType::DEFAULT;
+        let mut ext_first = None;
+        let mut ext_drop = None;
 
         let mut has_ext = imsg::has_flag(self.header, flag::Z);
         while has_ext {
@@ -110,6 +126,16 @@ where
                     ext_qos = q;
                     has_ext = ext;
                 }
+                ext::First::ID => {
+                    let (first, ext): (ext::First, bool) = eodec.read(&mut *reader)?;
+                    ext_first = Some(first);
+                    has_ext = ext;
+                }
+                ext::Drop::ID => {
+                    let (drop, ext): (ext::Drop, bool) = eodec.read(&mut *reader)?;
+                    ext_drop = Some(drop);
+                    has_ext = ext;
+                }
                 _ => {
                     has_ext = extension::skip(reader, "Fragment", ext)?;
                 }
@@ -121,6 +147,8 @@ where
             more,
             sn,
             ext_qos,
+            ext_first,
+            ext_drop,
         })
     }
 }
@@ -139,6 +167,8 @@ where
             sn,
             payload,
             ext_qos,
+            ext_first,
+            ext_drop,
         } = x;
 
         // Header
@@ -147,6 +177,8 @@ where
             more: *more,
             sn: *sn,
             ext_qos: *ext_qos,
+            ext_first: *ext_first,
+            ext_drop: *ext_drop,
         };
         self.write(&mut *writer, &header)?;
 
@@ -185,6 +217,8 @@ where
             more: header.more,
             sn: header.sn,
             ext_qos: header.ext_qos,
+            ext_first: header.ext_first,
+            ext_drop: header.ext_drop,
             payload,
         })
     }

diff --git a/commons/zenoh-codec/src/transport/init.rs b/commons/zenoh-codec/src/transport/init.rs
@@ -52,6 +52,7 @@ where
             ext_mlink,
             ext_lowlatency,
             ext_compression,
+            ext_patch,
         } = x;
 
         // Header
@@ -64,7 +65,8 @@ where
             + (ext_auth.is_some() as u8)
             + (ext_mlink.is_some() as u8)
             + (ext_lowlatency.is_some() as u8)
-            + (ext_compression.is_some() as u8);
+            + (ext_compression.is_some() as u8)
+            + (*ext_patch != ext::PatchType::NONE) as u8;
 
         #[cfg(feature = "shared-memory")]
         {
@@ -125,6 +127,10 @@ where
             n_exts -= 1;
             self.write(&mut *writer, (compression, n_exts != 0))?;
         }
+        if *ext_patch != ext::PatchType::NONE {
+            n_exts -= 1;
+            self.write(&mut *writer, (*ext_patch, n_exts != 0))?;
+        }
 
         Ok(())
     }
@@ -186,6 +192,7 @@ where
         let mut ext_mlink = None;
         let mut ext_lowlatency = None;
         let mut ext_compression = None;
+        let mut ext_patch = ext::PatchType::NONE;
 
         let mut has_ext = imsg::has_flag(self.header, flag::Z);
         while has_ext {
@@ -228,6 +235,11 @@ where
                     ext_compression = Some(q);
                     has_ext = ext;
                 }
+                ext::Patch::ID => {
+                    let (p, ext): (ext::PatchType, bool) = eodec.read(&mut *reader)?;
+                    ext_patch = p;
+                    has_ext = ext;
+                }
                 _ => {
                     has_ext = extension::skip(reader, "InitSyn", ext)?;
                 }
@@ -248,6 +260,7 @@ where
             ext_mlink,
             ext_lowlatency,
             ext_compression,
+            ext_patch,
         })
     }
 }
@@ -275,6 +288,7 @@ where
             ext_mlink,
             ext_lowlatency,
             ext_compression,
+            ext_patch,
         } = x;
 
         // Header
@@ -287,7 +301,8 @@ where
             + (ext_auth.is_some() as u8)
             + (ext_mlink.is_some() as u8)
             + (ext_lowlatency.is_some() as u8)
-            + (ext_compression.is_some() as u8);
+            + (ext_compression.is_some() as u8)
+            + (*ext_patch != ext::PatchType::NONE) as u8;
 
         #[cfg(feature = "shared-memory")]
         {
@@ -351,6 +366,10 @@ where
             n_exts -= 1;
             self.write(&mut *writer, (compression, n_exts != 0))?;
         }
+        if *ext_patch != ext::PatchType::NONE {
+            n_exts -= 1;
+            self.write(&mut *writer, (*ext_patch, n_exts != 0))?;
+        }
 
         Ok(())
     }
@@ -415,6 +434,7 @@ where
         let mut ext_mlink = None;
         let mut ext_lowlatency = None;
         let mut ext_compression = None;
+        let mut ext_patch = ext::PatchType::NONE;
 
         let mut has_ext = imsg::has_flag(self.header, flag::Z);
         while has_ext {
@@ -457,6 +477,11 @@ where
                     ext_compression = Some(q);
                     has_ext = ext;
                 }
+                ext::Patch::ID => {
+                    let (p, ext): (ext::PatchType, bool) = eodec.read(&mut *reader)?;
+                    ext_patch = p;
+                    has_ext = ext;
+                }
                 _ => {
                     has_ext = extension::skip(reader, "InitAck", ext)?;
                 }
@@ -478,6 +503,7 @@ where
             ext_mlink,
             ext_lowlatency,
             ext_compression,
+            ext_patch,
         })
     }
 }
diff --git a/commons/zenoh-codec/src/transport/join.rs b/commons/zenoh-codec/src/transport/join.rs
@@ -150,6 +150,7 @@ where
             next_sn,
             ext_qos,
             ext_shm,
+            ext_patch,
         } = x;
 
         // Header
@@ -160,7 +161,9 @@ where
         if resolution != &Resolution::default() || batch_size != &batch_size::MULTICAST {
             header |= flag::S;
         }
-        let mut n_exts = (ext_qos.is_some() as u8) + (ext_shm.is_some() as u8);
+        let mut n_exts = (ext_qos.is_some() as u8)
+            + (ext_shm.is_some() as u8)
+            + (*ext_patch != ext::PatchType::NONE) as u8;
         if n_exts != 0 {
             header |= flag::Z;
         }
@@ -201,6 +204,10 @@ where
             n_exts -= 1;
             self.write(&mut *writer, (shm, n_exts != 0))?;
         }
+        if *ext_patch != ext::PatchType::NONE {
+            n_exts -= 1;
+            self.write(&mut *writer, (*ext_patch, n_exts != 0))?;
+        }
 
         Ok(())
     }
@@ -264,6 +271,7 @@ where
         // Extensions
         let mut ext_qos = None;
         let mut ext_shm = None;
+        let mut ext_patch = ext::PatchType::NONE;
 
         let mut has_ext = imsg::has_flag(self.header, flag::Z);
         while has_ext {
@@ -280,6 +288,11 @@ where
                     ext_shm = Some(s);
                     has_ext = ext;
                 }
+                ext::Patch::ID => {
+                    let (p, ext): (ext::PatchType, bool) = eodec.read(&mut *reader)?;
+                    ext_patch = p;
+                    has_ext = ext;
+                }
                 _ => {
                     has_ext = extension::skip(reader, "Join", ext)?;
                 }
@@ -296,6 +309,7 @@ where
             next_sn,
             ext_qos,
             ext_shm,
+            ext_patch,
         })
     }
 }
diff --git a/commons/zenoh-codec/src/transport/mod.rs b/commons/zenoh-codec/src/transport/mod.rs
@@ -176,3 +176,43 @@ where
         Ok((ext.into(), more))
     }
 }
+
+// Extensions: Patch
+impl<W, const ID: u8> WCodec<(ext::PatchType<{ ID }>, bool), &mut W> for Zenoh080
+where
+    W: Writer,
+{
+    type Output = Result<(), DidntWrite>;
+
+    fn write(self, writer: &mut W, x: (ext::PatchType<{ ID }>, bool)) -> Self::Output {
+        let (x, more) = x;
+        let ext: ZExtZ64<{ ID }> = x.into();
+
+        self.write(&mut *writer, (&ext, more))
+    }
+}
+
+impl<R, const ID: u8> RCodec<(ext::PatchType<{ ID }>, bool), &mut R> for Zenoh080
+where
+    R: Reader,
+{
+    type Error = DidntRead;
+
+    fn read(self, reader: &mut R) -> Result<(ext::PatchType<{ ID }>, bool), Self::Error> {
+        let header: u8 = self.read(&mut *reader)?;
+        let codec = Zenoh080Header::new(header);
+        codec.read(reader)
+    }
+}
+
+impl<R, const ID: u8> RCodec<(ext::PatchType<{ ID }>, bool), &mut R> for Zenoh080Header
+where
+    R: Reader,
+{
+    type Error = DidntRead;
+
+    fn read(self, reader: &mut R) -> Result<(ext::PatchType<{ ID }>, bool), Self::Error> {
+        let (ext, more): (ZExtZ64<{ ID }>, bool) = self.read(&mut *reader)?;
+        Ok((ext.into(), more))
+    }
+}
diff --git a/commons/zenoh-protocol/src/transport/fragment.rs b/commons/zenoh-protocol/src/transport/fragment.rs
@@ -75,14 +75,26 @@ pub struct Fragment {
     pub sn: TransportSn,
     pub payload: ZSlice,
     pub ext_qos: ext::QoSType,
+    pub ext_first: Option<ext::First>,
+    pub ext_drop: Option<ext::Drop>,
 }
 
 // Extensions
 pub mod ext {
-    use crate::{common::ZExtZ64, zextz64};
+    use crate::{
+        common::{ZExtUnit, ZExtZ64},
+        zextunit, zextz64,
+    };
 
     pub type QoS = zextz64!(0x1, true);
     pub type QoSType = crate::transport::ext::QoSType<{ QoS::ID }>;
+
+    /// # Start extension
+    /// Mark the first fragment of a fragmented message
+    pub type First = zextunit!(0x2, false);
+    /// # Stop extension
+    /// Indicate that the remaining fragments has been dropped
+    pub type Drop = zextunit!(0x3, false);
 }
 
 impl Fragment {
@@ -97,13 +109,17 @@ impl Fragment {
         let sn: TransportSn = rng.gen();
         let payload = ZSlice::rand(rng.gen_range(8..128));
         let ext_qos = ext::QoSType::rand();
+        let ext_first = rng.gen_bool(0.5).then(ext::First::rand);
+        let ext_drop = rng.gen_bool(0.5).then(ext::Drop::rand);
 
         Fragment {
             reliability,
             sn,
             more,
             payload,
             ext_qos,
+            ext_first,
+            ext_drop,
         }
     }
 }
@@ -115,6 +131,8 @@ pub struct FragmentHeader {
     pub more: bool,
     pub sn: TransportSn,
     pub ext_qos: ext::QoSType,
+    pub ext_first: Option<ext::First>,
+    pub ext_drop: Option<ext::Drop>,
 }
 
 impl FragmentHeader {
@@ -128,12 +146,16 @@ impl FragmentHeader {
         let more = rng.gen_bool(0.5);
         let sn: TransportSn = rng.gen();
         let ext_qos = ext::QoSType::rand();
+        let ext_first = rng.gen_bool(0.5).then(ext::First::rand);
+        let ext_drop = rng.gen_bool(0.5).then(ext::Drop::rand);
 
         FragmentHeader {
             reliability,
             more,
             sn,
             ext_qos,
+            ext_first,
+            ext_drop,
         }
     }
 }