From 4bce413346cf77b22b2667895873ddc2580e71b6 Mon Sep 17 00:00:00 2001 From: Ian Dees Date: Wed, 26 Aug 2026 14:46:28 -0500 Subject: [PATCH] add guardrails for the aggregate endpoint's memory usage --- census_extractomatic/api.py | 22 ++++++++++++++--- census_extractomatic/user_geo.py | 42 ++++++++++++++++++++++++++++++++ 2 files changed, 60 insertions(+), 4 deletions(-) diff --git a/census_extractomatic/api.py b/census_extractomatic/api.py index 347d6ae..9c1a8ea 100644 --- a/census_extractomatic/api.py +++ b/census_extractomatic/api.py @@ -41,6 +41,7 @@ ) from .user_geo import ( COMPARISON_RELEASE_CODE, + GeographyTooLargeError, build_filename, create_block_xref_download, fetch_user_geodata, @@ -2383,10 +2384,23 @@ def aggregate(hash_digest, release, table_code): if url_exists(precomputed_url): return redirect(precomputed_url) - start = timer() - zf = create_aggregate_download(db, hash_digest, release, table_code) - end = timer() - return send_file(zf.name, 'application/zip', download_name=zipfile_name) + # Building this is expensive (full-geometry block-level query, held in memory), so + # only let one worker build a given (hash, release, table) at a time. Concurrent + # requests/retries for the same download get a 409 instead of piling on redundant work. + lock_key = f'aggregate-lock:{hash_digest}:{release}:{table_code}' + if not cache.add(lock_key, 1, timeout=300): + return jsonify(error="This geography is already being aggregated, please retry in a moment."), 409 + + try: + start = timer() + try: + zf = create_aggregate_download(db, hash_digest, release, table_code) + except GeographyTooLargeError as e: + return jsonify(error=str(e)), 413 + end = timer() + return send_file(zf.name, 'application/zip', download_name=zipfile_name) + finally: + cache.delete(lock_key) if __name__ == "__main__": diff --git a/census_extractomatic/user_geo.py b/census_extractomatic/user_geo.py index d0a7fc0..dfe0b74 100644 --- a/census_extractomatic/user_geo.py +++ b/census_extractomatic/user_geo.py @@ -41,6 +41,24 @@ def join_user_geo_to_blocks_task(user_geodata_id): COMPARISON_RELEASE_CODE = 'dec_pl94_compare_2020_2010' +# aggregate_decennial() pulls one row (with full geometry) per matching census block into +# memory via pandas, and the comparison release does this twice (2010 and 2020) and holds +# both dataframes at once. A user geography over this many blocks (e.g. a state's worth of +# counties, ~300k+ blocks) can balloon a single worker to multiple GB and take down the box. +# See the 2026-08-26 census-api memory incident. +MAX_SYNC_AGGREGATION_BLOCKS = 150_000 + + +class GeographyTooLargeError(Exception): + """Raised when a user geography has too many blocks to aggregate synchronously.""" + def __init__(self, block_count, limit): + self.block_count = block_count + self.limit = limit + super().__init__( + f"This geography has {block_count} census blocks, which is too many to " + f"aggregate on demand (limit {limit}). Contact us if you need this download." + ) + USER_GEODATA_INSERT_SQL = text(""" INSERT INTO aggregation.user_geodata (name, hash_digest, source_url, public, fields, bbox) VALUES (:name, :hash_digest, :source_url, :public, :fields, ST_MakeEnvelope(:xmin, :ymin, :xmax, :ymax, 4326)) @@ -351,6 +369,27 @@ def fetch_metadata(release=None, table_code=None): return None +def count_blocks_for_hash(db, hash_digest, release): + """Count the census blocks a user geography (identified by hash_digest) joins to + for the given release, without pulling any row data or geometry. + """ + try: + blocks_vintage_table = BLOCK_VINTAGE_TABLES[release] + except KeyError: + raise ValueError(f"No blocks vintage identified for given schema {release}") + sql = text(f""" + SELECT count(*) + FROM aggregation.user_geodata ug, + aggregation.user_geodata_geometry ugg, + aggregation.{blocks_vintage_table} ugb + WHERE ug.hash_digest = :hash_digest + AND ug.user_geodata_id = ugg.user_geodata_id + AND ugg.user_geodata_geometry_id = ugb.user_geodata_geometry_id + """) + with db.engine.begin() as con: + return con.execute(sql, dict(hash_digest=hash_digest)).scalar() + + def evaluateUserGeographySQLTemplate(schema, table_code): """Schemas and table names can't be handled as bindparams with SQLAlchemy, so this allows us to use a 'select *' syntax for multiple tables. @@ -370,6 +409,9 @@ def aggregate_decennial(db, hash_digest, release, table_code): """ if fetch_metadata(release=release, table_code=table_code): + block_count = count_blocks_for_hash(db, hash_digest, release) + if block_count > MAX_SYNC_AGGREGATION_BLOCKS: + raise GeographyTooLargeError(block_count, MAX_SYNC_AGGREGATION_BLOCKS) sql = evaluateUserGeographySQLTemplate(release, table_code) query = text(sql).bindparams(hash_digest=hash_digest) logger.info(f'aggregate_decennial: starting timer {hash_digest} {release} {table_code}')