From 61bb33ec9ccceb565e89e980991f6a58ef934a67 Mon Sep 17 00:00:00 2001 From: anatawa12 Date: Sun, 8 Oct 2023 21:39:54 +0900 Subject: [PATCH 1/2] perf: speedup H264Parser --- media/src/io/h264_reader/mod.rs | 96 +++++++++++++++++++-------------- 1 file changed, 56 insertions(+), 40 deletions(-) diff --git a/media/src/io/h264_reader/mod.rs b/media/src/io/h264_reader/mod.rs index 362ac1189..8fbb8b3eb 100644 --- a/media/src/io/h264_reader/mod.rs +++ b/media/src/io/h264_reader/mod.rs @@ -4,7 +4,7 @@ mod h264_reader_test; use std::fmt; use std::io::Read; -use bytes::{BufMut, Bytes, BytesMut}; +use bytes::{BufMut, BytesMut}; use crate::error::{Error, Result}; @@ -135,11 +135,14 @@ const NAL_PREFIX_4BYTES: [u8; 4] = [0, 0, 0, 1]; /// H264Reader reads data from stream and constructs h264 nal units pub struct H264Reader { reader: R, - nal_buffer: BytesMut, - count_of_consecutive_zero_bytes: usize, + // reading buffers + temp_buf: Box<[u8]>, + buf_read_end: usize, + buf_filled_end: usize, + // for reading nal_prefix_parsed: bool, - read_buffer: Vec, - temp_buf: Vec, + count_of_consecutive_zero_bytes: usize, + nal_buffer: BytesMut, } impl H264Reader { @@ -147,47 +150,63 @@ impl H264Reader { pub fn new(reader: R, capacity: usize) -> H264Reader { H264Reader { reader, - nal_buffer: BytesMut::new(), - count_of_consecutive_zero_bytes: 0, nal_prefix_parsed: false, - read_buffer: vec![], - temp_buf: vec![0u8; capacity], + temp_buf: vec![0u8; capacity].into_boxed_slice(), + buf_read_end: 0, + buf_filled_end: 0, + count_of_consecutive_zero_bytes: 0, + nal_buffer: BytesMut::new(), } } - fn read(&mut self, num_to_read: usize) -> Result { - let buf = &mut self.temp_buf; - while self.read_buffer.len() < num_to_read { - let n = match self.reader.read(buf) { - Ok(n) => { - if n == 0 { - break; - } - n - } - Err(e) => return Err(Error::Io(e.into())), - }; + fn read4(&mut self) -> Result<([u8; 4], usize)> { + let mut result = [0u8; 4]; + let mut result_filled = 0; + loop { + let in_buffer = self.buf_filled_end - self.buf_read_end; + + if in_buffer + result_filled >= 4 { + let consume = 4 - result_filled; + result[result_filled..].copy_from_slice(&self.temp_buf[self.buf_read_end..][..consume]); + self.buf_read_end += consume; + return Ok((result, 4)); + } + + result[result_filled..][..in_buffer].copy_from_slice(&self.temp_buf[self.buf_read_end..self.buf_filled_end]); + result_filled += in_buffer; - self.read_buffer.extend_from_slice(&buf[0..n]); + self.buf_read_end = 0; + self.buf_filled_end = self.reader.read(&mut self.temp_buf)?; + + if self.buf_filled_end == 0 { + return Ok((result, result_filled)); + } } + } - let num_should_read = if num_to_read <= self.read_buffer.len() { - num_to_read - } else { - self.read_buffer.len() - }; + fn read1(&mut self) -> Result> { + let in_buffer = self.buf_filled_end - self.buf_read_end; + if in_buffer != 0 { + let value = self.temp_buf[self.buf_read_end]; + self.buf_read_end += 1; + return Ok(Some(value)); + } + + self.buf_read_end = 0; + self.buf_filled_end = self.reader.read(&mut self.temp_buf)?; + + if self.buf_filled_end == 0 { + return Ok(None); + } - Ok(Bytes::from( - self.read_buffer - .drain(..num_should_read) - .collect::>(), - )) + let value = self.temp_buf[self.buf_read_end]; + self.buf_read_end += 1; + return Ok(Some(value)); } fn bit_stream_starts_with_h264prefix(&mut self) -> Result { - let prefix_buffer = self.read(4)?; + let (prefix_buffer, n) = self.read4()?; - let n = prefix_buffer.len(); if n == 0 { return Err(Error::ErrIoEOF); } @@ -229,13 +248,10 @@ impl H264Reader { } loop { - let buffer = self.read(1)?; - let n = buffer.len(); + let Some(read_byte) = self.read1()? else { + break + }; - if n != 1 { - break; - } - let read_byte = buffer[0]; let nal_found = self.process_byte(read_byte); if nal_found { let nal_unit_type = NalUnitType::from(self.nal_buffer[0] & 0x1F); From aff6b55f3e29460adf8a1a2ec3619752411dfa63 Mon Sep 17 00:00:00 2001 From: anatawa12 Date: Sun, 8 Oct 2023 21:53:04 +0900 Subject: [PATCH 2/2] refactor: split code for buffer management --- media/src/io/h264_reader/mod.rs | 78 +++++++++++++++++++++------------ 1 file changed, 51 insertions(+), 27 deletions(-) diff --git a/media/src/io/h264_reader/mod.rs b/media/src/io/h264_reader/mod.rs index 8fbb8b3eb..a489e1e7a 100644 --- a/media/src/io/h264_reader/mod.rs +++ b/media/src/io/h264_reader/mod.rs @@ -132,13 +132,49 @@ impl NAL { const NAL_PREFIX_3BYTES: [u8; 3] = [0, 0, 1]; const NAL_PREFIX_4BYTES: [u8; 4] = [0, 0, 0, 1]; +/// Wrapper class around reading buffer +struct ReadBuffer { + buffer: Box<[u8]>, + read_end: usize, + filled_end: usize, +} + +impl ReadBuffer { + fn new(capacity: usize) -> ReadBuffer { + Self { + buffer: vec![0u8; capacity].into_boxed_slice(), + read_end: 0, + filled_end: 0, + } + } + + #[inline] + fn in_buffer(&self) -> usize { + self.filled_end - self.read_end + } + + fn consume(&mut self, consume: usize) -> &[u8] { + debug_assert!(self.read_end + consume <= self.filled_end); + let result = &self.buffer[self.read_end..][..consume]; + self.read_end += consume; + result + } + + pub(crate) fn fill_buffer(&mut self, reader: &mut impl Read) -> Result<()> { + debug_assert_eq!(self.read_end, self.filled_end); + + self.read_end = 0; + self.filled_end = reader.read(&mut self.buffer)?; + + Ok(()) + } +} + /// H264Reader reads data from stream and constructs h264 nal units pub struct H264Reader { reader: R, // reading buffers - temp_buf: Box<[u8]>, - buf_read_end: usize, - buf_filled_end: usize, + buffer: ReadBuffer, // for reading nal_prefix_parsed: bool, count_of_consecutive_zero_bytes: usize, @@ -151,9 +187,7 @@ impl H264Reader { H264Reader { reader, nal_prefix_parsed: false, - temp_buf: vec![0u8; capacity].into_boxed_slice(), - buf_read_end: 0, - buf_filled_end: 0, + buffer: ReadBuffer::new(capacity), count_of_consecutive_zero_bytes: 0, nal_buffer: BytesMut::new(), } @@ -163,45 +197,35 @@ impl H264Reader { let mut result = [0u8; 4]; let mut result_filled = 0; loop { - let in_buffer = self.buf_filled_end - self.buf_read_end; + let in_buffer = self.buffer.in_buffer(); if in_buffer + result_filled >= 4 { let consume = 4 - result_filled; - result[result_filled..].copy_from_slice(&self.temp_buf[self.buf_read_end..][..consume]); - self.buf_read_end += consume; + result[result_filled..].copy_from_slice(self.buffer.consume(consume)); return Ok((result, 4)); } - result[result_filled..][..in_buffer].copy_from_slice(&self.temp_buf[self.buf_read_end..self.buf_filled_end]); + result[result_filled..][..in_buffer].copy_from_slice(self.buffer.consume(in_buffer)); result_filled += in_buffer; - self.buf_read_end = 0; - self.buf_filled_end = self.reader.read(&mut self.temp_buf)?; + self.buffer.fill_buffer(&mut self.reader)?; - if self.buf_filled_end == 0 { + if self.buffer.in_buffer() == 0 { return Ok((result, result_filled)); } } } fn read1(&mut self) -> Result> { - let in_buffer = self.buf_filled_end - self.buf_read_end; - if in_buffer != 0 { - let value = self.temp_buf[self.buf_read_end]; - self.buf_read_end += 1; - return Ok(Some(value)); - } + if self.buffer.in_buffer() == 0 { + self.buffer.fill_buffer(&mut self.reader)?; - self.buf_read_end = 0; - self.buf_filled_end = self.reader.read(&mut self.temp_buf)?; - - if self.buf_filled_end == 0 { - return Ok(None); + if self.buffer.in_buffer() == 0 { + return Ok(None); + } } - let value = self.temp_buf[self.buf_read_end]; - self.buf_read_end += 1; - return Ok(Some(value)); + Ok(Some(self.buffer.consume(1)[0])) } fn bit_stream_starts_with_h264prefix(&mut self) -> Result {