/src/h2/src/proto/go_away.rs

Source
use crate::codec::Codec;
use crate::frame::{self, Reason, StreamId};

use bytes::Buf;
use std::io;
use std::task::{Context, Poll};
use tokio::io::AsyncWrite;

/// Manages our sending of GOAWAY frames.
#[derive(Debug)]
pub(super) struct GoAway {
    /// Whether the connection should close now, or wait until idle.
    close_now: bool,
    /// Records if we've sent any GOAWAY before.
    going_away: Option<GoingAway>,
    /// Whether the user started the GOAWAY by calling `abrupt_shutdown`.
    is_user_initiated: bool,
    /// A GOAWAY frame that must be buffered in the Codec immediately.
    pending: Option<frame::GoAway>,
}

/// Keeps a memory of any GOAWAY frames we've sent before.
///
/// This looks very similar to a `frame::GoAway`, but is a separate type. Why?
/// Mostly for documentation purposes. This type is to record status. If it
/// were a `frame::GoAway`, it might appear like we eventually wanted to
/// serialize it. We **only** want to be able to look up these fields at a
/// later time.
#[derive(Debug)]
pub(crate) struct GoingAway {
    /// Stores the highest stream ID of a GOAWAY that has been sent.
    ///
    /// It's illegal to send a subsequent GOAWAY with a higher ID.
    last_processed_id: StreamId,

    /// Records the error code of any GOAWAY frame sent.
    reason: Reason,
}

impl GoAway {
    pub fn new() -> Self {
        GoAway {
            close_now: false,
            going_away: None,
            is_user_initiated: false,
            pending: None,
        }
    }

    /// Enqueue a GOAWAY frame to be written.
    ///
    /// The connection is expected to continue to run until idle.
    pub fn go_away(&mut self, f: frame::GoAway) {
        if let Some(ref going_away) = self.going_away {
            assert!(
                f.last_stream_id() <= going_away.last_processed_id,
                "GOAWAY stream IDs shouldn't be higher; \
                 last_processed_id = {:?}, f.last_stream_id() = {:?}",
                going_away.last_processed_id,
                f.last_stream_id(),
            );
        }

        self.going_away = Some(GoingAway {
            last_processed_id: f.last_stream_id(),
            reason: f.reason(),
        });
        self.pending = Some(f);
    }

    pub fn go_away_now(&mut self, f: frame::GoAway) {
        self.close_now = true;
        if let Some(ref going_away) = self.going_away {
            // Prevent sending the same GOAWAY twice.
            if going_away.last_processed_id == f.last_stream_id() && going_away.reason == f.reason()
            {
                return;
            }
        }
        self.go_away(f);
    }

    pub fn go_away_from_user(&mut self, f: frame::GoAway) {
        self.is_user_initiated = true;
        self.go_away_now(f);
    }

    /// Return if a GOAWAY has ever been scheduled.
    pub fn is_going_away(&self) -> bool {
        self.going_away.is_some()
    }

    pub fn is_user_initiated(&self) -> bool {
        self.is_user_initiated
    }

    /// Returns the going away info, if any.
    pub fn going_away(&self) -> Option<&GoingAway> {
        self.going_away.as_ref()
    }

    /// Returns if the connection should close now, or wait until idle.
    pub fn should_close_now(&self) -> bool {
        self.pending.is_none() && self.close_now
    }

    /// Returns if the connection should be closed when idle.
    pub fn should_close_on_idle(&self) -> bool {
        !self.close_now
            && self
                .going_away
                .as_ref()
                .map(|g| g.last_processed_id != StreamId::MAX)
                .unwrap_or(false)
    }

    /// Try to write a pending GOAWAY frame to the buffer.
    ///
    /// If a frame is written, the `Reason` of the GOAWAY is returned.
    pub fn send_pending_go_away<T, B>(
        &mut self,
        cx: &mut Context,
        dst: &mut Codec<T, B>,
    ) -> Poll<Option<io::Result<Reason>>>
    where
        T: AsyncWrite + Unpin,
        B: Buf,
    {
        if let Some(frame) = self.pending.take() {
            if !dst.poll_ready(cx)?.is_ready() {
                self.pending = Some(frame);
                return Poll::Pending;
            }

            let reason = frame.reason();
            dst.buffer(frame.into()).expect("invalid GOAWAY frame");

            return Poll::Ready(Some(Ok(reason)));
        } else if self.should_close_now() {
            return match self.going_away().map(|going_away| going_away.reason) {
                Some(reason) => Poll::Ready(Some(Ok(reason))),
                None => Poll::Ready(None),
            };
        }

        Poll::Ready(None)
    }
}

impl GoingAway {
    pub(crate) fn reason(&self) -> Reason {
        self.reason
    }
}

Line	Count	Source
1		use crate::codec::Codec;
2		use crate::frame::{self, Reason, StreamId};
3
4		use bytes::Buf;
5		use std::io;
6		use std::task::{Context, Poll};
7		use tokio::io::AsyncWrite;
8
9		/// Manages our sending of GOAWAY frames.
10		#[derive(Debug)]
11		pub(super) struct GoAway {
12		/// Whether the connection should close now, or wait until idle.
13		close_now: bool,
14		/// Records if we've sent any GOAWAY before.
15		going_away: Option<GoingAway>,
16		/// Whether the user started the GOAWAY by calling `abrupt_shutdown`.
17		is_user_initiated: bool,
18		/// A GOAWAY frame that must be buffered in the Codec immediately.
19		pending: Option<frame::GoAway>,
20		}
21
22		/// Keeps a memory of any GOAWAY frames we've sent before.
23		///
24		/// This looks very similar to a `frame::GoAway`, but is a separate type. Why?
25		/// Mostly for documentation purposes. This type is to record status. If it
26		/// were a `frame::GoAway`, it might appear like we eventually wanted to
27		/// serialize it. We only want to be able to look up these fields at a
28		/// later time.
29		#[derive(Debug)]
30		pub(crate) struct GoingAway {
31		/// Stores the highest stream ID of a GOAWAY that has been sent.
32		///
33		/// It's illegal to send a subsequent GOAWAY with a higher ID.
34		last_processed_id: StreamId,
35
36		/// Records the error code of any GOAWAY frame sent.
37		reason: Reason,
38		}
39
40		impl GoAway {
41	13.2k	pub fn new() -> Self {
42	13.2k	GoAway {
43	13.2k	close_now: false,
44	13.2k	going_away: None,
45	13.2k	is_user_initiated: false,
46	13.2k	pending: None,
47	13.2k	}
48	13.2k	}
49
50		/// Enqueue a GOAWAY frame to be written.
51		///
52		/// The connection is expected to continue to run until idle.
53	6.23k	pub fn go_away(&mut self, f: frame::GoAway) {
54	6.23k	if let Some(ref going_away) = self.going_away {
55	10	assert!(
56	10	f.last_stream_id() <= going_away.last_processed_id,
57	0	"GOAWAY stream IDs shouldn't be higher; \
58	0	last_processed_id = {:?}, f.last_stream_id() = {:?}",
59		going_away.last_processed_id,
60	0	f.last_stream_id(),
61		);
62	6.22k	}
63
64	6.23k	self.going_away = Some(GoingAway {
65	6.23k	last_processed_id: f.last_stream_id(),
66	6.23k	reason: f.reason(),
67	6.23k	});
68	6.23k	self.pending = Some(f);
69	6.23k	}
70
71	6.23k	pub fn go_away_now(&mut self, f: frame::GoAway) {
72	6.23k	self.close_now = true;
73	6.23k	if let Some(ref going_away) = self.going_away {
74		// Prevent sending the same GOAWAY twice.
75	10	if going_away.last_processed_id == f.last_stream_id() && going_away.reason == f.reason()
76		{
77	0	return;
78	10	}
79	6.22k	}
80	6.23k	self.go_away(f);
81	6.23k	}
82
83	0	pub fn go_away_from_user(&mut self, f: frame::GoAway) {
84	0	self.is_user_initiated = true;
85	0	self.go_away_now(f);
86	0	}
87
88		/// Return if a GOAWAY has ever been scheduled.
89	0	pub fn is_going_away(&self) -> bool {
90	0	self.going_away.is_some()
91	0	}
92
93	6.04k	pub fn is_user_initiated(&self) -> bool {
94	6.04k	self.is_user_initiated
95	6.04k	}
96
97		/// Returns the going away info, if any.
98	12.2k	pub fn going_away(&self) -> Option<&GoingAway> {
99	12.2k	self.going_away.as_ref()
100	12.2k	}
101
102		/// Returns if the connection should close now, or wait until idle.
103	911k	pub fn should_close_now(&self) -> bool {
104	911k	self.pending.is_none() && self.close_now
105	911k	}
106
107		/// Returns if the connection should be closed when idle.
108	5.58k	pub fn should_close_on_idle(&self) -> bool {
109	5.58k	!self.close_now
110	5.56k	&& self
111	5.56k	.going_away
112	5.56k	.as_ref()
113	5.56k	.map(\|g\| g.last_processed_id != StreamId::MAX)
114	5.56k	.unwrap_or(false)
115	5.58k	}
116
117		/// Try to write a pending GOAWAY frame to the buffer.
118		///
119		/// If a frame is written, the `Reason` of the GOAWAY is returned.
120	912k	pub fn send_pending_go_away<T, B>(
121	912k	&mut self,
122	912k	cx: &mut Context,
123	912k	dst: &mut Codec<T, B>,
124	912k	) -> Poll<Option<io::Result<Reason>>>
125	912k	where
126	912k	T: AsyncWrite + Unpin,
127	912k	B: Buf,
128		{
129	912k	if let Some(frame) = self.pending.take() {
130	6.30k	if !dst.poll_ready(cx)?.is_ready() {
131	137	self.pending = Some(frame);
132	137	return Poll::Pending;
133	6.04k	}
134
135	6.04k	let reason = frame.reason();
136	6.04k	dst.buffer(frame.into()).expect("invalid GOAWAY frame");
137
138	6.04k	return Poll::Ready(Some(Ok(reason)));
139	905k	} else if self.should_close_now() {
140	0	return match self.going_away().map(\|going_away\| going_away.reason) {
141	0	Some(reason) => Poll::Ready(Some(Ok(reason))),
142	0	None => Poll::Ready(None),
143		};
144	905k	}
145
146	905k	Poll::Ready(None)
147	912k	}
148		}
149
150		impl GoingAway {
151	6.04k	pub(crate) fn reason(&self) -> Reason {
152	6.04k	self.reason
153	6.04k	}
154		}

Coverage Report

Created: 2025-10-13 06:32