forked from mirrors/gstreamer-rs
513 lines
18 KiB
Rust
513 lines
18 KiB
Rust
|
use std::collections::HashMap;
|
||
|
use std::mem;
|
||
|
use std::sync::{atomic, Arc, Mutex};
|
||
|
|
||
|
use once_cell::sync::Lazy;
|
||
|
|
||
|
use gst::{glib, prelude::*};
|
||
|
|
||
|
static CAT: Lazy<gst::DebugCategory> = Lazy::new(|| {
|
||
|
gst::DebugCategory::new(
|
||
|
"utilsrs-stream-producer",
|
||
|
gst::DebugColorFlags::empty(),
|
||
|
Some("gst_app Stream Producer interface"),
|
||
|
)
|
||
|
});
|
||
|
|
||
|
/// The interface for transporting media data from one node
|
||
|
/// to another.
|
||
|
///
|
||
|
/// A producer is essentially a GStreamer `appsink` whose output
|
||
|
/// is sent to a set of consumers, who are essentially `appsrc` wrappers
|
||
|
#[derive(Debug, Clone)]
|
||
|
pub struct StreamProducer {
|
||
|
/// The appsink to dispatch data for
|
||
|
appsink: gst_app::AppSink,
|
||
|
/// The consumers to dispatch data to
|
||
|
consumers: Arc<Mutex<StreamConsumers>>,
|
||
|
}
|
||
|
|
||
|
impl PartialEq for StreamProducer {
|
||
|
fn eq(&self, other: &Self) -> bool {
|
||
|
self.appsink.eq(&other.appsink)
|
||
|
}
|
||
|
}
|
||
|
|
||
|
impl Eq for StreamProducer {}
|
||
|
|
||
|
impl StreamProducer {
|
||
|
/// Add an appsrc to dispatch data to
|
||
|
pub fn add_consumer(&self, consumer: &gst_app::AppSrc) {
|
||
|
let mut consumers = self.consumers.lock().unwrap();
|
||
|
if consumers
|
||
|
.consumers
|
||
|
.get(&(consumer.as_ptr() as usize))
|
||
|
.is_some()
|
||
|
{
|
||
|
gst::error!(CAT, obj: &self.appsink, "Consumer already added");
|
||
|
return;
|
||
|
}
|
||
|
|
||
|
gst::debug!(CAT, obj: &self.appsink, "Adding consumer");
|
||
|
|
||
|
consumer.set_property("max-buffers", 0u64);
|
||
|
consumer.set_property("max-bytes", 0u64);
|
||
|
consumer.set_property("max-time", gst::ClockTime::from_mseconds(500));
|
||
|
consumer.set_property("format", gst::Format::Time);
|
||
|
consumer.set_property_from_str("leaky-type", "downstream");
|
||
|
consumer
|
||
|
.try_set_property("handle-segment-change", true)
|
||
|
.ok();
|
||
|
consumer.set_automatic_eos(false);
|
||
|
|
||
|
// Forward force-keyunit events upstream to the appsink
|
||
|
let srcpad = consumer.static_pad("src").unwrap();
|
||
|
let appsink = &self.appsink;
|
||
|
let fku_probe_id = srcpad
|
||
|
.add_probe(
|
||
|
gst::PadProbeType::EVENT_UPSTREAM,
|
||
|
glib::clone!(@weak appsink, @weak consumer => @default-panic, move |_pad, info| {
|
||
|
if let Some(gst::PadProbeData::Event(ref ev)) = info.data {
|
||
|
if gst_video::UpstreamForceKeyUnitEvent::parse(ev).is_ok() {
|
||
|
gst::debug!(CAT, obj: &appsink, "Requesting keyframe");
|
||
|
let _ = appsink.send_event(ev.clone());
|
||
|
}
|
||
|
}
|
||
|
|
||
|
gst::PadProbeReturn::Ok
|
||
|
}),
|
||
|
)
|
||
|
.unwrap();
|
||
|
|
||
|
consumers.consumers.insert(
|
||
|
consumer.as_ptr() as usize,
|
||
|
StreamConsumer::new(consumer, fku_probe_id),
|
||
|
);
|
||
|
}
|
||
|
|
||
|
/// Remove a consumer appsrc by id
|
||
|
pub fn remove_consumer(&self, consumer: &gst_app::AppSrc) {
|
||
|
let name = consumer.name();
|
||
|
if self
|
||
|
.consumers
|
||
|
.lock()
|
||
|
.unwrap()
|
||
|
.consumers
|
||
|
.remove(&(consumer.as_ptr() as usize))
|
||
|
.is_some()
|
||
|
{
|
||
|
gst::debug!(CAT, obj: &self.appsink, "Removed consumer {}", name);
|
||
|
} else {
|
||
|
gst::debug!(CAT, obj: &self.appsink, "Consumer {} not found", name);
|
||
|
}
|
||
|
}
|
||
|
|
||
|
/// Stop discarding data samples and start forwarding them to the consumers.
|
||
|
///
|
||
|
/// This is useful for example for prerolling live sources.
|
||
|
pub fn forward(&self) {
|
||
|
self.consumers.lock().unwrap().discard = false;
|
||
|
}
|
||
|
|
||
|
/// Get the GStreamer `appsink` wrapped by this producer
|
||
|
pub fn appsink(&self) -> &gst_app::AppSink {
|
||
|
&self.appsink
|
||
|
}
|
||
|
}
|
||
|
|
||
|
impl<'a> From<&'a gst_app::AppSink> for StreamProducer {
|
||
|
fn from(appsink: &'a gst_app::AppSink) -> Self {
|
||
|
let consumers = Arc::new(Mutex::new(StreamConsumers {
|
||
|
current_latency: None,
|
||
|
latency_updated: false,
|
||
|
consumers: HashMap::new(),
|
||
|
discard: true,
|
||
|
}));
|
||
|
|
||
|
appsink.set_callbacks(
|
||
|
gst_app::AppSinkCallbacks::builder()
|
||
|
.new_sample(glib::clone!(@strong consumers => move |appsink| {
|
||
|
let mut consumers = consumers.lock().unwrap();
|
||
|
|
||
|
let sample = match appsink.pull_sample() {
|
||
|
Ok(sample) => sample,
|
||
|
Err(_err) => {
|
||
|
gst::debug!(CAT, obj: appsink, "Failed to pull sample");
|
||
|
return Err(gst::FlowError::Flushing);
|
||
|
}
|
||
|
};
|
||
|
|
||
|
if consumers.discard {
|
||
|
return Ok(gst::FlowSuccess::Ok);
|
||
|
}
|
||
|
|
||
|
let (is_discont, is_keyframe) = if let Some(buf) = sample.buffer() {
|
||
|
let flags = buf.flags();
|
||
|
|
||
|
(flags.contains(gst::BufferFlags::DISCONT),
|
||
|
!flags.contains(gst::BufferFlags::DELTA_UNIT))
|
||
|
} else {
|
||
|
(false, true)
|
||
|
};
|
||
|
|
||
|
gst::trace!(CAT, obj: appsink, "processing sample");
|
||
|
|
||
|
let latency = consumers.current_latency;
|
||
|
let latency_updated = mem::replace(&mut consumers.latency_updated, false);
|
||
|
let mut requested_keyframe = false;
|
||
|
|
||
|
let current_consumers = consumers
|
||
|
.consumers
|
||
|
.values()
|
||
|
.filter_map(|consumer| {
|
||
|
if let Some(latency) = latency {
|
||
|
if consumer.forwarded_latency
|
||
|
.compare_exchange(
|
||
|
false,
|
||
|
true,
|
||
|
atomic::Ordering::SeqCst,
|
||
|
atomic::Ordering::SeqCst,
|
||
|
)
|
||
|
.is_ok()
|
||
|
|| latency_updated
|
||
|
{
|
||
|
consumer.appsrc.set_latency(latency, gst::ClockTime::NONE);
|
||
|
}
|
||
|
}
|
||
|
|
||
|
if is_discont {
|
||
|
// Whenever we have a discontinuity, we need a new keyframe
|
||
|
consumer.needs_keyframe.store(
|
||
|
true,
|
||
|
atomic::Ordering::SeqCst,
|
||
|
);
|
||
|
}
|
||
|
|
||
|
if !is_keyframe && consumer.needs_keyframe.load(atomic::Ordering::SeqCst)
|
||
|
{
|
||
|
// If we need a keyframe (and this one isn't) request a keyframe upstream
|
||
|
if !requested_keyframe {
|
||
|
gst::debug!(CAT, obj: appsink, "Requesting keyframe for first buffer");
|
||
|
|
||
|
appsink.send_event(
|
||
|
gst_video::UpstreamForceKeyUnitEvent::builder()
|
||
|
.all_headers(true)
|
||
|
.build(),
|
||
|
);
|
||
|
requested_keyframe = true;
|
||
|
}
|
||
|
|
||
|
gst::debug!(CAT, obj: appsink, "Ignoring frame for {} while waiting for a keyframe",
|
||
|
consumer.appsrc.name());
|
||
|
None
|
||
|
} else {
|
||
|
consumer.needs_keyframe.store(false, atomic::Ordering::SeqCst);
|
||
|
|
||
|
Some(consumer.appsrc.clone())
|
||
|
}
|
||
|
})
|
||
|
.collect::<Vec<_>>();
|
||
|
drop(consumers);
|
||
|
|
||
|
for consumer in current_consumers {
|
||
|
if let Err(err) = consumer.push_sample(&sample) {
|
||
|
gst::warning!(CAT, obj: appsink, "Failed to push sample: {}", err);
|
||
|
}
|
||
|
}
|
||
|
|
||
|
Ok(gst::FlowSuccess::Ok)
|
||
|
}))
|
||
|
.eos(glib::clone!(@strong consumers => move |_| {
|
||
|
let current_consumers = consumers
|
||
|
.lock()
|
||
|
.unwrap()
|
||
|
.consumers
|
||
|
.values()
|
||
|
.map(|c| c.appsrc.clone())
|
||
|
.collect::<Vec<_>>();
|
||
|
|
||
|
for consumer in current_consumers {
|
||
|
let _ = consumer.end_of_stream();
|
||
|
}
|
||
|
}))
|
||
|
.build(),
|
||
|
);
|
||
|
|
||
|
let sinkpad = appsink.static_pad("sink").unwrap();
|
||
|
sinkpad.add_probe(
|
||
|
gst::PadProbeType::EVENT_UPSTREAM,
|
||
|
glib::clone!(@strong consumers => move |_pad, info| {
|
||
|
if let Some(gst::PadProbeData::Event(ref ev)) = info.data {
|
||
|
if let gst::EventView::Latency(ev) = ev.view() {
|
||
|
let latency = ev.latency();
|
||
|
let mut consumers = consumers.lock().unwrap();
|
||
|
consumers.current_latency = Some(latency);
|
||
|
consumers.latency_updated = true;
|
||
|
}
|
||
|
}
|
||
|
gst::PadProbeReturn::Ok
|
||
|
}),
|
||
|
);
|
||
|
|
||
|
StreamProducer {
|
||
|
appsink: appsink.clone(),
|
||
|
consumers,
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
|
||
|
/// Wrapper around a HashMap of consumers, exists for thread safety
|
||
|
/// and also protects some of the producer state
|
||
|
#[derive(Debug)]
|
||
|
struct StreamConsumers {
|
||
|
/// The currently-observed latency
|
||
|
current_latency: Option<gst::ClockTime>,
|
||
|
/// Whether the consumers' appsrc latency needs updating
|
||
|
latency_updated: bool,
|
||
|
/// The consumers, AppSrc pointer value -> consumer
|
||
|
consumers: HashMap<usize, StreamConsumer>,
|
||
|
/// Whether appsrc samples should be forwarded to consumers yet
|
||
|
discard: bool,
|
||
|
}
|
||
|
|
||
|
/// Wrapper around a consumer's `appsrc`
|
||
|
#[derive(Debug)]
|
||
|
struct StreamConsumer {
|
||
|
/// The GStreamer `appsrc` of the consumer
|
||
|
appsrc: gst_app::AppSrc,
|
||
|
/// The id of a pad probe that intercepts force-key-unit events
|
||
|
fku_probe_id: Option<gst::PadProbeId>,
|
||
|
/// Whether an initial latency was forwarded to the `appsrc`
|
||
|
forwarded_latency: atomic::AtomicBool,
|
||
|
/// Whether a first buffer has made it through, used to determine
|
||
|
/// whether a new key unit should be requested. Only useful for encoded
|
||
|
/// streams.
|
||
|
needs_keyframe: atomic::AtomicBool,
|
||
|
}
|
||
|
|
||
|
impl StreamConsumer {
|
||
|
/// Create a new consumer
|
||
|
fn new(appsrc: &gst_app::AppSrc, fku_probe_id: gst::PadProbeId) -> Self {
|
||
|
appsrc.set_callbacks(
|
||
|
gst_app::AppSrcCallbacks::builder()
|
||
|
.enough_data(move |appsrc| {
|
||
|
gst::debug!(
|
||
|
CAT,
|
||
|
obj: appsrc,
|
||
|
"consumer is not consuming fast enough, old samples are getting dropped",
|
||
|
);
|
||
|
})
|
||
|
.build(),
|
||
|
);
|
||
|
|
||
|
StreamConsumer {
|
||
|
appsrc: appsrc.clone(),
|
||
|
fku_probe_id: Some(fku_probe_id),
|
||
|
forwarded_latency: atomic::AtomicBool::new(false),
|
||
|
needs_keyframe: atomic::AtomicBool::new(true),
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
|
||
|
impl Drop for StreamConsumer {
|
||
|
fn drop(&mut self) {
|
||
|
if let Some(fku_probe_id) = self.fku_probe_id.take() {
|
||
|
let srcpad = self.appsrc.static_pad("src").unwrap();
|
||
|
srcpad.remove_probe(fku_probe_id);
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
|
||
|
impl PartialEq for StreamConsumer {
|
||
|
fn eq(&self, other: &Self) -> bool {
|
||
|
self.appsrc.eq(&other.appsrc)
|
||
|
}
|
||
|
}
|
||
|
|
||
|
impl Eq for StreamConsumer {}
|
||
|
|
||
|
impl std::hash::Hash for StreamConsumer {
|
||
|
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
|
||
|
std::hash::Hash::hash(&self.appsrc, state);
|
||
|
}
|
||
|
}
|
||
|
|
||
|
impl std::borrow::Borrow<gst_app::AppSrc> for StreamConsumer {
|
||
|
fn borrow(&self) -> &gst_app::AppSrc {
|
||
|
&self.appsrc
|
||
|
}
|
||
|
}
|
||
|
|
||
|
#[cfg(test)]
|
||
|
mod tests {
|
||
|
use std::{
|
||
|
str::FromStr,
|
||
|
sync::{Arc, Mutex},
|
||
|
};
|
||
|
|
||
|
use futures::channel::mpsc;
|
||
|
use futures::{channel::mpsc::Receiver, SinkExt, StreamExt};
|
||
|
use gst::prelude::*;
|
||
|
|
||
|
use crate::StreamProducer;
|
||
|
|
||
|
fn create_producer() -> (
|
||
|
gst::Pipeline,
|
||
|
gst_app::AppSrc,
|
||
|
gst_app::AppSink,
|
||
|
StreamProducer,
|
||
|
) {
|
||
|
let producer_pipe =
|
||
|
gst::parse_launch("appsrc name=producer_src ! appsink name=producer_sink")
|
||
|
.unwrap()
|
||
|
.downcast::<gst::Pipeline>()
|
||
|
.unwrap();
|
||
|
let producer_sink = producer_pipe
|
||
|
.by_name("producer_sink")
|
||
|
.unwrap()
|
||
|
.downcast::<gst_app::AppSink>()
|
||
|
.unwrap();
|
||
|
|
||
|
(
|
||
|
producer_pipe.clone(),
|
||
|
producer_pipe
|
||
|
.by_name("producer_src")
|
||
|
.unwrap()
|
||
|
.downcast::<gst_app::AppSrc>()
|
||
|
.unwrap(),
|
||
|
producer_sink.clone(),
|
||
|
StreamProducer::from(&producer_sink),
|
||
|
)
|
||
|
}
|
||
|
|
||
|
struct Consumer {
|
||
|
pipeline: gst::Pipeline,
|
||
|
src: gst_app::AppSrc,
|
||
|
sink: gst_app::AppSink,
|
||
|
receiver: Mutex<Receiver<gst::Sample>>,
|
||
|
connected: Mutex<bool>,
|
||
|
}
|
||
|
|
||
|
impl Consumer {
|
||
|
fn new(id: &str) -> Self {
|
||
|
let pipeline = gst::parse_launch(&format!("appsrc name={} ! appsink name=sink", id))
|
||
|
.unwrap()
|
||
|
.downcast::<gst::Pipeline>()
|
||
|
.unwrap();
|
||
|
|
||
|
let (sender, receiver) = mpsc::channel::<gst::Sample>(1000);
|
||
|
let sender = Arc::new(Mutex::new(sender));
|
||
|
let sink = pipeline
|
||
|
.by_name("sink")
|
||
|
.unwrap()
|
||
|
.downcast::<gst_app::AppSink>()
|
||
|
.unwrap();
|
||
|
|
||
|
sink.set_callbacks(
|
||
|
gst_app::AppSinkCallbacks::builder()
|
||
|
// Add a handler to the "new-sample" signal.
|
||
|
.new_sample(move |appsink| {
|
||
|
// Pull the sample in question out of the appsink's buffer.
|
||
|
let sender_clone = sender.clone();
|
||
|
futures::executor::block_on(
|
||
|
sender_clone
|
||
|
.lock()
|
||
|
.unwrap()
|
||
|
.send(appsink.pull_sample().unwrap()),
|
||
|
)
|
||
|
.unwrap();
|
||
|
|
||
|
Ok(gst::FlowSuccess::Ok)
|
||
|
})
|
||
|
.build(),
|
||
|
);
|
||
|
|
||
|
Self {
|
||
|
pipeline: pipeline.clone(),
|
||
|
src: pipeline
|
||
|
.by_name(id)
|
||
|
.unwrap()
|
||
|
.downcast::<gst_app::AppSrc>()
|
||
|
.unwrap(),
|
||
|
sink,
|
||
|
receiver: Mutex::new(receiver),
|
||
|
connected: Mutex::new(false),
|
||
|
}
|
||
|
}
|
||
|
|
||
|
fn connect(&self, producer: &StreamProducer) {
|
||
|
{
|
||
|
let mut connected = self.connected.lock().unwrap();
|
||
|
*connected = true;
|
||
|
}
|
||
|
|
||
|
producer.add_consumer(&self.src);
|
||
|
}
|
||
|
|
||
|
fn disconnect(&self, producer: &StreamProducer) {
|
||
|
{
|
||
|
let mut connected = self.connected.lock().unwrap();
|
||
|
*connected = false;
|
||
|
}
|
||
|
|
||
|
producer.remove_consumer(&self.src);
|
||
|
}
|
||
|
}
|
||
|
|
||
|
#[test]
|
||
|
fn simple() {
|
||
|
gst::init().unwrap();
|
||
|
|
||
|
let (producer_pipe, producer_src, _producer_sink, producer) = create_producer();
|
||
|
producer.forward();
|
||
|
producer_pipe
|
||
|
.set_state(gst::State::Playing)
|
||
|
.expect("Couldn't set producer pipeline state");
|
||
|
|
||
|
let mut consumers: Vec<Consumer> = Vec::new();
|
||
|
let consumer = Consumer::new("consumer1");
|
||
|
consumer.connect(&producer);
|
||
|
consumer
|
||
|
.pipeline
|
||
|
.set_state(gst::State::Playing)
|
||
|
.expect("Couldn't set producer pipeline state");
|
||
|
consumers.push(consumer);
|
||
|
|
||
|
let consumer = Consumer::new("consumer2");
|
||
|
consumer.connect(&producer);
|
||
|
consumer
|
||
|
.pipeline
|
||
|
.set_state(gst::State::Playing)
|
||
|
.expect("Couldn't set producer pipeline state");
|
||
|
consumers.push(consumer);
|
||
|
|
||
|
for i in 0..10 {
|
||
|
let caps = gst::Caps::from_str(&format!("test,n={}", i)).unwrap();
|
||
|
producer_src.set_caps(Some(&caps));
|
||
|
producer_src.push_buffer(gst::Buffer::new()).unwrap();
|
||
|
|
||
|
for consumer in &consumers {
|
||
|
if *consumer.connected.lock().unwrap() {
|
||
|
let sample =
|
||
|
futures::executor::block_on(consumer.receiver.lock().unwrap().next())
|
||
|
.expect("Received an empty buffer?");
|
||
|
sample.buffer().expect("No buffer on the sample?");
|
||
|
assert_eq!(sample.caps(), Some(caps.as_ref()));
|
||
|
} else {
|
||
|
debug_assert!(
|
||
|
consumer
|
||
|
.sink
|
||
|
.try_pull_sample(gst::ClockTime::from_nseconds(0))
|
||
|
.is_none(),
|
||
|
"Disconnected consumer got a new sample?!"
|
||
|
);
|
||
|
}
|
||
|
}
|
||
|
|
||
|
if i == 5 {
|
||
|
consumers.get(0).unwrap().disconnect(&producer);
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
}
|