Skip to main content

chorus_client/
grpc.rs

1use std::collections::{HashMap, HashSet};
2use std::sync::Arc;
3
4use arc_swap::{ArcSwap, ArcSwapOption};
5use async_trait::async_trait;
6use bytes::Bytes;
7use googleapis_tonic_google_storage_v2::google::storage::v2::{
8    bidi_write_object_request, bidi_write_object_response, storage_client::StorageClient,
9    write_object_request, write_object_response, AppendObjectSpec, BidiWriteHandle,
10    BidiWriteObjectRedirectedError, BidiWriteObjectRequest, BidiWriteObjectResponse,
11    ChecksummedData, DeleteObjectRequest, GetObjectRequest, ListObjectsRequest, Object,
12    ReadObjectRequest, UpdateObjectRequest, WriteObjectRequest, WriteObjectSpec,
13};
14use prost::Message as _;
15use tokio::sync::{mpsc, watch, Mutex as SessionMutex};
16use tokio_stream::wrappers::ReceiverStream;
17use tonic::metadata::MetadataValue;
18use tonic::transport::{Channel, ClientTlsConfig, Endpoint};
19use tonic::{Code, Request, Status, Streaming};
20
21use crate::auth::BearerAuth;
22use crate::error::Error;
23use crate::transport::{
24    AppendToken, LaneDurableChange, ListedObject, PackedAppend, PackedAppendMessage, Replica,
25    ReplicaFactory, ReplicaSnapshot, TransportCode, TransportError,
26};
27
28/// Cached zonal write routing token, learned from a `BidiWriteObjectRedirectedError`
29/// and replayed in `x-goog-request-params` to land on the bucket's location.
30/// Shared across every replica produced by one factory.
31type RoutingToken = Arc<ArcSwap<Option<Arc<String>>>>;
32
33#[derive(Clone)]
34struct GrpcReplica {
35    zone: usize,
36    bucket: String,
37    object: String,
38    auth: Option<BearerAuth>,
39    client: StorageClient<Channel>,
40    routing_token: RoutingToken,
41    /// The lane's persistent append session. The live service rate limits
42    /// per-object *mutations*, and every fresh append open is one — so a
43    /// lane opens its session once (at takeover) and streams every window
44    /// through it as continuations. Flushes are not rate limited.
45    /// Clones share the session; the hot send and
46    /// progress paths load its immutable handle without taking the replacement
47    /// mutex.
48    session: Arc<SessionSlot>,
49}
50
51#[derive(Clone)]
52/// Reusable Google Storage v2 channel, authentication handle, and zonal routing
53/// token for one Rapid bucket.
54pub struct GrpcReplicaFactory {
55    zone: usize,
56    bucket: String,
57    auth: Option<BearerAuth>,
58    client: StorageClient<Channel>,
59    routing_token: RoutingToken,
60}
61
62/// A live appendable write stream: an ordered request sender plus a
63/// background reader translating flush acknowledgments into a watchable
64/// durable tail. Send-ahead lanes write through `tx` while waiting on
65/// `state` — sends never block on acknowledgments.
66struct AppendSession {
67    handle: Arc<AppendSessionHandle>,
68    reader: tokio::task::JoinHandle<()>,
69}
70
71struct AppendSessionHandle {
72    tx: mpsc::Sender<BidiWriteObjectRequest>,
73    state: tokio::sync::watch::Receiver<LaneProgress>,
74}
75
76struct RedirectAwareStream {
77    tx: Option<mpsc::Sender<BidiWriteObjectRequest>>,
78    responses: Streaming<BidiWriteObjectResponse>,
79    first: BidiWriteObjectResponse,
80    attempt: u32,
81}
82
83struct SessionWait {
84    timeout: Option<std::time::Duration>,
85    clear_on_ready: bool,
86    inspect_before_error: bool,
87    reader_ended: &'static str,
88    stalled: &'static str,
89}
90
91struct SessionSlot {
92    current: ArcSwapOption<AppendSessionHandle>,
93    owned: SessionMutex<Option<AppendSession>>,
94    shutdown: watch::Sender<bool>,
95}
96
97impl SessionSlot {
98    fn new() -> Self {
99        let (shutdown, _) = watch::channel(false);
100        Self {
101            current: ArcSwapOption::empty(),
102            owned: SessionMutex::new(None),
103            shutdown,
104        }
105    }
106}
107
108impl Drop for AppendSession {
109    fn drop(&mut self) {
110        self.reader.abort();
111    }
112}
113
114impl AppendSession {
115    async fn shutdown(mut self) {
116        self.reader.abort();
117        let _ = (&mut self.reader).await;
118    }
119}
120
121#[derive(Clone, Debug, Default)]
122/// Durable tail plus any coalesced stream-termination error published by the
123/// session reader.
124struct LaneProgress {
125    durable: i64,
126    finalized: Option<ReplicaSnapshot>,
127    error: Option<TransportError>,
128}
129
130/// Target wire-message size for packed appends. Large enough to amortize
131/// per-message overhead (protobuf, CRC field, HTTP/2 framing, server
132/// per-message processing). The hard ceiling, measured live, is
133/// the service's 4 MiB gRPC inbound message cap — a 2 MiB + 1 byte chunk is
134/// accepted, 4 MiB is rejected with ResourceExhausted; the proto's only
135/// documented chunk constant (MaxReadChunkBytes = 2 MiB) is read-side.
136const WIRE_MESSAGE_TARGET_BYTES: usize = 262_144;
137
138pub(crate) fn pack_append(chunks: Vec<Bytes>) -> PackedAppend {
139    let total_len = chunks.iter().map(Bytes::len).sum::<usize>();
140    let mut packed =
141        bytes::BytesMut::with_capacity(total_len.min(WIRE_MESSAGE_TARGET_BYTES.saturating_mul(2)));
142    let mut messages = Vec::new();
143    let mut relative_offset = 0i64;
144    for data in &chunks {
145        if !packed.is_empty() && packed.len() + data.len() > WIRE_MESSAGE_TARGET_BYTES {
146            let content = packed.split().freeze();
147            let len = content.len() as i64;
148            messages.push(PackedAppendMessage {
149                relative_offset,
150                crc32c: crc32c::crc32c(&content),
151                content,
152            });
153            relative_offset += len;
154        }
155        packed.extend_from_slice(data);
156    }
157    if !packed.is_empty() {
158        let content = packed.freeze();
159        messages.push(PackedAppendMessage {
160            relative_offset,
161            crc32c: crc32c::crc32c(&content),
162            content,
163        });
164    }
165    PackedAppend::new(chunks, messages, total_len)
166}
167
168/// Per-response progress timeout for the persistent session. The session
169/// itself has no overall deadline (it lives for the segment), but a server
170/// that stops acknowledging flushed appends must fail the lane.
171const SESSION_PROGRESS_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
172const RPC_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
173const MAX_LIST_PAGES: usize = 10_000;
174
175impl AppendSessionHandle {
176    async fn send(
177        self: &Arc<Self>,
178        replica: &GrpcReplica,
179        request: BidiWriteObjectRequest,
180        disconnected: &'static str,
181    ) -> Result<(), TransportError> {
182        match tokio::time::timeout(SESSION_PROGRESS_TIMEOUT, self.tx.send(request)).await {
183            Ok(Ok(())) => Ok(()),
184            Ok(Err(_)) => {
185                replica.clear_session_if(self).await;
186                Err(replica.error(TransportCode::Unavailable, disconnected))
187            }
188            Err(_) => {
189                replica.clear_session_if(self).await;
190                Err(replica.error(
191                    TransportCode::DeadlineExceeded,
192                    "append session request channel made no progress",
193                ))
194            }
195        }
196    }
197
198    async fn wait_for<T>(
199        self: &Arc<Self>,
200        replica: &GrpcReplica,
201        wait: SessionWait,
202        mut inspect: impl FnMut(&LaneProgress) -> Option<Result<T, TransportError>>,
203    ) -> Result<T, TransportError> {
204        let mut state = self.state.clone();
205        loop {
206            let progress = state.borrow_and_update().clone();
207            if wait.inspect_before_error {
208                if let Some(result) = inspect(&progress) {
209                    if wait.clear_on_ready || result.is_err() {
210                        replica.clear_session_if(self).await;
211                    }
212                    return result;
213                }
214            }
215            if let Some(error) = progress.error {
216                replica.clear_session_if(self).await;
217                return Err(error);
218            }
219            if !wait.inspect_before_error {
220                if let Some(result) = inspect(&progress) {
221                    if wait.clear_on_ready || result.is_err() {
222                        replica.clear_session_if(self).await;
223                    }
224                    return result;
225                }
226            }
227            let changed = match wait.timeout {
228                Some(timeout) => match tokio::time::timeout(timeout, state.changed()).await {
229                    Ok(changed) => changed,
230                    Err(_) => {
231                        replica.clear_session_if(self).await;
232                        return Err(replica.error(TransportCode::DeadlineExceeded, wait.stalled));
233                    }
234                },
235                None => state.changed().await,
236            };
237            if changed.is_err() {
238                replica.clear_session_if(self).await;
239                return Err(replica.error(TransportCode::Unavailable, wait.reader_ended));
240            }
241        }
242    }
243}
244
245impl GrpcReplica {
246    fn live_session(&self) -> Result<Arc<AppendSessionHandle>, TransportError> {
247        self.session.current.load_full().ok_or_else(|| {
248            self.error(
249                TransportCode::Unavailable,
250                "no live append session (resume required)",
251            )
252        })
253    }
254
255    fn replace_session_locked(
256        &self,
257        owned: &mut Option<AppendSession>,
258        session: Option<AppendSession>,
259    ) -> Option<AppendSession> {
260        let previous = owned.take();
261        self.session
262            .current
263            .store(session.as_ref().map(|session| Arc::clone(&session.handle)));
264        *owned = session;
265        previous
266    }
267
268    async fn replace_session(&self, session: Option<AppendSession>) {
269        let previous = {
270            let mut owned = self.session.owned.lock().await;
271            self.replace_session_locked(&mut owned, session)
272        };
273        if let Some(previous) = previous {
274            previous.shutdown().await;
275        }
276    }
277
278    async fn clear_session_if(&self, expected: &Arc<AppendSessionHandle>) {
279        let previous = {
280            let mut owned = self.session.owned.lock().await;
281            if owned
282                .as_ref()
283                .is_some_and(|session| Arc::ptr_eq(&session.handle, expected))
284            {
285                self.replace_session_locked(&mut owned, None)
286            } else {
287                None
288            }
289        };
290        if let Some(previous) = previous {
291            previous.shutdown().await;
292        }
293    }
294
295    /// Like [`Self::request`], but without the per-RPC deadline: used for the
296    /// persistent append session, whose stream must outlive any fixed
297    /// deadline. Progress is bounded per response instead.
298    fn request_no_deadline<T>(&self, value: T) -> Result<Request<T>, TransportError> {
299        let mut request = self.request(value)?;
300        request.metadata_mut().remove("grpc-timeout");
301        Ok(request)
302    }
303
304    fn request<T>(&self, value: T) -> Result<Request<T>, TransportError> {
305        let mut request = Request::new(value);
306        request.set_timeout(RPC_TIMEOUT);
307        // GCS v2 gRPC routes every object RPC by the bucket resource path
308        // carried in x-goog-request-params; without it the service rejects the
309        // call with INVALID_ARGUMENT. Once a zonal write redirect has supplied a
310        // routing token, replay it here so the stream lands on the right location.
311        let guard = self.routing_token.load();
312        let value = match (**guard).as_ref() {
313            Some(token) => format!("bucket={}&routing_token={}", self.bucket, token),
314            None => format!("bucket={}", self.bucket),
315        };
316        let params = MetadataValue::try_from(value).map_err(|_| {
317            self.error(
318                TransportCode::Internal,
319                "request params are not valid gRPC metadata",
320            )
321        })?;
322        request
323            .metadata_mut()
324            .insert("x-goog-request-params", params);
325        if let Some(auth) = &self.auth {
326            let value = auth
327                .authorization_header()
328                .map_err(|error| self.error(TransportCode::Unauthenticated, error.to_string()))?;
329            request.metadata_mut().insert("authorization", value);
330        }
331        Ok(request)
332    }
333
334    fn error(&self, code: TransportCode, message: impl Into<String>) -> TransportError {
335        TransportError {
336            zone: self.zone,
337            code,
338            message: message.into(),
339        }
340    }
341
342    fn status(&self, status: Status) -> TransportError {
343        // Retry classification is a protocol contract, not a direct mirror of
344        // tonic's codes. This mapper is shared by append and non-append RPCs:
345        // RESOURCE_EXHAUSTED is GCS throttling, not a permanent rejection:
346        // per-object mutation-rate and per-project quotas reset over time, so
347        // retrying the valid request with backoff can succeed.
348        let code = match status.code() {
349            Code::NotFound => TransportCode::NotFound,
350            Code::AlreadyExists => TransportCode::AlreadyExists,
351            Code::InvalidArgument => TransportCode::InvalidArgument,
352            Code::FailedPrecondition => TransportCode::FailedPrecondition,
353            Code::Aborted => TransportCode::Aborted,
354            Code::OutOfRange => TransportCode::OutOfRange,
355            Code::ResourceExhausted => TransportCode::ResourceExhausted,
356            Code::Unimplemented => TransportCode::Unimplemented,
357            Code::DataLoss => TransportCode::DataLoss,
358            Code::Unauthenticated => TransportCode::Unauthenticated,
359            Code::PermissionDenied => TransportCode::PermissionDenied,
360            Code::Unavailable => TransportCode::Unavailable,
361            Code::DeadlineExceeded => TransportCode::DeadlineExceeded,
362            _ => TransportCode::Internal,
363        };
364        self.error(code, status.message())
365    }
366
367    fn snapshot_from_object(&self, object: Object, bytes: Vec<u8>) -> ReplicaSnapshot {
368        let crc32c = object
369            .checksums
370            .as_ref()
371            .and_then(|checksums| checksums.crc32c);
372        ReplicaSnapshot {
373            zone: self.zone,
374            generation: object.generation,
375            metageneration: object.metageneration,
376            persisted_size: bytes.len() as i64,
377            finalized: object.finalize_time.is_some(),
378            crc32c,
379            metadata: object.metadata,
380            bytes,
381        }
382    }
383
384    fn stat_from_object(&self, object: Object) -> ReplicaSnapshot {
385        let size = object.size;
386        let mut snapshot = self.snapshot_from_object(object, Vec::new());
387        if snapshot.finalized {
388            snapshot.persisted_size = size;
389        }
390        snapshot
391    }
392
393    /// If `status` is a zonal write redirect, cache its routing token for replay
394    /// and report `true` so the caller retries the same guarded open. The
395    /// redirect fields are intentionally ignored: conditional creates replay
396    /// their create precondition, takeovers replay their metageneration guard,
397    /// and resumes replay their existing handle.
398    fn try_capture_redirect(&self, status: &Status) -> bool {
399        match redirect_routing_token(status) {
400            Some(token) => {
401                self.routing_token.store(Arc::new(Some(Arc::new(token))));
402                true
403            }
404            None => false,
405        }
406    }
407
408    /// Open one write stream and obtain its first response, transparently
409    /// replaying redirects that occur before any response is observed.
410    ///
411    /// Both one-shot writes and persistent append sessions use this driver, so
412    /// routing-token capture, redirect limits, opening deadlines, and the
413    /// no-midstream-replay rule have one implementation.
414    async fn open_redirect_aware_stream(
415        &self,
416        requests: &[BidiWriteObjectRequest],
417        persistent: bool,
418    ) -> Result<Option<RedirectAwareStream>, TransportError> {
419        let mut attempt = 0u32;
420        let mut shutdown = self.session.shutdown.subscribe();
421        loop {
422            if *shutdown.borrow_and_update() {
423                return Err(self.error(
424                    TransportCode::Unavailable,
425                    "append session open cancelled by shutdown",
426                ));
427            }
428            attempt += 1;
429            let (tx, rx) = mpsc::channel(requests.len().max(64));
430            for request in requests {
431                if tx.send(request.clone()).await.is_err() {
432                    return Err(self.error(TransportCode::Internal, "write stream channel closed"));
433                }
434            }
435            let request = if persistent {
436                self.request_no_deadline(ReceiverStream::new(rx))?
437            } else {
438                self.request(ReceiverStream::new(rx))?
439            };
440            let mut tx = if persistent {
441                Some(tx)
442            } else {
443                drop(tx);
444                None
445            };
446            let mut client = self.client.clone();
447            let mut opening = Box::pin(client.bidi_write_object(request));
448            let opened = if persistent {
449                // The established session has no RPC deadline, so bound its
450                // opening handshake separately before returning the live stream.
451                tokio::select! {
452                    biased;
453                    changed = shutdown.changed() => {
454                        tx.take();
455                        if changed.is_ok() {
456                            if let Ok(Ok(response)) =
457                                tokio::time::timeout(SESSION_PROGRESS_TIMEOUT, &mut opening).await
458                            {
459                                let mut responses = response.into_inner();
460                                let _ = tokio::time::timeout(
461                                    SESSION_PROGRESS_TIMEOUT,
462                                    responses.message(),
463                                )
464                                .await;
465                            }
466                        }
467                        return Err(self.error(
468                            TransportCode::Unavailable,
469                            "append session open cancelled by shutdown",
470                        ));
471                    }
472                    opened = tokio::time::timeout(SESSION_PROGRESS_TIMEOUT, &mut opening) => {
473                        match opened {
474                            Ok(opened) => opened,
475                            Err(_) => {
476                                return Err(self.error(
477                                    TransportCode::DeadlineExceeded,
478                                    "append session open timed out",
479                                ));
480                            }
481                        }
482                    }
483                }
484            } else {
485                opening.await
486            };
487            let mut responses = match opened {
488                Ok(response) => response.into_inner(),
489                Err(status) => {
490                    if attempt <= MAX_WRITE_REDIRECTS && self.try_capture_redirect(&status) {
491                        continue;
492                    }
493                    return Err(self.status(status));
494                }
495            };
496            let first = if persistent {
497                tokio::select! {
498                    biased;
499                    changed = shutdown.changed() => {
500                        tx.take();
501                        if changed.is_ok() {
502                            let _ = tokio::time::timeout(
503                                SESSION_PROGRESS_TIMEOUT,
504                                responses.message(),
505                            )
506                            .await;
507                        }
508                        return Err(self.error(
509                            TransportCode::Unavailable,
510                            "append session open cancelled by shutdown",
511                        ));
512                    }
513                    response = tokio::time::timeout(
514                        SESSION_PROGRESS_TIMEOUT,
515                        responses.message(),
516                    ) => {
517                        match response {
518                            Ok(response) => response,
519                            Err(_) => {
520                                return Err(self.error(
521                                    TransportCode::DeadlineExceeded,
522                                    "append open made no progress",
523                                ));
524                            }
525                        }
526                    }
527                }
528            } else {
529                responses.message().await
530            };
531            match first {
532                Ok(Some(first)) => {
533                    return Ok(Some(RedirectAwareStream {
534                        tx,
535                        responses,
536                        first,
537                        attempt,
538                    }));
539                }
540                Ok(None) => return Ok(None),
541                Err(status)
542                    if attempt <= MAX_WRITE_REDIRECTS && self.try_capture_redirect(&status) =>
543                {
544                    continue;
545                }
546                Err(status) => return Err(self.status(status)),
547            }
548        }
549    }
550
551    /// Open (or resume) the persistent append session: send the first
552    /// message, wait for the opening acknowledgment, and return the live
553    /// session with the server-certified durable tail and session handle.
554    /// Redirects are captured and retried with the routing token.
555    async fn open_session(
556        &self,
557        first: BidiWriteObjectRequest,
558    ) -> Result<(AppendSession, i64, Option<Bytes>), TransportError> {
559        let (session, persisted_size, write_handle, _) = self.open_session_observed(first).await?;
560        Ok((session, persisted_size, write_handle))
561    }
562
563    async fn open_session_observed(
564        &self,
565        first: BidiWriteObjectRequest,
566    ) -> Result<(AppendSession, i64, Option<Bytes>, Option<i64>), TransportError> {
567        let Some(opened) = self
568            .open_redirect_aware_stream(std::slice::from_ref(&first), true)
569            .await?
570        else {
571            return Err(self.error(
572                TransportCode::Unavailable,
573                "append open closed without a response",
574            ));
575        };
576        let opening_has_status = opened.first.write_status.is_some();
577        let opening_resource_size = match opened.first.write_status.as_ref() {
578            Some(bidi_write_object_response::WriteStatus::Resource(resource)) => {
579                Some(resource.size)
580            }
581            _ => None,
582        };
583        let handle = opened
584            .first
585            .write_handle
586            .clone()
587            .map(|handle| handle.handle);
588        let mut progress = LaneProgress {
589            durable: if opening_has_status {
590                0
591            } else {
592                first.write_offset
593            },
594            finalized: None,
595            error: None,
596        };
597        self.fold_session_progress(&mut progress, opened.first);
598        let persisted = progress.durable;
599        let (state_tx, state_rx) = tokio::sync::watch::channel(progress);
600        let this = self.clone();
601        let reader = tokio::spawn(async move {
602            let mut responses = opened.responses;
603            loop {
604                match responses.message().await {
605                    Ok(Some(response)) => {
606                        state_tx.send_modify(|progress| {
607                            this.fold_session_progress(progress, response);
608                        });
609                    }
610                    Ok(None) => {
611                        state_tx.send_modify(|progress| {
612                            progress.error = Some(this.error(
613                                TransportCode::Unavailable,
614                                "append session closed by the service",
615                            ));
616                        });
617                        return;
618                    }
619                    Err(status) => {
620                        this.try_capture_redirect(&status);
621                        let error = this.status(status);
622                        state_tx.send_modify(|progress| {
623                            progress.error = Some(error.clone());
624                        });
625                        return;
626                    }
627                }
628            }
629        });
630        tracing::debug!(
631            zone = self.zone,
632            persisted_size = persisted,
633            has_write_handle = handle.is_some(),
634            attempt = opened.attempt,
635            "append session opened"
636        );
637        let tx = opened.tx.ok_or_else(|| {
638            self.error(
639                TransportCode::Internal,
640                "persistent write stream omitted its request sender",
641            )
642        })?;
643        Ok((
644            AppendSession {
645                handle: Arc::new(AppendSessionHandle {
646                    tx,
647                    state: state_rx,
648                }),
649                reader,
650            },
651            persisted,
652            handle,
653            opening_resource_size,
654        ))
655    }
656
657    fn fold_session_progress(
658        &self,
659        progress: &mut LaneProgress,
660        response: BidiWriteObjectResponse,
661    ) {
662        match response.write_status {
663            Some(bidi_write_object_response::WriteStatus::PersistedSize(size)) => {
664                progress.durable = progress.durable.max(size);
665            }
666            Some(bidi_write_object_response::WriteStatus::Resource(resource)) => {
667                let persisted_size = resource.size;
668                let snapshot = self.stat_from_object(resource);
669                progress.durable = progress.durable.max(persisted_size);
670                if snapshot.finalized {
671                    progress.finalized = Some(snapshot);
672                }
673            }
674            None => {}
675        }
676    }
677
678    /// The opening message for this object's append session.
679    fn session_open_request(
680        &self,
681        generation: i64,
682        metageneration: i64,
683        write_handle: Option<Bytes>,
684        write_offset: i64,
685    ) -> BidiWriteObjectRequest {
686        BidiWriteObjectRequest {
687            first_message: Some(bidi_write_object_request::FirstMessage::AppendObjectSpec(
688                AppendObjectSpec {
689                    bucket: self.bucket.clone(),
690                    object: self.object.clone(),
691                    generation,
692                    if_metageneration_match: Some(metageneration),
693                    write_handle: write_handle.map(|handle| BidiWriteHandle { handle }),
694                    ..Default::default()
695                },
696            )),
697            write_offset,
698            flush: true,
699            state_lookup: true,
700            ..Default::default()
701        }
702    }
703
704    /// Construct the candidate handle-free current-generation open. Generation
705    /// zero is wire-identical to an unset proto3 scalar; the live probe decides
706    /// whether GCS accepts it as a selector or rejects the missing generation.
707    #[cfg(feature = "probe-support")]
708    fn current_session_open_request(&self) -> BidiWriteObjectRequest {
709        BidiWriteObjectRequest {
710            first_message: Some(bidi_write_object_request::FirstMessage::AppendObjectSpec(
711                AppendObjectSpec {
712                    bucket: self.bucket.clone(),
713                    object: self.object.clone(),
714                    generation: 0,
715                    ..Default::default()
716                },
717            )),
718            write_offset: 0,
719            flush: true,
720            state_lookup: true,
721            ..Default::default()
722        }
723    }
724
725    #[cfg(feature = "probe-support")]
726    async fn takeover_current_generation(
727        &self,
728    ) -> Result<(AppendToken, Option<i64>), TransportError> {
729        let (session, persisted_size, write_handle, opening_resource_size) = self
730            .open_session_observed(self.current_session_open_request())
731            .await?;
732        self.replace_session(Some(session)).await;
733        tracing::debug!(
734            zone = self.zone,
735            persisted_size,
736            has_write_handle = write_handle.is_some(),
737            "current append session takeover completed"
738        );
739        Ok((
740            AppendToken {
741                zone: self.zone,
742                // The append-open response does not provide identity fields
743                // needed by later resume/finalize paths. Resolve them lazily.
744                generation: None,
745                metageneration: None,
746                persisted_size,
747                write_handle,
748            },
749            opening_resource_size,
750        ))
751    }
752
753    async fn resolve_token_identity(
754        &self,
755        token: &mut AppendToken,
756    ) -> Result<(i64, i64), TransportError> {
757        match (token.generation, token.metageneration) {
758            (Some(generation), Some(metageneration)) => Ok((generation, metageneration)),
759            (None, None) => {
760                let observed = self.stat().await?;
761                if observed.finalized {
762                    return Err(self.error(
763                        TransportCode::FailedPrecondition,
764                        "append session object is already finalized",
765                    ));
766                }
767                token.generation = Some(observed.generation);
768                token.metageneration = Some(observed.metageneration);
769                Ok((observed.generation, observed.metageneration))
770            }
771            _ => Err(self.error(
772                TransportCode::Internal,
773                "append token has incomplete generation identity",
774            )),
775        }
776    }
777
778    async fn finish_live_session(
779        &self,
780        write_offset: i64,
781        expected_generation: Option<i64>,
782    ) -> Result<ReplicaSnapshot, TransportError> {
783        let handle = self.live_session().map_err(|_| {
784            self.error(
785                TransportCode::Unavailable,
786                "no live append session to finalize",
787            )
788        })?;
789        let request = BidiWriteObjectRequest {
790            first_message: None,
791            write_offset,
792            finish_write: true,
793            ..Default::default()
794        };
795        handle
796            .send(
797                self,
798                request,
799                "append session disconnected while finalizing",
800            )
801            .await?;
802        handle
803            .wait_for(
804                self,
805                SessionWait {
806                    timeout: Some(SESSION_PROGRESS_TIMEOUT),
807                    clear_on_ready: true,
808                    inspect_before_error: true,
809                    reader_ended: "append session reader ended while finalizing",
810                    stalled: "append session finalization made no progress",
811                },
812                |progress| {
813                    let finalized = progress.finalized.clone()?;
814                    Some(
815                        if !finalized.finalized
816                            || expected_generation
817                                .is_some_and(|generation| finalized.generation != generation)
818                            || finalized.persisted_size != write_offset
819                        {
820                            tracing::warn!(
821                                expected_generation,
822                                actual_generation = finalized.generation,
823                                expected_size = write_offset,
824                                actual_size = finalized.persisted_size,
825                                finalized = finalized.finalized,
826                                "finalized append response did not match the requested prefix"
827                            );
828                            Err(self.error(
829                                TransportCode::DataLoss,
830                                "finalized segment does not match the committed prefix",
831                            ))
832                        } else {
833                            Ok(finalized)
834                        },
835                    )
836                },
837            )
838            .await
839    }
840
841    async fn drive_redirect_aware_stream<S>(
842        &self,
843        requests: Vec<BidiWriteObjectRequest>,
844        make_state: impl Fn() -> S,
845        mut step: impl FnMut(&mut S, BidiWriteObjectResponse),
846        complete: impl Fn(&S) -> bool,
847    ) -> (S, Option<TransportError>) {
848        let mut state = make_state();
849        let opened = match self.open_redirect_aware_stream(&requests, false).await {
850            Ok(Some(opened)) => opened,
851            Ok(None) => return (state, None),
852            Err(error) => return (state, Some(error)),
853        };
854        step(&mut state, opened.first);
855        if complete(&state) {
856            return (state, None);
857        }
858        let mut stream = opened.responses;
859        loop {
860            match stream.message().await {
861                Ok(Some(response)) => {
862                    step(&mut state, response);
863                    if complete(&state) {
864                        // Appendable sessions stay open server-side; once the
865                        // caller has what it needs, abandon the stream instead
866                        // of waiting for the live service's idle expiry.
867                        return (state, None);
868                    }
869                }
870                Ok(None) => return (state, None),
871                Err(status) => return (state, Some(self.status(status))),
872            }
873        }
874    }
875}
876
877impl GrpcReplicaFactory {
878    #[cfg(feature = "probe-support")]
879    fn probe_replica(&self, object: &str) -> GrpcReplica {
880        GrpcReplica {
881            zone: self.zone,
882            bucket: self.bucket.clone(),
883            object: object.to_string(),
884            auth: self.auth.clone(),
885            client: self.client.clone(),
886            routing_token: self.routing_token.clone(),
887            session: Arc::new(SessionSlot::new()),
888        }
889    }
890
891    /// Connect a zonal bucket with an optional static bearer token.
892    ///
893    /// The bucket must be a full v2 resource name such as
894    /// `projects/_/buckets/example-zone-a`. All replicas created by the factory
895    /// share the routing token learned from zonal write redirects. The normal
896    /// production endpoint is `https://storage.googleapis.com`; Cloud Storage
897    /// regional JSON/XML endpoints do not support this gRPC client.
898    pub async fn connect(
899        zone: usize,
900        endpoint: &str,
901        bucket: impl Into<String>,
902        bearer_token: Option<String>,
903    ) -> Result<Self, Error> {
904        let mut builder = Endpoint::from_shared(endpoint.to_string())
905            .map_err(|error| Error::Connection(error.to_string()))?;
906        builder = builder.connect_timeout(RPC_TIMEOUT);
907        if endpoint.starts_with("https://") {
908            builder = builder
909                .tls_config(ClientTlsConfig::new().with_webpki_roots())
910                .map_err(|error| Error::Connection(error.to_string()))?;
911        }
912        let channel = builder
913            .connect()
914            .await
915            .map_err(|error| Error::Connection(error.to_string()))?;
916        Ok(Self {
917            zone,
918            bucket: bucket.into(),
919            auth: bearer_token.map(BearerAuth::static_token),
920            client: StorageClient::new(channel),
921            routing_token: Arc::new(ArcSwap::from_pointee(None)),
922        })
923    }
924
925    /// Build a factory over an already-established channel. The deterministic
926    /// simulator uses this to route the production client over a simulated
927    /// transport (`Endpoint::connect_with_connector`); production callers use
928    /// `connect`/`connect_with_auth`.
929    pub fn from_channel(
930        zone: usize,
931        channel: Channel,
932        bucket: impl Into<String>,
933        bearer_token: Option<String>,
934    ) -> Self {
935        Self {
936            zone,
937            bucket: bucket.into(),
938            auth: bearer_token.map(BearerAuth::static_token),
939            client: StorageClient::new(channel),
940            routing_token: Arc::new(ArcSwap::from_pointee(None)),
941        }
942    }
943
944    /// Connect a zonal bucket using a shared static or refreshing auth handle.
945    /// Cloned refreshing handles update existing clients through `ArcSwap`.
946    /// Use full v2 bucket resource names and a gRPC-compatible endpoint such as
947    /// `https://storage.googleapis.com`.
948    pub async fn connect_with_auth(
949        zone: usize,
950        endpoint: &str,
951        bucket: impl Into<String>,
952        auth: BearerAuth,
953    ) -> Result<Self, Error> {
954        let mut builder = Endpoint::from_shared(endpoint.to_string())
955            .map_err(|error| Error::Connection(error.to_string()))?;
956        builder = builder.connect_timeout(RPC_TIMEOUT);
957        if endpoint.starts_with("https://") {
958            builder = builder
959                .tls_config(ClientTlsConfig::new().with_webpki_roots())
960                .map_err(|error| Error::Connection(error.to_string()))?;
961        }
962        let channel = builder
963            .connect()
964            .await
965            .map_err(|error| Error::Connection(error.to_string()))?;
966        Ok(Self {
967            zone,
968            bucket: bucket.into(),
969            auth: Some(auth),
970            client: StorageClient::new(channel),
971            routing_token: Arc::new(ArcSwap::from_pointee(None)),
972        })
973    }
974}
975
976#[cfg(feature = "probe-support")]
977#[derive(Clone, Debug)]
978/// One generation-zero append-open response observed by the live-GCS probe.
979pub struct GenerationZeroOpenObservation {
980    /// Authoritative tail derived from the opening response.
981    pub persisted_size: i64,
982    /// Object-resource size in the opening response, if GCS supplied a resource.
983    pub resource_size: Option<i64>,
984}
985
986#[cfg(feature = "probe-support")]
987#[derive(Clone, Debug)]
988/// Independent observations from the focused live-GCS generation-zero probe.
989pub struct GenerationZeroTakeoverProbeResult {
990    /// Number of bytes the probe requested GCS to persist.
991    pub expected_size: i64,
992    /// Result of creating an appendable object and flushing the payload.
993    pub append: Result<i64, TransportError>,
994    /// Present-object takeover result, absent only when creation or append failed.
995    pub takeover: Option<Result<GenerationZeroOpenObservation, TransportError>>,
996    /// Generation-zero open result for a never-created object.
997    pub absent: Result<GenerationZeroOpenObservation, TransportError>,
998    /// Cleanup result for the object that may have been created.
999    pub cleanup: Result<(), TransportError>,
1000}
1001
1002#[cfg(feature = "probe-support")]
1003/// Exercise the production gRPC append path against one present and one absent
1004/// object, then delete the present object before returning.
1005///
1006/// Both takeovers send `AppendObjectSpec.generation = 0` through the same
1007/// redirect-aware `open_session` implementation used by recovery. Operations
1008/// are reported independently so a provider rejection on the present object
1009/// does not suppress the absent-name observation.
1010pub async fn probe_generation_zero_takeover(
1011    factory: &GrpcReplicaFactory,
1012    present_object: &str,
1013    absent_object: &str,
1014    payload: Bytes,
1015) -> Result<GenerationZeroTakeoverProbeResult, TransportError> {
1016    let expected_size = i64::try_from(payload.len()).map_err(|_| TransportError {
1017        zone: factory.zone,
1018        code: TransportCode::InvalidArgument,
1019        message: "probe payload length does not fit in i64".into(),
1020    })?;
1021    if expected_size == 0 {
1022        return Err(TransportError {
1023            zone: factory.zone,
1024            code: TransportCode::InvalidArgument,
1025            message: "probe payload must be non-empty".into(),
1026        });
1027    }
1028
1029    let present = factory.probe_replica(present_object);
1030    let append = async {
1031        let token = present.create_append_session(HashMap::new()).await?;
1032        if token.persisted_size != 0 {
1033            return Err(present.error(
1034                TransportCode::DataLoss,
1035                format!(
1036                    "new appendable object opened at persisted_size={}, expected 0",
1037                    token.persisted_size
1038                ),
1039            ));
1040        }
1041        present
1042            .lane_send(token.persisted_size, std::slice::from_ref(&payload))
1043            .await?;
1044        let change = present.lane_durable_change(token.persisted_size).await?;
1045        if let Some(error) = change.error {
1046            return Err(error);
1047        }
1048        Ok(change.persisted_size)
1049    }
1050    .await;
1051    let takeover = if append.is_ok() {
1052        Some(
1053            present
1054                .takeover_current_generation()
1055                .await
1056                .map(|(token, resource_size)| GenerationZeroOpenObservation {
1057                    persisted_size: token.persisted_size,
1058                    resource_size,
1059                }),
1060        )
1061    } else {
1062        None
1063    };
1064
1065    present.shutdown().await;
1066    let cleanup = match present.stat().await {
1067        Ok(snapshot) => present.delete(snapshot.generation).await,
1068        Err(error) if error.code == TransportCode::NotFound => Ok(()),
1069        Err(error) => Err(error),
1070    };
1071
1072    let absent_replica = factory.probe_replica(absent_object);
1073    let absent =
1074        absent_replica
1075            .takeover_current_generation()
1076            .await
1077            .map(|(token, resource_size)| GenerationZeroOpenObservation {
1078                persisted_size: token.persisted_size,
1079                resource_size,
1080            });
1081    absent_replica.shutdown().await;
1082
1083    Ok(GenerationZeroTakeoverProbeResult {
1084        expected_size,
1085        append,
1086        takeover,
1087        absent,
1088        cleanup,
1089    })
1090}
1091
1092#[async_trait]
1093impl ReplicaFactory for GrpcReplicaFactory {
1094    fn bucket_name(&self) -> &str {
1095        self.bucket.rsplit('/').next().unwrap_or_default()
1096    }
1097
1098    fn replica(&self, object: &str) -> Arc<dyn Replica> {
1099        Arc::new(GrpcReplica {
1100            zone: self.zone,
1101            bucket: self.bucket.clone(),
1102            object: object.to_string(),
1103            auth: self.auth.clone(),
1104            client: self.client.clone(),
1105            routing_token: self.routing_token.clone(),
1106            session: Arc::new(SessionSlot::new()),
1107        })
1108    }
1109
1110    async fn list(&self, prefix: &str) -> Result<Vec<ListedObject>, TransportError> {
1111        let replica = GrpcReplica {
1112            zone: self.zone,
1113            bucket: self.bucket.clone(),
1114            object: String::new(),
1115            auth: self.auth.clone(),
1116            client: self.client.clone(),
1117            routing_token: self.routing_token.clone(),
1118            session: Arc::new(SessionSlot::new()),
1119        };
1120        let mut page_token = String::new();
1121        let mut seen_page_tokens = HashSet::new();
1122        let mut listed = Vec::new();
1123        for _ in 0..MAX_LIST_PAGES {
1124            let request = replica.request(ListObjectsRequest {
1125                parent: self.bucket.clone(),
1126                page_size: 1000,
1127                page_token,
1128                prefix: prefix.to_string(),
1129                ..Default::default()
1130            })?;
1131            let response = self
1132                .client
1133                .clone()
1134                .list_objects(request)
1135                .await
1136                .map_err(|status| replica.status(status))?
1137                .into_inner();
1138            listed.extend(response.objects.into_iter().map(|object| ListedObject {
1139                zone: self.zone,
1140                name: object.name,
1141                generation: object.generation,
1142                finalized: object.finalize_time.is_some(),
1143                metadata: object.metadata,
1144            }));
1145            if response.next_page_token.is_empty() {
1146                return Ok(listed);
1147            }
1148            if !seen_page_tokens.insert(response.next_page_token.clone()) {
1149                return Err(
1150                    replica.error(TransportCode::Internal, "ListObjects repeated a page token")
1151                );
1152            }
1153            page_token = response.next_page_token;
1154        }
1155        Err(replica.error(
1156            TransportCode::Internal,
1157            "ListObjects exceeded the pagination bound",
1158        ))
1159    }
1160}
1161
1162#[async_trait]
1163impl Replica for GrpcReplica {
1164    async fn stat(&self) -> Result<ReplicaSnapshot, TransportError> {
1165        let get = GetObjectRequest {
1166            bucket: self.bucket.clone(),
1167            object: self.object.clone(),
1168            ..Default::default()
1169        };
1170        let request = self.request(get)?;
1171        let object = self
1172            .client
1173            .clone()
1174            .get_object(request)
1175            .await
1176            .map_err(|status| self.status(status))?
1177            .into_inner();
1178        // A finalized object's size is frozen and authoritative; open
1179        // objects keep persisted_size = 0 here (stats are tail-blind).
1180        Ok(self.stat_from_object(object))
1181    }
1182
1183    async fn snapshot(&self) -> Result<ReplicaSnapshot, TransportError> {
1184        let get = GetObjectRequest {
1185            bucket: self.bucket.clone(),
1186            object: self.object.clone(),
1187            ..Default::default()
1188        };
1189        let request = self.request(get)?;
1190        let object = self
1191            .client
1192            .clone()
1193            .get_object(request)
1194            .await
1195            .map_err(|status| self.status(status))?
1196            .into_inner();
1197
1198        let read = ReadObjectRequest {
1199            bucket: self.bucket.clone(),
1200            object: self.object.clone(),
1201            generation: object.generation,
1202            if_generation_match: Some(object.generation),
1203            ..Default::default()
1204        };
1205        let request = self.request(read)?;
1206        let mut stream = self
1207            .client
1208            .clone()
1209            .read_object(request)
1210            .await
1211            .map_err(|status| self.status(status))?
1212            .into_inner();
1213        let mut bytes = Vec::new();
1214        while let Some(response) = stream
1215            .message()
1216            .await
1217            .map_err(|status| self.status(status))?
1218        {
1219            if let Some(data) = response.checksummed_data {
1220                if data
1221                    .crc32c
1222                    .is_some_and(|expected| expected != crc32c::crc32c(&data.content))
1223                {
1224                    return Err(self.error(
1225                        TransportCode::DataLoss,
1226                        "ReadObject response CRC32C mismatch",
1227                    ));
1228                }
1229                bytes.extend_from_slice(&data.content);
1230            }
1231        }
1232        Ok(self.snapshot_from_object(object, bytes))
1233    }
1234
1235    async fn create_appendable(
1236        &self,
1237        metadata: HashMap<String, String>,
1238    ) -> Result<ReplicaSnapshot, TransportError> {
1239        let object = Object {
1240            bucket: self.bucket.clone(),
1241            name: self.object.clone(),
1242            metadata: metadata.clone(),
1243            content_type: "application/vnd.chorus.records".into(),
1244            ..Default::default()
1245        };
1246        let request = BidiWriteObjectRequest {
1247            first_message: Some(bidi_write_object_request::FirstMessage::WriteObjectSpec(
1248                WriteObjectSpec {
1249                    resource: Some(object),
1250                    if_generation_match: Some(0),
1251                    appendable: Some(true),
1252                    ..Default::default()
1253                },
1254            )),
1255            write_offset: 0,
1256            flush: true,
1257            state_lookup: true,
1258            ..Default::default()
1259        };
1260        let (_, error) = self
1261            .drive_redirect_aware_stream(
1262                vec![request],
1263                || false,
1264                |seen, _| *seen = true,
1265                |seen| *seen,
1266            )
1267            .await;
1268        if let Some(error) = error {
1269            // The only precondition this request carries is
1270            // `if_generation_match=0`; the live service reports the conflict
1271            // as FAILED_PRECONDITION where the protocol expects AlreadyExists.
1272            if error.code == TransportCode::FailedPrecondition {
1273                return Err(TransportError {
1274                    code: TransportCode::AlreadyExists,
1275                    ..error
1276                });
1277            }
1278            return Err(error);
1279        }
1280        // Appendable create responses report only persisted size, not the
1281        // generation and metageneration required to guard the takeover open.
1282        // A metadata-only read supplies those fields without reading the
1283        // empty object body.
1284        self.stat().await
1285    }
1286
1287    async fn create_append_session(
1288        &self,
1289        metadata: HashMap<String, String>,
1290    ) -> Result<AppendToken, TransportError> {
1291        let object = Object {
1292            bucket: self.bucket.clone(),
1293            name: self.object.clone(),
1294            metadata,
1295            content_type: "application/vnd.chorus.records".into(),
1296            ..Default::default()
1297        };
1298        let request = BidiWriteObjectRequest {
1299            first_message: Some(bidi_write_object_request::FirstMessage::WriteObjectSpec(
1300                WriteObjectSpec {
1301                    resource: Some(object),
1302                    if_generation_match: Some(0),
1303                    appendable: Some(true),
1304                    ..Default::default()
1305                },
1306            )),
1307            write_offset: 0,
1308            flush: true,
1309            state_lookup: true,
1310            ..Default::default()
1311        };
1312        let (session, persisted_size, write_handle) =
1313            self.open_session(request).await.map_err(|error| {
1314                if error.code == TransportCode::FailedPrecondition {
1315                    TransportError {
1316                        code: TransportCode::AlreadyExists,
1317                        ..error
1318                    }
1319                } else {
1320                    error
1321                }
1322            })?;
1323        self.replace_session(Some(session)).await;
1324        Ok(AppendToken {
1325            zone: self.zone,
1326            generation: None,
1327            metageneration: None,
1328            persisted_size,
1329            write_handle,
1330        })
1331    }
1332
1333    async fn create_register(
1334        &self,
1335        metadata: HashMap<String, String>,
1336    ) -> Result<ReplicaSnapshot, TransportError> {
1337        let object = Object {
1338            bucket: self.bucket.clone(),
1339            name: self.object.clone(),
1340            metadata,
1341            content_type: "application/vnd.chorus.manifest".into(),
1342            ..Default::default()
1343        };
1344        // A plain one-shot WriteObject: regional buckets reject appendable
1345        // creates, and the register is never appended to anyway.
1346        let request = WriteObjectRequest {
1347            first_message: Some(write_object_request::FirstMessage::WriteObjectSpec(
1348                WriteObjectSpec {
1349                    resource: Some(object),
1350                    if_generation_match: Some(0),
1351                    ..Default::default()
1352                },
1353            )),
1354            write_offset: 0,
1355            finish_write: true,
1356            ..Default::default()
1357        };
1358        let request = self.request(tokio_stream::iter([request]))?;
1359        let response = self
1360            .client
1361            .clone()
1362            .write_object(request)
1363            .await
1364            .map_err(|status| {
1365                let error = self.status(status);
1366                // As with segment creates: the only precondition here is
1367                // `if_generation_match=0`, reported as FAILED_PRECONDITION.
1368                if error.code == TransportCode::FailedPrecondition {
1369                    TransportError {
1370                        code: TransportCode::AlreadyExists,
1371                        ..error
1372                    }
1373                } else {
1374                    error
1375                }
1376            })?
1377            .into_inner();
1378        let Some(write_object_response::WriteStatus::Resource(object)) = response.write_status
1379        else {
1380            return Err(self.error(
1381                TransportCode::Internal,
1382                "manifest create response omitted the object resource",
1383            ));
1384        };
1385        Ok(self.stat_from_object(object))
1386    }
1387
1388    async fn resume_tail(&self, token: &mut AppendToken) -> Result<i64, TransportError> {
1389        let (generation, metageneration) = self.resolve_token_identity(token).await?;
1390        let mut owned = self.session.owned.lock().await;
1391        let first = self.session_open_request(
1392            generation,
1393            metageneration,
1394            token.write_handle.clone(),
1395            token.persisted_size,
1396        );
1397        let (result, previous) = match self.open_session(first).await {
1398            Ok((session, persisted, _)) => {
1399                let previous = self.replace_session_locked(&mut owned, Some(session));
1400                tracing::debug!(
1401                    zone = self.zone,
1402                    persisted_size = persisted,
1403                    "append session resumed"
1404                );
1405                (Ok(persisted), previous)
1406            }
1407            Err(error) => {
1408                let previous = self.replace_session_locked(&mut owned, None);
1409                (Err(error), previous)
1410            }
1411        };
1412        drop(owned);
1413        if let Some(previous) = previous {
1414            previous.shutdown().await;
1415        }
1416        result
1417    }
1418
1419    async fn takeover(&self, observed: &ReplicaSnapshot) -> Result<AppendToken, TransportError> {
1420        let first = self.session_open_request(
1421            observed.generation,
1422            observed.metageneration,
1423            None, // handle-free: this open IS the fence
1424            observed.persisted_size,
1425        );
1426        let (session, persisted_size, write_handle) = self.open_session(first).await?;
1427        self.replace_session(Some(session)).await;
1428        tracing::debug!(
1429            zone = self.zone,
1430            persisted_size,
1431            has_write_handle = write_handle.is_some(),
1432            "append session takeover completed"
1433        );
1434        Ok(AppendToken {
1435            zone: self.zone,
1436            generation: Some(observed.generation),
1437            metageneration: Some(observed.metageneration),
1438            persisted_size,
1439            write_handle,
1440        })
1441    }
1442
1443    async fn update_register(
1444        &self,
1445        metageneration: i64,
1446        metadata: HashMap<String, String>,
1447    ) -> Result<ReplicaSnapshot, TransportError> {
1448        // Generation 0 addresses the live object; the register is never
1449        // deleted or recreated, so the metageneration precondition alone is
1450        // the CAS guard.
1451        let request = UpdateObjectRequest {
1452            object: Some(Object {
1453                bucket: self.bucket.clone(),
1454                name: self.object.clone(),
1455                metadata,
1456                ..Default::default()
1457            }),
1458            if_metageneration_match: Some(metageneration),
1459            update_mask: Some(prost_types::FieldMask {
1460                paths: vec!["metadata".into()],
1461            }),
1462            ..Default::default()
1463        };
1464        let object = self
1465            .client
1466            .clone()
1467            .update_object(self.request(request)?)
1468            .await
1469            .map_err(|status| self.status(status))?
1470            .into_inner();
1471        Ok(self.stat_from_object(object))
1472    }
1473
1474    async fn replace_appendable(
1475        &self,
1476        observed: &ReplicaSnapshot,
1477        data: Bytes,
1478        metadata: HashMap<String, String>,
1479    ) -> Result<AppendToken, TransportError> {
1480        let object = Object {
1481            bucket: self.bucket.clone(),
1482            name: self.object.clone(),
1483            metadata: metadata.clone(),
1484            content_type: "application/vnd.chorus.records".into(),
1485            ..Default::default()
1486        };
1487        let request = BidiWriteObjectRequest {
1488            first_message: Some(bidi_write_object_request::FirstMessage::WriteObjectSpec(
1489                WriteObjectSpec {
1490                    resource: Some(object),
1491                    if_generation_match: Some(observed.generation),
1492                    if_metageneration_match: Some(observed.metageneration),
1493                    appendable: Some(true),
1494                    ..Default::default()
1495                },
1496            )),
1497            write_offset: 0,
1498            data: Some(bidi_write_object_request::Data::ChecksummedData(
1499                ChecksummedData {
1500                    crc32c: Some(crc32c::crc32c(&data)),
1501                    content: data.clone(),
1502                },
1503            )),
1504            flush: true,
1505            state_lookup: true,
1506            ..Default::default()
1507        };
1508        let expected = data.len() as i64;
1509        let (persisted_size, error) = self
1510            .drive_redirect_aware_stream(
1511                vec![request],
1512                || None,
1513                |persisted_size, response| {
1514                    if let Some(bidi_write_object_response::WriteStatus::PersistedSize(size)) =
1515                        response.write_status
1516                    {
1517                        *persisted_size = Some(size);
1518                    }
1519                },
1520                |persisted_size| persisted_size.is_some_and(|size| size >= expected),
1521            )
1522            .await;
1523        if let Some(error) = error {
1524            return Err(error);
1525        }
1526        if persisted_size != Some(expected) {
1527            return Err(self.error(
1528                TransportCode::DataLoss,
1529                format!("replacement persisted {persisted_size:?}, expected {expected}"),
1530            ));
1531        }
1532        let snapshot = self.snapshot().await?;
1533        if snapshot.bytes != data[..] || snapshot.metadata != metadata {
1534            return Err(self.error(
1535                TransportCode::DataLoss,
1536                "replacement generation failed verification",
1537            ));
1538        }
1539        Ok(AppendToken {
1540            zone: self.zone,
1541            generation: Some(snapshot.generation),
1542            metageneration: Some(snapshot.metageneration),
1543            persisted_size: expected,
1544            write_handle: None,
1545        })
1546    }
1547
1548    async fn append(
1549        &self,
1550        token: &AppendToken,
1551        write_offset: i64,
1552        data: Vec<u8>,
1553    ) -> Result<i64, TransportError> {
1554        let Some(generation) = token.generation else {
1555            return Err(self.error(
1556                TransportCode::Internal,
1557                "one-shot append requires a generation-bound token",
1558            ));
1559        };
1560        let request = BidiWriteObjectRequest {
1561            first_message: Some(bidi_write_object_request::FirstMessage::AppendObjectSpec(
1562                AppendObjectSpec {
1563                    bucket: self.bucket.clone(),
1564                    object: self.object.clone(),
1565                    generation,
1566                    if_metageneration_match: token.metageneration,
1567                    write_handle: token
1568                        .write_handle
1569                        .clone()
1570                        .map(|handle| BidiWriteHandle { handle }),
1571                    ..Default::default()
1572                },
1573            )),
1574            write_offset,
1575            data: Some(bidi_write_object_request::Data::ChecksummedData(
1576                ChecksummedData {
1577                    content: Bytes::from(data.clone()),
1578                    crc32c: Some(crc32c::crc32c(&data)),
1579                },
1580            )),
1581            flush: true,
1582            state_lookup: true,
1583            ..Default::default()
1584        };
1585        let expected = write_offset + data.len() as i64;
1586        let (persisted_size, error) = self
1587            .drive_redirect_aware_stream(
1588                vec![request],
1589                || None,
1590                |persisted_size, response| {
1591                    if let Some(bidi_write_object_response::WriteStatus::PersistedSize(size)) =
1592                        response.write_status
1593                    {
1594                        *persisted_size = Some(size);
1595                    }
1596                },
1597                |persisted_size| persisted_size.is_some_and(|size| size >= expected),
1598            )
1599            .await;
1600        if let Some(error) = error {
1601            return Err(error);
1602        }
1603        match persisted_size {
1604            Some(size) if size >= expected => Ok(size),
1605            Some(size) => Err(self.error(
1606                TransportCode::DataLoss,
1607                format!("flush persisted {size}, expected at least {expected}"),
1608            )),
1609            None => Err(self.error(TransportCode::Internal, "missing persisted-size response")),
1610        }
1611    }
1612
1613    async fn lane_send(&self, write_offset: i64, chunks: &[Bytes]) -> Result<(), TransportError> {
1614        let packed = pack_append(chunks.to_vec());
1615        self.lane_send_packed(write_offset, &packed).await
1616    }
1617
1618    async fn lane_send_packed(
1619        &self,
1620        write_offset: i64,
1621        packed: &PackedAppend,
1622    ) -> Result<(), TransportError> {
1623        if packed.is_empty() {
1624            return Err(self.error(
1625                TransportCode::Internal,
1626                "append lane cannot send an empty flush group",
1627            ));
1628        }
1629        let handle = self.live_session()?;
1630        // The coalesced group was packed once before replica dispatch. Each
1631        // lane builds only its protobuf envelopes and shallow-clones the
1632        // refcounted message bytes; CRC32C and byte concatenation are shared.
1633        let messages = packed.messages();
1634        let last_index = messages.len() - 1;
1635        for (index, message) in messages.iter().enumerate() {
1636            let last = index == last_index;
1637            let request = BidiWriteObjectRequest {
1638                first_message: None,
1639                write_offset: write_offset + message.relative_offset,
1640                data: Some(bidi_write_object_request::Data::ChecksummedData(
1641                    ChecksummedData {
1642                        crc32c: Some(message.crc32c),
1643                        content: message.content.clone(),
1644                    },
1645                )),
1646                flush: last,
1647                state_lookup: last,
1648                ..Default::default()
1649            };
1650            handle
1651                .send(self, request, "append session disconnected while sending")
1652                .await?;
1653        }
1654        Ok(())
1655    }
1656
1657    async fn lane_durable_change(&self, seen: i64) -> Result<LaneDurableChange, TransportError> {
1658        let handle = self.live_session()?;
1659        let change = handle
1660            .wait_for(
1661                self,
1662                SessionWait {
1663                    timeout: None,
1664                    clear_on_ready: false,
1665                    // Preserve a coalesced response and stream error in one
1666                    // observation. Protocol code publishes the durable offset
1667                    // before deciding whether to recover or fence the writer.
1668                    inspect_before_error: true,
1669                    reader_ended: "append session reader ended",
1670                    stalled: "append session made no progress",
1671                },
1672                |progress| {
1673                    (progress.durable > seen).then(|| {
1674                        Ok(LaneDurableChange {
1675                            persisted_size: progress.durable,
1676                            error: progress.error.clone(),
1677                        })
1678                    })
1679                },
1680            )
1681            .await?;
1682        if change.error.is_some() {
1683            self.clear_session_if(&handle).await;
1684        }
1685        Ok(change)
1686    }
1687    async fn delete(&self, generation: i64) -> Result<(), TransportError> {
1688        let request = DeleteObjectRequest {
1689            bucket: self.bucket.clone(),
1690            object: self.object.clone(),
1691            generation,
1692            if_generation_match: Some(generation),
1693            ..Default::default()
1694        };
1695        let request = self.request(request)?;
1696        self.client
1697            .clone()
1698            .delete_object(request)
1699            .await
1700            .map_err(|status| self.status(status))?;
1701        Ok(())
1702    }
1703
1704    async fn finalize(
1705        &self,
1706        token: &mut AppendToken,
1707        write_offset: i64,
1708    ) -> Result<ReplicaSnapshot, TransportError> {
1709        // A healthy lane finishes on the same stream that conditionally
1710        // created the object. This needs neither another RPC nor a generation
1711        // lookup; the returned resource supplies the identity used by the
1712        // seal-time metadata CAS.
1713        if self.session.current.load().is_some() {
1714            let finalized = self
1715                .finish_live_session(write_offset, token.generation)
1716                .await?;
1717            token.generation = Some(finalized.generation);
1718            token.metageneration = Some(finalized.metageneration);
1719            return Ok(finalized);
1720        }
1721
1722        // If an idle or failed stream disappeared before finalization, stat
1723        // only to obtain a candidate identity, then require the original
1724        // write handle to resume that exact server-side session. A replacement
1725        // generation cannot be adopted through this path.
1726        let (generation, _) = self.resolve_token_identity(token).await?;
1727        let request = BidiWriteObjectRequest {
1728            first_message: Some(bidi_write_object_request::FirstMessage::AppendObjectSpec(
1729                AppendObjectSpec {
1730                    bucket: self.bucket.clone(),
1731                    object: self.object.clone(),
1732                    generation,
1733                    if_metageneration_match: None,
1734                    write_handle: token
1735                        .write_handle
1736                        .clone()
1737                        .map(|handle| BidiWriteHandle { handle }),
1738                    ..Default::default()
1739                },
1740            )),
1741            write_offset,
1742            finish_write: true,
1743            ..Default::default()
1744        };
1745        let (resource, error) = self
1746            .drive_redirect_aware_stream(
1747                vec![request],
1748                || None,
1749                |resource, response| {
1750                    if let Some(bidi_write_object_response::WriteStatus::Resource(object)) =
1751                        response.write_status
1752                    {
1753                        *resource = Some(object);
1754                    }
1755                },
1756                Option::is_some,
1757            )
1758            .await;
1759        if let Some(error) = error {
1760            return Err(error);
1761        }
1762        let Some(resource) = resource else {
1763            return Err(self.error(
1764                TransportCode::Internal,
1765                "missing finalized resource response",
1766            ));
1767        };
1768        let finalized = self.stat_from_object(resource);
1769        if !finalized.finalized
1770            || finalized.generation != generation
1771            || finalized.persisted_size != write_offset
1772        {
1773            return Err(self.error(
1774                TransportCode::DataLoss,
1775                "finalized segment does not match the recovered prefix",
1776            ));
1777        }
1778        Ok(finalized)
1779    }
1780
1781    async fn shutdown(&self) {
1782        self.session.shutdown.send_replace(true);
1783        self.replace_session(None).await;
1784    }
1785}
1786
1787/// Bounded number of consecutive zonal redirects honored per write open.
1788const MAX_WRITE_REDIRECTS: u32 = 5;
1789
1790/// Minimal view of `google.rpc.Status` for decoding the rich-error payload that
1791/// tonic exposes via `Status::details()` (the `grpc-status-details-bin` trailer).
1792#[derive(Clone, PartialEq, ::prost::Message)]
1793struct RichStatus {
1794    #[prost(int32, tag = "1")]
1795    code: i32,
1796    #[prost(string, tag = "2")]
1797    message: ::prost::alloc::string::String,
1798    #[prost(message, repeated, tag = "3")]
1799    details: ::prost::alloc::vec::Vec<::prost_types::Any>,
1800}
1801
1802/// Extract the routing token from a `BidiWriteObjectRedirectedError` attached to
1803/// an ABORTED status, if present. The feature-gated `probe_support` module
1804/// exports this for repository probes that must replay redirects exactly like
1805/// the production transport.
1806///
1807/// Normal users do not need this helper; the production transport consumes and
1808/// caches redirects automatically.
1809pub fn redirect_routing_token(status: &Status) -> Option<String> {
1810    let details = status.details();
1811    if details.is_empty() {
1812        return None;
1813    }
1814    let rich = RichStatus::decode(details).ok()?;
1815    for any in rich.details {
1816        if any
1817            .type_url
1818            .ends_with("google.storage.v2.BidiWriteObjectRedirectedError")
1819        {
1820            if let Ok(redirect) = BidiWriteObjectRedirectedError::decode(any.value.as_slice()) {
1821                if let Some(token) = redirect.routing_token {
1822                    return Some(token);
1823                }
1824            }
1825        }
1826    }
1827    None
1828}