Skip to content

Commit 63894d7

Browse files
Jessemoderakh
Jesse
authored and
moderakh
committed
Retry attempts that fail due to a connection timeout (databricks#24)
* Isolate delay bounding logic * Move error details scope up one-level. * Retry GetOperationStatus if an OSError was raised during execution. Add retry_delay_default to use in this case. * Log when a request is retried due to an OSError. Emit warnings for unexpected OSError codes * Update docstring for make_request * Nit: unit tests show the .warn message is deprecated. DeprecationWarning: The 'warn' function is deprecated, use 'warning' instead Signed-off-by: Jesse Whitehouse <jesse@whitehouse.dev>
1 parent e851ea4 commit 63894d7

File tree

2 files changed

+120
-15
lines changed

2 files changed

+120
-15
lines changed

‎src/databricks/sql/thrift_backend.py‎

Lines changed: 63 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,7 @@
11
from decimal import Decimal
2+
3+
import errno
4+
import logging
25
import math
36
import time
47
import threading
@@ -14,6 +17,9 @@
1417
from databricks.sql.auth.authenticators import CredentialsProvider
1518
from databricks.sql.thrift_api.TCLIService import TCLIService, ttypes
1619
from databricks.sql import *
20+
from databricks.sql.thrift_api.TCLIService.TCLIService import (
21+
Client as TCLIServiceClient,
22+
)
1723
from databricks.sql.utils import (
1824
ArrowQueue,
1925
ExecuteResponse,
@@ -38,6 +44,7 @@
3844
"_retry_delay_max": (float, 60, 5, 3600),
3945
"_retry_stop_after_attempts_count": (int, 30, 1, 60),
4046
"_retry_stop_after_attempts_duration": (float, 900, 1, 86400),
47+
"_retry_delay_default": (float, 5, 1, 60),
4148
}
4249

4350

@@ -70,6 +77,8 @@ def __init__(
7077
# _retry_delay_min (default: 1)
7178
# _retry_delay_max (default: 60)
7279
# {min,max} pre-retry delay bounds
80+
# _retry_delay_default (default: 5)
81+
# Only used when GetOperationStatus fails due to a TCP/OS Error.
7382
# _retry_stop_after_attempts_count (default: 30)
7483
# total max attempts during retry sequence
7584
# _retry_stop_after_attempts_duration (default: 900)
@@ -160,7 +169,7 @@ def _initialize_retry_args(self, kwargs):
160169
"retry parameter: {} given_or_default {}".format(key, given_or_default)
161170
)
162171
if bound != given_or_default:
163-
logger.warn(
172+
logger.warning(
164173
"Override out of policy retry parameter: "
165174
+ "{} given {}, restricted to {}".format(
166175
key, given_or_default, bound
@@ -245,7 +254,9 @@ def _handle_request_error(self, error_info, attempt, elapsed):
245254
# FUTURE: Consider moving to https://github.com/litl/backoff or
246255
# https://github.com/jd/tenacity for retry logic.
247256
def make_request(self, method, request):
248-
"""Execute given request, attempting retries when receiving HTTP 429/503.
257+
"""Execute given request, attempting retries when
258+
1. Receiving HTTP 429/503 from server
259+
2. OSError is raised during a GetOperationStatus
249260
250261
For delay between attempts, honor the given Retry-After header, but with bounds.
251262
Use lower bound of expontial-backoff based on _retry_delay_min,
@@ -262,17 +273,21 @@ def make_request(self, method, request):
262273
def get_elapsed():
263274
return time.time() - t0
264275

276+
def bound_retry_delay(attempt, proposed_delay):
277+
"""bound delay (seconds) by [min_delay*1.5^(attempt-1), max_delay]"""
278+
delay = int(proposed_delay)
279+
delay = max(delay, self._retry_delay_min * math.pow(1.5, attempt - 1))
280+
delay = min(delay, self._retry_delay_max)
281+
return delay
282+
265283
def extract_retry_delay(attempt):
266284
# encapsulate retry checks, returns None || delay-in-secs
267285
# Retry IFF 429/503 code + Retry-After header set
268286
http_code = getattr(self._transport, "code", None)
269287
retry_after = getattr(self._transport, "headers", {}).get("Retry-After")
270288
if http_code in [429, 503] and retry_after:
271289
# bound delay (seconds) by [min_delay*1.5^(attempt-1), max_delay]
272-
delay = int(retry_after)
273-
delay = max(delay, self._retry_delay_min * math.pow(1.5, attempt - 1))
274-
delay = min(delay, self._retry_delay_max)
275-
return delay
290+
return bound_retry_delay(attempt, int(retry_after))
276291
return None
277292

278293
def attempt_request(attempt):
@@ -281,24 +296,57 @@ def attempt_request(attempt):
281296
# - non-None method_return -> success, return and be done
282297
# - non-None retry_delay -> sleep delay before retry
283298
# - error, error_message always set when available
299+
300+
error, error_message, retry_delay = None, None, None
284301
try:
285302
logger.debug("Sending request: {}".format(request))
286303
response = method(request)
287304
logger.debug("Received response: {}".format(response))
288305
return response
289-
except Exception as error:
306+
except OSError as err:
307+
error = err
308+
error_message = str(err)
309+
310+
gos_name = TCLIServiceClient.GetOperationStatus.__name__
311+
if method.__name__ == gos_name:
312+
retry_delay = bound_retry_delay(attempt, self._retry_delay_default)
313+
314+
# fmt: off
315+
# The built-in errno package encapsulates OSError codes, which are OS-specific.
316+
# log.info for errors we believe are not unusual or unexpected. log.warn for
317+
# for others like EEXIST, EBADF, ERANGE which are not expected in this context.
318+
#
319+
# I manually tested this retry behaviour using mitmweb and confirmed that
320+
# GetOperationStatus requests are retried when I forced network connection
321+
# interruptions / timeouts / reconnects. See #24 for more info.
322+
# | Debian | Darwin |
323+
info_errs = [ # |--------|--------|
324+
errno.ESHUTDOWN, # | 32 | 32 |
325+
errno.EAFNOSUPPORT, # | 97 | 47 |
326+
errno.ECONNRESET, # | 104 | 54 |
327+
errno.ETIMEDOUT, # | 110 | 60 |
328+
]
329+
330+
# fmt: on
331+
log_string = f"{gos_name} failed with code {err.errno} and will attempt to retry"
332+
if err.errno in info_errs:
333+
logger.info(log_string)
334+
else:
335+
logger.warning(log_string)
336+
except Exception as err:
337+
error = err
290338
retry_delay = extract_retry_delay(attempt)
291339
error_message = ThriftBackend._extract_error_message_from_headers(
292340
getattr(self._transport, "headers", {})
293341
)
294-
return RequestErrorInfo(
295-
error=error,
296-
error_message=error_message,
297-
retry_delay=retry_delay,
298-
http_code=getattr(self._transport, "code", None),
299-
method=method.__name__,
300-
request=request,
301-
)
342+
return RequestErrorInfo(
343+
error=error,
344+
error_message=error_message,
345+
retry_delay=retry_delay,
346+
http_code=getattr(self._transport, "code", None),
347+
method=method.__name__,
348+
request=request,
349+
)
302350

303351
# The real work:
304352
# - for each available attempt:

‎tests/unit/test_thrift_backend.py‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ def retry_policy_factory():
1919
"_retry_delay_max": (float, 60, None, None),
2020
"_retry_stop_after_attempts_count": (int, 30, None, None),
2121
"_retry_stop_after_attempts_duration": (float, 900, None, None),
22+
"_retry_delay_default": (float, 5, 1, 60)
2223
}
2324

2425

@@ -968,6 +969,62 @@ def test_handle_execute_response_sets_active_op_handle(self):
968969

969970
self.assertEqual(mock_resp.operationHandle, mock_cursor.active_op_handle)
970971

972+
@patch("thrift.transport.THttpClient.THttpClient")
973+
@patch("databricks.sql.thrift_api.TCLIService.TCLIService.Client.GetOperationStatus")
974+
@patch("databricks.sql.thrift_backend._retry_policy", new_callable=retry_policy_factory)
975+
def test_make_request_will_retry_GetOperationStatus(
976+
self, mock_retry_policy, mock_GetOperationStatus, t_transport_class):
977+
978+
import thrift, errno
979+
from databricks.sql.thrift_api.TCLIService.TCLIService import Client
980+
from databricks.sql.exc import RequestError
981+
from databricks.sql.utils import NoRetryReason
982+
983+
this_gos_name = "GetOperationStatus"
984+
mock_GetOperationStatus.__name__ = this_gos_name
985+
mock_GetOperationStatus.side_effect = OSError(errno.ETIMEDOUT, "Connection timed out")
986+
987+
protocol = thrift.protocol.TBinaryProtocol.TBinaryProtocol(t_transport_class)
988+
client = Client(protocol)
989+
990+
req = ttypes.TGetOperationStatusReq(
991+
operationHandle=self.operation_handle,
992+
getProgressUpdate=False,
993+
)
994+
995+
EXPECTED_RETRIES = 2
996+
997+
thrift_backend = ThriftBackend(
998+
"foobar",
999+
443,
1000+
"path", [],
1001+
_retry_stop_after_attempts_count=EXPECTED_RETRIES,
1002+
_retry_delay_default=1)
1003+
1004+
1005+
with self.assertRaises(RequestError) as cm:
1006+
thrift_backend.make_request(client.GetOperationStatus, req)
1007+
1008+
self.assertEqual(NoRetryReason.OUT_OF_ATTEMPTS.value, cm.exception.context["no-retry-reason"])
1009+
self.assertEqual(f'{EXPECTED_RETRIES}/{EXPECTED_RETRIES}', cm.exception.context["attempt"])
1010+
1011+
# Unusual OSError code
1012+
mock_GetOperationStatus.side_effect = OSError(errno.EEXIST, "File does not exist")
1013+
1014+
with self.assertLogs("databricks.sql.thrift_backend", level=logging.WARNING) as cm:
1015+
with self.assertRaises(RequestError):
1016+
thrift_backend.make_request(client.GetOperationStatus, req)
1017+
1018+
# There should be two warning log messages: one for each retry
1019+
self.assertEqual(len(cm.output), EXPECTED_RETRIES)
1020+
1021+
# The warnings should be identical
1022+
self.assertEqual(cm.output[1], cm.output[0])
1023+
1024+
# The warnings should include this text
1025+
self.assertIn(f"{this_gos_name} failed with code {errno.EEXIST} and will attempt to retry", cm.output[0])
1026+
1027+
9711028
@patch("thrift.transport.THttpClient.THttpClient")
9721029
def test_make_request_wont_retry_if_headers_not_present(self, t_transport_class):
9731030
t_transport_instance = t_transport_class.return_value

0 commit comments

Comments
 (0)