diff options
Diffstat (limited to 'third_party/rust/h2/src/proto/go_away.rs')
-rw-r--r-- | third_party/rust/h2/src/proto/go_away.rs | 154 |
1 files changed, 154 insertions, 0 deletions
diff --git a/third_party/rust/h2/src/proto/go_away.rs b/third_party/rust/h2/src/proto/go_away.rs new file mode 100644 index 0000000000..d52252cd74 --- /dev/null +++ b/third_party/rust/h2/src/proto/go_away.rs @@ -0,0 +1,154 @@ +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 + } +} |