Limit the number of entities fetched per statement
Without this restriction, Cassandra can easily time out while
trying to fetch many rows from the entities table.
@@ -48,6 +48,9 @@ from google.appengine.datastore import entity_pb
# The directory Cassandra is installed to.
CASSANDRA_INSTALL_DIR = '/opt/cassandra'
+# The maximum amount of entities to fetch concurrently.
+ENTITY_FETCH_THRESHOLD = 100
+
# Full path for the nodetool binary.
NODE_TOOL = '{}/cassandra/bin/nodetool'.format(CASSANDRA_INSTALL_DIR)
@@ -195,23 +198,36 @@ class DatastoreProxy(AppDBInterface):
column=ThriftColumn.COLUMN_NAME,
)
query = SimpleStatement(statement, retry_policy=BASIC_RETRIES)
- parameters = (ValueSequence(row_keys_bytes), ValueSequence(column_names))
+ results = []
+ # Split the rows up into chunks to reduce the likelihood of timeouts.
+ chunk_indexes = [
+ (n, n + ENTITY_FETCH_THRESHOLD)
+ for n in xrange(0, len(row_keys_bytes), ENTITY_FETCH_THRESHOLD)]
+
+ # TODO: This can be made more efficient by maintaining a constant number
+ # of concurrent requests rather than waiting for each batch to complete.
+ for start, end in chunk_indexes:
+ parameters = (ValueSequence(row_keys_bytes[start:end]),
+ ValueSequence(column_names))
try:
- results = yield self.tornado_cassandra.execute(
+ batch_results = yield self.tornado_cassandra.execute(
query, parameters=parameters)
+ except dbconstants.TRANSIENT_CASSANDRA_ERRORS:
+ message = 'Exception during batch_get_entity'
+ logger.exception(message)
+ raise AppScaleDBConnectionError(message)
+
+ results.extend(list(batch_results))
results_dict = {row_key: {} for row_key in row_keys}
for (key, column, value) in results:
if key not in results_dict:
results_dict[key] = {}
+
results_dict[key][column] = value
raise gen.Return(results_dict)
- except dbconstants.TRANSIENT_CASSANDRA_ERRORS:
- message = 'Exception during batch_get_entity'
- logger.exception(message)
- raise AppScaleDBConnectionError(message)
@gen.coroutine
def batch_put_entity(self, table_name, row_keys, column_names, cell_values,
