feat: 截图线程 + 编码线程。

每个显示器分别使用一个截图线程和一个编码(png+base64)线程。避免程序不明原因崩溃。
This commit is contained in:
Ivan Li 2023-03-18 23:14:59 +08:00
parent 6ea8325b15
commit f92883199c
8 changed files with 478 additions and 28 deletions

107
src-tauri/Cargo.lock generated
View File

@ -546,6 +546,19 @@ dependencies = [
"cfg-if",
]
[[package]]
name = "env_logger"
version = "0.10.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "85cdab6a89accf66733ad5a1693a4dcced6aeff64602b634530dd73c1f3ee9f0"
dependencies = [
"humantime",
"is-terminal",
"log",
"regex",
"termcolor",
]
[[package]]
name = "errno"
version = "0.2.8"
@ -1027,6 +1040,12 @@ dependencies = [
"libc",
]
[[package]]
name = "hermit-abi"
version = "0.3.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fed44880c466736ef9a5c5b5facefb5ed0785676d0c02d612db14e54f0d84286"
[[package]]
name = "html5ever"
version = "0.25.2"
@ -1058,6 +1077,12 @@ version = "0.1.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "21dec9db110f5f872ed9699c3ecf50cf16f423502706ba5c72462e28d3157573"
[[package]]
name = "humantime"
version = "2.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9a3a5bfb195931eeb336b2a7b4d761daec841b97f947d34394601737a7bba5e4"
[[package]]
name = "ico"
version = "0.2.0"
@ -1153,6 +1178,18 @@ dependencies = [
"windows-sys 0.45.0",
]
[[package]]
name = "is-terminal"
version = "0.4.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "21b6b32576413a8e69b90e952e4a026476040d81017b80445deda5f2d3921857"
dependencies = [
"hermit-abi 0.3.1",
"io-lifetimes",
"rustix",
"windows-sys 0.45.0",
]
[[package]]
name = "itoa"
version = "0.4.8"
@ -1378,6 +1415,18 @@ dependencies = [
"adler",
]
[[package]]
name = "mio"
version = "0.8.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5b9d9a46eff5b4ff64b45a9e316a6d1e0bc719ef429cbec4dc630684212bfdf9"
dependencies = [
"libc",
"log",
"wasi 0.11.0+wasi-snapshot-preview1",
"windows-sys 0.45.0",
]
[[package]]
name = "ndk"
version = "0.6.0"
@ -1464,7 +1513,7 @@ version = "1.15.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0fac9e2da13b5eb447a6ce3d392f23a29d8694bff781bf03a16cd9ac8697593b"
dependencies = [
"hermit-abi",
"hermit-abi 0.2.6",
"libc",
]
@ -1563,6 +1612,12 @@ dependencies = [
"system-deps 6.0.3",
]
[[package]]
name = "paris"
version = "1.5.13"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2eaf2319cd71dd9ff38c72bebde61b9ea657134abcf26ae4205f54f772a32810"
[[package]]
name = "parking_lot"
version = "0.12.1"
@ -2174,6 +2229,15 @@ dependencies = [
"lazy_static",
]
[[package]]
name = "signal-hook-registry"
version = "1.4.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d8229b473baa5980ac72ef434c4415e70c4b5e71b423043adb4ba059f89c99a1"
dependencies = [
"libc",
]
[[package]]
name = "siphasher"
version = "0.3.10"
@ -2195,6 +2259,16 @@ version = "1.10.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a507befe795404456341dfab10cef66ead4c041f62b8b11bbb92bffe5d0953e0"
[[package]]
name = "socket2"
version = "0.4.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "64a4a911eed85daf18834cfaa86a79b7d266ff93ff5ba14005426219480ed662"
dependencies = [
"libc",
"winapi",
]
[[package]]
name = "soup2"
version = "0.2.1"
@ -2559,6 +2633,15 @@ dependencies = [
"utf-8",
]
[[package]]
name = "termcolor"
version = "1.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "be55cf8942feac5c765c2c993422806843c9a9a45d4d5c407ad6dd2ea95eb9b6"
dependencies = [
"winapi-util",
]
[[package]]
name = "test-demo"
version = "0.0.0"
@ -2567,12 +2650,15 @@ dependencies = [
"base64 0.21.0",
"core-graphics",
"display-info",
"env_logger",
"log",
"paris",
"png",
"serde",
"serde_json",
"tauri",
"tauri-build",
"tracing",
"tokio",
"webp",
]
@ -2662,12 +2748,29 @@ checksum = "03201d01c3c27a29c8a5cee5b55a93ddae1ccf6f08f65365c2c918f8c1b76f64"
dependencies = [
"autocfg",
"bytes",
"libc",
"memchr",
"mio",
"num_cpus",
"parking_lot",
"pin-project-lite",
"signal-hook-registry",
"socket2",
"tokio-macros",
"windows-sys 0.45.0",
]
[[package]]
name = "tokio-macros"
version = "1.8.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d266c00fde287f55d3f1c3e96c500c362a2b8c695076ec180f27918820bc6df8"
dependencies = [
"proc-macro2",
"quote",
"syn",
]
[[package]]
name = "toml"
version = "0.5.11"

View File

@ -22,7 +22,10 @@ core-graphics = "0.22.3"
display-info = "0.4.1"
png = "0.17.7"
anyhow = "1.0.69"
tracing = "0.1.37"
tokio = {version = "1.26.0", features = ["full"] }
paris = { version = "1.5", features = ["timestamps", "macros"] }
log = "0.4.17"
env_logger = "0.10.0"
[features]
# this feature is used for production builds or when `devPath` points to the filesystem

View File

@ -1,16 +1,18 @@
// Prevents additional console window on WiOk(ndows in release, DO NOT REMOVE!!
#![cfg_attr(not(debug_assertions), windows_subsystem = "windows")]
mod take_screenshot_loop;
pub mod screenshot;
mod screenshot_manager;
use base64::Engine;
use core_graphics::display::{
kCGNullWindowID, kCGWindowImageDefault, kCGWindowListOptionOnScreenOnly, CGDisplay,
};
use display_info::DisplayInfo;
use paris::error;
use screenshot_manager::ScreenshotManager;
use serde::{Deserialize, Serialize};
use serde_json::to_string;
use tracing::error;
#[derive(Serialize, Deserialize)]
#[serde(remote = "DisplayInfo")]
@ -63,7 +65,7 @@ fn take_screenshot(display_id: u32, scale_factor: f32) -> Result<String, String>
kCGWindowImageDefault,
)
.ok_or_else(|| anyhow::anyhow!("Display#{}: take screenshot failed", display_id))?;
println!("take screenshot took {}ms", start_at.elapsed().as_millis());
// println!("take screenshot took {}ms", start_at.elapsed().as_millis());
let buffer = cg_image.data();
let bytes_per_row = cg_image.bytes_per_row() as f32;
@ -74,10 +76,10 @@ fn take_screenshot(display_id: u32, scale_factor: f32) -> Result<String, String>
let image_height = (height as f32 / scale_factor) as u32;
let image_width = (width as f32 / scale_factor) as u32;
println!(
"raw image: {}x{}, output image: {}x{}",
width, height, image_width, image_height
);
// println!(
// "raw image: {}x{}, output image: {}x{}",
// width, height, image_width, image_height
// );
// // from bitmap vec
let mut image_buffer = vec![0u8; (image_width * image_height * 3) as usize];
@ -88,7 +90,6 @@ fn take_screenshot(display_id: u32, scale_factor: f32) -> Result<String, String>
let b = buffer[offset];
let g = buffer[offset + 1];
let r = buffer[offset + 2];
let a = buffer[offset + 3];
let offset = (y * image_width + x) as usize;
image_buffer[offset * 3] = r;
image_buffer[offset * 3 + 1] = g;
@ -136,12 +137,33 @@ fn take_screenshot(display_id: u32, scale_factor: f32) -> Result<String, String>
})
}
fn main() {
#[tauri::command]
async fn subscribe_encoded_screenshot_updated(
window: tauri::Window,
display_id: u32,
) -> Result<(), String> {
let screenshot_manager = ScreenshotManager::global().await;
screenshot_manager
.subscribe_encoded_screenshot_updated(window, display_id)
.await
.map_err(|err| {
error!("subscribe_encoded_screenshot_updated: {}", err);
err.to_string()
})
}
#[tokio::main]
async fn main() {
env_logger::init();
let screenshot_manager = ScreenshotManager::global().await;
screenshot_manager.start().unwrap();
tauri::Builder::default()
.invoke_handler(tauri::generate_handler![
greet,
take_screenshot,
list_display_info
list_display_info,
subscribe_encoded_screenshot_updated
])
.run(tauri::generate_context!())
.expect("error while running tauri application");

View File

@ -0,0 +1,42 @@
use std::sync::Arc;
use serde::Serialize;
use tauri::async_runtime::RwLock;
#[derive(Debug, Clone)]
pub struct Screenshot {
pub display_id: u32,
pub height: u32,
pub width: u32,
pub bytes_per_row: usize,
pub bytes: Arc<RwLock<Vec<u8>>>,
pub scale_factor: f32,
}
impl Screenshot {
pub fn new(
display_id: u32,
height: u32,
width: u32,
bytes_per_row: usize,
bytes: Vec<u8>,
scale_factor: f32,
) -> Self {
Self {
display_id,
height,
width,
bytes_per_row,
bytes: Arc::new(RwLock::new(bytes)),
scale_factor,
}
}
}
#[derive(Debug, Clone, Serialize)]
pub struct ScreenshotPayload {
pub display_id: u32,
pub height: u32,
pub width: u32,
pub base64_image: String,
}

View File

@ -0,0 +1,252 @@
use std::{collections::HashMap, sync::Arc};
use base64::Engine;
use core_graphics::display::{
kCGNullWindowID, kCGWindowImageDefault, kCGWindowListOptionOnScreenOnly, CGDisplay,
};
use paris::{error, info, warn};
use tauri::{async_runtime::RwLock, Window};
use tokio::sync::{watch, OnceCell};
use crate::screenshot::{Screenshot, ScreenshotPayload};
pub fn take_screenshot(display_id: u32, scale_factor: f32) -> anyhow::Result<Screenshot> {
log::debug!("take_screenshot");
// let start_at = std::time::Instant::now();
let cg_display = CGDisplay::new(display_id);
let cg_image = CGDisplay::screenshot(
cg_display.bounds(),
kCGWindowListOptionOnScreenOnly,
kCGNullWindowID,
kCGWindowImageDefault,
)
.ok_or_else(|| anyhow::anyhow!("Display#{}: take screenshot failed", display_id))?;
// println!("take screenshot took {}ms", start_at.elapsed().as_millis());
let buffer = cg_image.data();
let bytes_per_row = cg_image.bytes_per_row();
let height = cg_image.height();
let width = cg_image.width();
let mut bytes = vec![0u8; buffer.len() as usize];
bytes.copy_from_slice(&buffer);
Ok(Screenshot::new(
display_id,
height as u32,
width as u32,
bytes_per_row,
bytes,
scale_factor,
))
}
pub struct ScreenshotManager {
channels: Arc<RwLock<HashMap<u32, watch::Receiver<Screenshot>>>>,
encode_listeners: Arc<RwLock<HashMap<u32, Vec<Window>>>>,
}
impl ScreenshotManager {
pub async fn global() -> &'static Self {
static SCREENSHOT_MANAGER: OnceCell<ScreenshotManager> = OnceCell::const_new();
SCREENSHOT_MANAGER
.get_or_init(|| async {
let channels = Arc::new(RwLock::new(HashMap::new()));
let encode_listeners = Arc::new(RwLock::new(HashMap::new()));
Self {
channels,
encode_listeners,
}
})
.await
}
pub fn start(&self) -> anyhow::Result<()> {
let displays = display_info::DisplayInfo::all()?;
for display in displays {
self.start_one(display.id, display.scale_factor)?;
}
Ok(())
}
fn start_one(&self, display_id: u32, scale_factor: f32) -> anyhow::Result<()> {
let channels = self.channels.to_owned();
tokio::spawn(async move {
let screenshot = take_screenshot(display_id, scale_factor);
if screenshot.is_err() {
warn!("take_screenshot_loop: {}", screenshot.err().unwrap());
return;
}
let screenshot = screenshot.unwrap();
let (tx, rx) = watch::channel(screenshot);
{
let mut channels = channels.write().await;
channels.insert(display_id, rx);
}
loop {
Self::take_screenshot_loop(display_id, scale_factor, &tx).await;
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
}
});
Ok(())
}
pub async fn subscribe_encoded_screenshot_updated(
&self,
window: Window,
display_id: u32,
) -> anyhow::Result<()> {
let channels = self.channels.to_owned();
let encode_listeners = self.encode_listeners.to_owned();
log::info!("subscribe_encoded_screenshot_updated. {}", display_id);
{
let encode_listeners = encode_listeners.read().await;
let listening_windows = encode_listeners.get(&display_id);
if listening_windows.is_some() && listening_windows.unwrap().contains(&window) {
log::debug!("subscribe_encoded_screenshot_updated: already listening. display#{}, window#{}", display_id, window.label());
return Ok(());
}
}
{
encode_listeners
.write()
.await
.entry(display_id)
.or_default()
.push(window);
}
tokio::spawn(async move {
info!("subscribe_encoded_screenshot_updated: start");
let channels = channels.read().await;
let rx = channels.get(&display_id);
if rx.is_none() {
error!(
"subscribe_encoded_screenshot_updated: can not find display_id {}",
display_id
);
return;
}
let mut rx = rx.unwrap().clone();
loop {
if let Err(err) = rx.changed().await {
error!(
"subscribe_encoded_screenshot_updated: can not wait rx {}",
err
);
break;
}
let encode_listeners = encode_listeners.read().await;
let windows = encode_listeners.get(&display_id);
if windows.is_none() || windows.unwrap().is_empty() {
info!("subscribe_encoded_screenshot_updated: no listener, stop");
break;
}
let screenshot = rx.borrow().clone();
let base64_image = Self::encode_screenshot_to_base64(&screenshot).await;
let height = screenshot.height;
let width = screenshot.width;
if base64_image.is_err() {
error!(
"subscribe_encoded_screenshot_updated: encode_screenshot_to_base64 error {}",
base64_image.err().unwrap()
);
continue;
}
let base64_image = base64_image.unwrap();
for window in windows.unwrap().into_iter() {
let base64_image = base64_image.clone();
let payload = ScreenshotPayload {
display_id,
base64_image,
height,
width,
};
if let Err(err) = window.emit("encoded-screenshot-updated", payload) {
error!("subscribe_encoded_screenshot_updated: emit error {}", err)
} else {
info!(
"subscribe_encoded_screenshot_updated: emit success. display#{}",
display_id
)
}
}
}
});
Ok(())
}
async fn unsubscribe_encoded_screenshot_updated(&self, display_id: u32) -> anyhow::Result<()> {
let channels = self.channels.to_owned();
let mut channels = channels.write().await;
channels.remove(&display_id);
Ok(())
}
async fn take_screenshot_loop(
display_id: u32,
scale_factor: f32,
tx: &watch::Sender<Screenshot>,
) {
let screenshot = take_screenshot(display_id, scale_factor);
if let Ok(screenshot) = screenshot {
tx.send(screenshot).unwrap();
} else {
warn!("take_screenshot_loop: {}", screenshot.err().unwrap());
}
}
async fn encode_screenshot_to_base64(screenshot: &Screenshot) -> anyhow::Result<String> {
let bytes = screenshot.bytes.read().await;
let scale_factor = screenshot.scale_factor;
let image_height = (screenshot.height as f32 / scale_factor) as u32;
let image_width = (screenshot.width as f32 / scale_factor) as u32;
let mut image_buffer = vec![0u8; (image_width * image_height * 3) as usize];
for y in 0..image_height {
for x in 0..image_width {
let offset = (((y as f32) * screenshot.bytes_per_row as f32 + (x as f32) * 4.0)
* scale_factor) as usize;
let b = bytes[offset];
let g = bytes[offset + 1];
let r = bytes[offset + 2];
let offset = (y * image_width + x) as usize;
image_buffer[offset * 3] = r;
image_buffer[offset * 3 + 1] = g;
image_buffer[offset * 3 + 2] = b;
}
}
let mut image_png = Vec::new();
let mut encoder = png::Encoder::new(&mut image_png, image_width, image_height);
encoder.set_color(png::ColorType::Rgb);
encoder.set_depth(png::BitDepth::Eight);
let mut writer = encoder
.write_header()
.map_err(|e| anyhow::anyhow!("png: {}", anyhow::anyhow!(e.to_string())))?;
writer
.write_image_data(&image_buffer)
.map_err(|e| anyhow::anyhow!("png: {}", anyhow::anyhow!(e.to_string())))?;
writer
.finish()
.map_err(|e| anyhow::anyhow!("png: {}", anyhow::anyhow!(e.to_string())))?;
let mut base64_image = String::new();
base64_image.push_str("data:image/webp;base64,");
let encoded = base64::engine::general_purpose::STANDARD_NO_PAD.encode(&*image_png);
base64_image.push_str(encoded.as_str());
Ok(base64_image)
}
}

View File

@ -3,6 +3,7 @@ import { invoke } from "@tauri-apps/api/tauri";
import "./App.css";
import { DisplayInfo } from './models/display-info.model';
import { DisplayInfoPanel } from './components/display-info-panel';
import { ScreenView } from './components/screen-view';
function App() {
const [screenshots, setScreenshots] = createSignal<string[]>([]);
@ -14,9 +15,9 @@ function App() {
});
});
createEffect(() => {
take_all_display_screenshot();
}, [displays]);
// createEffect(() => {
// take_all_display_screenshot();
// }, [displays]);
async function take_all_display_screenshot() {
console.log('take_all_display_screenshot');
@ -37,18 +38,7 @@ function App() {
return (
<li>
<DisplayInfoPanel display={display} />
</li>
);
})}
</ol>
<ol>
{screenshots().map((screenshot) => {
return (
<li>
<img
style="object-fit: contain; height: 400px; width: 600px"
src={screenshot}
/>
<ScreenView displayId={display.id} />
</li>
);
})}

View File

@ -0,0 +1,38 @@
import { invoke } from '@tauri-apps/api';
import { listen } from '@tauri-apps/api/event';
import { Component, createEffect, createSignal, onCleanup } from 'solid-js';
type ScreenViewProps = {
displayId: number;
};
async function subscribeScreenshotUpdate(displayId: number) {
await invoke('subscribe_encoded_screenshot_updated', {
displayId,
});
}
export const ScreenView: Component<ScreenViewProps> = (props) => {
const [image, setImage] = createSignal<string>();
createEffect(() => {
const unlisten = listen<{ base64_image: string; display_id: number }>(
'encoded-screenshot-updated',
(event) => {
if (event.payload.display_id === props.displayId) {
setImage(event.payload.base64_image);
}
console.log(event.payload.display_id, props.displayId);
},
);
subscribeScreenshotUpdate(props.displayId);
onCleanup(() => {
unlisten.then((unlisten) => {
unlisten();
});
});
});
return <img src={image()} class="object-contain" />;
};