229 lines
7.4 KiB
Rust
229 lines
7.4 KiB
Rust
use crate::AppState;
|
|
use crate::error::AppError;
|
|
use crate::tiles::{DataId, TileId};
|
|
use anyhow::{anyhow, bail};
|
|
use axum::Json;
|
|
use axum::extract::{Path, State};
|
|
use axum::http::header;
|
|
use axum::response::IntoResponse;
|
|
use flate2::read::GzDecoder;
|
|
use image::codecs::png::PngEncoder;
|
|
use image::{Rgba, RgbaImage};
|
|
use moka::future::Cache;
|
|
use rayon::iter::IntoParallelIterator;
|
|
use std::collections::HashMap;
|
|
use std::f64::consts::PI;
|
|
use std::io;
|
|
use std::io::{Cursor, ErrorKind};
|
|
use std::sync::Arc;
|
|
use tracing::info_span;
|
|
use wxbox_grib2::GribMessage;
|
|
use wxbox_grib2::wgs84::LatLong;
|
|
use wxbox_pal::ColorPalette;
|
|
|
|
use rayon::iter::ParallelIterator;
|
|
use wxbox_common::{Grib2DataSource, GribMessageTimeMetadata, GribTileMetadata};
|
|
|
|
pub type Grib2DataCache = Cache<DataId, Arc<GribMessage>>;
|
|
pub type Grib2TileCache = Cache<TileId, Arc<Vec<u8>>>;
|
|
|
|
pub type Grib2DataConfig = HashMap<String, Grib2DataSource>;
|
|
|
|
#[tracing::instrument(level = "info")]
|
|
pub async fn grib2_metadata(
|
|
Path(source): Path<String>,
|
|
State(state): State<AppState>,
|
|
) -> Result<Json<wxbox_common::GribTileMetadata>, AppError> {
|
|
// is this even a valid data source?
|
|
let data_id = DataId { source };
|
|
let Some(ds) = state.config.data.grib2.get(&data_id.source) else {
|
|
return Err(anyhow!("invalid/unknown grib2 state").into());
|
|
};
|
|
|
|
// ok, so we don't have a tile image yet
|
|
// this means we are going to have to kick off a task to put that in the cache
|
|
// lets check if we have the raw data
|
|
let data = if !state.grib2_data_cache.contains_key(&data_id) {
|
|
// we don't, so let's start by starting a task for that
|
|
load_grib2_data(state.grib2_data_cache, data_id, ds.clone()).await?
|
|
} else {
|
|
state.grib2_data_cache.get(&data_id).await.unwrap()
|
|
};
|
|
|
|
Ok(Json(GribTileMetadata {
|
|
message_time: GribMessageTimeMetadata {
|
|
reference_time_significance: data.identification.reference_time_significance,
|
|
year: data.identification.year,
|
|
month: data.identification.month,
|
|
day: data.identification.day,
|
|
hour: data.identification.hour,
|
|
minute: data.identification.minute,
|
|
second: data.identification.second,
|
|
},
|
|
source_configuration: ds.clone(),
|
|
}))
|
|
}
|
|
|
|
#[tracing::instrument(level = "info")]
|
|
pub async fn grib2_handler(
|
|
Path((source, z, x, y)): Path<(String, usize, usize, String)>,
|
|
State(state): State<AppState>,
|
|
) -> Result<impl IntoResponse, AppError> {
|
|
let mut y = y
|
|
.strip_suffix(".png")
|
|
.ok_or(io::Error::new(ErrorKind::InvalidInput, "invalid"))?;
|
|
let mut size = 256;
|
|
if y.ends_with("@2x") {
|
|
size = 512;
|
|
y = y.strip_suffix("@2x").unwrap();
|
|
}
|
|
let y: usize = y.parse()?;
|
|
let tile_id = TileId {
|
|
source,
|
|
z,
|
|
x,
|
|
y,
|
|
size,
|
|
};
|
|
|
|
// do we have a pre-prepared tile? if so, return it immediately
|
|
if let Some(tile) = state.grib2_tile_cache.get(&tile_id).await {
|
|
return Ok(([(header::CONTENT_TYPE, "image/png")], tile.as_ref().clone()));
|
|
}
|
|
|
|
// is this even a valid data source?
|
|
let data_id = tile_id.data_id();
|
|
let Some(ds) = state.config.data.grib2.get(&data_id.source) else {
|
|
return Err(anyhow!("invalid/unknown grib2 state").into());
|
|
};
|
|
|
|
// ok, so we don't have a tile image yet
|
|
// this means we are going to have to kick off a task to put that in the cache
|
|
// lets check if we have the raw data
|
|
let data = if !state.grib2_data_cache.contains_key(&data_id) {
|
|
// we don't, so let's start by starting a task for that
|
|
load_grib2_data(state.grib2_data_cache, data_id, ds.clone()).await?
|
|
} else {
|
|
state.grib2_data_cache.get(&data_id).await.unwrap()
|
|
};
|
|
|
|
// we know we need to build the tile, so let's do that now
|
|
// it also returns it, so we can conveniently return it right now
|
|
let pixel_data =
|
|
render_to_png(state.grib2_tile_cache.clone(), data, tile_id, ds.clone()).await?;
|
|
|
|
Ok((
|
|
[(header::CONTENT_TYPE, "image/png")],
|
|
pixel_data.as_ref().clone(),
|
|
))
|
|
}
|
|
|
|
#[tracing::instrument(level = "info")]
|
|
async fn load_grib2_data(
|
|
cache: Grib2DataCache,
|
|
data_id: DataId,
|
|
data_source: Grib2DataSource,
|
|
) -> anyhow::Result<Arc<GribMessage>> {
|
|
let client = reqwest::Client::new();
|
|
let r = client.get(data_source.from.as_str()).send().await?;
|
|
|
|
if !r.status().is_success() {
|
|
bail!("grib2 data failed to load: {}", r.status());
|
|
}
|
|
|
|
let data = Cursor::new(r.bytes().await?.to_vec());
|
|
|
|
let msg = if data_source.needs_gzip {
|
|
let decoder = GzDecoder::new(data);
|
|
GribMessage::new(decoder)
|
|
} else {
|
|
GribMessage::new(data)
|
|
}?;
|
|
|
|
let data = Arc::new(msg);
|
|
cache.insert(data_id, data.clone()).await;
|
|
|
|
Ok(data)
|
|
}
|
|
|
|
const TWO_PI: f64 = PI * 2.0;
|
|
const HALF_PI: f64 = PI / 2.0;
|
|
|
|
async fn render_to_png(
|
|
cache: Grib2TileCache,
|
|
data: Arc<GribMessage>,
|
|
tile_id: TileId,
|
|
data_source: Grib2DataSource,
|
|
) -> anyhow::Result<Arc<Vec<u8>>> {
|
|
let span = info_span!("render_to_png");
|
|
let span = span.enter();
|
|
let mut image = RgbaImage::new(tile_id.size as u32, tile_id.size as u32);
|
|
|
|
let n = 2_usize.pow(tile_id.z as u32) as f64 * tile_id.size as f64;
|
|
let tile_x_times_tilesize = tile_id.x as f64 * tile_id.size as f64;
|
|
let tile_y_times_tilesize = tile_id.y as f64 * tile_id.size as f64;
|
|
|
|
let generate_pixels_span = info_span!("generate_pixels");
|
|
let pixels = (0..tile_id.size)
|
|
.into_par_iter()
|
|
.map(|x| {
|
|
(0..tile_id.size)
|
|
.into_par_iter()
|
|
.map(|y| {
|
|
let x_cartesian = (tile_x_times_tilesize + x as f64) / n;
|
|
let y_cartesian = (tile_y_times_tilesize + y as f64) / n;
|
|
|
|
let long = (TWO_PI * x_cartesian - PI).to_degrees();
|
|
let lat =
|
|
((PI - TWO_PI * y_cartesian).exp().atan() * 2.0_f64 - HALF_PI).to_degrees();
|
|
|
|
let nearest_value = data.value_for(LatLong { lat, long }).map(|u| u as f64);
|
|
|
|
colorize(nearest_value, &data_source).unwrap_or(Rgba::from([0, 0, 0, 0]))
|
|
})
|
|
.collect::<Vec<_>>()
|
|
})
|
|
.collect::<Vec<_>>();
|
|
drop(generate_pixels_span);
|
|
|
|
let put_pixels_span = info_span!("put_pixels");
|
|
|
|
pixels.iter().enumerate().for_each(|(x, r)| {
|
|
r.iter().enumerate().for_each(|(y, c)| {
|
|
image.put_pixel(x as u32, y as u32, *c);
|
|
});
|
|
});
|
|
|
|
drop(put_pixels_span);
|
|
|
|
// encode to png bytes and return
|
|
let mut result = vec![];
|
|
let encoder = PngEncoder::new(&mut result);
|
|
image.write_with_encoder(encoder)?;
|
|
|
|
let output = Arc::new(result);
|
|
|
|
drop(span);
|
|
|
|
cache.insert(tile_id, output.clone()).await;
|
|
|
|
Ok(output)
|
|
}
|
|
|
|
fn colorize(value: Option<f64>, data_source: &Grib2DataSource) -> anyhow::Result<Rgba<u8>> {
|
|
Ok(match value {
|
|
Some(c) if Some(c) == data_source.no_coverage => Rgba([0, 0, 0, 30]),
|
|
Some(c) if Some(c) == data_source.missing => Rgba([0, 0, 0, 0]),
|
|
Some(c) if Some(c) == data_source.range_folded => Rgba([119, 0, 125, 255]),
|
|
Some(value) => {
|
|
let color = wxbox_pal::parser::parse(&data_source.palette)?.colorize(value);
|
|
Rgba([
|
|
color.red,
|
|
color.green,
|
|
color.blue,
|
|
color.alpha, // guaranteed to be present afaict
|
|
])
|
|
}
|
|
None => Rgba([0, 0, 0, 30]),
|
|
})
|
|
}
|