#
google_cloud_pubsub/subscriber/
lease_state.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::leaser::Leaser;
    16use std::collections::HashSet;
    17use tokio::time::{Duration, Instant, Interval, interval_at};
    18
    19// An ack ID is less than 200 bytes. The limit for a request is 512kB. It should
    20// be safe to fit 2500 Ack IDs in a single RPC.
    21//
    22// https://docs.cloud.google.com/pubsub/quotas
    23const ACK_IDS_PER_RPC: usize = 2500;
    24
    25pub(super) struct LeaseOptions {
    26    /// How often we flush acks/nacks
    27    pub(super) flush_period: Duration,
    28    /// How long we wait for the initial flush
    29    pub(super) flush_start: Duration,
    30    /// How often we extend deadlines for messages under lease
    31    pub(super) extend_period: Duration,
    32    /// How long we wait for the initial extensions
    33    pub(super) extend_start: Duration,
    34}
    35
    36impl Default for LeaseOptions {
    37    fn default() -> Self {
    38        LeaseOptions {
    39            flush_period: Duration::from_secs(1),
    40            flush_start: Duration::from_secs(1),
    41            extend_period: Duration::from_secs(3),
    42            extend_start: Duration::from_millis(500),
    43        }
    44    }
    45}
    46
    47#[derive(Debug)]
    48pub(super) struct LeaseState<L>
    49where
    50    L: Leaser,
    51{
    52    // TODO(#3957) - support message expiry
    53    under_lease: HashSet<String>,
    54    to_ack: Vec<String>,
    55    to_nack: Vec<String>,
    56    // TODO(#3964) - support exactly once acks
    57    leaser: L,
    58
    59    // A timer for flushing acks/nacks
    60    flush_interval: Interval,
    61    // A timer for extending leases
    62    extend_interval: Interval,
    63}
    64
    65/// Actions taken by the `LeaseState` in the lease loop.
    66#[derive(Debug, PartialEq)]
    67pub(super) enum LeaseEvent {
    68    /// Flush acks/nacks
    69    Flush,
    70    /// Extend leases
    71    Extend,
    72}
    73
    74impl<L> LeaseState<L>
    75where
    76    L: Leaser,
    77{
    78    pub(super) fn new(leaser: L, options: LeaseOptions) -> Self {
    79        let flush_interval =
    80            interval_at(Instant::now() + options.flush_start, options.flush_period);
    81        let extend_interval =
    82            interval_at(Instant::now() + options.extend_start, options.extend_period);
    83        Self {
    84            under_lease: HashSet::new(),
    85            to_ack: Vec::new(),
    86            to_nack: Vec::new(),
    87            leaser,
    88            flush_interval,
    89            extend_interval,
    90        }
    91    }
    92
    93    /// A future that fires when it is time to either:
    94    /// - flush acks/nacks
    95    /// - extend leases
    96    ///
    97    /// We need to centralize event handling because the lease loop can only
    98    /// hold one mutable reference to `LeaseState` within its `select!`
    99    /// statement.
    100    pub(super) async fn next_event(&mut self) -> LeaseEvent {
    101        if self.to_ack.len() >= ACK_IDS_PER_RPC || self.to_nack.len() >= ACK_IDS_PER_RPC {
    102            // This is an OR because `Acknowledge` and `ModifyAckDeadline` are
    103            // separate RPCs, with separate limits.
    104            return LeaseEvent::Flush;
    105        }
    106
    107        tokio::select! {
    108            _ = self.flush_interval.tick() => LeaseEvent::Flush,
    109            _ = self.extend_interval.tick() => LeaseEvent::Extend,
    110        }
    111    }
    112
    113    /// Accept a new ack ID under lease management
    114    pub(super) fn add(&mut self, ack_id: String) {
    115        self.under_lease.insert(ack_id);
    116    }
    117
    118    /// Process an ack from the application
    119    pub(super) fn ack(&mut self, ack_id: String) {
    120        self.under_lease.remove(&ack_id);
    121        // Unconditionally add the ack ID to the next ack batch. It doesn't hurt
    122        // to optimistically add it, even if its lease has expired.
    123        self.to_ack.push(ack_id);
    124    }
    125
    126    /// Process a nack from the application
    127    pub(super) fn nack(&mut self, ack_id: String) {
    128        if self.under_lease.remove(&ack_id) {
    129            // Only add the ack ID to the nack batch if the message is under our
    130            // lease. If the message's lease has already expired, we do not need
    131            // to take any additional action.
    132            self.to_nack.push(ack_id);
    133        }
    134    }
    135
    136    /// Flush pending acks/nacks
    137    pub(super) async fn flush(&mut self) {
    138        let to_ack = std::mem::take(&mut self.to_ack);
    139        let to_nack = std::mem::take(&mut self.to_nack);
    140        // TODO(#3975) - await these concurrently.
    141        self.leaser.ack(to_ack).await;
    142        self.leaser.nack(to_nack).await;
    143    }
    144
    145    /// Extends leases for messages under lease management
    146    ///
    147    /// Drops messages whose lease deadline cannot be extended any further.
    148    pub(super) async fn extend(&mut self) {
    149        // TODO(#3957) - drop expired messages
    150        let all_ack_ids: Vec<&String> = self.under_lease.iter().collect();
    151        for chunk in all_ack_ids.chunks(ACK_IDS_PER_RPC) {
    152            // TODO(#3975) - await these concurrently.
    153            let ack_ids: Vec<String> = chunk.iter().map(|&s| s.clone()).collect();
    154            self.leaser.extend(ack_ids).await;
    155        }
    156    }
    157
    158    /// Shutdown the leaser
    159    ///
    160    /// This flushes all pending acks and nacks all other messages.
    161    pub(super) async fn shutdown(self) {
    162        let mut to_nack = self.to_nack;
    163        to_nack.extend(self.under_lease.into_iter());
    164        // TODO(#3975) - await these concurrently.
    165        self.leaser.ack(self.to_ack).await;
    166        self.leaser.nack(to_nack).await;
    167    }
    168}
    169
    170impl<L> PartialEq for LeaseState<L>
    171where
    172    L: Leaser,
    173{
    174    fn eq(&self, other: &Self) -> bool {
    175        self.under_lease == other.under_lease
    176            && self.to_ack == other.to_ack
    177            && self.to_nack == other.to_nack
    178    }
    179}
    180
    181#[cfg(test)]
    182pub(super) mod tests {
    183    use super::super::leaser::tests::MockLeaser;
    184    use super::*;
    185    use tokio::sync::mpsc::unbounded_channel;
    186    use tokio::time::interval;
    187
    188    // Cover the constant, converting it to an integer for convenience.
    189    const ACK_IDS_PER_RPC: i32 = super::ACK_IDS_PER_RPC as i32;
    190
    191    // Any valid `Interval` will do.
    192    fn test_interval() -> Interval {
    193        interval(Duration::from_secs(1))
    194    }
    195
    196    fn make_state<L, A, N>(under_lease: L, to_ack: A, to_nack: N) -> LeaseState<MockLeaser>
    197    where
    198        L: IntoIterator<Item = &'static str>,
    199        A: IntoIterator<Item = &'static str>,
    200        N: IntoIterator<Item = &'static str>,
    201    {
    202        LeaseState {
    203            under_lease: under_lease.into_iter().map(|s| s.to_string()).collect(),
    204            to_ack: to_ack.into_iter().map(|s| s.to_string()).collect(),
    205            to_nack: to_nack.into_iter().map(|s| s.to_string()).collect(),
    206            leaser: MockLeaser::new(),
    207            flush_interval: test_interval(),
    208            extend_interval: test_interval(),
    209        }
    210    }
    211
    212    pub(in super::super) fn test_id(v: i32) -> String {
    213        format!("{v:05}")
    214    }
    215
    216    pub(in super::super) fn test_ids(range: std::ops::Range<i32>) -> Vec<String> {
    217        range.map(test_id).collect()
    218    }
    219
    220    pub(in super::super) fn sorted(v: &[String]) -> Vec<String> {
    221        let mut s = v.to_owned();
    222        s.sort();
    223        s
    224    }
    225
    226    #[tokio::test]
    227    async fn basic_add_ack_nack() {
    228        let mock = MockLeaser::new();
    229        let mut state = LeaseState::new(mock, LeaseOptions::default());
    230        assert_eq!(state, make_state([], [], []));
    231
    232        state.add("1".to_string());
    233        assert_eq!(state, make_state(["1"], [], []));
    234
    235        state.add("2".to_string());
    236        assert_eq!(state, make_state(["1", "2"], [], []));
    237
    238        state.add("3".to_string());
    239        assert_eq!(state, make_state(["1", "2", "3"], [], []));
    240
    241        state.ack("1".to_string());
    242        assert_eq!(state, make_state(["2", "3"], ["1"], []));
    243
    244        state.nack("2".to_string());
    245        assert_eq!(state, make_state(["3"], ["1"], ["2"]));
    246
    247        state.add("4".to_string());
    248        assert_eq!(state, make_state(["3", "4"], ["1"], ["2"]));
    249
    250        state.ack("4".to_string());
    251        assert_eq!(state, make_state(["3"], ["1", "4"], ["2"]));
    252
    253        state.nack("3".to_string());
    254        assert_eq!(state, make_state([], ["1", "4"], ["2", "3"]));
    255    }
    256
    257    #[tokio::test]
    258    async fn flush() {
    259        let mut mock = MockLeaser::new();
    260        mock.expect_ack()
    261            .times(1)
    262            .withf(|v| sorted(v) == test_ids(0..10))
    263            .returning(|_| ());
    264        mock.expect_nack()
    265            .times(1)
    266            .withf(|v| sorted(v) == test_ids(10..20))
    267            .returning(|_| ());
    268
    269        let mut state = LeaseState::new(mock, LeaseOptions::default());
    270        for i in 0..100 {
    271            state.add(test_id(i));
    272        }
    273        for i in 0..10 {
    274            state.ack(test_id(i));
    275        }
    276        for i in 10..20 {
    277            state.nack(test_id(i));
    278        }
    279        let expected = LeaseState {
    280            under_lease: test_ids(20..100).into_iter().collect(),
    281            to_ack: test_ids(0..10),
    282            to_nack: test_ids(10..20),
    283            leaser: MockLeaser::new(),
    284            flush_interval: test_interval(),
    285            extend_interval: test_interval(),
    286        };
    287        assert_eq!(state, expected);
    288
    289        state.flush().await;
    290        let expected = LeaseState {
    291            under_lease: test_ids(20..100).into_iter().collect(),
    292            to_ack: Vec::new(),
    293            to_nack: Vec::new(),
    294            leaser: MockLeaser::new(),
    295            flush_interval: test_interval(),
    296            extend_interval: test_interval(),
    297        };
    298        assert_eq!(state, expected);
    299    }
    300
    301    #[tokio::test]
    302    async fn extend() {
    303        let mut seq = mockall::Sequence::new();
    304        let mut mock = MockLeaser::new();
    305        mock.expect_extend()
    306            .times(1)
    307            .in_sequence(&mut seq)
    308            .withf(|v| sorted(v) == test_ids(0..10))
    309            .returning(|_| ());
    310        mock.expect_extend()
    311            .times(1)
    312            .in_sequence(&mut seq)
    313            .withf(|v| sorted(v) == test_ids(0..20))
    314            .returning(|_| ());
    315        mock.expect_extend()
    316            .times(1)
    317            .in_sequence(&mut seq)
    318            .withf(|v| sorted(v) == test_ids(5..20))
    319            .returning(|_| ());
    320        mock.expect_extend()
    321            .times(1)
    322            .in_sequence(&mut seq)
    323            .withf(|v| sorted(v) == test_ids(10..20))
    324            .returning(|_| ());
    325
    326        let mut state = LeaseState::new(mock, LeaseOptions::default());
    327
    328        // Accept 10 messages. These are now under lease management.
    329        for i in 0..10 {
    330            state.add(test_id(i));
    331        }
    332        state.extend().await;
    333
    334        // Accept another 10 messages. These are now under lease management.
    335        for i in 10..20 {
    336            state.add(test_id(i));
    337        }
    338        state.extend().await;
    339
    340        // Ack the first 5 messages. We should not extend these leases.
    341        for i in 0..5 {
    342            state.ack(test_id(i));
    343        }
    344        state.extend().await;
    345
    346        // Nack the next 5 messages. We should not extend these leases.
    347        for i in 5..10 {
    348            state.nack(test_id(i));
    349        }
    350        state.extend().await;
    351    }
    352
    353    #[tokio::test]
    354    async fn shutdown() {
    355        let mut mock = MockLeaser::new();
    356        mock.expect_ack()
    357            .times(1)
    358            .withf(|v| sorted(v) == test_ids(0..10))
    359            .returning(|_| ());
    360        mock.expect_nack()
    361            .times(1)
    362            .withf(|v| sorted(v) == test_ids(10..30))
    363            .returning(|_| ());
    364
    365        let mut state = LeaseState::new(mock, LeaseOptions::default());
    366        for i in 0..30 {
    367            state.add(test_id(i));
    368        }
    369        for i in 0..10 {
    370            state.ack(test_id(i));
    371        }
    372        for i in 10..20 {
    373            state.nack(test_id(i));
    374        }
    375        state.shutdown().await;
    376    }
    377
    378    #[tokio::test]
    379    async fn ack_out_of_lease_included() {
    380        let mock = MockLeaser::new();
    381        let mut state = LeaseState::new(mock, LeaseOptions::default());
    382        assert_eq!(state, make_state([], [], []));
    383
    384        state.ack("1".to_string());
    385        assert_eq!(state, make_state([], ["1"], []));
    386    }
    387
    388    #[tokio::test]
    389    async fn nack_out_of_lease_ignored() {
    390        let mock = MockLeaser::new();
    391        let mut state = LeaseState::new(mock, LeaseOptions::default());
    392        assert_eq!(state, make_state([], [], []));
    393
    394        state.nack("1".to_string());
    395        assert_eq!(state, make_state([], [], []));
    396    }
    397
    398    #[tokio::test(start_paused = true)]
    399    async fn lease_events_timing() {
    400        let start = Instant::now();
    401
    402        // We expect flushes at time t=[1, 3, 5, 7, ...]
    403        const FLUSH_START: Duration = Duration::from_secs(1);
    404        const FLUSH_PERIOD: Duration = Duration::from_secs(2);
    405
    406        // We expect extensions at time t=[2, 6, 10, 14, ...]
    407        const EXTEND_START: Duration = Duration::from_secs(2);
    408        const EXTEND_PERIOD: Duration = Duration::from_secs(4);
    409
    410        let mock = MockLeaser::new();
    411        let options = LeaseOptions {
    412            flush_start: FLUSH_START,
    413            flush_period: FLUSH_PERIOD,
    414            extend_start: EXTEND_START,
    415            extend_period: EXTEND_PERIOD,
    416        };
    417        let mut state = LeaseState::new(mock, options);
    418
    419        assert_eq!(state.next_event().await, LeaseEvent::Flush);
    420        assert_eq!(start.elapsed(), Duration::from_secs(1));
    421
    422        assert_eq!(state.next_event().await, LeaseEvent::Extend);
    423        assert_eq!(start.elapsed(), Duration::from_secs(2));
    424
    425        assert_eq!(state.next_event().await, LeaseEvent::Flush);
    426        assert_eq!(start.elapsed(), Duration::from_secs(3));
    427
    428        assert_eq!(state.next_event().await, LeaseEvent::Flush);
    429        assert_eq!(start.elapsed(), Duration::from_secs(5));
    430
    431        assert_eq!(state.next_event().await, LeaseEvent::Extend);
    432        assert_eq!(start.elapsed(), Duration::from_secs(6));
    433
    434        assert_eq!(state.next_event().await, LeaseEvent::Flush);
    435        assert_eq!(start.elapsed(), Duration::from_secs(7));
    436    }
    437
    438    #[tokio::test(start_paused = true)]
    439    async fn limit_size_of_ack_batch() {
    440        let start = Instant::now();
    441
    442        const FLUSH_START: Duration = Duration::from_secs(1);
    443
    444        let mut mock = MockLeaser::new();
    445        mock.expect_ack()
    446            .times(1)
    447            .withf(|v| sorted(v) == test_ids(0..ACK_IDS_PER_RPC))
    448            .returning(|_| ());
    449        mock.expect_nack()
    450            .times(1)
    451            .withf(|v| v.is_empty())
    452            .returning(|_| ());
    453        let options = LeaseOptions {
    454            flush_start: FLUSH_START,
    455            flush_period: Duration::from_secs(100),
    456            extend_start: Duration::from_secs(100),
    457            extend_period: Duration::from_secs(100),
    458        };
    459        let mut state = LeaseState::new(mock, options);
    460
    461        for i in 0..ACK_IDS_PER_RPC {
    462            state.add(test_id(i));
    463            state.ack(test_id(i));
    464        }
    465        // With 2500 pending acks, the batch is full. We should flush it now.
    466        assert_eq!(state.next_event().await, LeaseEvent::Flush);
    467        assert_eq!(start.elapsed(), Duration::ZERO);
    468        state.flush().await;
    469
    470        // With 1000 pending acks, the batch is not full. The next event should
    471        // occur on the interval timer.
    472        for i in ACK_IDS_PER_RPC..ACK_IDS_PER_RPC + 1000 {
    473            state.add(test_id(i));
    474            state.ack(test_id(i));
    475        }
    476        assert_eq!(state.next_event().await, LeaseEvent::Flush);
    477        assert_eq!(start.elapsed(), FLUSH_START);
    478    }
    479
    480    #[tokio::test(start_paused = true)]
    481    async fn limit_size_of_nack_batch() {
    482        let start = Instant::now();
    483
    484        const FLUSH_START: Duration = Duration::from_secs(1);
    485
    486        let mut mock = MockLeaser::new();
    487        mock.expect_ack()
    488            .times(1)
    489            .withf(|v| v.is_empty())
    490            .returning(|_| ());
    491        mock.expect_nack()
    492            .times(1)
    493            .withf(|v| sorted(v) == test_ids(0..ACK_IDS_PER_RPC))
    494            .returning(|_| ());
    495        let options = LeaseOptions {
    496            flush_start: FLUSH_START,
    497            flush_period: Duration::from_secs(100),
    498            extend_start: Duration::from_secs(100),
    499            extend_period: Duration::from_secs(100),
    500        };
    501        let mut state = LeaseState::new(mock, options);
    502
    503        for i in 0..ACK_IDS_PER_RPC {
    504            state.add(test_id(i));
    505            state.nack(test_id(i));
    506        }
    507        // With 2500 pending nacks, the batch is full. We should flush it now.
    508        assert_eq!(state.next_event().await, LeaseEvent::Flush);
    509        assert_eq!(start.elapsed(), Duration::ZERO);
    510        state.flush().await;
    511
    512        // With 1000 pending nacks, the batch is not full. The next event should
    513        // occur on the interval timer.
    514        for i in ACK_IDS_PER_RPC..ACK_IDS_PER_RPC + 1000 {
    515            state.add(test_id(i));
    516            state.nack(test_id(i));
    517        }
    518        assert_eq!(state.next_event().await, LeaseEvent::Flush);
    519        assert_eq!(start.elapsed(), FLUSH_START);
    520    }
    521
    522    #[tokio::test(start_paused = true)]
    523    async fn ack_and_nack_batches_are_independent() {
    524        let start = Instant::now();
    525
    526        const FLUSH_START: Duration = Duration::from_secs(1);
    527
    528        let mock = MockLeaser::new();
    529        let options = LeaseOptions {
    530            flush_start: FLUSH_START,
    531            flush_period: Duration::from_secs(100),
    532            extend_start: Duration::from_secs(100),
    533            extend_period: Duration::from_secs(100),
    534        };
    535        let mut state = LeaseState::new(mock, options);
    536
    537        let over_half_full = ACK_IDS_PER_RPC / 2 + 100;
    538        for i in 0..over_half_full {
    539            state.add(test_id(i));
    540            state.ack(test_id(i));
    541
    542            state.add(test_id(over_half_full + i));
    543            state.nack(test_id(over_half_full + i));
    544        }
    545
    546        // While there are more than `ACK_IDS_PER_RPC` total messages under
    547        // lease management, neither the ack batch nor the nack batch are full.
    548        // The next event should occur on the interval timer.
    549        assert_eq!(state.next_event().await, LeaseEvent::Flush);
    550        assert_eq!(start.elapsed(), FLUSH_START);
    551    }
    552
    553    #[tokio::test(start_paused = true)]
    554    async fn limit_size_of_extends() -> anyhow::Result<()> {
    555        const NUM_BATCHES: i32 = 5;
    556
    557        // We use this channel to surface ack_ids from the mock expectation.
    558        let (ack_id_tx, mut ack_id_rx) = unbounded_channel();
    559
    560        let mut mock = MockLeaser::new();
    561        mock.expect_extend()
    562            .times(NUM_BATCHES as usize)
    563            .returning(move |ack_ids| {
    564                ack_id_tx
    565                    .send(ack_ids)
    566                    .expect("sending on channel always succeeds");
    567            });
    568        let mut state = LeaseState::new(mock, LeaseOptions::default());
    569
    570        let mut want = HashSet::new();
    571        for i in 0..NUM_BATCHES * ACK_IDS_PER_RPC {
    572            state.add(test_id(i));
    573
    574            // All ack IDs should be extended.
    575            want.insert(test_id(i));
    576        }
    577        state.extend().await;
    578
    579        let mut got = HashSet::new();
    580        for i in 0..NUM_BATCHES {
    581            let Some(ack_ids) = ack_id_rx.recv().await else {
    582                anyhow::bail!("expected batch {i}/{NUM_BATCHES}");
    583            };
    584            assert_eq!(ack_ids.len(), ACK_IDS_PER_RPC as usize);
    585            for ack_id in ack_ids {
    586                got.insert(ack_id);
    587            }
    588        }
    589        assert!(
    590            ack_id_rx.is_empty(),
    591            "There should be exactly {NUM_BATCHES} batches"
    592        );
    593
    594        // Make sure all ack IDs were extended.
    595        assert_eq!(got, want);
    596        Ok(())
    597    }
    598}