use std::fmt;
use std::io;
use windows_sys::Win32::Storage::FileSystem::{
Wow64DisableWow64FsRedirection, Wow64RevertWow64FsRedirection,
};
use windows_sys::Win32::System::Threading::{
GetCurrentThread, GetThreadInformation, MEMORY_PRIORITY, MEMORY_PRIORITY_BELOW_NORMAL,
MEMORY_PRIORITY_INFORMATION, MEMORY_PRIORITY_LOW, MEMORY_PRIORITY_MEDIUM,
MEMORY_PRIORITY_NORMAL, MEMORY_PRIORITY_VERY_LOW, SetThreadInformation, SetThreadPriority,
THREAD_MODE_BACKGROUND_BEGIN, THREAD_MODE_BACKGROUND_END, ThreadMemoryPriority,
};
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum MemoryPriority {
VeryLow,
Low,
Medium,
BelowNormal,
Normal,
}
impl MemoryPriority {
#[must_use]
pub const fn as_raw(self) -> MEMORY_PRIORITY {
match self {
Self::VeryLow => MEMORY_PRIORITY_VERY_LOW,
Self::Low => MEMORY_PRIORITY_LOW,
Self::Medium => MEMORY_PRIORITY_MEDIUM,
Self::BelowNormal => MEMORY_PRIORITY_BELOW_NORMAL,
Self::Normal => MEMORY_PRIORITY_NORMAL,
}
}
#[must_use]
pub const fn from_raw(raw: MEMORY_PRIORITY) -> Option<Self> {
match raw {
MEMORY_PRIORITY_VERY_LOW => Some(Self::VeryLow),
MEMORY_PRIORITY_LOW => Some(Self::Low),
MEMORY_PRIORITY_MEDIUM => Some(Self::Medium),
MEMORY_PRIORITY_BELOW_NORMAL => Some(Self::BelowNormal),
MEMORY_PRIORITY_NORMAL => Some(Self::Normal),
_ => None,
}
}
pub fn current() -> Result<Self, DeclaredError> {
let mut info = MEMORY_PRIORITY_INFORMATION {
MemoryPriority: MEMORY_PRIORITY_NORMAL,
};
let ok = unsafe {
GetThreadInformation(
GetCurrentThread(),
ThreadMemoryPriority,
std::ptr::from_mut(&mut info).cast(),
size_of::<MEMORY_PRIORITY_INFORMATION>() as u32,
)
};
if ok == 0 {
return Err(DeclaredError::new(DeclaredAspect::MemoryPriority));
}
Self::from_raw(info.MemoryPriority)
.ok_or_else(|| DeclaredError::without_os_error(DeclaredAspect::MemoryPriority))
}
fn install(self) -> Result<(), DeclaredError> {
let info = MEMORY_PRIORITY_INFORMATION {
MemoryPriority: self.as_raw(),
};
let ok = unsafe {
SetThreadInformation(
GetCurrentThread(),
ThreadMemoryPriority,
std::ptr::from_ref(&info).cast(),
size_of::<MEMORY_PRIORITY_INFORMATION>() as u32,
)
};
if ok == 0 {
return Err(DeclaredError::new(DeclaredAspect::MemoryPriority));
}
Ok(())
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum BackgroundMode {
Begin,
End,
}
impl BackgroundMode {
fn install(self) -> Result<(), DeclaredError> {
let value = match self {
Self::Begin => THREAD_MODE_BACKGROUND_BEGIN,
Self::End => THREAD_MODE_BACKGROUND_END,
};
let ok = unsafe { SetThreadPriority(GetCurrentThread(), value) };
if ok == 0 {
return Err(DeclaredError::new(DeclaredAspect::BackgroundMode));
}
Ok(())
}
const fn inverse(self) -> Self {
match self {
Self::Begin => Self::End,
Self::End => Self::Begin,
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum Wow64Redirection {
Disabled,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum DeclaredAspect {
Wow64Redirection,
MemoryPriority,
BackgroundMode,
}
#[derive(Debug)]
pub struct DeclaredError {
aspect: DeclaredAspect,
source: Option<io::Error>,
}
impl DeclaredError {
fn new(aspect: DeclaredAspect) -> Self {
Self {
aspect,
source: Some(io::Error::last_os_error()),
}
}
const fn without_os_error(aspect: DeclaredAspect) -> Self {
Self {
aspect,
source: None,
}
}
#[must_use]
pub const fn aspect(&self) -> DeclaredAspect {
self.aspect
}
#[must_use]
pub fn raw_os_error(&self) -> Option<i32> {
self.source.as_ref().and_then(io::Error::raw_os_error)
}
}
impl fmt::Display for DeclaredError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let what = match self.aspect {
DeclaredAspect::Wow64Redirection => {
"WOW64 filesystem redirection could not be changed (it exists only \
for a 32-bit process on 64-bit Windows)"
}
DeclaredAspect::MemoryPriority => "the thread memory priority could not be read or set",
DeclaredAspect::BackgroundMode => "background processing mode could not be changed",
};
match &self.source {
Some(source) => write!(f, "{what}: {source}"),
None => f.write_str(what),
}
}
}
impl std::error::Error for DeclaredError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
self.source
.as_ref()
.map(|source| source as &(dyn std::error::Error + 'static))
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
pub struct Declared {
pub wow64_redirection: Option<Wow64Redirection>,
pub memory_priority: Option<MemoryPriority>,
pub background_mode: Option<BackgroundMode>,
}
impl Declared {
#[must_use]
pub const fn none() -> Self {
Self {
wow64_redirection: None,
memory_priority: None,
background_mode: None,
}
}
#[must_use]
pub const fn is_empty(&self) -> bool {
self.wow64_redirection.is_none()
&& self.memory_priority.is_none()
&& self.background_mode.is_none()
}
#[must_use]
pub const fn with_wow64_redirection(mut self, value: Wow64Redirection) -> Self {
self.wow64_redirection = Some(value);
self
}
#[must_use]
pub const fn with_memory_priority(mut self, value: MemoryPriority) -> Self {
self.memory_priority = Some(value);
self
}
#[must_use]
pub const fn with_background_mode(mut self, value: BackgroundMode) -> Self {
self.background_mode = Some(value);
self
}
pub fn install(&self) -> Result<DeclaredGuard, DeclaredError> {
let background = match self.background_mode {
Some(mode) => {
mode.install()?;
Some(mode)
}
None => None,
};
let memory = match self.memory_priority {
Some(priority) => {
let previous = MemoryPriority::current();
match previous.and_then(|previous| priority.install().map(|()| previous)) {
Ok(previous) => Some(previous),
Err(error) => {
release_background(background);
return Err(error);
}
}
}
None => None,
};
let redirection = match self.wow64_redirection {
Some(Wow64Redirection::Disabled) => {
let mut old: *mut core::ffi::c_void = std::ptr::null_mut();
let ok = unsafe { Wow64DisableWow64FsRedirection(&mut old) };
if ok == 0 {
let error = DeclaredError::new(DeclaredAspect::Wow64Redirection);
release_memory(memory);
release_background(background);
return Err(error);
}
Some(old)
}
None => None,
};
Ok(DeclaredGuard {
background,
memory,
redirection,
released: false,
})
}
pub fn with_applied<F, T>(&self, operation: F) -> Result<T, DeclaredError>
where
F: FnOnce() -> T,
{
let guard = self.install()?;
let outcome = operation();
guard.release().map(|()| outcome)
}
}
#[must_use = "dropping the guard restores the aspects but discards any failure to do so"]
#[derive(Debug)]
pub struct DeclaredGuard {
background: Option<BackgroundMode>,
memory: Option<MemoryPriority>,
redirection: Option<*mut core::ffi::c_void>,
released: bool,
}
impl DeclaredGuard {
pub fn release(mut self) -> Result<(), DeclaredError> {
self.released = true;
Self::restore(self.background, self.memory, self.redirection)
}
fn restore(
background: Option<BackgroundMode>,
memory: Option<MemoryPriority>,
redirection: Option<*mut core::ffi::c_void>,
) -> Result<(), DeclaredError> {
let mut failure = None;
if let Some(old) = redirection {
let ok = unsafe { Wow64RevertWow64FsRedirection(old) };
if ok == 0 {
failure = Some(DeclaredError::new(DeclaredAspect::Wow64Redirection));
}
}
if let Some(previous) = memory
&& let Err(error) = previous.install()
{
failure = failure.or(Some(error));
}
if let Some(mode) = background
&& let Err(error) = mode.inverse().install()
{
failure = failure.or(Some(error));
}
match failure {
Some(error) => Err(error),
None => Ok(()),
}
}
}
impl Drop for DeclaredGuard {
fn drop(&mut self) {
if !self.released {
let _ = Self::restore(self.background, self.memory, self.redirection);
}
}
}
fn release_background(background: Option<BackgroundMode>) {
if let Some(mode) = background {
let _ = mode.inverse().install();
}
}
fn release_memory(memory: Option<MemoryPriority>) {
if let Some(previous) = memory {
let _ = previous.install();
}
}
#[cfg(test)]
mod tests;