#
google_cloud_pubsub/subscriber/
leaser.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::model::{AcknowledgeRequest, ModifyAckDeadlineRequest};
    17use gax::options::RequestOptions;
    18use gax::retry_policy::NeverRetry;
    19use std::sync::Arc;
    20
    21/// A trait representing leaser actions.
    22///
    23/// We stub out the interface, in order to test the lease management.
    24#[async_trait::async_trait]
    25pub(super) trait Leaser {
    26    /// Acknowledge a batch of messages.
    27    async fn ack(&self, ack_ids: Vec<String>);
    28    /// Negatively acknowledge a batch of messages.
    29    async fn nack(&self, ack_ids: Vec<String>);
    30    /// Extend lease deadlines for a batch of messages.
    31    async fn extend(&self, ack_ids: Vec<String>);
    32}
    33
    34pub(super) struct DefaultLeaser<T>
    35where
    36    T: Stub,
    37{
    38    inner: Arc<T>,
    39    subscription: String,
    40    ack_deadline_seconds: i32,
    41}
    42
    43impl<T> DefaultLeaser<T>
    44where
    45    T: Stub,
    46{
    47    pub(super) fn new(inner: Arc<T>, subscription: String, ack_deadline_seconds: i32) -> Self {
    48        DefaultLeaser {
    49            inner,
    50            subscription,
    51            ack_deadline_seconds,
    52        }
    53    }
    54}
    55
    56fn no_retry() -> RequestOptions {
    57    let mut o = RequestOptions::default();
    58    o.set_retry_policy(NeverRetry);
    59    o
    60}
    61
    62#[async_trait::async_trait]
    63impl<T> Leaser for DefaultLeaser<T>
    64where
    65    T: Stub,
    66{
    67    async fn ack(&self, ack_ids: Vec<String>) {
    68        if ack_ids.is_empty() {
    69            return;
    70        }
    71        let req = AcknowledgeRequest::new()
    72            .set_subscription(self.subscription.clone())
    73            .set_ack_ids(ack_ids);
    74        let _ = self.inner.acknowledge(req, no_retry()).await;
    75    }
    76    async fn nack(&self, ack_ids: Vec<String>) {
    77        if ack_ids.is_empty() {
    78            return;
    79        }
    80        let req = ModifyAckDeadlineRequest::new()
    81            .set_subscription(self.subscription.clone())
    82            .set_ack_ids(ack_ids)
    83            .set_ack_deadline_seconds(0);
    84        let _ = self.inner.modify_ack_deadline(req, no_retry()).await;
    85    }
    86    async fn extend(&self, ack_ids: Vec<String>) {
    87        if ack_ids.is_empty() {
    88            return;
    89        }
    90        let req = ModifyAckDeadlineRequest::new()
    91            .set_subscription(self.subscription.clone())
    92            .set_ack_ids(ack_ids)
    93            .set_ack_deadline_seconds(self.ack_deadline_seconds);
    94        let _ = self.inner.modify_ack_deadline(req, no_retry()).await;
    95    }
    96}
    97
    98#[cfg(test)]
    99pub(super) mod tests {
    100    use super::super::lease_state::tests::test_ids;
    101    use super::super::stub::tests::MockStub;
    102    use super::*;
    103    use gax::response::Response;
    104    use std::sync::Arc;
    105    use tokio::sync::Mutex;
    106
    107    mockall::mock! {
    108        #[derive(Debug)]
    109        pub(in super::super) Leaser {}
    110        #[async_trait::async_trait]
    111        impl Leaser for Leaser {
    112            async fn ack(&self, ack_ids: Vec<String>);
    113            async fn nack(&self, ack_ids: Vec<String>);
    114            async fn extend(&self, ack_ids: Vec<String>);
    115        }
    116    }
    117
    118    #[async_trait::async_trait]
    119    impl Leaser for Arc<Mutex<MockLeaser>> {
    120        async fn ack(&self, ack_ids: Vec<String>) {
    121            self.lock().await.ack(ack_ids).await
    122        }
    123        async fn nack(&self, ack_ids: Vec<String>) {
    124            self.lock().await.nack(ack_ids).await
    125        }
    126        async fn extend(&self, ack_ids: Vec<String>) {
    127            self.lock().await.extend(ack_ids).await
    128        }
    129    }
    130
    131    #[tokio::test]
    132    async fn ack() {
    133        let mut mock = MockStub::new();
    134        mock.expect_acknowledge().times(1).return_once(|r, o| {
    135            assert_eq!(
    136                r.subscription,
    137                "projects/my-project/subscriptions/my-subscription"
    138            );
    139            assert_eq!(r.ack_ids, test_ids(0..10));
    140            assert!(
    141                format!("{o:?}").contains("NeverRetry"),
    142                "Basic acks should not have a retry policy. o={o:?}"
    143            );
    144            Ok(Response::from(()))
    145        });
    146
    147        let leaser = DefaultLeaser::new(
    148            Arc::new(mock),
    149            "projects/my-project/subscriptions/my-subscription".to_string(),
    150            10,
    151        );
    152        leaser.ack(test_ids(0..10)).await;
    153    }
    154
    155    #[tokio::test]
    156    async fn nack() {
    157        let mut mock = MockStub::new();
    158        mock.expect_modify_ack_deadline()
    159            .times(1)
    160            .return_once(|r, o| {
    161                assert_eq!(r.ack_deadline_seconds, 0);
    162                assert_eq!(
    163                    r.subscription,
    164                    "projects/my-project/subscriptions/my-subscription"
    165                );
    166                assert_eq!(r.ack_ids, test_ids(0..10));
    167                assert!(
    168                    format!("{o:?}").contains("NeverRetry"),
    169                    "Basic modacks should not have a retry policy. o={o:?}"
    170                );
    171                Ok(Response::from(()))
    172            });
    173
    174        let leaser = DefaultLeaser::new(
    175            Arc::new(mock),
    176            "projects/my-project/subscriptions/my-subscription".to_string(),
    177            10,
    178        );
    179        leaser.nack(test_ids(0..10)).await;
    180    }
    181
    182    #[tokio::test]
    183    async fn extend() {
    184        let mut mock = MockStub::new();
    185        mock.expect_modify_ack_deadline()
    186            .times(1)
    187            .return_once(|r, o| {
    188                assert_eq!(r.ack_deadline_seconds, 10);
    189                assert_eq!(
    190                    r.subscription,
    191                    "projects/my-project/subscriptions/my-subscription"
    192                );
    193                assert_eq!(r.ack_ids, test_ids(0..10));
    194                assert!(
    195                    format!("{o:?}").contains("NeverRetry"),
    196                    "Basic acks should not have a retry policy. o={o:?}"
    197                );
    198                Ok(Response::from(()))
    199            });
    200
    201        let leaser = DefaultLeaser::new(
    202            Arc::new(mock),
    203            "projects/my-project/subscriptions/my-subscription".to_string(),
    204            10,
    205        );
    206        leaser.extend(test_ids(0..10)).await;
    207    }
    208
    209    #[tokio::test]
    210    async fn empty_is_noop() {
    211        let mock = MockStub::new();
    212        // We expect no calls on the mock stub, as there is no reason to send an
    213        // RPC with an empty list of ack IDs.
    214
    215        let leaser = DefaultLeaser::new(
    216            Arc::new(mock),
    217            "projects/my-project/subscriptions/my-subscription".to_string(),
    218            10,
    219        );
    220        leaser.ack(Vec::new()).await;
    221        leaser.nack(Vec::new()).await;
    222        leaser.extend(Vec::new()).await;
    223    }
    224}