refactor: cleanup bandwidth tracking functionality

This commit is contained in:
zhom
2025-11-30 16:55:23 +04:00
parent cdba9aac33
commit f098128988
10 changed files with 379 additions and 62 deletions
+79 -27
View File
@@ -267,6 +267,18 @@ export async function startCamoufoxProcess(
});
}
/**
* Check if a process is running by PID
*/
function isProcessRunning(pid: number): boolean {
try {
process.kill(pid, 0);
return true;
} catch {
return false;
}
}
/**
* Stop a Camoufox process
* @param id The Camoufox ID to stop
@@ -279,45 +291,85 @@ export async function stopCamoufoxProcess(id: string): Promise<boolean> {
return false;
}
const pid = config.processId;
try {
// Method 1: If we have a process ID, kill by PID with proper signal sequence
if (config.processId) {
if (pid && isProcessRunning(pid)) {
try {
// First try SIGTERM for graceful shutdown
process.kill(config.processId, "SIGTERM");
// Give it more time to terminate gracefully (increased from 2s to 5s)
await new Promise((resolve) => setTimeout(resolve, 5000));
process.kill(pid, "SIGTERM");
// Check if process is still running
try {
process.kill(config.processId, 0); // Signal 0 checks if process exists
process.kill(config.processId, "SIGKILL");
} catch {}
} catch {}
// Wait up to 3 seconds for graceful shutdown
for (let i = 0; i < 30; i++) {
await new Promise((resolve) => setTimeout(resolve, 100));
if (!isProcessRunning(pid)) {
break;
}
}
// If still running, force kill
if (isProcessRunning(pid)) {
process.kill(pid, "SIGKILL");
// Wait for SIGKILL to take effect
for (let i = 0; i < 20; i++) {
await new Promise((resolve) => setTimeout(resolve, 100));
if (!isProcessRunning(pid)) {
break;
}
}
}
} catch {
// Process might have already exited
}
}
// Method 2: Pattern-based kill as fallback
const killByPattern = spawn(
"pkill",
["-TERM", "-f", `camoufox-worker.*${id}`],
{
stdio: "ignore",
},
);
// Wait for pattern-based kill command to complete
// Method 2: Pattern-based kill as fallback (kills any child processes)
await new Promise<void>((resolve) => {
const killByPattern = spawn(
"pkill",
["-TERM", "-f", `camoufox-worker.*${id}`],
{ stdio: "ignore" },
);
killByPattern.on("exit", () => resolve());
// Timeout after 3 seconds
setTimeout(() => resolve(), 3000);
setTimeout(() => resolve(), 1000);
});
// Final cleanup with SIGKILL if needed
setTimeout(() => {
spawn("pkill", ["-KILL", "-f", `camoufox-worker.*${id}`], {
stdio: "ignore",
// Wait a moment then force kill any remaining
await new Promise((resolve) => setTimeout(resolve, 500));
await new Promise<void>((resolve) => {
const killByPatternForce = spawn(
"pkill",
["-KILL", "-f", `camoufox-worker.*${id}`],
{ stdio: "ignore" },
);
killByPatternForce.on("exit", () => resolve());
setTimeout(() => resolve(), 1000);
});
// Also kill any Firefox processes associated with this profile
if (config.profilePath) {
await new Promise<void>((resolve) => {
const killFirefox = spawn(
"pkill",
["-KILL", "-f", config.profilePath!],
{ stdio: "ignore" },
);
killFirefox.on("exit", () => resolve());
setTimeout(() => resolve(), 1000);
});
}, 1000);
}
// Verify process is actually dead
if (pid && isProcessRunning(pid)) {
// Last resort: SIGKILL again
try {
process.kill(pid, "SIGKILL");
} catch {
// Ignore
}
}
// Delete the configuration
deleteCamoufoxConfig(id);
+1
View File
@@ -66,6 +66,7 @@ export function parseProxyString(proxyString: LaunchOptions["proxy"] | string) {
// Try parsing as URL first (handles protocol://username:password@host:port)
if (trimmed.includes("://")) {
const url = new URL(trimmed);
// Playwright accepts short form "host:port" for HTTP proxies
server = `${url.hostname}:${url.port}`;
if (url.username) {
+83 -8
View File
@@ -1064,6 +1064,19 @@ impl BrowserRunner {
profile.id
);
// Stop the proxy associated with this profile first
let profile_id_str = profile.id.to_string();
if let Err(e) = PROXY_MANAGER
.stop_proxy_by_profile_id(app_handle.clone(), &profile_id_str)
.await
{
log::warn!(
"Warning: Failed to stop proxy for profile {}: {e}",
profile_id_str
);
}
let mut process_actually_stopped = false;
match self
.camoufox_manager
.find_camoufox_by_profile(&profile_path_str)
@@ -1083,13 +1096,69 @@ impl BrowserRunner {
{
Ok(stopped) => {
if stopped {
log::info!(
"Successfully stopped Camoufox process: {} (PID: {:?})",
camoufox_process.id,
camoufox_process.processId
);
// Verify the process actually died by checking after a short delay
if let Some(pid) = camoufox_process.processId {
use tokio::time::{sleep, Duration};
sleep(Duration::from_millis(500)).await;
use sysinfo::{Pid, System};
let system = System::new_all();
process_actually_stopped = system.process(Pid::from(pid as usize)).is_none();
if process_actually_stopped {
log::info!(
"Successfully stopped Camoufox process: {} (PID: {:?}) - verified process is dead",
camoufox_process.id,
pid
);
} else {
log::warn!(
"Camoufox stop command returned success but process {} (PID: {:?}) is still running - forcing kill",
camoufox_process.id,
pid
);
// Force kill the process
#[cfg(target_os = "macos")]
{
use crate::platform_browser;
if let Err(e) = platform_browser::macos::kill_browser_process_impl(
pid,
Some(&profile_path_str),
)
.await
{
log::error!("Failed to force kill Camoufox process {}: {}", pid, e);
} else {
process_actually_stopped = true;
}
}
#[cfg(target_os = "linux")]
{
use crate::platform_browser;
if let Err(e) = platform_browser::linux::kill_browser_process_impl(pid).await
{
log::error!("Failed to force kill Camoufox process {}: {}", pid, e);
} else {
process_actually_stopped = true;
}
}
#[cfg(target_os = "windows")]
{
use crate::platform_browser;
if let Err(e) =
platform_browser::windows::kill_browser_process_impl(pid).await
{
log::error!("Failed to force kill Camoufox process {}: {}", pid, e);
} else {
process_actually_stopped = true;
}
}
}
} else {
process_actually_stopped = true; // No PID to verify, assume stopped
}
} else {
log::info!(
log::warn!(
"Failed to stop Camoufox process: {} (PID: {:?})",
camoufox_process.id,
camoufox_process.processId
@@ -1097,7 +1166,7 @@ impl BrowserRunner {
}
}
Err(e) => {
log::info!(
log::error!(
"Error stopping Camoufox process {}: {}",
camoufox_process.id,
e
@@ -1111,9 +1180,10 @@ impl BrowserRunner {
profile.name,
profile.id
);
process_actually_stopped = true; // No process found, consider it stopped
}
Err(e) => {
log::info!(
log::error!(
"Error finding Camoufox process for profile {}: {}",
profile.name,
e
@@ -1121,6 +1191,11 @@ impl BrowserRunner {
}
}
// Log warning if process wasn't confirmed stopped, but continue with cleanup
if !process_actually_stopped {
log::warn!("Camoufox process may still be running, but proceeding with cleanup");
}
// Clear the process ID from the profile
let mut updated_profile = profile.clone();
updated_profile.process_id = None;
+89 -8
View File
@@ -664,14 +664,32 @@ impl ProxyManager {
&& existing.upstream_port == desired_port;
if is_same_upstream {
// Reuse existing local proxy
return Ok(ProxySettings {
proxy_type: "http".to_string(),
host: "127.0.0.1".to_string(),
port: existing.local_port,
username: None,
password: None,
});
// Check if profile_id matches - if not, we need to restart to update tracking
let profile_id_matches = match (profile_id, &existing.profile_id) {
(Some(ref new_id), Some(ref old_id)) => new_id == old_id,
(None, None) => true,
_ => false,
};
if profile_id_matches {
// Reuse existing local proxy (profile_id matches)
return Ok(ProxySettings {
proxy_type: "http".to_string(),
host: "127.0.0.1".to_string(),
port: existing.local_port,
username: None,
password: None,
});
} else {
// Profile ID changed - need to restart proxy to update tracking
log::info!(
"Profile ID changed for proxy {}: {:?} -> {:?}, restarting proxy",
existing.id,
existing.profile_id,
profile_id
);
needs_restart = true;
}
} else {
// Upstream changed; we must restart the local proxy so that traffic is routed correctly
needs_restart = true;
@@ -864,6 +882,69 @@ impl ProxyManager {
Ok(())
}
// Stop the proxy associated with a profile ID
pub async fn stop_proxy_by_profile_id(
&self,
app_handle: tauri::AppHandle,
profile_id: &str,
) -> Result<(), String> {
// Find the proxy ID for this profile
let proxy_id = {
let map = self.profile_active_proxy_ids.lock().unwrap();
map.get(profile_id).cloned()
};
if let Some(proxy_id) = proxy_id {
// Find the PID for this proxy
let pid = {
let proxies = self.active_proxies.lock().unwrap();
proxies.iter().find_map(|(pid, proxy)| {
if proxy.id == proxy_id {
Some(*pid)
} else {
None
}
})
};
if let Some(pid) = pid {
// Use the existing stop_proxy method
self.stop_proxy(app_handle, pid).await
} else {
// Proxy not found in active_proxies, try to stop it directly by ID
let proxy_cmd = app_handle
.shell()
.sidecar("donut-proxy")
.map_err(|e| format!("Failed to create sidecar: {e}"))?
.arg("proxy")
.arg("stop")
.arg("--id")
.arg(&proxy_id);
let output = proxy_cmd.output().await.unwrap();
if !output.status.success() {
let stderr = String::from_utf8_lossy(&output.stderr);
log::warn!("Proxy stop error: {stderr}");
}
// Clear profile-to-proxy mapping
let mut map = self.profile_active_proxy_ids.lock().unwrap();
map.remove(profile_id);
// Emit event for reactive UI updates
if let Err(e) = app_handle.emit("proxies-changed", ()) {
log::error!("Failed to emit proxies-changed event: {e}");
}
Ok(())
}
} else {
// No proxy found for this profile
Ok(())
}
}
// Update the PID mapping for an existing proxy
pub fn update_proxy_pid(&self, old_pid: u32, new_pid: u32) -> Result<(), String> {
let mut proxies = self.active_proxies.lock().unwrap();
+9 -1
View File
@@ -30,9 +30,17 @@ pub async fn start_proxy_process_with_profile(
listener.local_addr().unwrap().port()
});
let config = ProxyConfig::new(id.clone(), upstream, Some(local_port)).with_profile_id(profile_id);
let config =
ProxyConfig::new(id.clone(), upstream, Some(local_port)).with_profile_id(profile_id.clone());
save_proxy_config(&config)?;
// Log profile_id for debugging
if let Some(ref pid) = profile_id {
log::info!("Saved proxy config {} with profile_id: {}", id, pid);
} else {
log::info!("Saved proxy config {} without profile_id", id);
}
// Spawn proxy worker process in the background using std::process::Command
// This ensures proper process detachment on Unix systems
let exe = std::env::current_exe()?;
+24 -9
View File
@@ -45,12 +45,14 @@ impl<S: AsyncRead + Unpin> AsyncRead for CountingStream<S> {
let result = Pin::new(&mut self.inner).poll_read(cx, buf);
if let Poll::Ready(Ok(())) = &result {
let bytes_read = buf.filled().len() - filled_before;
self
.bytes_read
.fetch_add(bytes_read as u64, Ordering::Relaxed);
// Update global tracker
if let Some(tracker) = get_traffic_tracker() {
tracker.add_bytes_received(bytes_read as u64);
if bytes_read > 0 {
self
.bytes_read
.fetch_add(bytes_read as u64, Ordering::Relaxed);
// Update global tracker - count as received (data coming into proxy)
if let Some(tracker) = get_traffic_tracker() {
tracker.add_bytes_received(bytes_read as u64);
}
}
}
result
@@ -66,7 +68,7 @@ impl<S: AsyncWrite + Unpin> AsyncWrite for CountingStream<S> {
let result = Pin::new(&mut self.inner).poll_write(cx, buf);
if let Poll::Ready(Ok(n)) = &result {
self.bytes_written.fetch_add(*n as u64, Ordering::Relaxed);
// Update global tracker
// Update global tracker - count as sent (data going out of proxy)
if let Some(tracker) = get_traffic_tracker() {
tracker.add_bytes_sent(*n as u64);
}
@@ -522,15 +524,17 @@ pub async fn run_proxy_server(config: ProxyConfig) -> Result<(), Box<dyn std::er
};
log::error!(
"Found config: id={}, port={:?}, upstream={}",
"Found config: id={}, port={:?}, upstream={}, profile_id={:?}",
config.id,
config.local_port,
config.upstream_url
config.upstream_url,
config.profile_id
);
log::error!("Starting proxy server for config id: {}", config.id);
// Initialize traffic tracker with profile ID if available
// This can now be called multiple times to update the tracker
init_traffic_tracker(config.id.clone(), config.profile_id.clone());
log::error!(
"Traffic tracker initialized for proxy: {} (profile_id: {:?})",
@@ -538,6 +542,17 @@ pub async fn run_proxy_server(config: ProxyConfig) -> Result<(), Box<dyn std::er
config.profile_id
);
// Verify tracker was initialized correctly
if let Some(tracker) = crate::traffic_stats::get_traffic_tracker() {
log::error!(
"Tracker verified: proxy_id={}, profile_id={:?}",
tracker.proxy_id,
tracker.profile_id
);
} else {
log::error!("WARNING: Tracker was not initialized!");
}
// Determine the bind address
let bind_addr = SocketAddr::from(([127, 0, 0, 1], config.local_port.unwrap_or(0)));
+15 -5
View File
@@ -150,7 +150,7 @@ impl TrafficStats {
}
}
// Add new data point
// Add new data point (even if bytes are zero, to ensure chart has continuous data)
self.bandwidth_history.push(BandwidthDataPoint {
timestamp: now,
bytes_sent,
@@ -392,6 +392,11 @@ impl LiveTrafficTracker {
let mut stats = load_traffic_stats(&self.proxy_id)
.unwrap_or_else(|| TrafficStats::new(self.proxy_id.clone(), self.profile_id.clone()));
// Ensure profile_id is set (in case stats were loaded from disk without it)
if stats.profile_id.is_none() && self.profile_id.is_some() {
stats.profile_id = self.profile_id.clone();
}
// Update bandwidth history
stats.record_bandwidth(bytes_sent, bytes_received);
@@ -419,17 +424,22 @@ impl LiveTrafficTracker {
}
/// Global traffic tracker that can be accessed from connection handlers
pub static TRAFFIC_TRACKER: std::sync::OnceLock<Arc<LiveTrafficTracker>> =
std::sync::OnceLock::new();
/// Using RwLock to allow reinitialization when proxy config changes
static TRAFFIC_TRACKER: std::sync::RwLock<Option<Arc<LiveTrafficTracker>>> =
std::sync::RwLock::new(None);
/// Initialize the global traffic tracker
/// This can be called multiple times to update the tracker when proxy config changes
pub fn init_traffic_tracker(proxy_id: String, profile_id: Option<String>) {
let _ = TRAFFIC_TRACKER.set(Arc::new(LiveTrafficTracker::new(proxy_id, profile_id)));
let tracker = Arc::new(LiveTrafficTracker::new(proxy_id, profile_id));
if let Ok(mut guard) = TRAFFIC_TRACKER.write() {
*guard = Some(tracker);
}
}
/// Get the global traffic tracker
pub fn get_traffic_tracker() -> Option<Arc<LiveTrafficTracker>> {
TRAFFIC_TRACKER.get().cloned()
TRAFFIC_TRACKER.read().ok().and_then(|guard| guard.clone())
}
#[cfg(test)]
+9 -2
View File
@@ -884,7 +884,10 @@ export function ProfilesDataTable({
const newSnapshots: Record<string, TrafficSnapshot> = {};
for (const snapshot of allSnapshots) {
if (snapshot.profile_id) {
newSnapshots[snapshot.profile_id] = snapshot;
const existing = newSnapshots[snapshot.profile_id];
if (!existing || snapshot.last_update > existing.last_update) {
newSnapshots[snapshot.profile_id] = snapshot;
}
}
}
setTrafficSnapshots(newSnapshots);
@@ -1693,13 +1696,17 @@ export function ProfilesDataTable({
if (isRunning && meta.trafficSnapshots) {
// Find the traffic snapshot for this profile by matching profile_id
const snapshot = meta.trafficSnapshots[profile.id];
const bandwidthData = snapshot?.recent_bandwidth || [];
// Create a new array reference to ensure React detects changes
const bandwidthData = snapshot?.recent_bandwidth
? [...snapshot.recent_bandwidth]
: [];
const currentBandwidth =
(snapshot?.current_bytes_sent || 0) +
(snapshot?.current_bytes_received || 0);
return (
<BandwidthMiniChart
key={`${profile.id}-${snapshot?.last_update || 0}-${bandwidthData.length}`}
data={bandwidthData}
currentBandwidth={currentBandwidth}
onClick={() => meta.onOpenTrafficDialog?.(profile.id)}
+10 -2
View File
@@ -83,8 +83,16 @@ export function TrafficDetailsDialog({
const fetchStats = async () => {
try {
const allStats = await invoke<TrafficStats[]>("get_all_traffic_stats");
const profileStats = allStats.find((s) => s.profile_id === profileId);
setStats(profileStats || null);
const matchingStats = allStats.filter(
(s) => s.profile_id === profileId,
);
const profileStats =
matchingStats.length > 0
? matchingStats.reduce((latest, current) =>
current.last_update > latest.last_update ? current : latest,
)
: null;
setStats(profileStats);
} catch (error) {
console.error("Failed to fetch traffic stats:", error);
}
+60
View File
@@ -16,6 +16,11 @@ export interface ThemeColors extends Record<string, string> {
"--destructive": string;
"--destructive-foreground": string;
"--border": string;
"--chart-1": string;
"--chart-2": string;
"--chart-3": string;
"--chart-4": string;
"--chart-5": string;
}
export interface Theme {
@@ -46,6 +51,11 @@ export const THEMES: Theme[] = [
"--destructive": "#f7768e",
"--destructive-foreground": "#1a1b26",
"--border": "#3b4261",
"--chart-1": "#7aa2f7",
"--chart-2": "#9ece6a",
"--chart-3": "#bb9af7",
"--chart-4": "#2ac3de",
"--chart-5": "#ff9e64",
},
},
{
@@ -69,6 +79,11 @@ export const THEMES: Theme[] = [
"--destructive": "#ff5555",
"--destructive-foreground": "#f8f8f2",
"--border": "#6272a4",
"--chart-1": "#bd93f9",
"--chart-2": "#50fa7b",
"--chart-3": "#ff79c6",
"--chart-4": "#8be9fd",
"--chart-5": "#ffb86c",
},
},
{
@@ -92,6 +107,11 @@ export const THEMES: Theme[] = [
"--destructive": "#ff819f",
"--destructive-foreground": "#273136",
"--border": "#304e37",
"--chart-1": "#7eb08a",
"--chart-2": "#d2b48c",
"--chart-3": "#7ea4b0",
"--chart-4": "#a8c97f",
"--chart-5": "#e6c07b",
},
},
{
@@ -115,6 +135,11 @@ export const THEMES: Theme[] = [
"--destructive": "#ef4444",
"--destructive-foreground": "#f7f7f8",
"--border": "#2a2e39",
"--chart-1": "#5755d9",
"--chart-2": "#0ea5e9",
"--chart-3": "#f25f4c",
"--chart-4": "#22c55e",
"--chart-5": "#f59e0b",
},
},
{
@@ -138,6 +163,11 @@ export const THEMES: Theme[] = [
"--destructive": "#f07178",
"--destructive-foreground": "#b3b1ad",
"--border": "#1f2430",
"--chart-1": "#39bae6",
"--chart-2": "#c2d94c",
"--chart-3": "#d2a6ff",
"--chart-4": "#ffb454",
"--chart-5": "#f07178",
},
},
{
@@ -161,6 +191,11 @@ export const THEMES: Theme[] = [
"--destructive": "#f07178",
"--destructive-foreground": "#fafafa",
"--border": "#e7eaed",
"--chart-1": "#399ee6",
"--chart-2": "#86b300",
"--chart-3": "#a37acc",
"--chart-4": "#fa8d3e",
"--chart-5": "#f07178",
},
},
{
@@ -184,6 +219,11 @@ export const THEMES: Theme[] = [
"--destructive": "#d20f39",
"--destructive-foreground": "#eff1f5",
"--border": "#9ca0b0",
"--chart-1": "#1e66f5",
"--chart-2": "#40a02b",
"--chart-3": "#8839ef",
"--chart-4": "#04a5e5",
"--chart-5": "#df8e1d",
},
},
{
@@ -207,6 +247,11 @@ export const THEMES: Theme[] = [
"--destructive": "#e78284",
"--destructive-foreground": "#303446",
"--border": "#737994",
"--chart-1": "#8caaee",
"--chart-2": "#a6d189",
"--chart-3": "#ca9ee6",
"--chart-4": "#99d1db",
"--chart-5": "#e5c890",
},
},
{
@@ -230,6 +275,11 @@ export const THEMES: Theme[] = [
"--destructive": "#ed8796",
"--destructive-foreground": "#24273a",
"--border": "#6e738d",
"--chart-1": "#8aadf4",
"--chart-2": "#a6da95",
"--chart-3": "#c6a0f6",
"--chart-4": "#91d7e3",
"--chart-5": "#eed49f",
},
},
{
@@ -253,6 +303,11 @@ export const THEMES: Theme[] = [
"--destructive": "#f38ba8",
"--destructive-foreground": "#1e1e2e",
"--border": "#585b70",
"--chart-1": "#89b4fa",
"--chart-2": "#a6e3a1",
"--chart-3": "#cba6f7",
"--chart-4": "#89dceb",
"--chart-5": "#f9e2af",
},
},
];
@@ -276,6 +331,11 @@ export const THEME_VARIABLES: Array<{ key: keyof ThemeColors; label: string }> =
{ key: "--destructive", label: "Destructive" },
{ key: "--destructive-foreground", label: "Destructive FG" },
{ key: "--border", label: "Border" },
{ key: "--chart-1", label: "Chart 1" },
{ key: "--chart-2", label: "Chart 2" },
{ key: "--chart-3", label: "Chart 3" },
{ key: "--chart-4", label: "Chart 4" },
{ key: "--chart-5", label: "Chart 5" },
];
export function getThemeById(id: string): Theme | undefined {