#
google_cloud_pubsub/subscriber/
keepalive.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 crate::google::pubsub::v1::StreamingPullRequest;
    16use tokio::sync::mpsc::Sender;
    17use tokio::task::JoinHandle;
    18use tokio::time::{Duration, Instant, interval_at};
    19use tokio_util::sync::CancellationToken;
    20
    21pub(super) const KEEPALIVE_PERIOD: Duration = Duration::from_secs(30);
    22
    23/// Spawns a task to keepalive a stream
    24///
    25/// This task periodically writes requests into a channel. The receiver of this
    26/// channel is the request stream for a StreamingPull bidi RPC.
    27///
    28/// Callers may signal a graceful shutdown of this task by cancelling the
    29/// `CancellationToken` and `await`ing the returned handle.
    30///
    31/// Callers can also just drop the returned handle to shutdown.
    32pub(super) fn spawn(
    33    request_tx: Sender<StreamingPullRequest>,
    34    shutdown: CancellationToken,
    35) -> JoinHandle<()> {
    36    tokio::spawn(async move {
    37        let mut keepalive = interval_at(Instant::now() + KEEPALIVE_PERIOD, KEEPALIVE_PERIOD);
    38        loop {
    39            tokio::select! {
    40                _ = shutdown.cancelled() => break,
    41                _ = keepalive.tick() => {
    42                    let _ = request_tx.send(StreamingPullRequest::default()).await;
    43                }
    44            }
    45        }
    46    })
    47}
    48
    49#[cfg(test)]
    50mod tests {
    51    use super::*;
    52    use tokio::sync::mpsc::channel;
    53
    54    #[tokio::test(start_paused = true)]
    55    async fn keepalive_interval() {
    56        let start = Instant::now();
    57        let (request_tx, mut request_rx) = channel(1);
    58        let shutdown = CancellationToken::new();
    59        let _handle = spawn(request_tx, shutdown);
    60
    61        // Wait for the first keepalive
    62        let r = request_rx.recv().await.unwrap();
    63        assert_eq!(r, StreamingPullRequest::default());
    64        assert_eq!(start.elapsed(), KEEPALIVE_PERIOD);
    65
    66        // Wait for the second keepalive
    67        let r = request_rx.recv().await.unwrap();
    68        assert_eq!(r, StreamingPullRequest::default());
    69        assert_eq!(start.elapsed(), KEEPALIVE_PERIOD * 2);
    70
    71        // Wait for the third keepalive
    72        let r = request_rx.recv().await.unwrap();
    73        assert_eq!(r, StreamingPullRequest::default());
    74        assert_eq!(start.elapsed(), KEEPALIVE_PERIOD * 3);
    75    }
    76
    77    #[tokio::test(start_paused = true)]
    78    async fn shutdown_immediately() -> anyhow::Result<()> {
    79        let start = Instant::now();
    80        let (request_tx, mut request_rx) = channel(1);
    81        let shutdown = CancellationToken::new();
    82        let handle = spawn(request_tx, shutdown.clone());
    83
    84        // Wait for the first keepalive
    85        let _ = request_rx.recv().await.unwrap();
    86        assert_eq!(start.elapsed(), KEEPALIVE_PERIOD);
    87
    88        // Simulate the loop running for a bit.
    89        const DELTA: Duration = Duration::from_secs(10);
    90        tokio::time::advance(DELTA).await;
    91
    92        // Shutdown the task
    93        shutdown.cancel();
    94        handle.await?;
    95
    96        // Verify that we did not wait for the full keepalive interval.
    97        assert_eq!(start.elapsed(), KEEPALIVE_PERIOD + DELTA);
    98        Ok(())
    99    }
    100}