#
google_cloud_pubsub/subscriber/
stream.rs

    1// Copyright 2025 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::stub::Stub;
    16use crate::google::pubsub::v1::StreamingPullRequest;
    17use crate::{Error, Result};
    18use gax::options::RequestOptions;
    19use std::sync::Arc;
    20use tokio::sync::mpsc;
    21
    22/// Open a stream for the `StreamingPull` RPC.
    23///
    24/// This returns the stream and a Sender for feeding the stream writes.
    25pub(super) async fn open_stream<T>(
    26    inner: Arc<T>,
    27    initial_req: StreamingPullRequest,
    28) -> Result<(<T as Stub>::Stream, mpsc::Sender<StreamingPullRequest>)>
    29where
    30    T: Stub,
    31{
    32    // The only writes we perform are keepalives, which are sent so infrequently
    33    // that we don't fear any back pressure on this channel.
    34    let (request_tx, request_rx) = mpsc::channel(1);
    35    request_tx.send(initial_req).await.map_err(Error::io)?;
    36    let stream = inner
    37        .streaming_pull(request_rx, RequestOptions::default())
    38        .await?
    39        .into_inner();
    40
    41    Ok((stream, request_tx))
    42}
    43
    44#[cfg(test)]
    45mod tests {
    46    use super::super::lease_state::tests::test_ids;
    47    use super::super::stub::TonicStreaming;
    48    use super::super::stub::tests::MockStub;
    49    use super::*;
    50    use crate::google::pubsub::v1::{ReceivedMessage, StreamingPullResponse};
    51
    52    fn test_response(range: std::ops::Range<i32>) -> StreamingPullResponse {
    53        StreamingPullResponse {
    54            received_messages: test_ids(range)
    55                .into_iter()
    56                .map(|ack_id| ReceivedMessage {
    57                    ack_id,
    58                    ..Default::default()
    59                })
    60                .collect(),
    61            ..Default::default()
    62        }
    63    }
    64
    65    fn initial_request() -> StreamingPullRequest {
    66        StreamingPullRequest {
    67            subscription: "projects/my-project/subscriptions/my-subscription".to_string(),
    68            stream_ack_deadline_seconds: 10,
    69            ..Default::default()
    70        }
    71    }
    72
    73    fn keepalive_request() -> StreamingPullRequest {
    74        StreamingPullRequest {
    75            subscription: "projects/my-project/subscriptions/my-subscription".to_string(),
    76            ..Default::default()
    77        }
    78    }
    79
    80    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    81    async fn success() -> anyhow::Result<()> {
    82        let (response_tx, response_rx) = mpsc::channel(10);
    83        // We use this channel to surface writes (requests) from outside our
    84        // mock stream.
    85        let (recover_writes_tx, mut recover_writes_rx) = mpsc::channel(1);
    86
    87        let mut mock = MockStub::new();
    88        mock.expect_streaming_pull()
    89            .times(1)
    90            .return_once(move |mut request_rx, _o| {
    91                tokio::spawn(async move {
    92                    // Note that this task stays alive as long as we hold
    93                    // `recover_writes_rx`.
    94                    while let Some(request) = request_rx.recv().await {
    95                        recover_writes_tx
    96                            .send(request)
    97                            .await
    98                            .expect("forwarding writes always succeeds");
    99                    }
    100                });
    101                Ok(tonic::Response::from(response_rx))
    102            });
    103
    104        let (mut stream, request_tx) = open_stream(Arc::new(mock), initial_request()).await?;
    105
    106        // Verify the stream is seeded with the initial request.
    107        assert_eq!(recover_writes_rx.recv().await, Some(initial_request()));
    108
    109        // Verify we can read from the stream.
    110        response_tx.send(Ok(test_response(1..10))).await?;
    111        response_tx.send(Ok(test_response(11..20))).await?;
    112        response_tx.send(Ok(test_response(21..30))).await?;
    113        assert_eq!(stream.next_message().await?, Some(test_response(1..10)));
    114        assert_eq!(stream.next_message().await?, Some(test_response(11..20)));
    115
    116        // Verify we can write to the stream from `request_tx`.
    117        request_tx.send(keepalive_request()).await?;
    118        assert_eq!(recover_writes_rx.recv().await, Some(keepalive_request()));
    119
    120        // Read the last batch of messages (verifying bidi nature of stream).
    121        assert_eq!(stream.next_message().await?, Some(test_response(21..30)));
    122
    123        // Drop the sender
    124        drop(request_tx);
    125        assert_eq!(recover_writes_rx.recv().await, None);
    126
    127        Ok(())
    128    }
    129
    130    #[tokio::test]
    131    async fn error() -> anyhow::Result<()> {
    132        let mut mock = MockStub::new();
    133        mock.expect_streaming_pull()
    134            .times(1)
    135            .return_once(|_, _| Err(Error::io("fail")));
    136
    137        let err = open_stream(Arc::new(mock), initial_request())
    138            .await
    139            .expect_err("open_stream should fail");
    140        assert!(err.is_io(), "{err:?}");
    141
    142        Ok(())
    143    }
    144}