modeco80/vncstream
1
0
Fork 0
Code Issues Pull requests Actions Packages Projects Releases Wiki Activity

decouple encoder thread from VNC client

Browse source
This commit is contained in:
Lily Tsuru 2024-09-09 00:08:08 -04:00
parent c8f9c9d9ea
commit 86490ed80b
3 changed files with 175 additions and 186 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

40
server/src/encoder_thread.rs
View file

@ -7,7 +7,6 @@ use tokio::sync::mpsc::{self, error::TryRecvError};
pub enum EncodeThreadInput {
Init { size: crate::types::Size },
ForceKeyframe,
Shutdown,
SendFrame,
}
@ -20,7 +19,6 @@ pub enum EncodeThreadOutput {
fn set_frame_flags(frame: &mut ffmpeg_the_third::Frame, force_keyframe: bool) {
unsafe {
if force_keyframe {
//println!("frame {frame_number} will be a keyframe");
(*frame.as_mut_ptr()).pict_type =
ffmpeg_the_third::sys::AVPictureType::AV_PICTURE_TYPE_I;
(*frame.as_mut_ptr()).flags = ffmpeg_the_third::sys::AV_FRAME_FLAG_KEY;
@ -49,8 +47,6 @@ fn encoder_thread_main(
let mut frame_number = 0usize;
let mut force_keyframe = false;
println!("encoder thread spawned");
loop {
match rx.try_recv() {
Ok(msg) => match msg {
@ -82,26 +78,6 @@ fn encoder_thread_main(
);
}
EncodeThreadInput::Shutdown => {
if encoder.is_some() {
let enc = encoder.as_mut().unwrap();
enc.send_eof();
enc.receive_packet(&mut packet)
.expect("failed to recv eof packet");
unsafe {
if !packet.is_empty() {
let _ = tx.blocking_send(EncodeThreadOutput::Frame {
packet: packet.clone(),
});
}
}
}
break;
}
EncodeThreadInput::ForceKeyframe => {
force_keyframe = true;
}
@ -157,6 +133,22 @@ fn encoder_thread_main(
}
}
}
if encoder.is_some() {
let enc = encoder.as_mut().unwrap();
enc.send_eof();
enc.receive_packet(&mut packet)
.expect("failed to recv eof packet");
unsafe {
if !packet.is_empty() {
let _ = tx.blocking_send(EncodeThreadOutput::Frame {
packet: packet.clone(),
});
}
}
}
}
pub fn encoder_thread_spawn(

215
server/src/main.rs
View file

@ -4,10 +4,16 @@ mod surface;
mod types;
mod vnc_engine;
use std::sync::{Arc, Mutex};
use std::{
sync::{Arc, Mutex},
time::Duration,
};
use rand::distributions::DistString;
use tokio::sync::{broadcast, mpsc};
use tokio::sync::{
broadcast,
mpsc::{self, error::TryRecvError},
};
use vnc_engine::VncMessageOutput;
use std::net::SocketAddr;
@ -15,7 +21,7 @@ use std::net::SocketAddr;
use axum::{
extract::{
connect_info::ConnectInfo,
ws::{Message, WebSocket, WebSocketUpgrade},
ws::{self, Message, WebSocket, WebSocketUpgrade},
State,
},
response::IntoResponse,
@ -25,39 +31,149 @@ use axum::{
use futures::{sink::SinkExt, stream::StreamExt};
#[derive(Clone, Debug)]
enum WsMessage {
Text(String),
Buffer(Vec<u8>),
}
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>,
websocket_broadcast_tx: broadcast::Sender<WsMessage>,
websocket_broadcast_tx: broadcast::Sender<ws::Message>,
}
impl AppState {
fn new(engine_tx: mpsc::Sender<vnc_engine::VncMessageInput>) -> Self {
fn new(
engine_tx: mpsc::Sender<vnc_engine::VncMessageInput>,
encoder_tx: mpsc::Sender<encoder_thread::EncodeThreadInput>,
) -> Self {
let (chat_tx, _chat_rx) = broadcast::channel(10);
Self {
engine_tx: engine_tx,
encoder_tx: encoder_tx,
websocket_broadcast_tx: chat_tx,
}
}
}
#[tokio::main(flavor = "multi_thread", worker_threads = 4)]
#[tokio::main(flavor = "multi_thread", worker_threads = 8)]
async fn main() -> anyhow::Result<()> {
let surface = Arc::new(Mutex::new(surface::Surface::new()));
let (engine_output_tx, mut engine_output_rx) = mpsc::channel(32);
let (engine_input_tx, engine_input_rx) = mpsc::channel(16);
let state = Arc::new(AppState::new(engine_input_tx));
// H.264 encoder related
let frame: Arc<Mutex<Option<ffmpeg_the_third::frame::Video>>> = Arc::new(Mutex::new(None));
let (mut encoder_rx, encoder_tx) = encoder_thread::encoder_thread_spawn(&frame);
let state = Arc::new(AppState::new(engine_input_tx, encoder_tx));
// VNC related
let mut vnc_client =
vnc_engine::Client::new(engine_output_tx, engine_input_rx, surface.clone());
tokio::spawn(async move {
let addr = vnc_engine::Address::Tcp("10.16.0.1:5930".parse().expect("its over"));
//let addr = vnc_engine::Address::Tcp("127.0.0.1:6930".parse().expect("its over"));
vnc_client.connect_and_run(addr).await
});
let state_clone = state.clone();
let vnc_recv_handle = tokio::spawn(async move {
let frame_clone = frame.clone();
let encoder_tx_clone = state_clone.encoder_tx.clone();
let surface_clone = surface.clone();
loop {
match engine_output_rx.try_recv() {
Ok(msg) => match msg {
VncMessageOutput::Connect => {
println!("connected")
}
VncMessageOutput::Disconnect => {
println!("disconnect")
}
VncMessageOutput::FramebufferUpdate => {
// let's encode a frame
// First we copy the current VNC framebuffer to the shared
// frame between the encoder thread and ffmpeg
{
let mut frame_locked = frame
.lock()
.expect("Couldn't lock frame on our end. Did the VNC end panic?");
let mut_frame = frame_locked.as_mut().expect("it's None? why?");
let width = mut_frame.width();
let height = mut_frame.height();
let mut surf = surface_clone.lock().expect(
"locking the surface to paint it to the ffmpeg frame failed",
);
let surf_buf = surf.get_buffer();
let buf_ptr =
unsafe { (*(*mut_frame.as_mut_ptr()).buf[0]).data as *mut u32 };
for y in 0..height {
for x in 0..width {
unsafe {
let ofs = (y * width + x) as usize;
*buf_ptr.add(ofs) = surf_buf[ofs];
}
}
}
}
let _ = encoder_tx_clone
.send(crate::encoder_thread::EncodeThreadInput::SendFrame)
.await;
}
VncMessageOutput::FramebufferResized(size) => {
// make a new frame for the encoder
{
let mut lk_frame = frame_clone.lock().expect("Couldn't lock frame");
*lk_frame = Some(ffmpeg_the_third::frame::Video::new(
ffmpeg_the_third::format::Pixel::BGRA,
size.clone().width,
size.clone().height,
));
}
let _ = encoder_tx_clone
.send(crate::encoder_thread::EncodeThreadInput::Init { size: size })
.await;
}
},
Err(TryRecvError::Disconnected) => break,
Err(TryRecvError::Empty) => {}
}
match encoder_rx.try_recv() {
Ok(msg) => match msg {
encoder_thread::EncodeThreadOutput::Frame { mut packet } => {
let vec = {
let data = packet.data_mut().expect("packet is empty somehow");
data.to_vec()
};
let _ = state_clone
.websocket_broadcast_tx
.send(ws::Message::Binary(vec));
}
},
Err(TryRecvError::Empty) => {}
_ => break,
}
tokio::time::sleep(Duration::from_millis(2)).await;
}
});
// Axum websocket server
let app: Router<()> = Router::new()
.route(
"/",
@ -71,48 +187,26 @@ async fn main() -> anyhow::Result<()> {
)
.with_state(state.clone());
let mut engine = vnc_engine::Client::new(engine_output_tx, engine_input_rx, surface);
let tcp_listener = tokio::net::TcpListener::bind("0.0.0.0:4940")
.await
.expect("failed to listen");
let vnc_engine_handle = tokio::spawn(async move {
let addr = vnc_engine::Address::Tcp("10.16.0.1:5930".parse().expect("its over"));
//let addr = vnc_engine::Address::Tcp("127.0.0.1:6930".parse().expect("its over"));
engine.connect_and_run(addr).await
});
let state_clone = Arc::clone(&state);
tokio::spawn(async move {
while let Some(msg) = engine_output_rx.recv().await {
match msg {
VncMessageOutput::Connect => {
println!("connected")
}
VncMessageOutput::Disconnect => {
println!("disconnect")
}
VncMessageOutput::FramebufferUpdate(vec) => {
let _ = state_clone
.websocket_broadcast_tx
.send(WsMessage::Buffer(vec));
}
_ => {}
}
}
});
let (res1, res2) = tokio::join!(
axum::serve(
tcp_listener,
app.into_make_service_with_connect_info::<SocketAddr>(),
),
vnc_engine_handle
let axum_future = axum::serve(
tcp_listener,
app.into_make_service_with_connect_info::<SocketAddr>(),
);
// If the VNC client disconnects we should exit.
tokio::select! {
_ = axum_future => {
println!("axum died");
}
_ = vnc_recv_handle => {
println!("VNC client disconnected, exiting");
}
}
Ok(())
}
@ -133,8 +227,13 @@ async fn handle_socket(socket: WebSocket, who: SocketAddr, state: Arc<AppState>)
// Force a ws connection to mean a keyframe
let _ = state
.engine_tx
.send(vnc_engine::VncMessageInput::ForceKeyframe)
.encoder_tx
.send(crate::encoder_thread::EncodeThreadInput::ForceKeyframe)
.await;
let _ = state
.encoder_tx
.send(crate::encoder_thread::EncodeThreadInput::SendFrame)
.await;
// random username
@ -148,14 +247,8 @@ async fn handle_socket(socket: WebSocket, who: SocketAddr, state: Arc<AppState>)
let mut sub = send_clone.websocket_broadcast_tx.subscribe();
while let Ok(msg) = sub.recv().await {
match msg {
WsMessage::Text(c) => {
let _ = sender.send(Message::Text(c)).await;
}
WsMessage::Buffer(buffer) => {
let _ = sender.send(Message::Binary(buffer)).await;
}
if sender.send(msg).await.is_err() {
break;
}
}
});
@ -186,8 +279,8 @@ async fn handle_socket(socket: WebSocket, who: SocketAddr, state: Arc<AppState>)
state
.websocket_broadcast_tx
.send(WsMessage::Text(
serde_json::to_string(&send).expect("penis"),
.send(ws::Message::Text(
serde_json::to_string(&send).expect("oh well"),
))
.expect("boom");

106
server/src/vnc_engine.rs
View file

@ -1,7 +1,5 @@
//! Native-side VNC client. This is usually run in another OS thread.
use crate::encoder_thread::{encoder_thread_spawn, EncodeThreadOutput};
use super::surface::Surface;
use super::types::*;
@ -29,24 +27,14 @@ pub enum VncMessageOutput {
Connect,
Disconnect,
// this will contain a single annex B packet
FramebufferUpdate(Vec<u8>),
FramebufferUpdate,
FramebufferResized(Size),
}
#[derive(Debug)]
pub enum VncMessageInput {
KeyEvent {
keysym: u32,
pressed: bool,
},
MouseEvent {
pt: Point,
buttons: u8,
},
Disconnect,
/// Forces a keyframe to occur in the video stream.
ForceKeyframe,
KeyEvent { keysym: u32, pressed: bool },
MouseEvent { pt: Point, buttons: u8 },
}
pub struct Client {
@ -110,10 +98,6 @@ impl Client {
self.out_tx.send(VncMessageOutput::Connect).await?;
// h.264 encoder related
let frame: Arc<Mutex<Option<ffmpeg_the_third::frame::Video>>> = Arc::new(Mutex::new(None));
let (mut encoder_rx, encoder_tx) = encoder_thread_spawn(&frame);
loop {
// Pull a event and act on it. If none are there, it's fine and we can just move on to
// advancing the vnc client, but if the channel is closed, that means we are to disconnect
@ -137,16 +121,6 @@ impl Client {
}))
.await?;
}
VncMessageInput::Disconnect => break,
VncMessageInput::ForceKeyframe => {
encoder_tx
.send(crate::encoder_thread::EncodeThreadInput::ForceKeyframe)
.await?;
encoder_tx
.send(crate::encoder_thread::EncodeThreadInput::SendFrame)
.await?;
}
},
Err(TryRecvError::Empty) => {}
@ -176,22 +150,6 @@ impl Client {
height: res.height as u32,
};
// make a new frame for the encoder
{
let mut lk_frame = frame.lock().expect("oh FUCK");
*lk_frame = Some(ffmpeg_the_third::frame::Video::new(
ffmpeg_the_third::format::Pixel::BGRA,
cvm_size.clone().width,
cvm_size.clone().height,
));
}
encoder_tx
.send(crate::encoder_thread::EncodeThreadInput::Init {
size: cvm_size.clone(),
})
.await?;
self.out_tx
.send(VncMessageOutput::FramebufferResized(cvm_size))
.await?;
@ -230,44 +188,13 @@ impl Client {
if !self.rects_in_frame.is_empty() {
// We don't care what rects ARE there, but
// if none were pressent then we probably need not bother
//println!("vnc engine frame");
// let's encode a frame
// First we copy the current VNC framebuffer to the shared
// frame between the encoder thread and ffmpeg
{
let mut frame_locked =
frame.lock().expect("Couldn't lock the damn frame. FUCK");
let mut_frame = frame_locked.as_mut().expect("NOOOO");
let width = mut_frame.width();
let height = mut_frame.height();
let mut surf = self.surf.lock().expect("frame lock fail");
let surf_buf = surf.get_buffer();
let buf_ptr =
unsafe { (*(*mut_frame.as_mut_ptr()).buf[0]).data as *mut u32 };
for y in 0..height {
for x in 0..width {
unsafe {
let ofs = (y * width + x) as usize;
*buf_ptr.add(ofs) = surf_buf[ofs];
}
}
}
}
encoder_tx
.send(crate::encoder_thread::EncodeThreadInput::SendFrame)
self.out_tx
.send(VncMessageOutput::FramebufferUpdate)
.await?;
self.rects_in_frame.clear();
}
// send frame to encoder thread, pull encoded data back
}
// TODO: we might want to pass this to js at some point
@ -276,29 +203,6 @@ impl Client {
}
}
match encoder_rx.try_recv() {
Ok(msg) => {
match msg {
EncodeThreadOutput::Frame { mut packet } => {
let vec = {
//f !packet.is_empty()
let data = packet.data_mut().expect("packet is not empty");
data.to_vec()
};
self.out_tx
.send(VncMessageOutput::FramebufferUpdate(vec))
.await?;
}
}
}
Err(TryRecvError::Empty) => {}
_ => break,
}
// Sleep to give CPU time
tokio::time::sleep(Duration::from_millis(2)).await;
}