From f92883199c60f011550891f1c11e915814eb5e94 Mon Sep 17 00:00:00 2001 From: Ivan Li Date: Sat, 18 Mar 2023 23:14:59 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=88=AA=E5=9B=BE=E7=BA=BF=E7=A8=8B=20?= =?UTF-8?q?+=20=E7=BC=96=E7=A0=81=E7=BA=BF=E7=A8=8B=E3=80=82?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 每个显示器分别使用一个截图线程和一个编码（png+base64）线程。避免程序不明原因崩溃。 --- src-tauri/Cargo.lock | 107 +++++++- src-tauri/Cargo.toml | 5 +- .../{take_screenshot_loop.rs => logger.rs} | 0 src-tauri/src/main.rs | 42 ++- src-tauri/src/screenshot.rs | 42 +++ src-tauri/src/screenshot_manager.rs | 252 ++++++++++++++++++ src/App.tsx | 20 +- src/components/screen-view.tsx | 38 +++ 8 files changed, 478 insertions(+), 28 deletions(-) rename src-tauri/src/{take_screenshot_loop.rs => logger.rs} (100%) create mode 100644 src-tauri/src/screenshot.rs create mode 100644 src-tauri/src/screenshot_manager.rs create mode 100644 src/components/screen-view.tsx diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock index 7c3c95c..293bcc8 100644 --- a/src-tauri/Cargo.lock +++ b/src-tauri/Cargo.lock @@ -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" diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index eebbe5b..686cfb8 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -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 diff --git a/src-tauri/src/take_screenshot_loop.rs b/src-tauri/src/logger.rs similarity index 100% rename from src-tauri/src/take_screenshot_loop.rs rename to src-tauri/src/logger.rs diff --git a/src-tauri/src/main.rs b/src-tauri/src/main.rs index 87945df..ba4b579 100644 --- a/src-tauri/src/main.rs +++ b/src-tauri/src/main.rs @@ -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 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 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 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 }) } -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"); diff --git a/src-tauri/src/screenshot.rs b/src-tauri/src/screenshot.rs new file mode 100644 index 0000000..148fa02 --- /dev/null +++ b/src-tauri/src/screenshot.rs @@ -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>>, + pub scale_factor: f32, +} + +impl Screenshot { + pub fn new( + display_id: u32, + height: u32, + width: u32, + bytes_per_row: usize, + bytes: Vec, + 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, +} diff --git a/src-tauri/src/screenshot_manager.rs b/src-tauri/src/screenshot_manager.rs new file mode 100644 index 0000000..f6e5daf --- /dev/null +++ b/src-tauri/src/screenshot_manager.rs @@ -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 { + 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>>>, + encode_listeners: Arc>>>, +} + +impl ScreenshotManager { + pub async fn global() -> &'static Self { + static SCREENSHOT_MANAGER: OnceCell = 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, + ) { + 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 { + 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) + } +} diff --git a/src/App.tsx b/src/App.tsx index 3c03d7e..3c89d91 100644 --- a/src/App.tsx +++ b/src/App.tsx @@ -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([]); @@ -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 (
  • -
    • - ); - })} - -
      - {screenshots().map((screenshot) => { - return ( -
    1. - +
      2. ); })} diff --git a/src/components/screen-view.tsx b/src/components/screen-view.tsx new file mode 100644 index 0000000..7746a6c --- /dev/null +++ b/src/components/screen-view.tsx @@ -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 = (props) => { + const [image, setImage] = createSignal(); + 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 ; +};