Skip to content

Commit 90bf654

Browse files
jacobtreadjacobtread
committed
refactor: separate thread suspend logic, take suspended threads allocation when resuming
Resuming will now swap the allocated suspended threads vec with an empty non allocated vec, preventing the capacity allocation from staying for the life of the program (Even though its an super tiny allocation)
1 parent eb63b3b commit 90bf654

3 files changed

Lines changed: 133 additions & 101 deletions

File tree

‎src/lib.rs‎

Lines changed: 4 additions & 94 deletions
Original file line numberDiff line numberDiff line change
@@ -7,26 +7,15 @@ use core::{
77
};
88
use log::error;
99
use pocket_relay_client_shared as core;
10-
use std::{path::Path, sync::Mutex};
10+
use std::path::Path;
1111
use ui::{confirm_message, error_message};
12-
use windows_sys::Win32::{
13-
Foundation::{CloseHandle, INVALID_HANDLE_VALUE},
14-
System::{
15-
Diagnostics::ToolHelp::{
16-
CreateToolhelp32Snapshot, Thread32First, Thread32Next, TH32CS_SNAPTHREAD, THREADENTRY32,
17-
},
18-
SystemServices::{DLL_PROCESS_ATTACH, DLL_PROCESS_DETACH},
19-
Threading::{
20-
GetCurrentProcessId, GetCurrentThreadId, OpenThread, ResumeThread, SuspendThread,
21-
THREAD_QUERY_INFORMATION, THREAD_SUSPEND_RESUME,
22-
},
23-
},
24-
};
12+
use windows_sys::Win32::System::SystemServices::{DLL_PROCESS_ATTACH, DLL_PROCESS_DETACH};
2513

2614
pub mod config;
2715
pub mod game;
2816
pub mod hooks;
2917
pub mod servers;
18+
pub mod threads;
3019
pub mod ui;
3120
pub mod update;
3221

@@ -36,7 +25,7 @@ pub const APP_VERSION: &str = env!("CARGO_PKG_VERSION");
3625
/// Handles the plugin being attached to the game
3726
fn attach() {
3827
// Suspend all game threads so the user has a chance to connect to a server
39-
unsafe { suspend_all_threads() };
28+
unsafe { threads::suspend_all_threads() };
4029

4130
// Debug allocates a console window to display output
4231
#[cfg(debug_assertions)]
@@ -79,85 +68,6 @@ fn detach() {
7968
}
8069
}
8170

82-
// Threads that were suspended
83-
static SUSPENDED_THREADS: Mutex<Vec<u32>> = Mutex::new(Vec::new());
84-
85-
unsafe fn suspend_all_threads() {
86-
let current_thread_id = GetCurrentThreadId();
87-
let target_process_id = GetCurrentProcessId();
88-
89-
let snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
90-
if snapshot == INVALID_HANDLE_VALUE {
91-
return;
92-
}
93-
94-
let mut thread_entry: THREADENTRY32 = unsafe { std::mem::zeroed() };
95-
thread_entry.dwSize = std::mem::size_of::<THREADENTRY32>() as u32;
96-
97-
// Read the first thread entry
98-
if Thread32First(snapshot, &mut thread_entry) == 0 {
99-
return;
100-
}
101-
102-
let mut suspended_threads = Vec::new();
103-
104-
loop {
105-
// Suspend threads that aren't the current thread
106-
if thread_entry.th32OwnerProcessID == target_process_id
107-
&& thread_entry.th32ThreadID != current_thread_id
108-
{
109-
let thread_handle = unsafe {
110-
OpenThread(
111-
THREAD_SUSPEND_RESUME | THREAD_QUERY_INFORMATION,
112-
0,
113-
thread_entry.th32ThreadID,
114-
)
115-
};
116-
117-
if thread_handle != 0 {
118-
SuspendThread(thread_handle);
119-
CloseHandle(thread_handle);
120-
121-
suspended_threads.push(thread_entry.th32ThreadID);
122-
}
123-
}
124-
125-
// Read the next thread
126-
if Thread32Next(snapshot, &mut thread_entry) == 0 {
127-
break;
128-
}
129-
}
130-
131-
CloseHandle(snapshot);
132-
133-
// Store the threads we suspended
134-
if let Ok(mut value) = SUSPENDED_THREADS.lock() {
135-
*value = suspended_threads;
136-
}
137-
}
138-
139-
unsafe fn resume_all_threads() {
140-
// Get the suspended threads
141-
let suspended_threads = match SUSPENDED_THREADS.lock() {
142-
Ok(mut value) => value.split_off(0),
143-
Err(_) => return,
144-
};
145-
146-
// Resume the threads that were suspended
147-
for thread_id in suspended_threads {
148-
let thread_handle = OpenThread(
149-
THREAD_SUSPEND_RESUME | THREAD_QUERY_INFORMATION,
150-
0,
151-
thread_id,
152-
);
153-
154-
if thread_handle != 0 {
155-
ResumeThread(thread_handle);
156-
CloseHandle(thread_handle);
157-
}
158-
}
159-
}
160-
16171
/// Attempts to load an identity file if one is present
16272
fn load_identity() -> Option<Identity> {
16373
// Load the client identity

‎src/threads.rs‎

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
//! # Threads
2+
//!
3+
//! Logic for managing threads, in this case it handles pausing and
4+
//! resuming process threads on startup. This is what allows the user
5+
//! to connect to a server before the game properly starts
6+
7+
use std::{mem::swap, sync::Mutex};
8+
use windows_sys::Win32::{
9+
Foundation::{CloseHandle, FALSE, INVALID_HANDLE_VALUE},
10+
System::{
11+
Diagnostics::ToolHelp::{
12+
CreateToolhelp32Snapshot, Thread32First, Thread32Next, TH32CS_SNAPTHREAD, THREADENTRY32,
13+
},
14+
Threading::{
15+
GetCurrentProcessId, GetCurrentThreadId, OpenThread, ResumeThread, SuspendThread,
16+
THREAD_QUERY_INFORMATION, THREAD_SUSPEND_RESUME,
17+
},
18+
},
19+
};
20+
21+
// Threads that were suspended
22+
static SUSPENDED_THREADS: Mutex<Vec<u32>> = Mutex::new(Vec::new());
23+
24+
/// Suspends all threads on the process excluding the current thread. Suspended
25+
/// threads are stored in [SUSPENDED_THREADS] and can be later resumed with
26+
/// [resume_all_threads].
27+
///
28+
/// Should only be called on initial startup to prevent interrupting any network
29+
/// connection threads.
30+
pub fn suspend_all_threads() {
31+
let (current_thread_id, target_process_id) =
32+
unsafe { (GetCurrentThreadId(), GetCurrentProcessId()) };
33+
34+
let snapshot = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0) };
35+
if snapshot == INVALID_HANDLE_VALUE {
36+
return;
37+
}
38+
39+
let mut thread_entry: THREADENTRY32 = unsafe { std::mem::zeroed() };
40+
thread_entry.dwSize = std::mem::size_of::<THREADENTRY32>() as u32;
41+
42+
// Read the first thread entry
43+
if unsafe { Thread32First(snapshot, &mut thread_entry) } == FALSE {
44+
return;
45+
}
46+
47+
let mut suspended_threads = Vec::new();
48+
49+
loop {
50+
// Suspend threads that aren't the current thread
51+
if thread_entry.th32OwnerProcessID == target_process_id
52+
&& thread_entry.th32ThreadID != current_thread_id
53+
{
54+
let thread_handle = unsafe {
55+
OpenThread(
56+
THREAD_SUSPEND_RESUME | THREAD_QUERY_INFORMATION,
57+
FALSE,
58+
thread_entry.th32ThreadID,
59+
)
60+
};
61+
62+
if thread_handle != 0 {
63+
unsafe {
64+
SuspendThread(thread_handle);
65+
CloseHandle(thread_handle);
66+
}
67+
68+
suspended_threads.push(thread_entry.th32ThreadID);
69+
}
70+
}
71+
72+
// Read the next thread
73+
if unsafe { Thread32Next(snapshot, &mut thread_entry) } == FALSE {
74+
break;
75+
}
76+
}
77+
78+
unsafe {
79+
CloseHandle(snapshot);
80+
}
81+
82+
// Store the threads we suspended
83+
if let Ok(mut value) = SUSPENDED_THREADS.lock() {
84+
*value = suspended_threads;
85+
}
86+
}
87+
88+
/// Resumes all suspended threads
89+
pub fn resume_all_threads() {
90+
// Get the suspended threads
91+
let suspended_threads = match SUSPENDED_THREADS.lock() {
92+
// Take the collection of locked threads
93+
Ok(mut value) => {
94+
// Swap the allocated threads list with an unallocated vec
95+
//
96+
// Reason: Allows us to take the allocated capacity so it doesn't
97+
// maintain it for the life of the program like split_off would
98+
let mut threads = Vec::new();
99+
swap(value.as_mut(), &mut threads);
100+
101+
threads
102+
}
103+
104+
// Lock is poisoned, shouldn't have reached a reusable point if
105+
// the main thread crashed
106+
Err(_) => return,
107+
};
108+
109+
// Resume the threads that were suspended
110+
for thread_id in suspended_threads {
111+
let thread_handle = unsafe {
112+
OpenThread(
113+
THREAD_SUSPEND_RESUME | THREAD_QUERY_INFORMATION,
114+
0,
115+
thread_id,
116+
)
117+
};
118+
119+
if thread_handle != 0 {
120+
unsafe {
121+
ResumeThread(thread_handle);
122+
CloseHandle(thread_handle);
123+
}
124+
}
125+
}
126+
}

‎src/ui.rs‎

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,8 @@ use crate::{
55
reqwest::Client,
66
servers::{has_server_tasks, stop_server_tasks},
77
},
8-
resume_all_threads,
98
servers::start_all_servers,
9+
threads::resume_all_threads,
1010
update,
1111
};
1212
use futures::FutureExt;
@@ -172,9 +172,7 @@ impl App {
172172
self.set_button.set_text("Disconnect");
173173

174174
// Resume game threads
175-
unsafe {
176-
resume_all_threads();
177-
}
175+
resume_all_threads();
178176
}
179177
}
180178

@@ -223,9 +221,7 @@ pub fn init(config: Option<ClientConfig>, client: Client) {
223221
dispatch_thread_events();
224222

225223
// Resume the game threads if we close the UI
226-
unsafe {
227-
resume_all_threads();
228-
}
224+
resume_all_threads();
229225

230226
let shutdown_signal = tokio::signal::ctrl_c();
231227
let _ = runtime.block_on(shutdown_signal);

0 commit comments

Comments
 (0)