Implement exponential backoff for batch writes (#728)

The old code failed to respect configured backoff policy, this PR changes it so it will no longer retry in infinite tight loop but instead apply the same backoff/retry policy as the rest of the api


* Implement exponential backoff for batch writes

* Review fix

Update pynamodb/connection/base.py

Co-Authored-By: Ilya Konstantinov <ilya.konstantinov@gmail.com>

* Add test for throwing PutError when max retries exceeded

* Fix format string type

Co-Authored-By: Ilya Konstantinov <ilya.konstantinov@gmail.com>

Author: edholland

pynamodb/connection/base.py

@@ -61,6 +61,7 @@
 from pynamodb.types import HASH, RANGE
 
 BOTOCORE_EXCEPTIONS = (BotoCoreError, ClientError)
+RATE_LIMITING_ERROR_CODES = ['ProvisionedThroughputExceededException', 'ThrottlingException']
 
 log = logging.getLogger(__name__)
 log.addHandler(logging.NullHandler())
@@ -442,7 +443,7 @@ def _make_api_call(self, operation_name, operation_kwargs):
                     if is_last_attempt_for_exceptions:
                         log.debug('Reached the maximum number of retry attempts: %s', attempt_number)
                         raise
-                    elif status_code < 500 and code != 'ProvisionedThroughputExceededException':
+                    elif status_code < 500 and code not in RATE_LIMITING_ERROR_CODES:
                         # We don't retry on a ConditionalCheckFailedException or other 4xx (except for
                         # throughput related errors) because we assume they will fail in perpetuity.
                         # Retrying when there is already contention could cause other problems

pynamodb/models.py

@@ -2,6 +2,7 @@
 DynamoDB Models for PynamoDB
 """
 import json
+import random
 import time
 import six
 import logging
@@ -11,7 +12,7 @@
 from six import add_metaclass
 
 from pynamodb.expressions.condition import NotExists, Comparison
-from pynamodb.exceptions import DoesNotExist, TableDoesNotExist, TableError, InvalidStateError
+from pynamodb.exceptions import DoesNotExist, TableDoesNotExist, TableError, InvalidStateError, PutError
 from pynamodb.attributes import (
     Attribute, AttributeContainer, AttributeContainerMeta, MapAttribute, TTLAttribute, VersionAttribute
 )
@@ -52,6 +53,7 @@ def __init__(self, model, auto_commit=True):
         self.auto_commit = auto_commit
         self.max_operations = BATCH_WRITE_PAGE_LIMIT
         self.pending_operations = []
+        self.failed_operations = []
 
     def __enter__(self):
         return self
@@ -130,16 +132,24 @@ def commit(self):
         )
         if data is None:
             return
+        retries = 0
         unprocessed_items = data.get(UNPROCESSED_ITEMS, {}).get(self.model.Meta.table_name)
         while unprocessed_items:
+            sleep_time = random.randint(0, self.model.Meta.base_backoff_ms * (2 ** retries)) / 1000
+            time.sleep(sleep_time)
+            retries += 1
+            if retries >= self.model.Meta.max_retry_attempts:
+                self.failed_operations = unprocessed_items
+                raise PutError("Failed to batch write items: max_retry_attempts exceeded")
             put_items = []
             delete_items = []
             for item in unprocessed_items:
                 if PUT_REQUEST in item:
                     put_items.append(item.get(PUT_REQUEST).get(ITEM))
                 elif DELETE_REQUEST in item:
                     delete_items.append(item.get(DELETE_REQUEST).get(KEY))
-            log.info("Resending %s unprocessed keys for batch operation", len(unprocessed_items))
+            log.info("Resending %d unprocessed keys for batch operation after %d seconds sleep",
+                     len(unprocessed_items), sleep_time)
             data = self.model._get_connection().batch_write_item(
                 put_items=put_items,
                 delete_items=delete_items

tests/test_model.py

@@ -15,7 +15,7 @@
 
 from .deep_eq import deep_eq
 from pynamodb.connection.util import pythonic
-from pynamodb.exceptions import DoesNotExist, TableError
+from pynamodb.exceptions import DoesNotExist, TableError, PutError
 from pynamodb.types import RANGE
 from pynamodb.constants import (
     ITEM, STRING_SHORT, ALL, KEYS_ONLY, INCLUDE, REQUEST_ITEMS, UNPROCESSED_KEYS, CAMEL_COUNT,
@@ -236,6 +236,20 @@ class Meta:
     ttl = TTLAttribute(null=True)
 
 
+
+class BatchModel(Model):
+    """
+    A testing model
+    """
+
+    class Meta:
+        table_name = 'BatchModel'
+        max_retry_attempts = 0
+
+    user_name = UnicodeAttribute(hash_key=True)
+
+
+
 class HostSpecificModel(Model):
     """
     A testing model
@@ -1940,6 +1954,7 @@ def test_batch_write(self):
                 for item in items:
                     batch.save(item)
 
+
     def test_batch_write_with_unprocessed(self):
         picture_blob = b'FFD8FFD8'
 
@@ -1984,6 +1999,43 @@ def test_batch_write_with_unprocessed(self):
 
             self.assertEqual(len(req.mock_calls), 3)
 
+    def test_batch_write_raises_put_error(self):
+        items = []
+        for idx in range(10):
+            items.append(BatchModel(
+                '{}'.format(idx)
+            ))
+
+        unprocessed_items = []
+        for idx in range(5, 10):
+            unprocessed_items.append({
+                'PutRequest': {
+                    'Item': {
+                        'user_name': {STRING_SHORT: 'daniel'},
+                    }
+                }
+            })
+
+        with patch(PATCH_METHOD) as req:
+            req.side_effect = [
+                {
+                    UNPROCESSED_ITEMS: {
+                        BatchModel.Meta.table_name: unprocessed_items[:2],
+                    }
+                },
+                {
+                    UNPROCESSED_ITEMS: {
+                        BatchModel.Meta.table_name: unprocessed_items[2:],
+                    }
+                },
+                {}
+            ]
+            with self.assertRaises(PutError):
+                with BatchModel.batch_write() as batch:
+                    for item in items:
+                        batch.save(item)
+            self.assertEqual(len(batch.failed_operations), 3)
+
     def test_index_queries(self):
         """
         Model.Index.Query