Improve robustness of S3 reading

CHANGELOG.md

@@ -1,5 +1,7 @@
 # Unreleased
 
+- Improve robustness of S3 reading (PR [#552](https://github.com/RaRe-Technologies/smart_open/pull/552), [@mpenkov](https://github.com/mpenkov))
+
 # 4.1.0, 30 Dec 2020
 
 - Refactor `s3` submodule to minimize resource usage (PR [#569](https://github.com/RaRe-Technologies/smart_open/pull/569), [@mpenkov](https://github.com/mpenkov))

howto.md

@@ -224,6 +224,47 @@ To access such buckets, you need to pass some special transport parameters:
 
 This works only when reading and writing via S3.
 
+## How to Make S3 I/O Robust to Network Errors
+
+Boto3 has a [built-in mechanism](https://boto3.amazonaws.com/v1/documentation/api/latest/guide/retries.html) for retrying after a recoverable error.
+You can fine-tune it using several ways:
+
+### Pre-configuring a boto3 resource and then passing the resource to smart_open
+
+```python
+>>> import boto3
+>>> import botocore.config
+>>> import smart_open
+>>> config = botocore.config.Config(retries={'mode': 'standard'})
+>>> resource = boto3.resource('s3', config=config)
+>>> tp = {'resource': resource}
+>>> with smart_open.open('s3://commoncrawl/robots.txt', transport_params=tp) as fin:
+...     print(fin.readline())
+User-Agent: *
+```
+
+### Directly passing configuration as transport parameters to smart_open
+
+```python
+>>> import boto3
+>>> import botocore.config
+>>> import smart_open
+>>> config = botocore.config.Config(retries={'mode': 'standard'})
+>>> tp = {'resource_kwargs': {'config': config}}
+>>> with smart_open.open('s3://commoncrawl/robots.txt', transport_params=tp) as fin:
+...     print(fin.readline())
+User-Agent: *
+```
+
+To verify your settings have effect:
+
+```python
+import logging
+logging.getLogger('smart_open.s3').setLevel(logging.DEBUG)
+```
+
+and check the log output of your code.
+
 ## How to Read/Write from localstack
 
 [localstack](https://github.com/localstack/localstack) is a convenient test framework for developing cloud apps.

smart_open/s3.py

@@ -11,6 +11,7 @@
 import functools
 import logging
 import time
+import urllib3.exceptions
 
 try:
     import boto3
@@ -307,7 +308,12 @@ class _SeekableRawReader(object):
     This class is internal to the S3 submodule.
     """
 
-    def __init__(self, s3_object, version_id=None, object_kwargs=None):
+    def __init__(
+        self,
+        s3_object,
+        version_id=None,
+        object_kwargs=None,
+    ):
         self._object = s3_object
         self._content_length = None
         self._version_id = version_id
@@ -397,18 +403,22 @@ def _open_body(self, start=None, stop=None):
             self._position = self._content_length = int(error_response['ActualObjectSize'])
             self._body = io.BytesIO()
         else:
+            #
+            # Keep track of how many times boto3's built-in retry mechanism
+            # activated.
+            #
+            # https://boto3.amazonaws.com/v1/documentation/api/latest/guide/retries.html#checking-retry-attempts-in-an-aws-service-response
+            #
+            logger.debug(
+                '%s: RetryAttempts: %d',
+                self,
+                response['ResponseMetadata']['RetryAttempts'],
+            )
             units, start, stop, length = smart_open.utils.parse_content_range(response['ContentRange'])
             self._content_length = length
             self._position = start
             self._body = response['Body']
 
-    def _read_from_body(self, size=-1):
-        if size == -1:
-            binary = self._body.read()
-        else:
-            binary = self._body.read(size)
-        return binary
-
     def read(self, size=-1):
         """Read from the continuous connection with the remote peer."""
         if self._body is None:
@@ -417,14 +427,50 @@ def read(self, size=-1):
         if self._position >= self._content_length:
             return b''
 
-        try:
-            binary = self._read_from_body(size)
-        except botocore.exceptions.IncompleteReadError:
-            # The underlying connection of the self._body was closed by the remote peer.
-            self._open_body()
-            binary = self._read_from_body(size)
-        self._position += len(binary)
-        return binary
+        #
+        # Boto3 has built-in error handling and retry mechanisms:
+        #
+        # https://boto3.amazonaws.com/v1/documentation/api/latest/guide/error-handling.html
+        # https://boto3.amazonaws.com/v1/documentation/api/latest/guide/retries.html
+        #
+        # Unfortunately, it isn't always enough. There is still a non-zero
+        # possibility that an exception will slip past these mechanisms and
+        # terminate the read prematurely.  Luckily, at this stage, it's very
+        # simple to recover from the problem: wait a little bit, reopen the
+        # HTTP connection and try again.  Usually, a single retry attempt is
+        # enough to recover, but we try multiple times "just in case".
+        #
+        for attempt, seconds in enumerate([1, 2, 4, 8, 16], 1):
+            try:
+                if size == -1:
+                    binary = self._body.read()
+                else:
+                    binary = self._body.read(size)
+            except (
+                ConnectionResetError,
+                botocore.exceptions.BotoCoreError,
+                urllib3.exceptions.HTTPError,
+            ) as err:
+                logger.warning(
+                    '%s: caught %r while reading %d bytes, sleeping %ds before retry',
+                    self,
+                    err,
+                    size,
+                    seconds,
+                )
+                time.sleep(seconds)
+                self._open_body()
+            else:
+                self._position += len(binary)
+                return binary
+
+        raise IOError('%s: failed to read %d bytes after %d attempts' % (self, size, attempt))
+
+    def __str__(self):
+        return 'smart_open.s3._SeekableReader(%r, %r)' % (
+            self._object.bucket_name,
+            self._object.key,
+        )
 
 
 def _initialize_boto3(rw, session, resource, resource_kwargs):

smart_open/tests/test_s3.py

@@ -155,6 +155,72 @@ def test_read_from_a_closed_body(self):
         self.assertEqual(reader.read(2), b'23')
 
 
+class CrapStream(io.BytesIO):
+    """Raises an exception on every second read call."""
+    def __init__(self, *args, modulus=2, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._count = 0
+        self._modulus = modulus
+
+    def read(self, size=-1):
+        self._count += 1
+        if self._count % self._modulus == 0:
+            raise botocore.exceptions.BotoCoreError()
+        the_bytes = super().read(size)
+        return the_bytes
+
+
+class CrapObject:
+    def __init__(self, data, modulus=2):
+        self._datasize = len(data)
+        self._body = CrapStream(data, modulus=modulus)
+        self.bucket_name, self.key = 'crap', 'object'
+
+    def get(self, *args, **kwargs):
+        return {
+            'ActualObjectSize': self._datasize,
+            'ContentLength': self._datasize,
+            'ContentRange': 'bytes 0-%d/%d' % (self._datasize, self._datasize),
+            'Body': self._body,
+            'ResponseMetadata': {'RetryAttempts': 1},
+        }
+
+
+class IncrementalBackoffTest(unittest.TestCase):
+    def test_every_read_fails(self):
+        reader = smart_open.s3._SeekableRawReader(CrapObject(b'hello', 1))
+        with mock.patch('time.sleep') as mock_sleep:
+            with self.assertRaises(IOError):
+                reader.read()
+
+            #
+            # Make sure our incremental backoff is actually happening here.
+            #
+            mock_sleep.assert_has_calls([mock.call(s) for s in (1, 2, 4, 8, 16)])
+
+    def test_every_second_read_fails(self):
+        """Can we read from a stream that raises exceptions from time to time?"""
+        reader = smart_open.s3._SeekableRawReader(CrapObject(b'hello'))
+        with mock.patch('time.sleep') as mock_sleep:
+            assert reader.read(1) == b'h'
+            mock_sleep.assert_not_called()
+
+            assert reader.read(1) == b'e'
+            mock_sleep.assert_called_with(1)
+            mock_sleep.reset_mock()
+
+            assert reader.read(1) == b'l'
+            mock_sleep.reset_mock()
+
+            assert reader.read(1) == b'l'
+            mock_sleep.assert_called_with(1)
+            mock_sleep.reset_mock()
+
+            assert reader.read(1) == b'o'
+            mock_sleep.assert_called_with(1)
+            mock_sleep.reset_mock()
+
+
 @moto.mock_s3
 class SeekableBufferedInputBaseTest(BaseTest):
     def setUp(self):