Notify not working when triggered by a button

Specific Domains Visualization
1

My code for updating a heatmap is not working when triggered by a button.

I have tried to boil things down to a minimum working example. The expected behaviour is that at each evaluation of the cost function, one pixel of the heatmap is updated. In the example this works when run from code, but not when the update is induced by pressing a button.

Any ideas on what might be going wrong?

module ObservableHeatmaps

export ObservableHeatmap, update_observable_heatmap!

using Base.Iterators
using Makie


struct ObservableHeatmap{H}
    xs::Observable{Vector{Float64}}
    ys::Observable{Vector{Float64}}
    zs::Observable{Matrix{Float64}}
    ax::Axis
    hm::H
    function ObservableHeatmap(ax)
        xs, ys, zs = Observable(zeros(1)), Observable(zeros(1)),  Observable(zeros(1, 1))
        hm = heatmap!(ax, xs, ys, zs)    
        new{typeof(hm)}(xs, ys, zs, ax, hm)
    end
end

function scheduled_n_tasks(n_schedule_max, tasks, ondone)
    unscheduled_task_idxs = collect(eachindex(tasks))
    scheduled_task_idxs = empty(unscheduled_task_idxs)
    while length(unscheduled_task_idxs) + length(scheduled_task_idxs) > 0
        while length(scheduled_task_idxs) < n_schedule_max && length(unscheduled_task_idxs) > 0
            i = popfirst!(unscheduled_task_idxs)
            task = tasks[i]
            schedule(task)
            # Threads.@spawn :default () -> task.code()
            push!(scheduled_task_idxs, i)
        end
        scheduled_task_idxs = filter(scheduled_task_idxs) do i
            task = tasks[i]
            # If task has just finished
            isdone = istaskdone(task)
            if isdone
                ondone(fetch(task)...)
            end
            ~isdone
        end
        yield()
    end
end

function update_observable_heatmap!(
        hm::ObservableHeatmap,
        cost_func::Function,
        new_xs::AbstractVector{Float64},
        new_ys::AbstractVector{Float64})
    (;xs, ys, zs, ax) = hm
    
    # Reset heatmap to zero
    xs[], ys[], zs[] = new_xs, new_ys, zeros(length(new_xs), length(new_ys))
    notify.( (xs, ys, zs) )
    notify(ax.limits)
    sleep(0.1)

    # Create task for each pixel
    xy_iterator = product(enumerate(xs[]), enumerate(ys[]))
    tasks = Vector{Task}()
    for ((i, x), (j, y)) in xy_iterator
        task = Task() do
            C::Float64 = cost_func(x, y)
            i, j, C
        end
        push!(tasks, task)
    end

    # Callback to update plot (in main thread) when done
    function ondone(i, j, C)
        zs[][i, j] = C
        notify(zs)
        yield()
    end
    scheduled_n_tasks(Threads.nthreads()-1, tasks, ondone)
    nothing
end

end

using GLMakie
using .ObservableHeatmaps

fig = Figure()
display(fig)
ax = Axis(fig[1,1])
hmplot = ObservableHeatmap(ax)
# cb = Colorbar(fig[1, 2], hmplot.hm)

function expensive_cost_function(x, y)
    val = sin(x) + cos(y)
    for i = 1:1000
        val += val/i
    end
    return val
end

xs, ys = LinRange(0., 10., 50), LinRange(0., 10., 50)
update_observable_heatmap!(hmplot, expensive_cost_function, xs, ys)

button = Button(fig[2, 1], label="Recompute", tellwidth=false)
on(button.clicks) do n
    println("Recomputing...")
    update_observable_heatmap!(hmplot, expensive_cost_function, xs, ys)
    println("Done.")
end