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

Source
// Copyright 2024 The NativeLink Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// 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>,
    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_msg16 = stream
            .next()
            .await
            .err_tip(|| "Error receiving first message in stream")?0
            .err_tip(|| "Expected WriteRequest struct in stream")?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_fn25(|cx| Pin::new(&mut *self).poll_next(cx)).await25
    }

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

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 {
            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_message28 = if let Some(first_msg16) = 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!("Stream error at byte {}"0, self.bytes_received0)),
                Poll::Ready(None) => Err(make_input_err!("Expected WriteRequest struct in stream")),
            }
        };

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

            // Check that we haven't read past the expected end.
            if 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_message3 {
            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: 2025-06-24 08:57

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