Skip to content

surface_distance Dijkstra heap is sized height*width and overflows on dense target rasters, corrupting memory #3723

Description

@brendancol

Description

surface_distance(), surface_allocation() and surface_direction() can write past the end of their Dijkstra heap arrays and corrupt the process heap. On a 48x48 raster with 40% target pixels over rugged terrain the call aborts the interpreter with double free or corruption (!prev).

_dijkstra() and _dijkstra_geodesic() in xrspatial/surface_distance.py size the heap at height * width:

max_heap = height * width
h_keys = np.empty(max_heap, dtype=np.float64)
h_rows = np.empty(max_heap, dtype=np.int64)
h_cols = np.empty(max_heap, dtype=np.int64)

That is not the right bound. Both kernels use a lazy-deletion min-heap: a pixel is pushed again every time its tentative distance improves, and stale entries stay in the heap until they are popped and skipped by the visited check. Live heap occupancy is therefore bounded by the seed count plus the number of improving relaxations, not by the pixel count.

_heap_push is @ngjit with no bounds checking, so once h_size reaches height * width the next push writes into whatever allocation follows h_keys / h_rows / h_cols.

This is the same defect that was fixed in cost_distance (xrspatial/cost_distance.py:161), which now sizes its heap height * width * (n_neighbors + 1) with a comment explaining that the old height * width sizing "underflows that bound and let _heap_push write past the end of the arrays, corrupting memory". surface_distance was never updated.

Affected backends

  • numpy: _surface_distance_numpy calls _dijkstra / _dijkstra_geodesic directly.
  • dask+numpy: the bounded map_overlap path calls _surface_distance_numpy per chunk, and _run_tile in the iterative path calls _dijkstra per tile.
  • cupy and dask+cupy bounded are not affected; they use the parallel relaxation kernel, which has no heap.

Reproduction

import numpy as np
import xarray as xr
from xrspatial.surface_distance import surface_distance

N = 48
rng = np.random.default_rng(1)
src = np.zeros((N, N), dtype=np.float64)
src.flat[rng.choice(N * N, size=921, replace=False)] = 1.0   # 40% targets
elev = rng.random((N, N)) * 200.0                            # rugged relief

coords = {'y': np.arange(N, dtype=np.float64),
          'x': np.arange(N, dtype=np.float64)}
raster = xr.DataArray(src, dims=['y', 'x'], coords=coords,
                      attrs={'res': (1.0, 1.0)})
elevation = xr.DataArray(elev, dims=['y', 'x'], coords=coords,
                         attrs={'res': (1.0, 1.0)})

surface_distance(raster, elevation, connectivity=8)

Observed on this host (numba 0.62, Python 3.14):

$ python repro.py
double free or corruption (!prev)
Aborted

With bounds checking on, the out-of-bounds write is reported instead of corrupting memory:

$ NUMBA_BOUNDSCHECK=1 python repro.py
  File "xrspatial/surface_distance.py", line 474, in _surface_distance_numpy
    _dijkstra(elev_data, H, W, max_distance,
IndexError: index is out of bounds

How close normal inputs get

Instrumenting _dijkstra to record peak heap occupancy against the height * width cap, over random scenes varying grid size, relief amplitude and target density:

N relief targets cap (H*W) peak heap peak / cap
48 200 46 (2%) 2304 1314 0.57
48 200 115 (5%) 2304 1697 0.74
48 200 230 (10%) 2304 1940 0.84
48 200 460 (20%) 2304 2259 0.98
48 200 921 (40%) 2304 2519 1.09 -- overflow

Total pushes are already over the cap at 2% target density (1.6x at the top of this table). Only the interleaved pops keep live occupancy under it in the lower rows, and that margin disappears as target density or relief goes up. Dense target rasters are not exotic here: a landcover raster where roads or built-up cells are the targets easily reaches 20-40%.

Suggested fix

Grow the heap on demand instead of pre-allocating a fixed cap. Starting at height * width and doubling keeps the common case at today's footprint and removes the cap entirely.

The static bound height * width * (n_neighbors + 1) that cost_distance uses would also be correct, but for 8-connectivity it multiplies the heap allocation by 9, from 24 to 216 bytes per pixel. surface_distance documents an 80 bytes-per-pixel working set in _BYTES_PER_PIXEL and enforces it in _check_memory, so that constant would have to rise to roughly 272 and the guard would start refusing rasters about 3.4x smaller than it does today.

(Related, not fixed here: cost_distance._BYTES_PER_PIXEL is still 40, which was the pre-fix figure and no longer covers its own 9x heap.)

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:proximityArea: proximitybugSomething isn't workingdaskDask backend / chunked arraysoomOut-of-memory risk with large datasetsperformancePR touches performance-sensitive codeseverity:criticalSweep finding: CRITICALsweep-performanceFound by /sweep-performance

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions