#
google_cloud_pubsub/subscriber/
lease_loop.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::handler::AckResult;
    16use super::lease_state::{LeaseEvent, LeaseOptions, LeaseState};
    17use super::leaser::Leaser;
    18use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel};
    19use tokio::task::JoinHandle;
    20
    21/// A convenience struct that groups the components of the lease loop.
    22pub(super) struct LeaseLoop {
    23    /// A handle to the task running the lease loop.
    24    pub(super) handle: JoinHandle<()>,
    25    /// For sending messages from the stream to the lease loop.
    26    pub(super) message_tx: UnboundedSender<String>,
    27    /// For sending acks/nacks from the application to the lease loop.
    28    pub(super) ack_tx: UnboundedSender<AckResult>,
    29}
    30
    31impl LeaseLoop {
    32    pub(super) fn new<L>(leaser: L, options: LeaseOptions) -> Self
    33    where
    34        L: Leaser + Send + 'static,
    35    {
    36        let (message_tx, mut message_rx) = unbounded_channel();
    37        let (ack_tx, mut ack_rx) = unbounded_channel();
    38        let mut state = LeaseState::new(leaser, options);
    39
    40        let handle = tokio::spawn(async move {
    41            loop {
    42                tokio::select! {
    43                    biased;
    44                    event = state.next_event() => {
    45                        match event {
    46                            LeaseEvent::Flush => state.flush().await,
    47                            LeaseEvent::Extend => state.extend().await,
    48                        }
    49                    },
    50                    message = message_rx.recv() => {
    51                        match message {
    52                            None => break shutdown(state, ack_rx).await,
    53                            Some(ack_id) => state.add(ack_id),
    54                        }
    55                    },
    56                    ack_id = ack_rx.recv() => {
    57                        match ack_id {
    58                            None => break,
    59                            Some(AckResult::Ack(ack_id)) => state.ack(ack_id),
    60                            Some(AckResult::Nack(ack_id)) => state.nack(ack_id),
    61                        }
    62                    },
    63                }
    64            }
    65        });
    66        LeaseLoop {
    67            handle,
    68            message_tx,
    69            ack_tx,
    70        }
    71    }
    72}
    73
    74// Shuts down lease management.
    75//
    76// Processes any acks from the application that we already know about and
    77// triggers a shutdown of the lease state.
    78async fn shutdown<L>(mut state: LeaseState<L>, mut ack_rx: UnboundedReceiver<AckResult>)
    79where
    80    L: Leaser + Send + 'static,
    81{
    82    while let Ok(r) = ack_rx.try_recv() {
    83        if let AckResult::Ack(ack_id) = r {
    84            state.ack(ack_id);
    85        }
    86    }
    87    state.shutdown().await;
    88}
    89
    90#[cfg(test)]
    91mod tests {
    92    use super::super::lease_state::tests::{sorted, test_id, test_ids};
    93    use super::super::leaser::tests::MockLeaser;
    94    use super::*;
    95    use std::sync::Arc;
    96    use tokio::sync::Mutex;
    97    use tokio::time::{Duration, Instant};
    98
    99    #[tokio::test(start_paused = true)]
    100    async fn flush_acks_nacks_on_interval() -> anyhow::Result<()> {
    101        const FLUSH_PERIOD: Duration = Duration::from_secs(1);
    102        const FLUSH_START: Duration = Duration::from_millis(200);
    103
    104        let mock = Arc::new(Mutex::new(MockLeaser::new()));
    105
    106        let options = LeaseOptions {
    107            flush_period: FLUSH_PERIOD,
    108            flush_start: FLUSH_START,
    109            // effectively disable extensions to simplify this test.
    110            extend_start: Duration::from_secs(900),
    111            ..Default::default()
    112        };
    113        let lease_loop = LeaseLoop::new(mock.clone(), options);
    114        // Yield execution, so tokio can actually start the lease loop.
    115        tokio::task::yield_now().await;
    116
    117        // Seed the lease loop with some messages
    118        for i in 0..30 {
    119            lease_loop.message_tx.send(test_id(i))?;
    120        }
    121
    122        // Ack 10 messages
    123        for i in 0..10 {
    124            lease_loop.ack_tx.send(AckResult::Ack(test_id(i)))?;
    125        }
    126
    127        // Confirm initial state
    128        mock.lock().await.checkpoint();
    129
    130        // Advance to and validate the first flush
    131        {
    132            mock.lock()
    133                .await
    134                .expect_ack()
    135                .times(1)
    136                .withf(|v| sorted(v) == test_ids(0..10))
    137                .returning(move |_| ());
    138            mock.lock()
    139                .await
    140                .expect_nack()
    141                .times(1)
    142                .withf(|v| v.is_empty())
    143                .returning(|_| ());
    144            tokio::time::advance(FLUSH_START).await;
    145
    146            // Yield the current task, so tokio can execute the flush().
    147            tokio::task::yield_now().await;
    148            mock.lock().await.checkpoint();
    149        }
    150
    151        // Nack 10 messages
    152        for i in 10..20 {
    153            lease_loop.ack_tx.send(AckResult::Nack(test_id(i)))?;
    154        }
    155
    156        // Advance to and validate the second flush
    157        {
    158            mock.lock()
    159                .await
    160                .expect_ack()
    161                .times(1)
    162                .withf(|v| v.is_empty())
    163                .returning(move |_| ());
    164            mock.lock()
    165                .await
    166                .expect_nack()
    167                .times(1)
    168                .withf(|v| sorted(v) == test_ids(10..20))
    169                .returning(|_| ());
    170            tokio::time::advance(FLUSH_PERIOD).await;
    171
    172            // Yield the current task, so tokio can execute the flush().
    173            tokio::task::yield_now().await;
    174            mock.lock().await.checkpoint();
    175        }
    176
    177        // Ack 5 messages
    178        for i in 20..25 {
    179            lease_loop.ack_tx.send(AckResult::Ack(test_id(i)))?;
    180        }
    181        // Nack 5 messages
    182        for i in 25..30 {
    183            lease_loop.ack_tx.send(AckResult::Nack(test_id(i)))?;
    184        }
    185
    186        // Advance to the third flush
    187        {
    188            mock.lock()
    189                .await
    190                .expect_ack()
    191                .times(1)
    192                .withf(|v| sorted(v) == test_ids(20..25))
    193                .returning(move |_| ());
    194            mock.lock()
    195                .await
    196                .expect_nack()
    197                .times(1)
    198                .withf(|v| sorted(v) == test_ids(25..30))
    199                .returning(|_| ());
    200            tokio::time::advance(FLUSH_PERIOD).await;
    201
    202            // Yield the current task, so tokio can execute the flush().
    203            tokio::task::yield_now().await;
    204            mock.lock().await.checkpoint();
    205        }
    206
    207        Ok(())
    208    }
    209
    210    #[tokio::test(start_paused = true)]
    211    async fn deadline_interval() -> anyhow::Result<()> {
    212        const EXTEND_PERIOD: Duration = Duration::from_secs(1);
    213        const EXTEND_START: Duration = Duration::from_millis(200);
    214
    215        let mock = Arc::new(Mutex::new(MockLeaser::new()));
    216
    217        let options = LeaseOptions {
    218            // effectively disable ack/nack flushes to simplify this test.
    219            flush_start: Duration::from_secs(900),
    220            extend_period: EXTEND_PERIOD,
    221            extend_start: EXTEND_START,
    222            ..Default::default()
    223        };
    224        let lease_loop = LeaseLoop::new(mock.clone(), options);
    225        // Yield execution, so tokio can actually start the lease loop.
    226        tokio::task::yield_now().await;
    227
    228        // Seed the lease loop with some messages
    229        for i in 0..30 {
    230            lease_loop.message_tx.send(test_id(i))?;
    231        }
    232
    233        // Confirm initial state
    234        mock.lock().await.checkpoint();
    235
    236        // Advance to and validate the first extension
    237        {
    238            mock.lock()
    239                .await
    240                .expect_extend()
    241                .times(1)
    242                .withf(|v| sorted(v) == test_ids(0..30))
    243                .returning(move |_| ());
    244            tokio::time::advance(EXTEND_START).await;
    245
    246            // Yield the current task, so tokio can execute the flush().
    247            tokio::task::yield_now().await;
    248            mock.lock().await.checkpoint();
    249        }
    250
    251        // Ack 10 messages
    252        for i in 0..10 {
    253            lease_loop.ack_tx.send(AckResult::Ack(test_id(i)))?;
    254        }
    255
    256        // Advance to and validate the second extension
    257        {
    258            mock.lock()
    259                .await
    260                .expect_extend()
    261                .times(1)
    262                .withf(|v| sorted(v) == test_ids(10..30))
    263                .returning(|_| ());
    264            tokio::time::advance(EXTEND_PERIOD).await;
    265
    266            // Yield the current task, so tokio can execute the flush().
    267            tokio::task::yield_now().await;
    268            mock.lock().await.checkpoint();
    269        }
    270
    271        Ok(())
    272    }
    273
    274    #[tokio::test(start_paused = true)]
    275    async fn drop_does_not_wait_for_pending_operations() -> anyhow::Result<()> {
    276        let start = Instant::now();
    277        let mock = MockLeaser::new();
    278        let lease_loop = LeaseLoop::new(mock, LeaseOptions::default());
    279        // Yield execution, so tokio can actually start the lease loop.
    280        tokio::task::yield_now().await;
    281
    282        // Seed the lease loop with some messages
    283        for i in 0..30 {
    284            lease_loop.message_tx.send(test_id(i))?;
    285        }
    286
    287        // Ack 10 messages
    288        for i in 0..10 {
    289            lease_loop.ack_tx.send(AckResult::Ack(test_id(i)))?;
    290        }
    291
    292        // Drop the lease_loop.
    293        drop(lease_loop);
    294
    295        // Verify no time has passed.
    296        assert_eq!(start.elapsed(), Duration::ZERO);
    297
    298        Ok(())
    299    }
    300
    301    #[tokio::test(start_paused = true)]
    302    async fn close_waits_for_flush() -> anyhow::Result<()> {
    303        const EXPECTED_SLEEP: Duration = Duration::from_millis(100);
    304
    305        let start = Instant::now();
    306
    307        struct FakeLeaser;
    308        #[async_trait::async_trait]
    309        impl Leaser for FakeLeaser {
    310            async fn ack(&self, mut ack_ids: Vec<String>) {
    311                ack_ids.sort();
    312                assert_eq!(ack_ids, test_ids(0..10));
    313                tokio::time::sleep(EXPECTED_SLEEP).await;
    314            }
    315            async fn nack(&self, mut ack_ids: Vec<String>) {
    316                ack_ids.sort();
    317                assert_eq!(ack_ids, test_ids(10..30));
    318            }
    319            async fn extend(&self, _ack_ids: Vec<String>) {}
    320        }
    321
    322        let lease_loop = LeaseLoop::new(FakeLeaser, LeaseOptions::default());
    323
    324        // Seed the lease loop with some messages
    325        for i in 0..30 {
    326            lease_loop.message_tx.send(test_id(i))?;
    327        }
    328
    329        // Ack 10 messages
    330        for i in 0..10 {
    331            lease_loop.ack_tx.send(AckResult::Ack(test_id(i)))?;
    332        }
    333
    334        // Shutdown the lease_loop.
    335        drop(lease_loop.message_tx);
    336        lease_loop.handle.await?;
    337
    338        // Verify that we flushed the acks immediately, and waited for them to
    339        // complete.
    340        assert_eq!(start.elapsed(), EXPECTED_SLEEP);
    341
    342        Ok(())
    343    }
    344
    345    #[tokio::test(start_paused = true)]
    346    async fn no_add_and_ack_race() -> anyhow::Result<()> {
    347        // This test validates the use of `biased` in the select statement.
    348        //
    349        // Specifically, we want incoming messages to be processed by the event
    350        // loop before any incoming acks/nacks.
    351        //
    352        // If these channels can race, we might accept an application's
    353        // acknowledgement, and then immediately after, accept that message under
    354        // lease management.
    355
    356        for _ in 0..1000 {
    357            // Run this test enough times to trigger a race, if one existed.
    358
    359            let mock = Arc::new(Mutex::new(MockLeaser::new()));
    360            let options = LeaseOptions {
    361                flush_start: Duration::from_millis(100),
    362                extend_start: Duration::from_millis(200),
    363                ..Default::default()
    364            };
    365            let lease_loop = LeaseLoop::new(mock.clone(), options);
    366            // Yield execution, so tokio can actually start the lease loop.
    367            tokio::task::yield_now().await;
    368
    369            // Seed the lease loop with a message
    370            lease_loop.message_tx.send(test_id(1))?;
    371            // Immediately ack the message
    372            lease_loop.ack_tx.send(AckResult::Ack(test_id(1)))?;
    373
    374            // Advance to and validate the first flush
    375            {
    376                mock.lock()
    377                    .await
    378                    .expect_ack()
    379                    .times(1)
    380                    .withf(|v| *v == vec![test_id(1)])
    381                    .returning(|_| ());
    382                mock.lock()
    383                    .await
    384                    .expect_nack()
    385                    .times(1)
    386                    .withf(|v| v.is_empty())
    387                    .returning(|_| ());
    388                tokio::time::advance(Duration::from_millis(100)).await;
    389
    390                // Yield the current task, so tokio can execute the flush().
    391                tokio::task::yield_now().await;
    392                mock.lock().await.checkpoint();
    393            }
    394
    395            // Confirm that no messages are under lease management.
    396            {
    397                mock.lock().await.expect_extend().times(0);
    398                tokio::time::advance(Duration::from_millis(100)).await;
    399
    400                // Yield the current task, so tokio can execute the flush().
    401                tokio::task::yield_now().await;
    402                mock.lock().await.checkpoint();
    403            }
    404        }
    405        Ok(())
    406    }
    407}