This commit is contained in:
Lily Tsuru 2024-09-11 21:48:09 -04:00
parent 2da9b38974
commit f388eb32b2
4 changed files with 35 additions and 23 deletions

View file

@ -1,6 +1,6 @@
mod video;
mod surface;
mod types;
mod video;
mod vnc_engine;
use video::encoder_thread;
@ -15,6 +15,7 @@ use rand::distributions::DistString;
use tokio::sync::{
broadcast,
mpsc::{self, error::TryRecvError},
Mutex as TokioMutex,
};
use vnc_engine::VncMessageOutput;
@ -37,7 +38,7 @@ struct AppState {
/// Channel for sending things to the VNC engine
/// should later be used for control
engine_tx: mpsc::Sender<vnc_engine::VncMessageInput>,
encoder_tx: mpsc::Sender<encoder_thread::EncodeThreadInput>,
encoder_tx: Arc<TokioMutex<mpsc::Sender<encoder_thread::EncodeThreadInput>>>,
websocket_broadcast_tx: broadcast::Sender<ws::Message>,
}
@ -50,7 +51,7 @@ impl AppState {
let (chat_tx, _chat_rx) = broadcast::channel(10);
Self {
engine_tx: engine_tx,
encoder_tx: encoder_tx,
encoder_tx: Arc::new(TokioMutex::new(encoder_tx)),
websocket_broadcast_tx: chat_tx,
}
}
@ -101,6 +102,8 @@ async fn main() -> anyhow::Result<()> {
// let's encode a frame
// First we copy the current VNC framebuffer to the shared
// frame between the encoder thread and ffmpeg
// TODO: Do this on the encoder thread
{
let mut frame_locked = frame
.lock()
@ -130,6 +133,8 @@ async fn main() -> anyhow::Result<()> {
}
let _ = encoder_tx_clone
.lock()
.await
.send(crate::encoder_thread::EncodeThreadInput::SendFrame)
.await;
}
@ -146,6 +151,8 @@ async fn main() -> anyhow::Result<()> {
}
let _ = encoder_tx_clone
.lock()
.await
.send(crate::encoder_thread::EncodeThreadInput::Init { size: size })
.await;
}
@ -227,16 +234,18 @@ async fn ws_handler(
async fn handle_socket(socket: WebSocket, who: SocketAddr, state: Arc<AppState>) {
let (mut sender, mut receiver) = socket.split();
// Force a ws connection to mean a keyframe
let _ = state
.encoder_tx
.send(crate::encoder_thread::EncodeThreadInput::ForceKeyframe)
.await;
{
let locked = state.encoder_tx.lock().await;
let _ = state
.encoder_tx
.send(crate::encoder_thread::EncodeThreadInput::SendFrame)
.await;
// Force a ws connection to mean a keyframe
let _ = locked
.send(crate::encoder_thread::EncodeThreadInput::ForceKeyframe)
.await;
let _ = locked
.send(crate::encoder_thread::EncodeThreadInput::SendFrame)
.await;
}
// random username
let username: Arc<String> =

View file

@ -40,6 +40,7 @@ fn encoder_thread_main(
) {
let mut packet = ffmpeg::Packet::empty();
let mut encoder: Option<H264Encoder> = None;
let mut sws = None;
@ -81,6 +82,7 @@ fn encoder_thread_main(
}
EncodeThreadInput::ForceKeyframe => {
println!("got force keyframe request");
force_keyframe = true;
}

View file

@ -30,14 +30,7 @@ fn create_context_and_set_common_parameters(
let mut video_encoder_context = create_context_from_codec(encoder)?.encoder().video()?;
// TODO: Either no GOP, or a fairly large one.
// idk
let gop = /*if max_framerate / 2 != 0 {
max_framerate / 2
} else {
max_framerate
} */
i32::MAX as u32;
video_encoder_context.set_width(size.width);
video_encoder_context.set_height(size.height);
@ -53,7 +46,11 @@ fn create_context_and_set_common_parameters(
video_encoder_context.set_time_base(ffmpeg::Rational(1, max_framerate as i32).invert());
video_encoder_context.set_format(ffmpeg::format::Pixel::YUV420P);
video_encoder_context.set_gop(gop);
// We manually control the GOP of the stream so we only need one stream
// that can be broadcast to many users.
//
// TODO: Either no GOP, or a fairly large one.
video_encoder_context.set_gop(i32::MAX as u32);
video_encoder_context.set_max_b_frames(0);
unsafe {
@ -104,6 +101,8 @@ impl H264Encoder {
dict.set("crf", "43");
dict.set("crf_max", "48");
dict.set("forced-idr", "1");
let encoder = video_encoder_context
.open_as_with(encoder, dict)
.with_context(|| "While opening x264 video codec")?;

View file

@ -1,6 +1,8 @@
pub mod h264_encoder;
pub mod encoder_thread;
pub mod h264_encoder;
//pub mod lc_muxer;
/// Re-export of `ffmpeg_the_third` crate in an infinitely less obtuse name.
pub use ffmpeg_the_third as ffmpeg;
pub use encoder_thread::*;
pub use encoder_thread::*;