Speed up ERA5 batch pipeline: parallel CDS + stream downloads + hot index

Four independent wins from a perf audit of the era5_batch worker path:

1. Parallel CDS submits per job. fetch_single_level_month and
   fetch_pressure_level_month have no ordering dependency — each blocks
   30+ min on CDS queue/download — so running them in two linked Tasks
   halves wall time for every job in the backlog.

2. Stream GRIB2 downloads directly to disk via Req `into: File.stream!`.
   A month-tile GRIB is 50–200 MB and was previously slurped into the
   BEAM heap and then written back out to a temp file before wgrib2 ran.
   New Era5Client.submit_and_download_to_file/3 skips both copies. Saves
   100–200 MB heap per worker × 12 concurrent ≈ 1.2–2.4 GB peak under
   full parallelism. Partial files are removed on HTTP error.

3. Wgrib2.extract_grid_messages_from_file/3 lets batch decoding read the
   streamed temp file directly. The binary variant still exists for the
   single-point path and now delegates through the same helper.

4. Composite index (valid_time, lat, lon) on era5_profiles, matching the
   three-range scan used by Weather.find_nearest_era5 (contact detail +
   path profiles) and StatusLive.count_era5_done (status page EXISTS
   subquery). The unique index on (lat, lon, valid_time) couldn't serve
   it efficiently.

Also bumps the bulk insert chunk size 2k → 4k to halve insert round-trips
per month-tile (still well under Postgres' 65k-parameter cap).

lib/microwaveprop/weather/era5_batch_client.ex

@@ -81,14 +81,39 @@ defmodule Microwaveprop.Weather.Era5BatchClient do
 
     Logger.info("ERA5Batch: fetching #{year}-#{pad(month)} tile #{tile_lat},#{tile_lon} (#{days} days)")
 
-    with {:ok, single_grib} <- fetch_single_level_month(year, month, days, area),
+    tmp_dir = System.tmp_dir!()
-         {:ok, pressure_grib} <- fetch_pressure_level_month(year, month, days, area),
+    job_tag = "#{year}_#{pad(month)}_#{tile_lat}_#{tile_lon}_#{System.unique_integer([:positive])}"
-         {:ok, profiles} <- build_month_profiles(tile_lat, tile_lon, single_grib, pressure_grib) do
+    single_path = Path.join(tmp_dir, "era5_single_#{job_tag}.grib2")
-      inserted = bulk_insert_profiles(profiles)
+    pressure_path = Path.join(tmp_dir, "era5_pressure_#{job_tag}.grib2")
 
-      Logger.info("ERA5Batch: stored #{inserted} profiles for #{year}-#{pad(month)} tile #{tile_lat},#{tile_lon}")
+    # Fire both CDS submits in parallel — they have no ordering dependency and
+    # each spends 30+ min blocking on CDS queue/download. Running them serially
+    # doubles wall time for zero benefit. Task.async links to the Oban worker,
+    # so if the worker exits (deploy, crash) both tasks die with it. We still
+    # shut down any unfinished task on early error so we don't waste CDS quota.
+    # Both downloads stream directly to disk (Req `into: File.stream!`), so
+    # peak heap stays O(chunk) rather than O(month-tile GRIB size ~50-200 MB).
+    single_task =
+      Task.async(fn -> fetch_single_level_month(year, month, days, area, single_path) end)
 
-      {:ok, inserted}
+    pressure_task =
+      Task.async(fn -> fetch_pressure_level_month(year, month, days, area, pressure_path) end)
+
+    try do
+      with :ok <- Task.await(single_task, :infinity),
+           :ok <- Task.await(pressure_task, :infinity),
+           {:ok, profiles} <- build_month_profiles(tile_lat, tile_lon, single_path, pressure_path) do
+        inserted = bulk_insert_profiles(profiles)
+
+        Logger.info("ERA5Batch: stored #{inserted} profiles for #{year}-#{pad(month)} tile #{tile_lat},#{tile_lon}")
+
+        {:ok, inserted}
+      end
+    after
+      Task.shutdown(single_task, :brutal_kill)
+      Task.shutdown(pressure_task, :brutal_kill)
+      File.rm(single_path)
+      File.rm(pressure_path)
     end
   end
 
@@ -123,7 +148,7 @@ defmodule Microwaveprop.Weather.Era5BatchClient do
 
   # -- CDS requests ----------------------------------------------------------
 
-  defp fetch_single_level_month(year, month, days, area) do
+  defp fetch_single_level_month(year, month, days, area, dest_path) do
     request = %{
       "product_type" => ["reanalysis"],
       "variable" => Era5Client.single_level_vars(),
@@ -135,10 +160,10 @@ defmodule Microwaveprop.Weather.Era5BatchClient do
       "data_format" => "grib"
     }
 
-    Era5Client.submit_and_download("reanalysis-era5-single-levels", request)
+    Era5Client.submit_and_download_to_file("reanalysis-era5-single-levels", request, dest_path)
   end
 
-  defp fetch_pressure_level_month(year, month, days, area) do
+  defp fetch_pressure_level_month(year, month, days, area, dest_path) do
     request = %{
       "product_type" => ["reanalysis"],
       "variable" => Era5Client.pressure_level_vars(),
@@ -151,7 +176,7 @@ defmodule Microwaveprop.Weather.Era5BatchClient do
       "data_format" => "grib"
     }
 
-    Era5Client.submit_and_download("reanalysis-era5-pressure-levels", request)
+    Era5Client.submit_and_download_to_file("reanalysis-era5-pressure-levels", request, dest_path)
   end
 
   # -- GRIB extraction -------------------------------------------------------
@@ -169,12 +194,21 @@ defmodule Microwaveprop.Weather.Era5BatchClient do
     }
   end
 
-  defp build_month_profiles(tile_lat, tile_lon, single_grib, pressure_grib) do
+  defp build_month_profiles(tile_lat, tile_lon, single_path, pressure_path) do
     grid_spec = grid_spec_for_tile(tile_lat, tile_lon)
 
     with {:ok, single_msgs} <-
-           Wgrib2.extract_grid_messages(single_grib, ":(TMP|DPT|PRES|UGRD|VGRD|PWAT|HPBL):", grid_spec),
+           Wgrib2.extract_grid_messages_from_file(
-         {:ok, pressure_msgs} <- Wgrib2.extract_grid_messages(pressure_grib, ":(TMP|SPFH|HGT):", grid_spec) do
+             single_path,
+             ":(TMP|DPT|PRES|UGRD|VGRD|PWAT|HPBL):",
+             grid_spec
+           ),
+         {:ok, pressure_msgs} <-
+           Wgrib2.extract_grid_messages_from_file(
+             pressure_path,
+             ":(TMP|SPFH|HGT):",
+             grid_spec
+           ) do
       # Index both message sets by (datetime, lat, lon) → var:level map.
       points_single = index_messages(single_msgs)
       points_pressure = index_messages(pressure_msgs)
@@ -291,7 +325,7 @@ defmodule Microwaveprop.Weather.Era5BatchClient do
 
   # Insert in chunks to keep each query under Postgres' parameter limit
   # (65,535 params / ~15 fields ≈ 4,000 rows per statement).
-  @chunk_size 2_000
+  @chunk_size 4_000
 
   defp bulk_insert_profiles(profiles) do
     now = DateTime.truncate(DateTime.utc_now(), :second)

lib/microwaveprop/weather/era5_client.ex

@@ -127,15 +127,39 @@ defmodule Microwaveprop.Weather.Era5Client do
     do_submit_and_download(dataset, request)
   end
 
+  @doc """
+  Like `submit_and_download/2` but streams the GRIB2 response directly to
+  `path` on disk rather than loading the whole body (50-200 MB for month
+  tiles) into the BEAM heap.
+  """
+  @spec submit_and_download_to_file(String.t(), map(), Path.t()) :: :ok | {:error, term()}
+  def submit_and_download_to_file(dataset, request, path) do
+    with {:ok, source, api_key} <- submit_and_poll(dataset, request) do
+      case source do
+        {:body, body} -> File.write(path, body)
+        {:url, href} -> download_href_to_file(href, api_key, path)
+      end
+    end
+  end
+
   defp do_submit_and_download(dataset, request) do
+    with {:ok, source, api_key} <- submit_and_poll(dataset, request) do
+      case source do
+        {:body, body} -> {:ok, body}
+        {:url, href} -> download_file(href, api_key)
+      end
+    end
+  end
+
+  defp submit_and_poll(dataset, request) do
     api_key = api_key()
 
     if is_nil(api_key) do
       {:error, "ERA5_CDS_API_KEY not configured"}
     else
-      case submit_request(dataset, request, api_key) do
+      with {:ok, job_id} <- submit_request(dataset, request, api_key),
-        {:ok, job_id} -> poll_and_download(job_id, api_key)
+           {:ok, source} <- poll_until_complete(job_id, api_key) do
-        {:error, reason} -> {:error, reason}
+        {:ok, source, api_key}
       end
     end
   end
@@ -166,24 +190,26 @@ defmodule Microwaveprop.Weather.Era5Client do
     end
   end
 
-  defp poll_and_download(job_id, api_key, attempt \\ 0)
+  # Returns {:ok, {:url, href}} when CDS returns a JSON asset pointer (new CDS)
+  # or {:ok, {:body, binary}} for direct binary responses (legacy CDS).
+  defp poll_until_complete(job_id, api_key, attempt \\ 0)
 
-  defp poll_and_download(_job_id, _api_key, attempt) when attempt >= @max_poll_attempts do
+  defp poll_until_complete(_job_id, _api_key, attempt) when attempt >= @max_poll_attempts do
     {:error, "ERA5: poll timeout after #{@max_poll_attempts} attempts"}
   end
 
-  defp poll_and_download(job_id, api_key, attempt) do
+  defp poll_until_complete(job_id, api_key, attempt) do
     url = "#{@cds_url}/retrieve/v1/jobs/#{job_id}"
 
     case Req.get(url, headers: [{"PRIVATE-TOKEN", api_key}], receive_timeout: 10_000) do
       {:ok, %{status: 200, body: %{"status" => "successful"}}} ->
         results_url = "#{url}/results"
-        download_result(results_url, api_key)
+        fetch_results_source(results_url, api_key)
 
       {:ok, %{status: 200, body: %{"status" => status}}}
       when status in ["accepted", "running"] ->
         Process.sleep(@poll_interval_ms)
-        poll_and_download(job_id, api_key, attempt + 1)
+        poll_until_complete(job_id, api_key, attempt + 1)
 
       {:ok, %{status: 200, body: %{"status" => "failed"} = body}} ->
         # CDS often omits "message" — fall back to dumping the whole body so the
@@ -199,18 +225,21 @@ defmodule Microwaveprop.Weather.Era5Client do
     end
   end
 
-  defp download_result(url, api_key) do
+  defp fetch_results_source(url, api_key) do
     case Req.get(url, headers: [{"PRIVATE-TOKEN", api_key}], receive_timeout: 120_000) do
       {:ok, %{status: 200, body: %{"asset" => %{"value" => %{"href" => href}}}}} ->
         # New CDS returns a JSON with download link
-        download_file(href, api_key)
+        {:ok, {:url, href}}
 
       {:ok, %{status: 200, headers: headers, body: body}} ->
         # Direct binary download
-        content_type = %Req.Response{headers: headers} |> Req.Response.get_header("content-type") |> List.first("")
+        content_type =
+          %Req.Response{headers: headers}
+          |> Req.Response.get_header("content-type")
+          |> List.first("")
 
         if String.contains?(content_type, "grib") or is_binary(body) do
-          {:ok, body}
+          {:ok, {:body, body}}
         else
           {:error, "ERA5: unexpected response format"}
         end
@@ -231,6 +260,32 @@ defmodule Microwaveprop.Weather.Era5Client do
     end
   end
 
+  # Streams the GRIB2 response directly to `path`. Uses Req's `into:` with a
+  # `File.stream!` collectable, so chunks land on disk as they arrive and peak
+  # heap stays O(chunk size) instead of O(full response). Target path is
+  # removed on HTTP error so a failed download never leaves a partial tmp file.
+  defp download_href_to_file(href, api_key, path) do
+    result =
+      Req.get(href,
+        headers: [{"PRIVATE-TOKEN", api_key}],
+        receive_timeout: 300_000,
+        into: File.stream!(path)
+      )
+
+    case result do
+      {:ok, %{status: 200}} ->
+        :ok
+
+      {:ok, %{status: status}} ->
+        File.rm(path)
+        {:error, "ERA5 file download HTTP #{status}"}
+
+      {:error, reason} ->
+        File.rm(path)
+        {:error, "ERA5 file download failed: #{inspect(reason)}"}
+    end
+  end
+
   defp build_profile(lat, lon, valid_time, single_grib, pressure_grib) do
     # Decode GRIB2 data using existing infrastructure
     with {:ok, single_points} <- extract_point(single_grib, lat, lon),

lib/microwaveprop/weather/grib2/wgrib2.ex

@@ -354,13 +354,44 @@ defmodule Microwaveprop.Weather.Grib2.Wgrib2 do
           | {:error, term()}
   def extract_grid_messages(grib_binary, match_pattern, grid_spec) do
     if available?() do
-      extract_messages_with_wgrib2(grib_binary, match_pattern, grid_spec)
+      tmp_grib = Path.join(System.tmp_dir!(), "era5_#{System.unique_integer([:positive])}.grib2")
+
+      try do
+        File.write!(tmp_grib, grib_binary)
+        extract_messages_with_wgrib2(tmp_grib, match_pattern, grid_spec)
+      after
+        File.rm(tmp_grib)
+      end
     else
       {:error, :wgrib2_not_available}
     end
   end
 
-  defp extract_messages_with_wgrib2(grib_binary, match_pattern, grid_spec) do
+  @doc """
+  Like `extract_grid_messages/3`, but reads GRIB2 from a file on disk instead
+  of a binary in memory. Saves 50-200 MB of heap per call for month-tile
+  fetches by avoiding the intermediate `File.write!(tmp, binary)` step.
+  """
+  @spec extract_grid_messages_from_file(Path.t(), String.t(), grid_spec()) ::
+          {:ok,
+           [
+             %{
+               datetime: DateTime.t() | nil,
+               var: String.t(),
+               level: String.t(),
+               values: %{{float(), float()} => float()}
+             }
+           ]}
+          | {:error, term()}
+  def extract_grid_messages_from_file(grib_path, match_pattern, grid_spec) do
+    if available?() do
+      extract_messages_with_wgrib2(grib_path, match_pattern, grid_spec)
+    else
+      {:error, :wgrib2_not_available}
+    end
+  end
+
+  defp extract_messages_with_wgrib2(grib_path, match_pattern, grid_spec) do
     %{
       lon_start: lon_start,
       lon_count: lon_count,
@@ -374,13 +405,10 @@ defmodule Microwaveprop.Weather.Grib2.Wgrib2 do
     lon_spec = "#{wgrib2_lon_start}:#{lon_count}:#{lon_step}"
     lat_spec = "#{lat_start}:#{lat_count}:#{lat_step}"
 
-    tmp_grib = Path.join(System.tmp_dir!(), "era5_#{System.unique_integer([:positive])}.grib2")
+    tmp_bin = grib_path <> ".lola.#{System.unique_integer([:positive])}.bin"
-    tmp_bin = tmp_grib <> ".lola.bin"
 
     try do
-      File.write!(tmp_grib, grib_binary)
+      args = [grib_path, "-match", match_pattern, "-lola", lon_spec, lat_spec, tmp_bin, "bin"]
-
-      args = [tmp_grib, "-match", match_pattern, "-lola", lon_spec, lat_spec, tmp_bin, "bin"]
 
       case System.cmd(wgrib2_path(), args, stderr_to_stdout: true) do
         {output, 0} ->
@@ -401,7 +429,6 @@ defmodule Microwaveprop.Weather.Grib2.Wgrib2 do
           {:error, "wgrib2 failed (exit #{exit_code}): #{String.slice(output, 0, 200)}"}
       end
     after
-      File.rm(tmp_grib)
       File.rm(tmp_bin)
     end
   end

priv/repo/migrations/20260413205533_add_era5_profiles_valid_time_lat_lon_index.exs

@@ -0,0 +1,17 @@
+defmodule Microwaveprop.Repo.Migrations.AddEra5ProfilesValidTimeLatLonIndex do
+  use Ecto.Migration
+
+  # Matches the hot range-scan access pattern used by both
+  # `Weather.find_nearest_era5/3` (contact detail + path profiles) and
+  # `StatusLive.count_era5_done/0` (status page EXISTS correlated subquery):
+  #   WHERE valid_time BETWEEN ? AND ?
+  #     AND lat BETWEEN ? AND ?
+  #     AND lon BETWEEN ? AND ?
+  # `valid_time` leads because it collapses the search to a ~1-hour bucket
+  # first; lat/lon then filter the ~9×9 tile grid within that bucket.
+  # The existing unique index (lat, lon, valid_time) handles upsert dedup
+  # but can't serve three-range scans efficiently.
+  def change do
+    create index(:era5_profiles, [:valid_time, :lat, :lon])
+  end
+end