def __init__(self, schema):
self._read_queues = dict([(t.name_in_db, RequestQueue(t.name, False, '%s reads' % (t.name), t.read_units)) \
for t in schema.GetTables()])
self._write_queues = dict([(t.name_in_db, RequestQueue(t.name, True, '%s writes' % (t.name), t.write_units)) \
for t in schema.GetTables()])
self._cp_read_only_queue = RequestQueue('ControlPlane', False, 'Control Plane R/O', 100)
self._cp_mutate_queue = RequestQueue('ControlPlane', True, 'Control Plane Mutate', 1)
self._paused = False
self._asyncdynamo = AsyncDynamoDB(secrets.GetSecret('aws_access_key_id'),
secrets.GetSecret('aws_secret_access_key'))
class RequestScheduler(object):
"""Prioritizes and schedules competing requests to the DynamoDB
backend. Requests are organized by tables. Each table has its own
provisioning for read and writes per second. Depending on failures
indicating that provisioned throughput is being exceeded, requests
are placed into priority queues and throttled to just under the
maximum sustainable rate.
"""
_READ_ONLY_METHODS = ('ListTables', 'DescribeTable', 'GetItem', 'Query', 'Scan', 'BatchGetItem')
def __init__(self, schema):
self._read_queues = dict([(t.name_in_db, RequestQueue(t.name, False, '%s reads' % (t.name), t.read_units)) \
for t in schema.GetTables()])
self._write_queues = dict([(t.name_in_db, RequestQueue(t.name, True, '%s writes' % (t.name), t.write_units)) \
for t in schema.GetTables()])
self._cp_read_only_queue = RequestQueue('ControlPlane', False, 'Control Plane R/O', 100)
self._cp_mutate_queue = RequestQueue('ControlPlane', True, 'Control Plane Mutate', 1)
self._paused = False
self._asyncdynamo = AsyncDynamoDB(secrets.GetSecret('aws_access_key_id'),
secrets.GetSecret('aws_secret_access_key'))
def Schedule(self, method, request, callback):
"""Creates a DynamoDB request to API call 'method' with JSON
encoded arguments 'request'. Invokes 'callback' with JSON decoded
response as an argument.
"""
if method in ('ListTables', 'DescribeTable'):
queue = self._cp_read_only_queue
elif method in ('CreateTable', 'DeleteTable'):
queue = self._cp_mutate_queue
elif method in ('GetItem', 'Query', 'Scan'):
queue = self._read_queues[request['TableName']]
elif method in ('BatchGetItem',):
table_names = request['RequestItems'].keys()
assert len(table_names) == 1, table_names
queue = self._read_queues[table_names[0]]
else:
assert method in ('DeleteItem', 'PutItem', 'UpdateItem'), method
queue = self._write_queues[request['TableName']]
# The execution callback that we initialize the dynamodb request with is wrapped
# so that on execution, errors will be handled in the context of this method's caller.
dyn_req = DynDBRequest(method=method, request=request, op=None, finish_cb=callback,
execute_cb=stack_context.wrap(partial(self._ExecuteRequest, queue)))
queue.Push(dyn_req)
self._ProcessQueue(queue)
def _ExecuteRequest(self, queue, dyn_req):
"""Helper function to execute a DynamoDB request within the context
in which is was scheduled. This way, if an unrecoverable exception is
thrown during execution, it can be re-raised to the appropriate caller.
"""
def _OnResponse(start_time, json_response):
if dyn_req.method in ('BatchGetItem',):
consumed_units = next(json_response.get('Responses').itervalues()).get('ConsumedCapacityUnits', 1)
else:
consumed_units = json_response.get('ConsumedCapacityUnits', 1)
logging.debug('%s response: %d bytes, %d units, %.3fs elapsed' %
(dyn_req.method, len(json_response), consumed_units,
time.time() - start_time))
queue.Report(True, consumed_units)
dyn_req.finish_cb(json_response)
def _OnException(type, value, tb):
if type in (DBProvisioningExceededError, DBLimitExceededError):
# Retry on DynamoDB throttling errors. Report the failure to the queue so that it will backoff the
# requests/sec rate.
queue.Report(False)
elif type == DynamoDBResponseError and value.status in [500, 599] and \
dyn_req.method in RequestScheduler._READ_ONLY_METHODS:
# DynamoDB returns 500 when the service is unavailable for some reason.
# Curl returns 599 when something goes wrong with the connection, such as a timeout or connection reset.
# Only retry if this is a read-only request, since otherwise an update may be applied twice.
pass
else:
# Re-raise the exception now that we're in the stack context of original caller.
logging.warning('error calling "%s" with this request: %s' % (dyn_req.method, dyn_req.request))
raise type, value, tb
if dyn_req.method in ('BatchGetItem',):
table_name = next(dyn_req.request['RequestItems'].iterkeys())
else:
table_name = dyn_req.request.get('TableName', None)
logging.warning('%s against %s table failed: %s' % (dyn_req.method, table_name, value))
queue.Push(dyn_req)
self._ProcessQueue(queue)
logging.debug('sending %s (%d bytes) dynamodb request' % (dyn_req.method, len(dyn_req.request)))
with util.MonoBarrier(partial(_OnResponse, time.time()), on_exception=partial(_OnException)) as b:
self._asyncdynamo.make_request(dyn_req.method, json.dumps(dyn_req.request), b.Callback())
def _ProcessQueue(self, queue):
"""If the queue is not empty and adequate provisioning is expected,
sends the highest priority queue item(s) to DynamoDB.
When all items have been sent, resets the queue processing timeout.
"""
if self._paused:
return
while not queue.IsEmpty() and not queue.NeedsBackoff():
dyn_req = queue.Pop()
dyn_req.execute_cb(dyn_req)
queue.ResetTimeout(partial(self._ProcessQueue, queue))
def _Pause(self):
"""Pauses all queue processing. No requests will be sent until
_Resume() is invoked.
NOTE: intended for testing.
"""
self._paused = True
def _Resume(self):
"""Resume the scheduler if paused."""
if self._paused:
self._paused = False
[self._ProcessQueue(q) for q in self._read_queues.values()]
[self._ProcessQueue(q) for q in self._write_queues.values()]
[self._ProcessQueue(q) for q in (self._cp_read_only_queue, self._cp_mutate_queue)]
