diff --git a/README.md b/README.md index 1d7e7cf..ce5983f 100644 --- a/README.md +++ b/README.md @@ -321,8 +321,12 @@ normalization, `offset_gain`, `linear` and `dual_mono`. `gs://` and `s3://` objects are fetched with streaming reads and concurrent range requests on seek, buffered in memory (`AVTENSOR_MAX_CLOUD_OBJECT_BYTES` -caps the buffer; default 16 GiB). HTTP(S) inputs go through FFmpeg's own -protocol layer. +caps the buffer; default 16 GiB). A seek that lands more than +`AVTENSOR_CLOUD_SEEK_THRESHOLD_BYTES` (default 64 MiB) past the +already-fetched data opens a new range request instead of streaming through +the gap, and stops readers that would keep downloading the skipped bytes; +lower it to skip more aggressively when decoding a subset of streams. +HTTP(S) inputs go through FFmpeg's own protocol layer. ### Google Cloud Storage (`gs://`) diff --git a/src/decoder/io.rs b/src/decoder/io.rs index e78be6b..761d2d7 100644 --- a/src/decoder/io.rs +++ b/src/decoder/io.rs @@ -9,7 +9,10 @@ use std::{ ffi::CStr, fs::File, io::{Read, Seek, SeekFrom}, - sync::{Arc, Mutex}, + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, Mutex, + }, }; use tokio::task::JoinHandle; @@ -203,6 +206,33 @@ fn max_cloud_object_size() -> usize { } } +/// Default distance (bytes) between the fetched frontier and a forward-seek +/// target beyond which the seek callback opens a new range request instead +/// of waiting for an existing reader to stream through the gap. +const DEFAULT_CLOUD_SEEK_THRESHOLD: usize = 64 * 1024 * 1024; // 64 MiB + +/// The forward-seek threshold: `AVTENSOR_CLOUD_SEEK_THRESHOLD_BYTES` (bytes) +/// when set to a valid integer, otherwise [`DEFAULT_CLOUD_SEEK_THRESHOLD`]. +/// Lowering it makes seeks over cloud objects switch to a new range request +/// (skipping the bytes in between) more eagerly — useful when decoding a +/// subset of streams from finely-interleaved containers, where the skipped +/// gaps are much smaller than 64 MiB. +fn cloud_seek_threshold() -> usize { + match std::env::var("AVTENSOR_CLOUD_SEEK_THRESHOLD_BYTES") { + Ok(value) => match value.parse::() { + Ok(bytes) => bytes, + Err(_) => { + log::warn!( + "Ignoring invalid AVTENSOR_CLOUD_SEEK_THRESHOLD_BYTES value {value:?}; \ + using the default of {DEFAULT_CLOUD_SEEK_THRESHOLD} bytes" + ); + DEFAULT_CLOUD_SEEK_THRESHOLD + } + }, + Err(_) => DEFAULT_CLOUD_SEEK_THRESHOLD, + } +} + /// A parsed cloud-storage URI. #[derive(Clone)] enum CloudUri { @@ -269,10 +299,14 @@ async fn open_cloud_stream( /// Arguments: /// - `stream`: The byte stream to read data from. /// - `start`: The index of the first byte in the buffer to start writing data to. +/// - `range_end`: One past the last byte this reader should fetch. Starts at +/// the object size and may be truncated concurrently by the seek callback +/// when FFmpeg skips forward past this reader's position. /// - `context`: A handle to the storage buffer to which we should write data as it comes in. async fn cloud_asset_reader( mut stream: CloudChunkStream, start: usize, + range_end: Arc, context: Arc>, ) -> Result<(), anyhow::Error> { let mut current_position = start; @@ -306,6 +340,19 @@ async fn cloud_asset_reader( // Update the write position current_position = new_position; + // Stop when a forward seek truncated this reader's range: + // the bytes past `range_end` were skipped by FFmpeg, and if + // they are ever needed another reader fetches them. + let end = range_end.load(Ordering::Relaxed); + if current_position >= end { + log::debug!( + "cloud_asset_reader: range was truncated to end at {} - stopping at {}", + end, + current_position + ); + break; + } + // Check whether we should continue to consume the stream. let next_position = buffer .indicators @@ -335,6 +382,37 @@ async fn cloud_asset_reader( Ok(()) } +/// Decides which readers to truncate when a forward seek to `seek_target` +/// spawns a new reader. +/// +/// `ranges` holds the `(start, end)` byte ranges of the existing readers and +/// `frontier` is one past the last fetched byte below the seek target +/// (`closest_available_byte + 1`). A reader whose range covers the frontier +/// would keep streaming through the `[frontier, seek_target)` gap FFmpeg +/// just skipped, so it is truncated to end at the frontier — its useful +/// extent. Returns `(reader_index, new_end)` pairs. +/// +/// A later backward seek into the truncated gap still works: the truncated +/// reader no longer covers the gap, so the seek (and the read callback's +/// stall recovery) sees no live reader for those bytes and spawns a fresh +/// one at the requested position. +fn reader_truncations( + ranges: &[(usize, usize)], + frontier: usize, + seek_target: usize, +) -> Vec<(usize, usize)> { + if frontier >= seek_target { + // No gap is being skipped. + return Vec::new(); + } + ranges + .iter() + .enumerate() + .filter(|(_, (start, end))| *start < frontier && frontier < *end) + .map(|(index, _)| (index, frontier)) + .collect() +} + /// Returns an [`AVFormatContextInput`] capable of streaming data from cloud storage. /// /// `s3_config` is the explicit S3 client configuration for `s3://` URIs (see @@ -422,21 +500,21 @@ pub fn cloud_storage_avio_reader( })); // Spawn a future to sequentially read data from the stream into the StorageBuffer. - let read_sequential_future = cloud_asset_reader(stream, 0, context.clone()); + let range_end = Arc::new(AtomicUsize::new(obj_size)); + let read_sequential_future = cloud_asset_reader(stream, 0, range_end.clone(), context.clone()); runtime.block_on(async { let mut context = context.lock().await; context.readers.push(DataReader { range: ReadRange { start: 0, - end: obj_size, + end: range_end, }, handle: runtime.spawn(read_sequential_future), }); }); - // TODO (rikheijdens): we could make these configurable and tunable. let ffmpeg_buffer_size = 32768; - let data_segment_size = 64 * 1024 * 1024; // 64 MiB + let data_segment_size = cloud_seek_threshold(); let read_context = context.clone(); let seek_context = context.clone(); @@ -520,7 +598,7 @@ pub fn cloud_storage_avio_reader( // waiting; otherwise surface it for a respawn. let covered = readers.iter().any(|reader| { reader.range.start <= first_missing - && first_missing < reader.range.end + && first_missing < reader.range.end.load(Ordering::Relaxed) && !reader.handle.is_finished() }); (!covered).then_some(first_missing) @@ -552,15 +630,20 @@ pub fn cloud_storage_avio_reader( let (_, stream) = open_cloud_stream(&read_uri, pos) .await .context("reopening cloud stream after a reader stall")?; - let read_future = - cloud_asset_reader(stream, pos, read_context.clone()); + let range_end = Arc::new(AtomicUsize::new(obj_size)); + let read_future = cloud_asset_reader( + stream, + pos, + range_end.clone(), + read_context.clone(), + ); let CloudStorageReaderContext { ref mut readers, .. } = *read_context.lock().await; readers.push(DataReader { range: ReadRange { start: pos, - end: obj_size, + end: range_end, }, handle: runtime.spawn(read_future), }); @@ -665,7 +748,7 @@ pub fn cloud_storage_avio_reader( // then we'll assume the data will come in from that data loader. for reader in readers.iter() { if reader.range.start <= closest_available_byte - && closest_available_byte < reader.range.end + && closest_available_byte < reader.range.end.load(Ordering::Relaxed) && !reader.handle.is_finished() { // We have an active reader that is already fetching the data we need, @@ -678,16 +761,47 @@ pub fn cloud_storage_avio_reader( // There isn't a reader that is reading any data that we're after, or that is close to it // spawn a new one by making a range request, to minimize the number of requests that we're making // we'll request a read range starting from the requested offset to the end of the file. + // + // Truncate readers that would otherwise keep streaming + // through the `[frontier, position)` gap FFmpeg just + // skipped, so the skipped bytes are not downloaded. + let frontier = closest_available_byte + 1; + let ranges: Vec<(usize, usize)> = readers + .iter() + .map(|reader| { + ( + reader.range.start, + reader.range.end.load(Ordering::Relaxed), + ) + }) + .collect(); + for (index, new_end) in reader_truncations(&ranges, frontier, position) { + log::debug!( + "Truncating DataReader [{}..{}] to end at {} after seek to {}", + ranges[index].0, + ranges[index].1, + new_end, + position + ); + readers[index].range.end.store(new_end, Ordering::Relaxed); + } + let (_, stream) = open_cloud_stream(&seek_uri, position).await?; - let read_future = cloud_asset_reader(stream, position, seek_context.clone()); + let range_end = Arc::new(AtomicUsize::new(obj_size)); + let read_future = cloud_asset_reader( + stream, + position, + range_end.clone(), + seek_context.clone(), + ); // Spawn the future and keep track of the reader. log::debug!("Spawned new DataReader to read range [{}..{}]", position, obj_size); readers.push(DataReader { range: ReadRange { start: position, - end: obj_size + end: range_end }, handle: runtime.spawn(read_future) }); @@ -714,8 +828,11 @@ struct CloudStorageBuffer { struct ReadRange { // Start of the range in bytes. start: usize, - // End of the range in bytes. - end: usize, + // End of the range in bytes (exclusive). Shared with the reader task, + // which stops once its write position reaches it, so a later forward + // seek can truncate the range and stop the reader from downloading + // bytes FFmpeg skipped over. + end: Arc, } struct DataReader { @@ -755,4 +872,45 @@ mod tests { let filename = CString::new(test_video.path().to_str().unwrap()).unwrap(); avio_reader(filename.as_c_str()).unwrap(); } + + /// Exercises the default, a valid override, and an invalid override in a + /// single test: environment variables are process-global, so splitting + /// the cases across tests would race under the parallel test runner. + #[test] + fn test_cloud_seek_threshold_env_override() { + init_logger(); + + std::env::remove_var("AVTENSOR_CLOUD_SEEK_THRESHOLD_BYTES"); + assert_eq!(cloud_seek_threshold(), DEFAULT_CLOUD_SEEK_THRESHOLD); + + std::env::set_var("AVTENSOR_CLOUD_SEEK_THRESHOLD_BYTES", "1048576"); + assert_eq!(cloud_seek_threshold(), 1048576); + + std::env::set_var("AVTENSOR_CLOUD_SEEK_THRESHOLD_BYTES", "not-a-number"); + assert_eq!(cloud_seek_threshold(), DEFAULT_CLOUD_SEEK_THRESHOLD); + + std::env::remove_var("AVTENSOR_CLOUD_SEEK_THRESHOLD_BYTES"); + } + + #[test] + fn test_reader_truncations() { + // The eager reader covering the frontier is truncated to it. + assert_eq!(reader_truncations(&[(0, 1000)], 10, 500), vec![(0, 10)]); + // A reader starting past the seek target is untouched. + assert_eq!(reader_truncations(&[(600, 1000)], 10, 500), vec![]); + // A reader that already stops at (or below) the frontier is untouched. + assert_eq!(reader_truncations(&[(0, 10)], 10, 500), vec![]); + assert_eq!(reader_truncations(&[(0, 5)], 10, 500), vec![]); + // No gap between the frontier and the seek target -> nothing to do. + assert_eq!(reader_truncations(&[(0, 1000)], 500, 500), vec![]); + assert_eq!(reader_truncations(&[(0, 1000)], 600, 500), vec![]); + // No readers at all. + assert_eq!(reader_truncations(&[], 10, 500), vec![]); + // Multiple readers: only the ones covering the frontier are + // truncated, and the returned indices identify them. + assert_eq!( + reader_truncations(&[(0, 1000), (20, 1000), (600, 1000)], 100, 500), + vec![(0, 100), (1, 100)] + ); + } } diff --git a/src/decoder/mod.rs b/src/decoder/mod.rs index 07a4221..f90bf08 100644 --- a/src/decoder/mod.rs +++ b/src/decoder/mod.rs @@ -216,6 +216,13 @@ pub fn decode_media( let video_stream = select_video_stream(&input_format_context, video_stream_request)?; let audio_streams = select_audio_streams(&input_format_context, audio_streams)?; + // Tell the demuxer to skip every stream we are not going to decode. + let mut selected_indices: Vec = video_stream.iter().filter_map(|s| s.index).collect(); + if let Some(audio_streams) = &audio_streams { + selected_indices.extend(audio_streams.iter().filter_map(|s| s.index)); + } + discard_unselected_streams(&mut input_format_context, &selected_indices); + // Count the number of streams that we'll be decoding. let num_streams_to_decode = if video_stream.is_some() { 1 } else { 0 } + audio_streams @@ -552,6 +559,29 @@ pub fn probe_media( Ok(probed) } +/// Marks every stream whose index is not in `selected_indices` with +/// `AVDISCARD_ALL`, so the demuxer skips their packets instead of parsing +/// them. Index-based demuxers (notably mp4/mov) additionally skip *reading* +/// discarded streams' chunks on seekable inputs, which turns into real +/// bandwidth savings on remote (http(s)/gs/s3) sources — e.g. an audio-only +/// decode no longer downloads the video track where the container layout +/// allows skipping it. +/// +/// The packet loop still drops packets without a matching filter context; +/// that stays as defense in depth for demuxers that ignore `discard`. +fn discard_unselected_streams( + input_format_context: &mut AVFormatContextInput, + selected_indices: &[usize], +) { + for stream in input_format_context.streams_mut() { + let index = stream.index as usize; + if !selected_indices.contains(&index) { + log::debug!("Marking unselected stream #{index} as AVDISCARD_ALL"); + stream.set_discard(ffi::AVDISCARD_ALL); + } + } +} + /// Selects video streams to decode /// /// Arguments: @@ -4433,6 +4463,107 @@ mod tests { Ok(()) } + /// Streams that were not selected for decoding must be marked + /// `AVDISCARD_ALL`, so the demuxer skips their packets (and, on seekable + /// inputs, avoids reading their chunks); selected streams must be left + /// untouched. + #[test] + fn test_discard_unselected_streams() -> anyhow::Result<()> { + init_logger(); + + let test_video = generate_test_video_file(&TestVideoParameters::default())?; + let source = MediaSource::Uri(test_video.path().to_str().unwrap().into()); + let mut input_format_context = source.open(None)?; + + let audio_index = input_format_context + .find_best_stream(ffi::AVMEDIA_TYPE_AUDIO)? + .map(|(index, _)| index) + .context("test video should contain an audio stream")?; + + discard_unselected_streams(&mut input_format_context, &[audio_index]); + + assert!( + input_format_context.streams().len() > 1, + "test video should contain more than one stream" + ); + for stream in input_format_context.streams() { + if stream.index as usize == audio_index { + assert_ne!( + stream.discard, + ffi::AVDISCARD_ALL, + "selected stream #{audio_index} must not be discarded" + ); + } else { + assert_eq!( + stream.discard, + ffi::AVDISCARD_ALL, + "unselected stream #{} must be discarded", + stream.index + ); + } + } + Ok(()) + } + + /// An audio-only decode — with the video stream discarded at the demuxer + /// — must produce exactly the same samples as the audio stream of a full + /// A/V decode. + #[test] + fn test_decode_media_audio_only_matches_full_decode() -> anyhow::Result<()> { + init_logger(); + + let test_video = generate_test_video_file(&TestVideoParameters { + channel_layout: ChannelLayout::Stereo, + ..Default::default() + })?; + let path: String = test_video.path().to_str().unwrap().into(); + + let decode = + |video_stream: Option| -> anyhow::Result> { + decode_media( + MediaDecodeRequest { + source: MediaSource::Uri(path.clone()), + start_time: None, + end_time: None, + video_stream, + audio_streams: Some(vec![AudioStreamRequest::default()]), + }, + None, + ) + }; + + let mut audio_only = decode(None)?; + assert_eq!( + audio_only.len(), + 1, + "audio-only decode must yield exactly one stream" + ); + let audio_only = audio_only.remove(0); + assert_eq!(audio_only.stream_type(), StreamType::Audio); + assert!(!audio_only.decoded_frames.is_empty()); + let audio_only_data = audio_only.data.context("audio stream should have data")?; + + let full_audio_data = decode(Some(VideoStreamRequest::default()))? + .into_iter() + .find(|s| s.stream_type() == StreamType::Audio) + .context("full decode should contain an audio stream")? + .data + .context("audio stream should have data")?; + + assert_eq!(audio_only_data.size(), full_audio_data.size()); + let max_diff: f64 = audio_only_data + .f_sub(&full_audio_data)? + .abs() + .max() + .to_kind(tch::Kind::Float) + .try_into()?; + assert_eq!( + max_diff, 0.0, + "audio-only decode must match the full decode's audio samples" + ); + Ok(()) + } + #[test] fn test_decode_media_from_bytes() -> anyhow::Result<()> { init_logger();