/build/source/nativelink-util/src/proto_stream_utils.rs

Source
// Copyright 2024 The NativeLink Authors. All rights reserved.
//
// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    See LICENSE file for details
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use core::fmt::Debug;
use core::mem;
use core::pin::Pin;
use core::task::{Context, Poll};
use std::borrow::Cow;
use std::sync::Arc;

use futures::{Stream, StreamExt};
use nativelink_error::{Error, ResultExt, error_if, make_input_err};
use nativelink_proto::google::bytestream::{ReadResponse, WriteRequest};
use parking_lot::Mutex;
use tonic::{Status, Streaming};

use crate::resource_info::ResourceInfo;

pub struct WriteRequestStreamWrapper<T> {
    pub resource_info: ResourceInfo<'static>,
    pub bytes_received: usize,
    stream: T,
    first_msg: Option<WriteRequest>,
    pub write_finished: bool,
}

impl<T> Debug for WriteRequestStreamWrapper<T> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("WriteRequestStreamWrapper")
            .field("resource_info", &self.resource_info)
            .field("bytes_received", &self.bytes_received)
            .field("first_msg", &self.first_msg)
            .field("write_finished", &self.write_finished)
            .finish()
    }
}

impl<T, E> WriteRequestStreamWrapper<T>
where
    T: Stream<Item = Result<WriteRequest, E>> + Unpin,
    E: Into<Error>,
{
    pub async fn from(mut stream: T) -> Result<Self, Error> {
        let first_msg28 = stream
            .next()
            .await
            .err_tip(|| "Error receiving first message in stream")?0
            .err_tip(|| "Expected WriteRequest struct in stream (from)")?1;

        let resource_info = ResourceInfo::new(&first_msg.resource_name, true)
            .err_tip(|| {0
                format!(
                    "Could not extract resource info from first message of stream: {}",
                    first_msg.resource_name
                )
            })?
            .to_owned();

        Ok(Self {
            resource_info,
            bytes_received: 0,
            stream,
            first_msg: Some(first_msg),
            write_finished: false,
        })
    }

    pub async fn next(&mut self) -> Option<Result<WriteRequest, Error>> {
        futures::future::poll_fn41(|cx| Pin::new(&mut *self).poll_next(cx)).await41
    }

    pub const fn is_first_msg(&self) -> bool {
        self.first_msg.is_some()
    }

    /// Returns whether the first message has `finish_write` set to true.
    /// This indicates a single-shot upload where all data is in one message.
    pub fn is_first_msg_complete(&self) -> bool {
        self.first_msg.as_ref().is_some_and(|msg| msg.finish_write)
    }

    fn enforce_wire_size_matches_digest_size(&self) -> bool {
        matches!(
            self.resource_info.compressor.as_deref(),
            None | Some("identity")
        )
    }
}

impl<T, E> Stream for WriteRequestStreamWrapper<T>
where
    E: Into<Error>,
    T: Stream<Item = Result<WriteRequest, E>> + Unpin,
{
    type Item = Result<WriteRequest, Error>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        // If the stream said that the previous message was the last one, then
        // return a stream EOF (i.e. None).
        if self.write_finished {
            if self.enforce_wire_size_matches_digest_size() {
                error_if!(
                    self.bytes_received != self.resource_info.expected_size,
                    "Did not send enough data. Expected {}, but so far received {}",
                    self.resource_info.expected_size,
                    self.bytes_received
                );
            }
            return Poll::Ready(None);
        }

        // Gets the next message, this is either the cached first or a
        // subsequent message from the wrapped Stream.
        let maybe_message47 = if let Some(first_msg27) = self.first_msg.take() {
            Ok(first_msg)
        } else {
            match Pin::new(&mut self.stream).poll_next(cx) {
                Poll::Pending => return Poll::Pending,
                Poll::Ready(Some(maybe_message)) => maybe_message
                    .err_tip(|| format!0("Stream error at byte {}", self.bytes_received0)),
                Poll::Ready(None) => Err(make_input_err!(
                    "Expected WriteRequest struct in stream (got None)"
                )),
            }
        };

        // If we successfully got a message, update our internal state with the
        // message meta data.
        Poll::Ready(Some(maybe_message.and_then(|message| {44
            self.write_finished = message.finish_write;
            self.bytes_received += message.data.len();

            // Check that we haven't read past the expected end.
            if self.enforce_wire_size_matches_digest_size()
                && self.bytes_received > self.resource_info.expected_size
            {
                Err(make_input_err!(
                    "Sent too much data. Expected {}, but so far received {}",
                    self.resource_info.expected_size,
                    self.bytes_received
                ))
            } else {
                Ok(message)
            }
        })))
    }
}

/// Represents the state of the first response in a `FirstStream`.
#[derive(Debug)]
pub enum FirstResponseState {
    /// Contains an optional first response that hasn't been consumed yet.
    /// A `None` value indicates the first response was EOF.
    Unused(Option<ReadResponse>),
    /// Indicates the first response has been consumed and future reads should
    /// come from the underlying stream.
    Used,
}

/// This provides a buffer for the first response from GrpcStore.read in order
/// to allow the first read to occur within the retry loop.  That means that if
/// the connection establishes fine, but reading the first byte of the file
/// fails we have the ability to retry before returning to the caller.
#[derive(Debug)]
pub struct FirstStream {
    /// The current state of the first response. When in the `Unused` state,
    /// contains an optional response which could be `None` or an EOF.
    /// Once consumed, transitions to the `Used` state.
    state: FirstResponseState,
    /// The stream to get responses from after the first response is consumed.
    stream: Streaming<ReadResponse>,
}

impl FirstStream {
    /// Creates a new `FirstStream` with the given first response and underlying
    /// stream.
    pub const fn new(
        first_response: Option<ReadResponse>,
        stream: Streaming<ReadResponse>,
    ) -> Self {
        Self {
            state: FirstResponseState::Unused(first_response),
            stream,
        }
    }
}

impl Stream for FirstStream {
    type Item = Result<ReadResponse, Status>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        match mem::replace(&mut self.state, FirstResponseState::Used) {
            FirstResponseState::Unused(first_response) => Poll::Ready(first_response.map(Ok)),
            FirstResponseState::Used => Pin::new(&mut self.stream).poll_next(cx),
        }
    }
}

/// This structure wraps all of the information required to perform a write
/// request on the `GrpcStore`.  It stores the last message retrieved which allows
/// the write to resume since the UUID allows upload resume at the server.
#[derive(Debug)]
pub struct WriteState<T, E>
where
    T: Stream<Item = Result<WriteRequest, E>> + Unpin + Send + 'static,
    E: Into<Error> + 'static,
{
    instance_name: String,
    read_stream_error: Option<Error>,
    read_stream: WriteRequestStreamWrapper<T>,
    // Tonic doesn't appear to report an error until it has taken two messages,
    // therefore we are required to buffer the last two messages.
    cached_messages: [Option<WriteRequest>; 2],
    // When resuming after an error, the previous messages are cloned into this
    // queue upfront to allow them to be served back.
    resume_queue: [Option<WriteRequest>; 2],
    // An optimisation to avoid having to manage resume_queue when it's empty.
    is_resumed: bool,
}

impl<T, E> WriteState<T, E>
where
    T: Stream<Item = Result<WriteRequest, E>> + Unpin + Send + 'static,
    E: Into<Error> + 'static,
{
    pub const fn new(instance_name: String, read_stream: WriteRequestStreamWrapper<T>) -> Self {
        Self {
            instance_name,
            read_stream_error: None,
            read_stream,
            cached_messages: [None, None],
            resume_queue: [None, None],
            is_resumed: false,
        }
    }

    fn push_message(&mut self, message: WriteRequest) {
        self.cached_messages.swap(0, 1);
        self.cached_messages[0] = Some(message);
    }

    const fn resumed_message(&mut self) -> Option<WriteRequest> {
        if self.is_resumed {
            // The resume_queue is a circular buffer, that we have to shift,
            // since its only got two elements its a trivial swap.
            self.resume_queue.swap(0, 1);
            let message = self.resume_queue[0].take();
            if message.is_none() {
                self.is_resumed = false;
            }
            message
        } else {
            None
        }
    }

    pub const fn can_resume(&self) -> bool {
        self.read_stream_error.is_none()
            && (self.cached_messages[0].is_some() || self.read_stream.is_first_msg())
    }

    pub fn resume(&mut self) {
        self.resume_queue.clone_from(&self.cached_messages);
        self.is_resumed = true;
    }

    pub const fn take_read_stream_error(&mut self) -> Option<Error> {
        self.read_stream_error.take()
    }
}

/// A wrapper around `WriteState` to allow it to be reclaimed from the underlying
/// write call in the case of failure.
#[derive(Debug)]
pub struct WriteStateWrapper<T, E>
where
    T: Stream<Item = Result<WriteRequest, E>> + Unpin + Send + 'static,
    E: Into<Error> + 'static,
{
    shared_state: Arc<Mutex<WriteState<T, E>>>,
}

impl<T, E> WriteStateWrapper<T, E>
where
    T: Stream<Item = Result<WriteRequest, E>> + Unpin + Send + 'static,
    E: Into<Error> + 'static,
{
    pub const fn new(shared_state: Arc<Mutex<WriteState<T, E>>>) -> Self {
        Self { shared_state }
    }
}

impl<T, E> Stream for WriteStateWrapper<T, E>
where
    T: Stream<Item = Result<WriteRequest, E>> + Unpin + Send + 'static,
    E: Into<Error> + 'static,
{
    type Item = WriteRequest;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        const IS_UPLOAD_TRUE: bool = true;

        // This should be an uncontended lock since write was called.
        let mut local_state = self.shared_state.lock();
        // If this is the first or second call after a failure and we have
        // cached messages, then use the cached write requests.
        let cached_message = local_state.resumed_message();
        if cached_message.is_some() {
            return Poll::Ready(cached_message);
        }
        // Read a new write request from the downstream.
        let Poll::Ready(maybe_message) = Pin::new(&mut local_state.read_stream).poll_next(cx)
        else {
            return Poll::Pending;
        };
        // Update the instance name in the write request and forward it on.
        let result = match maybe_message7 {
            Some(Ok(mut message)) => {
                if !message.resource_name.is_empty() {
                    // Replace the instance name in the resource name if it is
                    // different from the instance name in the write state.
                    match ResourceInfo::new(&message.resource_name, IS_UPLOAD_TRUE) {
                        Ok(mut resource_name) => {
                            if resource_name.instance_name != local_state.instance_name {
                                resource_name.instance_name =
                                    Cow::Borrowed(&local_state.instance_name);
                                message.resource_name = resource_name.to_string(IS_UPLOAD_TRUE);
                            }
                        }
                        Err(err) => {
                            local_state.read_stream_error = Some(err);
                            return Poll::Ready(None);
                        }
                    }
                }
                // Cache the last request in case there is an error to allow
                // the upload to be resumed.
                local_state.push_message(message.clone());
                Some(message)
            }
            Some(Err(err)) => {
                local_state.read_stream_error = Some(err);
                None
            }
            None => None,
        };
        Poll::Ready(result)
    }
}

Coverage Report

Created: 2026-07-17 15:58

Line	Count	Source
1		// Copyright 2024 The NativeLink Authors. All rights reserved.
2		//
3		// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License");
4		// you may not use this file except in compliance with the License.
5		// You may obtain a copy of the License at
6		//
7		// See LICENSE file for details
8		//
9		// Unless required by applicable law or agreed to in writing, software
10		// distributed under the License is distributed on an "AS IS" BASIS,
11		// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12		// See the License for the specific language governing permissions and
13		// limitations under the License.
14
15		use core::fmt::Debug;
16		use core::mem;
17		use core::pin::Pin;
18		use core::task::{Context, Poll};
19		use std::borrow::Cow;
20		use std::sync::Arc;
21
22		use futures::{Stream, StreamExt};
23		use nativelink_error::{Error, ResultExt, error_if, make_input_err};
24		use nativelink_proto::google::bytestream::{ReadResponse, WriteRequest};
25		use parking_lot::Mutex;
26		use tonic::{Status, Streaming};
27
28		use crate::resource_info::ResourceInfo;
29
30		pub struct WriteRequestStreamWrapper<T> {
31		pub resource_info: ResourceInfo<'static>,
32		pub bytes_received: usize,
33		stream: T,
34		first_msg: Option<WriteRequest>,
35		pub write_finished: bool,
36		}
37
38		impl<T> Debug for WriteRequestStreamWrapper<T> {
39	13	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
40	13	f.debug_struct("WriteRequestStreamWrapper")
41	13	.field("resource_info", &self.resource_info)
42	13	.field("bytes_received", &self.bytes_received)
43	13	.field("first_msg", &self.first_msg)
44	13	.field("write_finished", &self.write_finished)
45	13	.finish()
46	13	}
47		}
48
49		impl<T, E> WriteRequestStreamWrapper<T>
50		where
51		T: Stream<Item = Result<WriteRequest, E>> + Unpin,
52		E: Into<Error>,
53		{
54	29	pub async fn from(mut stream: T) -> Result<Self, Error> {
55	29	let first_msg28 = stream
56	29	.next()
57	29	.await
58	29	.err_tip(\|\| "Error receiving first message in stream") ?0
59	29	.err_tip(\|\| "Expected WriteRequest struct in stream (from)") ?1 ;
60
61	28	let resource_info = ResourceInfo::new(&first_msg.resource_name, true)
62	28	.err_tip(\|\| {0
63	0	format!(
64		"Could not extract resource info from first message of stream: {}",
65		first_msg.resource_name
66		)
67	0	})?
68	28	.to_owned();
69
70	28	Ok(Self {
71	28	resource_info,
72	28	bytes_received: 0,
73	28	stream,
74	28	first_msg: Some(first_msg),
75	28	write_finished: false,
76	28	})
77	29	}
78
79	41	pub async fn next(&mut self) -> Option<Result<WriteRequest, Error>> {
80	143	futures::future::poll_fn41 (\|cx\| Pin::new(&mut *self).poll_next(cx)). await41
81	40	}
82
83	0	pub const fn is_first_msg(&self) -> bool {
84	0	self.first_msg.is_some()
85	0	}
86
87		/// Returns whether the first message has `finish_write` set to true.
88		/// This indicates a single-shot upload where all data is in one message.
89	18	pub fn is_first_msg_complete(&self) -> bool {
90	18	self.first_msg.as_ref().is_some_and(\|msg\| msg.finish_write)
91	18	}
92
93	47	fn enforce_wire_size_matches_digest_size(&self) -> bool {
94	12	matches!(
95	47	self.resource_info.compressor.as_deref(),
96	13	None \| Some("identity")
97		)
98	47	}
99		}
100
101		impl<T, E> Stream for WriteRequestStreamWrapper<T>
102		where
103		E: Into<Error>,
104		T: Stream<Item = Result<WriteRequest, E>> + Unpin,
105		{
106		type Item = Result<WriteRequest, Error>;
107
108	153	fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
109		// If the stream said that the previous message was the last one, then
110		// return a stream EOF (i.e. None).
111	153	if self.write_finished {
112	3	if self.enforce_wire_size_matches_digest_size() {
113	0	error_if!(
114	3	self.bytes_received != self.resource_info.expected_size,
115		"Did not send enough data. Expected {}, but so far received {}",
116	0	self.resource_info.expected_size,
117	0	self.bytes_received
118		);
119	0	}
120	3	return Poll::Ready(None);
121	150	}
122
123		// Gets the next message, this is either the cached first or a
124		// subsequent message from the wrapped Stream.
125	150	let maybe_message47 = if let Some( first_msg27 ) = self.first_msg.take() {
126	27	Ok(first_msg)
127		} else {
128	123	match Pin::new(&mut self.stream).poll_next(cx) {
129	103	Poll::Pending => return Poll::Pending,
130	17	Poll::Ready(Some(maybe_message)) => maybe_message
131	17	.err_tip(\|\| format!0 ("Stream error at byte {}", self.bytes_received0 )),
132	3	Poll::Ready(None) => Err(make_input_err!(
133	3	"Expected WriteRequest struct in stream (got None)"
134	3	)),
135		}
136		};
137
138		// If we successfully got a message, update our internal state with the
139		// message meta data.
140	47	Poll::Ready(Some(maybe_message.and_then(\|message\| {44
141	44	self.write_finished = message.finish_write;
142	44	self.bytes_received += message.data.len();
143
144		// Check that we haven't read past the expected end.
145	44	if self.enforce_wire_size_matches_digest_size()
146	32	&& self.bytes_received > self.resource_info.expected_size
147		{
148	2	Err(make_input_err!(
149	2	"Sent too much data. Expected {}, but so far received {}",
150	2	self.resource_info.expected_size,
151	2	self.bytes_received
152	2	))
153		} else {
154	42	Ok(message)
155		}
156	44	})))
157	153	}
158		}
159
160		/// Represents the state of the first response in a `FirstStream`.
161		#[derive(Debug)]
162		pub enum FirstResponseState {
163		/// Contains an optional first response that hasn't been consumed yet.
164		/// A `None` value indicates the first response was EOF.
165		Unused(Option<ReadResponse>),
166		/// Indicates the first response has been consumed and future reads should
167		/// come from the underlying stream.
168		Used,
169		}
170
171		/// This provides a buffer for the first response from GrpcStore.read in order
172		/// to allow the first read to occur within the retry loop. That means that if
173		/// the connection establishes fine, but reading the first byte of the file
174		/// fails we have the ability to retry before returning to the caller.
175		#[derive(Debug)]
176		pub struct FirstStream {
177		/// The current state of the first response. When in the `Unused` state,
178		/// contains an optional response which could be `None` or an EOF.
179		/// Once consumed, transitions to the `Used` state.
180		state: FirstResponseState,
181		/// The stream to get responses from after the first response is consumed.
182		stream: Streaming<ReadResponse>,
183		}
184
185		impl FirstStream {
186		/// Creates a new `FirstStream` with the given first response and underlying
187		/// stream.
188	6	pub const fn new(
189	6	first_response: Option<ReadResponse>,
190	6	stream: Streaming<ReadResponse>,
191	6	) -> Self {
192	6	Self {
193	6	state: FirstResponseState::Unused(first_response),
194	6	stream,
195	6	}
196	6	}
197		}
198
199		impl Stream for FirstStream {
200		type Item = Result<ReadResponse, Status>;
201
202	14	fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
203	14	match mem::replace(&mut self.state, FirstResponseState::Used) {
204	6	FirstResponseState::Unused(first_response) => Poll::Ready(first_response.map(Ok)),
205	8	FirstResponseState::Used => Pin::new(&mut self.stream).poll_next(cx),
206		}
207	14	}
208		}
209
210		/// This structure wraps all of the information required to perform a write
211		/// request on the `GrpcStore`. It stores the last message retrieved which allows
212		/// the write to resume since the UUID allows upload resume at the server.
213		#[derive(Debug)]
214		pub struct WriteState<T, E>
215		where
216		T: Stream<Item = Result<WriteRequest, E>> + Unpin + Send + 'static,
217		E: Into<Error> + 'static,
218		{
219		instance_name: String,
220		read_stream_error: Option<Error>,
221		read_stream: WriteRequestStreamWrapper<T>,
222		// Tonic doesn't appear to report an error until it has taken two messages,
223		// therefore we are required to buffer the last two messages.
224		cached_messages: [Option<WriteRequest>; 2],
225		// When resuming after an error, the previous messages are cloned into this
226		// queue upfront to allow them to be served back.
227		resume_queue: [Option<WriteRequest>; 2],
228		// An optimisation to avoid having to manage resume_queue when it's empty.
229		is_resumed: bool,
230		}
231
232		impl<T, E> WriteState<T, E>
233		where
234		T: Stream<Item = Result<WriteRequest, E>> + Unpin + Send + 'static,
235		E: Into<Error> + 'static,
236		{
237	3	pub const fn new(instance_name: String, read_stream: WriteRequestStreamWrapper<T>) -> Self {
238	3	Self {
239	3	instance_name,
240	3	read_stream_error: None,
241	3	read_stream,
242	3	cached_messages: [None, None],
243	3	resume_queue: [None, None],
244	3	is_resumed: false,
245	3	}
246	3	}
247
248	7	fn push_message(&mut self, message: WriteRequest) {
249	7	self.cached_messages.swap(0, 1);
250	7	self.cached_messages[0] = Some(message);
251	7	}
252
253	10	const fn resumed_message(&mut self) -> Option<WriteRequest> {
254	10	if self.is_resumed {
255		// The resume_queue is a circular buffer, that we have to shift,
256		// since its only got two elements its a trivial swap.
257	0	self.resume_queue.swap(0, 1);
258	0	let message = self.resume_queue[0].take();
259	0	if message.is_none() {
260	0	self.is_resumed = false;
261	0	}
262	0	message
263		} else {
264	10	None
265		}
266	10	}
267
268	0	pub const fn can_resume(&self) -> bool {
269	0	self.read_stream_error.is_none()
270	0	&& (self.cached_messages[0].is_some() \|\| self.read_stream.is_first_msg())
271	0	}
272
273	0	pub fn resume(&mut self) {
274	0	self.resume_queue.clone_from(&self.cached_messages);
275	0	self.is_resumed = true;
276	0	}
277
278	3	pub const fn take_read_stream_error(&mut self) -> Option<Error> {
279	3	self.read_stream_error.take()
280	3	}
281		}
282
283		/// A wrapper around `WriteState` to allow it to be reclaimed from the underlying
284		/// write call in the case of failure.
285		#[derive(Debug)]
286		pub struct WriteStateWrapper<T, E>
287		where
288		T: Stream<Item = Result<WriteRequest, E>> + Unpin + Send + 'static,
289		E: Into<Error> + 'static,
290		{
291		shared_state: Arc<Mutex<WriteState<T, E>>>,
292		}
293
294		impl<T, E> WriteStateWrapper<T, E>
295		where
296		T: Stream<Item = Result<WriteRequest, E>> + Unpin + Send + 'static,
297		E: Into<Error> + 'static,
298		{
299	3	pub const fn new(shared_state: Arc<Mutex<WriteState<T, E>>>) -> Self {
300	3	Self { shared_state }
301	3	}
302		}
303
304		impl<T, E> Stream for WriteStateWrapper<T, E>
305		where
306		T: Stream<Item = Result<WriteRequest, E>> + Unpin + Send + 'static,
307		E: Into<Error> + 'static,
308		{
309		type Item = WriteRequest;
310
311	10	fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
312		const IS_UPLOAD_TRUE: bool = true;
313
314		// This should be an uncontended lock since write was called.
315	10	let mut local_state = self.shared_state.lock();
316		// If this is the first or second call after a failure and we have
317		// cached messages, then use the cached write requests.
318	10	let cached_message = local_state.resumed_message();
319	10	if cached_message.is_some() {
320	0	return Poll::Ready(cached_message);
321	10	}
322		// Read a new write request from the downstream.
323	10	let Poll::Ready(maybe_message) = Pin::new(&mut local_state.read_stream).poll_next(cx)
324		else {
325	0	return Poll::Pending;
326		};
327		// Update the instance name in the write request and forward it on.
328	10	let result = match maybe_message7 {
329	7	Some(Ok(mut message)) => {
330	7	if !message.resource_name.is_empty() {
331		// Replace the instance name in the resource name if it is
332		// different from the instance name in the write state.
333	5	match ResourceInfo::new(&message.resource_name, IS_UPLOAD_TRUE) {
334	5	Ok(mut resource_name) => {
335	5	if resource_name.instance_name != local_state.instance_name {
336	0	resource_name.instance_name =
337	0	Cow::Borrowed(&local_state.instance_name);
338	0	message.resource_name = resource_name.to_string(IS_UPLOAD_TRUE);
339	5	}
340		}
341	0	Err(err) => {
342	0	local_state.read_stream_error = Some(err);
343	0	return Poll::Ready(None);
344		}
345		}
346	2	}
347		// Cache the last request in case there is an error to allow
348		// the upload to be resumed.
349	7	local_state.push_message(message.clone());
350	7	Some(message)
351		}
352	0	Some(Err(err)) => {
353	0	local_state.read_stream_error = Some(err);
354	0	None
355		}
356	3	None => None,
357		};
358	10	Poll::Ready(result)
359	10	}
360		}