#
google_cloud_pubsub/subscriber/
session.rs

    1// Copyright 2026 Google LLC
    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//     https://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
    15use super::builder::StreamingPull;
    16use super::handler::{AckResult, AtLeastOnce, Handler};
    17use super::keepalive;
    18use super::lease_loop::LeaseLoop;
    19use super::lease_state::LeaseOptions;
    20use super::leaser::DefaultLeaser;
    21use super::stream::open_stream;
    22use super::stub::{Stub, TonicStreaming};
    23use super::transport::Transport;
    24use crate::google::pubsub::v1::StreamingPullRequest;
    25use crate::model::PubsubMessage;
    26use crate::{Error, Result};
    27use gaxi::grpc::from_status::to_gax_error;
    28use gaxi::prost::FromProto as _;
    29use std::collections::VecDeque;
    30use tokio::sync::mpsc::UnboundedSender;
    31use tokio_util::sync::{CancellationToken, DropGuard};
    32
    33/// Represents an open subscribe session.
    34///
    35/// This is a stream-like struct for serving messages to an application.
    36///
    37/// # Example
    38/// ```
    39/// # use google_cloud_pubsub::client::Subscriber;
    40/// # async fn sample(client: Subscriber) -> anyhow::Result<()> {
    41/// let mut session = client
    42///     .streaming_pull("projects/my-project/subscriptions/my-subscription")
    43///     .start()
    44///     .await?;
    45/// while let Some((m, h)) = session.next().await.transpose()? {
    46///     println!("Received message m={m:?}");
    47///     h.ack();
    48/// }
    49/// # Ok(()) }
    50/// ```
    51#[derive(Debug)]
    52pub struct Session {
    53    /// The bidirectional stream.
    54    stream: <Transport as Stub>::Stream,
    55
    56    /// Applications ask for messages one at a time. Individual stream responses
    57    /// can contain multiple messages. We use `pool` to hold the extra messages
    58    /// while we wait to serve them to applications.
    59    ///
    60    /// A FIFO queue is necessary to preserve ordering.
    61    pool: VecDeque<(PubsubMessage, Handler)>,
    62
    63    /// A sender for sending new messages from the stream into the lease
    64    /// management task.
    65    message_tx: UnboundedSender<String>,
    66
    67    /// A sender for forwarding acks/nacks from the application to the lease
    68    /// management task. Each `Handler` holds a clone of this.
    69    ack_tx: UnboundedSender<AckResult>,
    70
    71    /// A guard which signals a shutdown to the task sending keepalive pings
    72    /// when it is dropped.
    73    _keepalive_guard: DropGuard,
    74
    75    /// A handle on the lease loop task.
    76    ///
    77    /// We hold onto this handle so we can await pending lease operations. While awaiting pending
    78    /// lease operations is useful for setting expectations in our unit tests, it is not that
    79    /// helpful to applications in practice.
    80    _lease_loop: tokio::task::JoinHandle<()>,
    81}
    82
    83impl Session {
    84    pub(super) async fn new(builder: StreamingPull) -> Result<Self> {
    85        let shutdown = CancellationToken::new();
    86        let inner = builder.inner;
    87        let subscription = builder.subscription;
    88
    89        let leaser = DefaultLeaser::new(
    90            inner.clone(),
    91            subscription.clone(),
    92            builder.ack_deadline_seconds,
    93        );
    94        let LeaseLoop {
    95            handle: _lease_loop,
    96            message_tx,
    97            ack_tx,
    98        } = LeaseLoop::new(leaser, LeaseOptions::default());
    99
    100        let initial_req = StreamingPullRequest {
    101            subscription,
    102            stream_ack_deadline_seconds: builder.ack_deadline_seconds,
    103            max_outstanding_messages: builder.max_outstanding_messages,
    104            max_outstanding_bytes: builder.max_outstanding_bytes,
    105            ..Default::default()
    106        };
    107        let (stream, request_tx) = open_stream(inner, initial_req).await?;
    108        keepalive::spawn(request_tx, shutdown.clone());
    109
    110        Ok(Self {
    111            stream,
    112            pool: VecDeque::new(),
    113            message_tx,
    114            ack_tx,
    115            _keepalive_guard: shutdown.drop_guard(),
    116            _lease_loop,
    117        })
    118    }
    119
    120    /// Returns the next message received on this subscription.
    121    ///
    122    /// The message data is returned along with a [Handler] for acknowledging
    123    /// (ack) or rejecting (nack) the message.
    124    ///
    125    /// If the underlying stream encounters a permanent error, an `Error` is
    126    /// returned instead.
    127    ///
    128    /// `None` represents the end of a stream, but in practice, the stream stays
    129    /// open until it is cancelled or encounters a permanent error.
    130    ///
    131    /// # Example
    132    /// ```
    133    /// # use google_cloud_pubsub::subscriber::session::Session;
    134    /// # async fn sample(mut session: Session) -> anyhow::Result<()> {
    135    /// while let Some((m, h)) = session.next().await.transpose()? {
    136    ///     println!("Received message m={m:?}");
    137    ///     h.ack();
    138    /// }
    139    /// # Ok(()) }
    140    /// ```
    141    pub async fn next(&mut self) -> Option<Result<(PubsubMessage, Handler)>> {
    142        loop {
    143            // Serve a message if we have one ready.
    144            if let Some(item) = self.pool.pop_front() {
    145                return Some(Ok(item));
    146            }
    147            // Otherwise, read more messages from the stream.
    148            if let Err(e) = self.stream_next().await? {
    149                return Some(Err(e));
    150            }
    151        }
    152    }
    153
    154    async fn stream_next(&mut self) -> Option<Result<()>> {
    155        let resp = match self.stream.next_message().await.transpose()? {
    156            Ok(resp) => resp,
    157            Err(e) => return Some(Err(to_gax_error(e))),
    158        };
    159        for rm in resp.received_messages {
    160            let Some(message) = rm.message else {
    161                // The message field should always be present. If not, the proto
    162                // message was corrupted while in transit, or there is a bug in
    163                // the service.
    164                //
    165                // The client can just ignore an ack ID without an associated
    166                // message.
    167                continue;
    168            };
    169            let _ = self.message_tx.send(rm.ack_id.clone());
    170            let message = match message.cnv().map_err(Error::deser) {
    171                Ok(message) => message,
    172                Err(e) => return Some(Err(e)),
    173            };
    174            self.pool.push_back((
    175                message,
    176                Handler::AtLeastOnce(AtLeastOnce {
    177                    ack_id: rm.ack_id,
    178                    ack_tx: self.ack_tx.clone(),
    179                }),
    180            ));
    181        }
    182        Some(Ok(()))
    183    }
    184
    185    #[cfg(test)]
    186    /// Close the session, awaiting all pending acks and nacks.
    187    ///
    188    /// This is a useful method for setting clean test expectations.
    189    async fn close(self) -> anyhow::Result<()> {
    190        // Signal a shutdown to the keepalive task.
    191        drop(self._keepalive_guard);
    192
    193        // Signal a shutdown to the lease management background task.
    194        drop(self.message_tx);
    195
    196        // Wait for the lease management task to complete.
    197        self._lease_loop.await?;
    198
    199        Ok(())
    200    }
    201}
    202
    203#[cfg(test)]
    204mod tests {
    205    use super::super::keepalive::KEEPALIVE_PERIOD;
    206    use super::super::lease_state::tests::{test_id, test_ids};
    207    use super::super::transport::tests::test_transport;
    208    use super::*;
    209    use pubsub_grpc_mock::google::pubsub::v1;
    210    use pubsub_grpc_mock::{MockSubscriber, start};
    211    use std::sync::Arc;
    212    use tokio::sync::mpsc::{channel, unbounded_channel};
    213    use tokio::task::JoinHandle;
    214    use tokio::time::Duration;
    215
    216    fn sorted(mut v: Vec<String>) -> Vec<String> {
    217        v.sort();
    218        v
    219    }
    220
    221    fn test_data(v: i32) -> bytes::Bytes {
    222        bytes::Bytes::from(format!("data-{}", test_id(v)))
    223    }
    224
    225    fn test_response(range: std::ops::Range<i32>) -> v1::StreamingPullResponse {
    226        v1::StreamingPullResponse {
    227            received_messages: range
    228                .into_iter()
    229                .map(|i| v1::ReceivedMessage {
    230                    ack_id: test_id(i),
    231                    message: Some(v1::PubsubMessage {
    232                        data: test_data(i).to_vec(),
    233                        ..Default::default()
    234                    }),
    235                    ..Default::default()
    236                })
    237                .collect(),
    238            ..Default::default()
    239        }
    240    }
    241
    242    #[tokio::test]
    243    async fn error_starting_stream() -> anyhow::Result<()> {
    244        let mut mock = MockSubscriber::new();
    245        mock.expect_streaming_pull()
    246            .return_once(|_| Err(tonic::Status::internal("fail")));
    247        let (endpoint, _server) = start("0.0.0.0:0", mock).await?;
    248        let transport = test_transport(endpoint).await?;
    249        let builder = StreamingPull::new(
    250            Arc::new(transport),
    251            "projects/p/subscriptions/s".to_string(),
    252        );
    253        let err = Session::new(builder)
    254            .await
    255            .expect_err("Session should not be created.");
    256        assert!(err.status().is_some(), "{err:?}");
    257        let status = err.status().unwrap();
    258        assert_eq!(status.code, gax::error::rpc::Code::Internal);
    259        assert_eq!(status.message, "fail");
    260
    261        Ok(())
    262    }
    263
    264    #[tokio::test]
    265    async fn initial_request() -> anyhow::Result<()> {
    266        const MIB: i64 = 1024 * 1024;
    267
    268        // We use this channel to surface writes (requests) from outside our
    269        // mock expectation.
    270        let (recover_writes_tx, mut recover_writes_rx) = channel(1);
    271
    272        let mut mock = MockSubscriber::new();
    273        mock.expect_streaming_pull().return_once(move |request| {
    274            tokio::spawn(async move {
    275                // Note that this task stays alive as long as we hold
    276                // `recover_writes_rx`.
    277                let mut request_rx = request.into_inner();
    278                while let Some(request) = request_rx.recv().await {
    279                    recover_writes_tx
    280                        .send(request)
    281                        .await
    282                        .expect("forwarding writes always succeeds");
    283                }
    284            });
    285            Err(tonic::Status::internal("fail"))
    286        });
    287
    288        let (endpoint, _server) = start("0.0.0.0:0", mock).await?;
    289        let transport = test_transport(endpoint).await?;
    290        let builder = StreamingPull::new(
    291            Arc::new(transport),
    292            "projects/p/subscriptions/s".to_string(),
    293        )
    294        .set_ack_deadline_seconds(20)
    295        .set_max_outstanding_messages(2000)
    296        .set_max_outstanding_bytes(200 * MIB);
    297        let _ = Session::new(builder).await;
    298
    299        let initial_req = recover_writes_rx
    300            .recv()
    301            .await
    302            .expect("should receive a request")?;
    303        assert_eq!(initial_req.subscription, "projects/p/subscriptions/s");
    304        assert_eq!(initial_req.stream_ack_deadline_seconds, 20);
    305        assert_eq!(initial_req.max_outstanding_messages, 2000);
    306        assert_eq!(initial_req.max_outstanding_bytes, 200 * MIB);
    307
    308        Ok(())
    309    }
    310
    311    #[tokio::test(start_paused = true)]
    312    async fn basic_success() -> anyhow::Result<()> {
    313        let (response_tx, response_rx) = channel(10);
    314        let (ack_tx, mut ack_rx) = unbounded_channel();
    315
    316        let mut mock = MockSubscriber::new();
    317        mock.expect_streaming_pull()
    318            .return_once(|_| Ok(tonic::Response::from(response_rx)));
    319        mock.expect_acknowledge().returning(move |r| {
    320            ack_tx
    321                .send(r.into_inner())
    322                .expect("sending on channel always succeeds");
    323            Ok(tonic::Response::from(()))
    324        });
    325        let (endpoint, _server) = start("0.0.0.0:0", mock).await?;
    326        let transport = test_transport(endpoint).await?;
    327        let builder = StreamingPull::new(
    328            Arc::new(transport),
    329            "projects/p/subscriptions/s".to_string(),
    330        );
    331        let mut session = Session::new(builder).await?;
    332
    333        response_tx.send(Ok(test_response(1..2))).await?;
    334        response_tx.send(Ok(test_response(2..4))).await?;
    335        response_tx.send(Ok(test_response(4..7))).await?;
    336        drop(response_tx);
    337
    338        for i in 1..7 {
    339            let (m, Handler::AtLeastOnce(h)) =
    340                session.next().await.transpose()?.expect("message {i}/6");
    341            assert_eq!(m.data, test_data(i));
    342            assert_eq!(h.ack_id, test_id(i));
    343            h.ack();
    344        }
    345        let end = session.next().await.transpose()?;
    346        assert!(end.is_none(), "Received extra message: {end:?}");
    347
    348        // Wait for the session to join its background tasks.
    349        session.close().await?;
    350
    351        // Verify the acks went through.
    352        let ack_req = ack_rx.try_recv()?;
    353        assert_eq!(ack_req.subscription, "projects/p/subscriptions/s");
    354        assert_eq!(sorted(ack_req.ack_ids), test_ids(1..7));
    355
    356        Ok(())
    357    }
    358
    359    #[tokio::test(start_paused = true)]
    360    async fn basic_lease_management() -> anyhow::Result<()> {
    361        let (response_tx, response_rx) = channel(10);
    362        let (ack_tx, mut ack_rx) = unbounded_channel();
    363        let (nack_tx, mut nack_rx) = unbounded_channel();
    364        let (extend_tx, mut extend_rx) = unbounded_channel();
    365
    366        let mut mock = MockSubscriber::new();
    367        mock.expect_streaming_pull()
    368            .return_once(|_| Ok(tonic::Response::from(response_rx)));
    369        mock.expect_acknowledge().returning(move |r| {
    370            ack_tx
    371                .send(r.into_inner())
    372                .expect("sending on channel always succeeds");
    373            Ok(tonic::Response::from(()))
    374        });
    375        mock.expect_modify_ack_deadline().returning(move |r| {
    376            let r = r.into_inner();
    377            if r.ack_deadline_seconds == 0 {
    378                nack_tx.send(r).expect("sending on channel always succeeds");
    379            } else {
    380                extend_tx
    381                    .send(r)
    382                    .expect("sending on channel always succeeds");
    383            }
    384            Ok(tonic::Response::from(()))
    385        });
    386        let (endpoint, _server) = start("0.0.0.0:0", mock).await?;
    387        let transport = test_transport(endpoint).await?;
    388        let builder = StreamingPull::new(
    389            Arc::new(transport),
    390            "projects/p/subscriptions/s".to_string(),
    391        );
    392        let mut session = Session::new(builder).await?;
    393
    394        response_tx.send(Ok(test_response(0..30))).await?;
    395        drop(response_tx);
    396
    397        // Ack some messages
    398        for i in 0..10 {
    399            let Some((_, Handler::AtLeastOnce(h))) = session.next().await.transpose()? else {
    400                anyhow::bail!("expected message {i}")
    401            };
    402            h.ack();
    403        }
    404        // Nack some messages
    405        for i in 10..20 {
    406            let Some((_, Handler::AtLeastOnce(h))) = session.next().await.transpose()? else {
    407                anyhow::bail!("expected message {i}")
    408            };
    409            h.nack();
    410        }
    411        // Take a long time to process some messages
    412        let mut hold = Vec::new();
    413        for i in 20..30 {
    414            let Some((_, Handler::AtLeastOnce(h))) = session.next().await.transpose()? else {
    415                anyhow::bail!("expected message {i}")
    416            };
    417            hold.push(h);
    418        }
    419
    420        // Advance the clock 10s, which is the default stream ack deadline. In
    421        // this time, we should attempt at least one lease extension RPC.
    422        tokio::time::advance(Duration::from_secs(10)).await;
    423
    424        // Close the session, to make sure pending operations complete.
    425        session.close().await?;
    426
    427        // Verify the acks went through.
    428        let ack_req = ack_rx.try_recv()?;
    429        assert_eq!(ack_req.subscription, "projects/p/subscriptions/s");
    430        assert_eq!(sorted(ack_req.ack_ids), test_ids(0..10));
    431        assert!(ack_rx.is_empty());
    432
    433        // Verify the initial nacks went through.
    434        let nack_req = nack_rx.try_recv()?;
    435        assert_eq!(nack_req.subscription, "projects/p/subscriptions/s");
    436        assert_eq!(nack_req.ack_deadline_seconds, 0);
    437        assert_eq!(sorted(nack_req.ack_ids), test_ids(10..20));
    438
    439        // Verify that we nack the leftover messages when the stream shuts down.
    440        let nack_req = nack_rx.try_recv()?;
    441        assert_eq!(nack_req.subscription, "projects/p/subscriptions/s");
    442        assert_eq!(nack_req.ack_deadline_seconds, 0);
    443        assert_eq!(sorted(nack_req.ack_ids), test_ids(20..30));
    444        assert!(nack_rx.is_empty());
    445
    446        // Verify at least one lease extension attempt was made.
    447        let extend_req = extend_rx.try_recv()?;
    448        assert_eq!(extend_req.subscription, "projects/p/subscriptions/s");
    449        assert_eq!(extend_req.ack_deadline_seconds, 10);
    450        assert_eq!(sorted(extend_req.ack_ids), test_ids(20..30));
    451
    452        Ok(())
    453    }
    454
    455    #[tokio::test(start_paused = true)]
    456    async fn delayed_responses() -> anyhow::Result<()> {
    457        // In this test, we verify the case where an application asks for a
    458        // message, but a response is not immediately available on the stream.
    459
    460        let (response_tx, response_rx) = channel(10);
    461        let handle: JoinHandle<anyhow::Result<()>> = tokio::spawn(async move {
    462            tokio::time::sleep(Duration::from_millis(20)).await;
    463            response_tx.send(Ok(test_response(1..2))).await?;
    464            Ok(())
    465        });
    466
    467        let mut mock = MockSubscriber::new();
    468        mock.expect_streaming_pull()
    469            .return_once(|_| Ok(tonic::Response::from(response_rx)));
    470        let (endpoint, _server) = start("0.0.0.0:0", mock).await?;
    471        let transport = test_transport(endpoint).await?;
    472        let builder = StreamingPull::new(
    473            Arc::new(transport),
    474            "projects/p/subscriptions/s".to_string(),
    475        );
    476        let mut session = Session::new(builder).await?;
    477        let (m, Handler::AtLeastOnce(h)) = session
    478            .next()
    479            .await
    480            .transpose()?
    481            .expect("stream should wait for a message");
    482        assert_eq!(m.data, test_data(1));
    483        assert_eq!(h.ack_id, test_id(1));
    484
    485        handle.await??;
    486
    487        Ok(())
    488    }
    489
    490    #[tokio::test]
    491    async fn serves_messages_immediately() -> anyhow::Result<()> {
    492        // This test verifies we do not do something crazy like draining the
    493        // stream (which would never end) before serving messages to the
    494        // application.
    495
    496        let (response_tx, response_rx) = channel(10);
    497
    498        let mut mock = MockSubscriber::new();
    499        mock.expect_streaming_pull()
    500            .return_once(|_| Ok(tonic::Response::from(response_rx)));
    501        let (endpoint, _server) = start("0.0.0.0:0", mock).await?;
    502        let transport = test_transport(endpoint).await?;
    503        let builder = StreamingPull::new(
    504            Arc::new(transport),
    505            "projects/p/subscriptions/s".to_string(),
    506        );
    507        let mut session = Session::new(builder).await?;
    508
    509        for i in 1..7 {
    510            response_tx.send(Ok(test_response(i..i + 1))).await?;
    511
    512            let (m, Handler::AtLeastOnce(h)) =
    513                session.next().await.transpose()?.expect("message {i}/6");
    514            assert_eq!(m.data, test_data(i));
    515            assert_eq!(h.ack_id, test_id(i));
    516        }
    517        drop(response_tx);
    518        let end = session.next().await.transpose()?;
    519        assert!(end.is_none(), "Received extra message: {end:?}");
    520
    521        Ok(())
    522    }
    523
    524    #[tokio::test]
    525    async fn handles_empty_response() -> anyhow::Result<()> {
    526        let (response_tx, response_rx) = channel(10);
    527
    528        let mut mock = MockSubscriber::new();
    529        mock.expect_streaming_pull()
    530            .return_once(|_| Ok(tonic::Response::from(response_rx)));
    531        let (endpoint, _server) = start("0.0.0.0:0", mock).await?;
    532        let transport = test_transport(endpoint).await?;
    533        let builder = StreamingPull::new(
    534            Arc::new(transport),
    535            "projects/p/subscriptions/s".to_string(),
    536        );
    537        let mut session = Session::new(builder).await?;
    538
    539        response_tx.send(Ok(test_response(1..2))).await?;
    540        // See if we can handle an empty range
    541        response_tx.send(Ok(test_response(2..2))).await?;
    542        response_tx.send(Ok(test_response(2..3))).await?;
    543        drop(response_tx);
    544
    545        for i in 1..3 {
    546            let (m, Handler::AtLeastOnce(h)) =
    547                session.next().await.transpose()?.expect("message {i}/2");
    548            assert_eq!(m.data, test_data(i));
    549            assert_eq!(h.ack_id, test_id(i));
    550        }
    551        let end = session.next().await.transpose()?;
    552        assert!(end.is_none(), "Received extra message: {end:?}");
    553
    554        Ok(())
    555    }
    556
    557    #[tokio::test(start_paused = true)]
    558    async fn handles_missing_message_field() -> anyhow::Result<()> {
    559        let (response_tx, response_rx) = channel(10);
    560        let (extend_tx, mut extend_rx) = unbounded_channel();
    561
    562        let bad = v1::StreamingPullResponse {
    563            received_messages: vec![v1::ReceivedMessage {
    564                ack_id: "ignored-ack-id".to_string(),
    565                message: None,
    566                ..Default::default()
    567            }],
    568            ..Default::default()
    569        };
    570
    571        let mut mock = MockSubscriber::new();
    572        mock.expect_streaming_pull()
    573            .return_once(|_| Ok(tonic::Response::from(response_rx)));
    574        mock.expect_acknowledge()
    575            .returning(|_| Ok(tonic::Response::from(())));
    576        mock.expect_modify_ack_deadline().returning(move |r| {
    577            extend_tx
    578                .send(r.into_inner())
    579                .expect("sending on channel always succeeds");
    580            Ok(tonic::Response::from(()))
    581        });
    582        let (endpoint, _server) = start("0.0.0.0:0", mock).await?;
    583        let transport = test_transport(endpoint).await?;
    584        let builder = StreamingPull::new(
    585            Arc::new(transport),
    586            "projects/p/subscriptions/s".to_string(),
    587        );
    588        let mut session = Session::new(builder).await?;
    589
    590        response_tx.send(Ok(test_response(1..4))).await?;
    591        // See if we can handle an empty range
    592        response_tx.send(Ok(bad)).await?;
    593        response_tx.send(Ok(test_response(4..7))).await?;
    594        drop(response_tx);
    595
    596        for i in 1..7 {
    597            let (m, Handler::AtLeastOnce(h)) =
    598                session.next().await.transpose()?.expect("message {i}/6");
    599            assert_eq!(m.data, test_data(i));
    600            assert_eq!(h.ack_id, test_id(i));
    601        }
    602        let end = session.next().await.transpose()?;
    603        assert!(end.is_none(), "Received extra message: {end:?}");
    604
    605        // Advance the clock 10s, which is the default stream ack deadline. In
    606        // this time, we should attempt at least one lease extension RPC.
    607        tokio::time::advance(Duration::from_secs(10)).await;
    608
    609        // Close the session, to make sure pending operations complete.
    610        session.close().await?;
    611
    612        // Verify at least one lease extension attempt was made.
    613        let extend_req = extend_rx.try_recv()?;
    614        assert_eq!(extend_req.subscription, "projects/p/subscriptions/s");
    615        assert_eq!(extend_req.ack_deadline_seconds, 10);
    616        // Note that we do not expect to see "ignored-ack-id".
    617        assert_eq!(sorted(extend_req.ack_ids), test_ids(1..7));
    618
    619        Ok(())
    620    }
    621
    622    #[tokio::test]
    623    async fn permanent_error_midstream() -> anyhow::Result<()> {
    624        let (response_tx, response_rx) = channel(10);
    625
    626        let mut mock = MockSubscriber::new();
    627        mock.expect_streaming_pull()
    628            .return_once(|_| Ok(tonic::Response::from(response_rx)));
    629        let (endpoint, _server) = start("0.0.0.0:0", mock).await?;
    630        let transport = test_transport(endpoint).await?;
    631        let builder = StreamingPull::new(
    632            Arc::new(transport),
    633            "projects/p/subscriptions/s".to_string(),
    634        );
    635        let mut session = Session::new(builder).await?;
    636
    637        response_tx.send(Ok(test_response(1..4))).await?;
    638        response_tx
    639            .send(Err(tonic::Status::internal("fail")))
    640            .await?;
    641        drop(response_tx);
    642
    643        for i in 1..4 {
    644            let (m, Handler::AtLeastOnce(h)) =
    645                session.next().await.transpose()?.expect("message {i}/3");
    646            assert_eq!(m.data, test_data(i));
    647            assert_eq!(h.ack_id, test_id(i));
    648        }
    649        let err = session
    650            .next()
    651            .await
    652            .transpose()
    653            .expect_err("expected an error from stream");
    654        assert!(err.status().is_some(), "{err:?}");
    655        let status = err.status().unwrap();
    656        assert_eq!(status.code, gax::error::rpc::Code::Internal);
    657        assert_eq!(status.message, "fail");
    658
    659        Ok(())
    660    }
    661
    662    #[tokio::test(start_paused = true)]
    663    async fn keepalives() -> anyhow::Result<()> {
    664        // We use this channel to surface writes (requests) from outside our
    665        // mock expectation.
    666        let (recover_writes_tx, mut recover_writes_rx) = channel(1);
    667        let (_response_tx, response_rx) = channel(10);
    668
    669        let mut mock = MockSubscriber::new();
    670        mock.expect_streaming_pull().return_once(move |request| {
    671            tokio::spawn(async move {
    672                // Note that this task stays alive as long as we hold
    673                // `recover_writes_rx`.
    674                let mut request_rx = request.into_inner();
    675                while let Some(request) = request_rx.recv().await {
    676                    recover_writes_tx
    677                        .send(request)
    678                        .await
    679                        .expect("forwarding writes always succeeds");
    680                }
    681            });
    682            Ok(tonic::Response::from(response_rx))
    683        });
    684
    685        let (endpoint, _server) = start("0.0.0.0:0", mock).await?;
    686        let transport = test_transport(endpoint).await?;
    687        let builder = StreamingPull::new(
    688            Arc::new(transport),
    689            "projects/p/subscriptions/s".to_string(),
    690        );
    691        let session = Session::new(builder).await;
    692
    693        let initial_req = recover_writes_rx
    694            .recv()
    695            .await
    696            .expect("should receive an initial request")?;
    697        assert_eq!(initial_req.subscription, "projects/p/subscriptions/s");
    698
    699        // Verify that we receive at least one keepalive request on the stream.
    700        tokio::time::advance(KEEPALIVE_PERIOD).await;
    701        let keepalive_req = recover_writes_rx
    702            .recv()
    703            .await
    704            .expect("should receive a keepalive request")?;
    705        assert_eq!(keepalive_req, v1::StreamingPullRequest::default());
    706
    707        // Drop the session, which should signal a shutdown of the keepalive
    708        // task.
    709        drop(session);
    710
    711        // Advance the time far enough to expect a keepalive ping, if the
    712        // keepalive task was still running.
    713        tokio::time::advance(4 * KEEPALIVE_PERIOD).await;
    714        assert!(recover_writes_rx.is_empty());
    715
    716        Ok(())
    717    }
    718}