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)]