Exactly-once changes

google/cloud/pubsub_v1/proto/pubsub.proto

@@ -1164,6 +1164,7 @@ message StreamingPullRequest {
 message StreamingPullResponse {
   // Subscription properties sent as part of the response.
   message SubscriptionProperties {
+    bool exactly_once_delivery_enabled = 1;
     // True iff message ordering is enabled for this subscription.
     bool message_ordering_enabled = 2;
   }

google/cloud/pubsub_v1/subscriber/_protocol/dispatcher.py

@@ -22,6 +22,7 @@
 import typing
 from typing import List, Optional, Sequence, Union
 import warnings
+from google.api_core.retry import exponential_sleep_generator
 
 from google.cloud.pubsub_v1.subscriber._protocol import helper_threads
 from google.cloud.pubsub_v1.subscriber._protocol import requests
@@ -66,6 +67,13 @@
 IDs at a time.
 """
 
+_MIN_ACK_MODACK_RETRY_DURATION_SECS = 1
+"""The time to wait for the first retry of failed acks and modacks when exactly-once
+is enabled."""
+
+_MAX_ACK_MODACK_RETRY_DURATION_SECS = 20 * 60
+"""The maximum amount of time in seconds to retry failed acks and modacks when
+exactly-once is enabled."""
 
 class Dispatcher(object):
     def __init__(self, manager: "StreamingPullManager", queue: "queue.Queue"):
@@ -168,17 +176,35 @@ def ack(self, items: Sequence[requests.AckRequest]) -> None:
 
         # We must potentially split the request into multiple smaller requests
         # to avoid the server-side max request size limit.
-        ack_ids = (item.ack_id for item in items)
+        items_gen = iter(items)
+        ack_ids_gen = (item.ack_id for item in items)
         total_chunks = int(math.ceil(len(items) / _ACK_IDS_BATCH_SIZE))
 
         for _ in range(total_chunks):
-            request = gapic_types.StreamingPullRequest(
-                ack_ids=itertools.islice(ack_ids, _ACK_IDS_BATCH_SIZE)
-            )
-            self._manager.send(request)
-
-        # Remove the message from lease management.
-        self.drop(items)
+            future_reqs_dict = {req.ack_id: req for req in itertools.islice(items_gen, _ACK_IDS_BATCH_SIZE) if req.future}
+            requests_completed, requests_to_retry =  self._manager.send_unary_ack(
+                ack_ids=list(itertools.islice(ack_ids_gen, _ACK_IDS_BATCH_SIZE)),
+                future_reqs_dict=future_reqs_dict)
+
+            # Remove the completed messages from lease management.
+            self.drop(requests_completed)
+
+            # retry acks
+            retry_delay_gen = exponential_sleep_generator(initial=_MIN_ACK_MODACK_RETRY_DURATION_SECS,
+                                                          maximum=_MAX_ACK_MODACK_RETRY_DURATION_SECS)
+            while requests_to_retry:
+                time_to_wait = retry_delay_gen()
+                _LOGGER.debug("Retrying {len(requests_to_retry)} ack(s) after delay of " +
+                              str(time_to_wait) + " seconds")
+                time.sleep(time_to_wait)
+
+                future_reqs_dict = {req.ack_id: req for req in requests_to_retry if req.future}
+                requests_completed, requests_to_retry = self._manager.send_unary_ack(
+                    ack_ids=[req.ack_id for req in requests_to_retry],
+                    future_reqs_dict=future_reqs_dict)
+                assert len(requests_to_retry) <= _ACK_IDS_BATCH_SIZE, "Too many requests to be retried."
+                # Remove the completed messages from lease management.
+                self.drop(requests_completed)
 
     def drop(
         self,
@@ -215,16 +241,35 @@ def modify_ack_deadline(self, items: Sequence[requests.ModAckRequest]) -> None:
         """
         # We must potentially split the request into multiple smaller requests
         # to avoid the server-side max request size limit.
-        ack_ids = (item.ack_id for item in items)
-        seconds = (item.seconds for item in items)
+        items_gen = iter(items)
+        ack_ids_gen = (item.ack_id for item in items)
+        deadline_seconds_gen = (item.seconds for item in items)
         total_chunks = int(math.ceil(len(items) / _ACK_IDS_BATCH_SIZE))
 
         for _ in range(total_chunks):
-            request = gapic_types.StreamingPullRequest(
-                modify_deadline_ack_ids=itertools.islice(ack_ids, _ACK_IDS_BATCH_SIZE),
-                modify_deadline_seconds=itertools.islice(seconds, _ACK_IDS_BATCH_SIZE),
-            )
-            self._manager.send(request)
+            future_reqs_dict = {req.ack_id: req for req in itertools.islice(items_gen, _ACK_IDS_BATCH_SIZE) if req.future}
+            # no further work needs to be done for `requests_to_retry`
+            requests_completed, requests_to_retry = self._manager.send_unary_modack(
+                modify_deadline_ack_ids=list(itertools.islice(ack_ids_gen, _ACK_IDS_BATCH_SIZE)),
+                modify_deadline_seconds=list(itertools.islice(deadline_seconds_gen, _ACK_IDS_BATCH_SIZE)),
+                future_reqs_dict=future_reqs_dict)
+            assert len(requests_to_retry) <= _ACK_IDS_BATCH_SIZE, "Too many requests to be retried."
+
+            # retry modacks
+            retry_delay_gen = exponential_sleep_generator(initial=_MIN_ACK_MODACK_RETRY_DURATION_SECS,
+                                                          maximum=_MAX_ACK_MODACK_RETRY_DURATION_SECS)
+            while requests_to_retry:
+                time_to_wait = retry_delay_gen()
+                _LOGGER.debug("Retrying {len(requests_to_retry)} modack(s) after delay of " +
+                              str(time_to_wait) + " seconds")
+                time.sleep(time_to_wait)
+
+                print(requests_to_retry)
+                future_reqs_dict = {req.ack_id: req for req in requests_to_retry if req.future}
+                requests_to_retry = self._manager.send_unary_modack(
+                    modify_deadline_ack_ids=[req.ack_id for req in requests_to_retry],
+                    modify_deadline_seconds=[req.seconds for req in requests_to_retry],
+                    future_reqs_dict=future_reqs_dict)
 
     def nack(self, items: Sequence[requests.NackRequest]) -> None:
         """Explicitly deny receipt of messages.
@@ -233,6 +278,6 @@ def nack(self, items: Sequence[requests.NackRequest]) -> None:
             items: The items to deny.
         """
         self.modify_ack_deadline(
-            [requests.ModAckRequest(ack_id=item.ack_id, seconds=0) for item in items]
+            [requests.ModAckRequest(ack_id=item.ack_id, seconds=0, future=item.future) for item in items]
         )
-        self.drop([requests.DropRequest(*item) for item in items])
+        self.drop([requests.DropRequest(ack_id=item.ack_id, byte_size=item.byte_size, ordering_key=item.ordering_key) for item in items])

google/cloud/pubsub_v1/subscriber/_protocol/leaser.py

@@ -195,7 +195,7 @@ def maintain_leases(self) -> None:
                 #       is inactive.
                 assert self._manager.dispatcher is not None
                 self._manager.dispatcher.modify_ack_deadline(
-                    [requests.ModAckRequest(ack_id, deadline) for ack_id in ack_ids]
+                    [requests.ModAckRequest(ack_id, deadline, None) for ack_id in ack_ids]
                 )
 
             # Now wait an appropriate period of time and do this again.

google/cloud/pubsub_v1/subscriber/_protocol/requests.py

@@ -12,8 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import typing
 from typing import NamedTuple, Optional
 
+if typing.TYPE_CHECKING:  # pragma: NO COVER
+    from google.cloud.pubsub_v1.subscriber import futures
 
 # Namedtuples for management requests. Used by the Message class to communicate
 # items of work back to the policy.
@@ -22,6 +25,7 @@ class AckRequest(NamedTuple):
     byte_size: int
     time_to_ack: float
     ordering_key: Optional[str]
+    future: Optional["pubsub_v1.subscriber.futures.Future"]
 
 
 class DropRequest(NamedTuple):
@@ -39,9 +43,11 @@ class LeaseRequest(NamedTuple):
 class ModAckRequest(NamedTuple):
     ack_id: str
     seconds: float
+    future: Optional["pubsub_v1.subscriber.futures.Future"]
 
 
 class NackRequest(NamedTuple):
     ack_id: str
     byte_size: int
     ordering_key: Optional[str]
+    future: Optional["pubsub_v1.subscriber.futures.Future"]