use crate::{GitError, Result};
use std::error::Error as StdError;
use std::fmt;
use std::io::{self, Read};
use std::process::Child;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum StreamControl {
#[default]
Continue,
Stop,
}
impl StreamControl {
#[inline]
pub const fn is_stop(self) -> bool {
matches!(self, Self::Stop)
}
#[inline]
pub const fn is_continue(self) -> bool {
matches!(self, Self::Continue)
}
}
#[derive(Debug, Default)]
pub struct AtomicCancel {
cancelled: AtomicBool,
}
impl AtomicCancel {
#[inline]
pub const fn new() -> Self {
Self {
cancelled: AtomicBool::new(false),
}
}
#[inline]
pub fn cancel(&self) {
self.cancelled.store(true, Ordering::Release);
}
#[inline]
pub fn clear(&self) {
self.cancelled.store(false, Ordering::Release);
}
#[inline]
pub fn is_cancelled(&self) -> bool {
self.cancelled.load(Ordering::Acquire)
}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct CancelFlag<'a> {
source: Option<&'a AtomicCancel>,
}
impl CancelFlag<'static> {
#[inline]
pub const fn never() -> Self {
Self { source: None }
}
}
impl<'a> CancelFlag<'a> {
#[inline]
pub const fn new(source: &'a AtomicCancel) -> Self {
Self {
source: Some(source),
}
}
#[inline]
pub const fn never_dyn() -> CancelFlag<'static> {
CancelFlag::never()
}
#[inline]
pub fn is_cancelled(self) -> bool {
self.source.is_some_and(AtomicCancel::is_cancelled)
}
#[inline]
pub fn check(self) -> Result<()> {
if self.is_cancelled() {
Err(GitError::Cancelled)
} else {
Ok(())
}
}
#[inline]
pub fn control(self) -> StreamControl {
if self.is_cancelled() {
StreamControl::Stop
} else {
StreamControl::Continue
}
}
#[inline]
pub fn as_ref(self) -> CancelFlag<'a> {
self
}
#[inline]
pub fn source(self) -> Option<&'a AtomicCancel> {
self.source
}
}
pub type DynCancelFlag<'a> = CancelFlag<'a>;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct OperationCancelled;
impl fmt::Display for OperationCancelled {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("operation cancelled")
}
}
impl StdError for OperationCancelled {}
#[derive(Debug)]
pub struct CancellableRead<'a, R> {
inner: R,
cancel: CancelFlag<'a>,
}
impl<'a, R> CancellableRead<'a, R> {
#[inline]
pub fn new(inner: R, cancel: CancelFlag<'a>) -> Self {
Self { inner, cancel }
}
#[inline]
pub fn get_ref(&self) -> &R {
&self.inner
}
#[inline]
pub fn get_mut(&mut self) -> &mut R {
&mut self.inner
}
#[inline]
pub fn into_inner(self) -> R {
self.inner
}
#[inline]
pub fn cancel(&self) -> CancelFlag<'a> {
self.cancel
}
}
impl<R: Read> Read for CancellableRead<'_, R> {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
if self.cancel.is_cancelled() {
return Err(cancelled_io_error());
}
self.inner.read(buf)
}
}
#[inline]
pub fn cancelled_io_error() -> io::Error {
io::Error::other(OperationCancelled)
}
#[inline]
pub fn is_cancelled_io(err: &io::Error) -> bool {
err.get_ref()
.is_some_and(|inner| inner.downcast_ref::<OperationCancelled>().is_some())
}
#[inline]
pub fn map_cancel_io(err: io::Error) -> GitError {
if is_cancelled_io(&err) {
GitError::Cancelled
} else {
GitError::from(err)
}
}
#[inline]
pub fn is_cancelled_error(err: &GitError) -> bool {
matches!(err, GitError::Cancelled)
}
#[inline]
pub fn kill_child_if_cancelled(child: &mut Child, cancel: CancelFlag<'_>) {
if cancel.is_cancelled() {
let _ = child.kill();
}
}
#[inline]
pub fn cancel_flag_from_arc(source: &Arc<AtomicCancel>) -> CancelFlag<'_> {
CancelFlag::new(source.as_ref())
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::{Cursor, Read};
#[test]
fn never_flag_is_never_cancelled() {
let flag = CancelFlag::never();
assert!(!flag.is_cancelled());
assert!(flag.check().is_ok());
assert_eq!(flag.control(), StreamControl::Continue);
}
#[test]
fn atomic_cancel_trips_flag() {
let source = AtomicCancel::new();
let flag = CancelFlag::new(&source);
assert!(!flag.is_cancelled());
source.cancel();
assert!(flag.is_cancelled());
assert_eq!(flag.check(), Err(GitError::Cancelled));
assert_eq!(flag.control(), StreamControl::Stop);
source.clear();
assert!(!flag.is_cancelled());
}
#[test]
fn cancellable_read_fails_with_operation_cancelled_not_interrupted() {
let source = AtomicCancel::new();
let data = b"hello world";
let mut reader = CancellableRead::new(Cursor::new(&data[..]), CancelFlag::new(&source));
let mut buf = [0u8; 5];
assert_eq!(reader.read(&mut buf).expect("read"), 5);
source.cancel();
let err = reader.read(&mut buf).expect_err("cancelled");
assert_ne!(
err.kind(),
io::ErrorKind::Interrupted,
"Interrupted is retried by read_exact / pkt-line"
);
assert!(is_cancelled_io(&err));
assert_eq!(map_cancel_io(err), GitError::Cancelled);
}
#[test]
fn read_exact_does_not_spin_on_cancel() {
let source = AtomicCancel::new();
source.cancel();
let mut reader = CancellableRead::new(Cursor::new(&b"abcd"[..]), CancelFlag::new(&source));
let mut buf = [0u8; 4];
let err = reader.read_exact(&mut buf).expect_err("cancelled");
assert!(is_cancelled_io(&err));
}
#[test]
fn never_dyn_alias_matches_never() {
assert!(!CancelFlag::never().is_cancelled());
}
#[test]
fn kill_child_if_cancelled_is_noop_when_not_cancelled() {
let mut child = std::process::Command::new("sleep")
.arg("60")
.stdout(std::process::Stdio::null())
.stderr(std::process::Stdio::null())
.spawn()
.expect("spawn sleep");
let source = AtomicCancel::new();
kill_child_if_cancelled(&mut child, CancelFlag::new(&source));
assert!(child.try_wait().expect("try_wait").is_none());
let _ = child.kill();
let _ = child.wait();
}
#[test]
fn kill_child_if_cancelled_kills_when_flag_set() {
let mut child = std::process::Command::new("sleep")
.arg("60")
.stdout(std::process::Stdio::null())
.stderr(std::process::Stdio::null())
.spawn()
.expect("spawn sleep");
let source = AtomicCancel::new();
source.cancel();
kill_child_if_cancelled(&mut child, CancelFlag::new(&source));
let status = child.wait().expect("wait after kill");
assert!(!status.success());
}
}