engine refactor
This commit is contained in:
parent
9e663db9dc
commit
81b8e763bd
37 changed files with 3532 additions and 1639 deletions
151
internal/datasets/gefs/source.go
Normal file
151
internal/datasets/gefs/source.go
Normal file
|
|
@ -0,0 +1,151 @@
|
|||
// Package gefs implements datasets.Source for NOAA GEFS (Global Ensemble
|
||||
// Forecast System) forecasts.
|
||||
//
|
||||
// Each ensemble member is treated as its own dataset, selected via
|
||||
// DatasetID.Subset.Members. The download skeleton (HTTP, idx parsing,
|
||||
// parallel blit) lives in internal/datasets/grib; this package only
|
||||
// supplies GEFS-specific URL templating and member resolution.
|
||||
package gefs
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
|
||||
"predictor-refactored/internal/datasets"
|
||||
"predictor-refactored/internal/datasets/grib"
|
||||
"predictor-refactored/internal/weather"
|
||||
wgfs "predictor-refactored/internal/weather/gfs"
|
||||
)
|
||||
|
||||
// Source is the GEFS implementation of datasets.Source.
|
||||
type Source struct {
|
||||
Variant *wgfs.Variant
|
||||
Parallel int
|
||||
Client *http.Client
|
||||
Log *zap.Logger
|
||||
}
|
||||
|
||||
// NewSource returns a default Source over variant. If variant is nil,
|
||||
// GEFS 0.5° 3-hour is used.
|
||||
func NewSource(variant *wgfs.Variant, log *zap.Logger) *Source {
|
||||
if variant == nil {
|
||||
variant = wgfs.GEFS0p50_3h
|
||||
}
|
||||
return &Source{
|
||||
Variant: variant,
|
||||
Parallel: 8,
|
||||
Client: &http.Client{Timeout: 2 * time.Minute},
|
||||
Log: log,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Source) ID() string { return s.Variant.ID }
|
||||
|
||||
func (s *Source) downloader() *grib.Downloader {
|
||||
return &grib.Downloader{
|
||||
Variant: s.Variant,
|
||||
URLs: s.url,
|
||||
Parallel: s.Parallel,
|
||||
Client: s.Client,
|
||||
Log: s.Log,
|
||||
}
|
||||
}
|
||||
|
||||
// url generates the GEFS URL for (date, runHour, member, step, levelSet).
|
||||
func (s *Source) url(date string, runHour, member, step int, ls wgfs.LevelSet) string {
|
||||
if ls == wgfs.LevelSetB {
|
||||
return wgfs.GefsGribURLB(date, runHour, member, step, s.Variant.ResToken)
|
||||
}
|
||||
return wgfs.GefsGribURL(date, runHour, member, step, s.Variant.ResToken)
|
||||
}
|
||||
|
||||
// LatestEpoch HEAD-checks the control member's final forecast hour.
|
||||
func (s *Source) LatestEpoch(ctx context.Context) (time.Time, error) {
|
||||
now := time.Now().UTC()
|
||||
hour := now.Hour() - (now.Hour() % 6)
|
||||
current := time.Date(now.Year(), now.Month(), now.Day(), hour, 0, 0, 0, time.UTC)
|
||||
|
||||
client := s.Client
|
||||
if client == nil {
|
||||
client = &http.Client{Timeout: 2 * time.Minute}
|
||||
}
|
||||
log := s.Log
|
||||
if log == nil {
|
||||
log = zap.NewNop()
|
||||
}
|
||||
|
||||
for range 8 {
|
||||
date := current.Format("20060102")
|
||||
url := wgfs.GefsGribURL(date, current.Hour(), 0, s.Variant.MaxHour, s.Variant.ResToken) + ".idx"
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodHead, url, nil)
|
||||
if err == nil {
|
||||
resp, err := client.Do(req)
|
||||
if err == nil {
|
||||
resp.Body.Close()
|
||||
if resp.StatusCode == http.StatusOK {
|
||||
log.Info("latest GEFS run discovered",
|
||||
zap.Time("run", current),
|
||||
zap.String("verified_url", url))
|
||||
return current, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
current = current.Add(-6 * time.Hour)
|
||||
}
|
||||
return time.Time{}, fmt.Errorf("no recent GEFS run found")
|
||||
}
|
||||
|
||||
// Coverage returns the extent of id.
|
||||
func (s *Source) Coverage(id datasets.DatasetID) datasets.Coverage {
|
||||
v := s.Variant
|
||||
cov := datasets.Coverage{
|
||||
Region: datasets.Region{MinLat: -90, MaxLat: 90, MinLng: 0, MaxLng: 360},
|
||||
StartTime: id.Epoch,
|
||||
EndTime: id.Epoch.Add(time.Duration(v.MaxHour) * time.Hour),
|
||||
}
|
||||
if r := id.Subset.Region; r != nil {
|
||||
cov.Region = *r
|
||||
}
|
||||
if h := id.Subset.HourRange; h != nil {
|
||||
cov.StartTime = id.Epoch.Add(time.Duration(h.MinHour) * time.Hour)
|
||||
cov.EndTime = id.Epoch.Add(time.Duration(h.MaxHour) * time.Hour)
|
||||
}
|
||||
return cov
|
||||
}
|
||||
|
||||
// Open loads a stored GEFS dataset as a WindField.
|
||||
func (s *Source) Open(_ context.Context, id datasets.DatasetID, store datasets.Storage) (weather.WindField, error) {
|
||||
if !store.Exists(id) {
|
||||
return nil, fmt.Errorf("dataset %s not found", id.Filename())
|
||||
}
|
||||
file, err := wgfs.Open(store.Path(id), s.Variant, id.Epoch.UTC())
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return wgfs.NewWind(file), nil
|
||||
}
|
||||
|
||||
// memberOf extracts the single member index encoded by id.Subset.Members.
|
||||
func memberOf(id datasets.DatasetID) (int, error) {
|
||||
if len(id.Subset.Members) != 1 {
|
||||
return 0, fmt.Errorf("gefs dataset id must specify exactly one member (got %v)", id.Subset.Members)
|
||||
}
|
||||
m := id.Subset.Members[0]
|
||||
if m < 0 || m >= wgfs.GEFSMembers {
|
||||
return 0, fmt.Errorf("gefs member %d out of range [0, %d)", m, wgfs.GEFSMembers)
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
|
||||
// Download fetches one ensemble member's dataset.
|
||||
func (s *Source) Download(ctx context.Context, id datasets.DatasetID, store datasets.Storage, prog datasets.ProgressSink, throttle datasets.Throttle) error {
|
||||
member, err := memberOf(id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return s.downloader().Run(ctx, id, member, store, prog, throttle)
|
||||
}
|
||||
|
|
@ -1,85 +1,96 @@
|
|||
// Package gfs implements datasets.Source for NOAA GFS 0.5-degree forecasts.
|
||||
// Package gfs implements datasets.Source for NOAA GFS forecasts.
|
||||
//
|
||||
// The package serves multiple GFS variants (0.5° 3-hour, 0.25° 3-hour,
|
||||
// 0.25° 1-hour); the variant is selected at construction time. The
|
||||
// download skeleton (HTTP, idx parsing, parallel blit) lives in
|
||||
// internal/datasets/grib; this package only supplies URL templating and
|
||||
// the Source-interface plumbing.
|
||||
package gfs
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"math"
|
||||
"net/http"
|
||||
"os"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/nilsmagnus/grib/griblib"
|
||||
"go.uber.org/zap"
|
||||
"golang.org/x/sync/errgroup"
|
||||
|
||||
"predictor-refactored/internal/datasets"
|
||||
"predictor-refactored/internal/datasets/grib"
|
||||
"predictor-refactored/internal/weather"
|
||||
wgfs "predictor-refactored/internal/weather/gfs"
|
||||
)
|
||||
|
||||
// Source is the GFS implementation of datasets.Source.
|
||||
type Source struct {
|
||||
Parallel int // max concurrent step downloads
|
||||
Client *http.Client // optional; defaults to a 2-minute-timeout client
|
||||
Variant *wgfs.Variant
|
||||
Parallel int
|
||||
Client *http.Client
|
||||
Log *zap.Logger
|
||||
}
|
||||
|
||||
// NewSource returns a default Source.
|
||||
func NewSource(log *zap.Logger) *Source {
|
||||
// NewSource returns a default Source over variant. If variant is nil,
|
||||
// GFS 0.5° 3-hour is used (the historical Tawhiri default).
|
||||
func NewSource(variant *wgfs.Variant, log *zap.Logger) *Source {
|
||||
if variant == nil {
|
||||
variant = wgfs.GFS0p50_3h
|
||||
}
|
||||
return &Source{
|
||||
Variant: variant,
|
||||
Parallel: 8,
|
||||
Client: &http.Client{Timeout: 2 * time.Minute},
|
||||
Log: log,
|
||||
}
|
||||
}
|
||||
|
||||
// ID returns the source identifier.
|
||||
func (s *Source) ID() string { return "noaa-gfs-0p50" }
|
||||
// ID returns the variant's ID.
|
||||
func (s *Source) ID() string { return s.Variant.ID }
|
||||
|
||||
func (s *Source) log() *zap.Logger {
|
||||
if s.Log == nil {
|
||||
return zap.NewNop()
|
||||
func (s *Source) downloader() *grib.Downloader {
|
||||
return &grib.Downloader{
|
||||
Variant: s.Variant,
|
||||
URLs: s.url,
|
||||
Parallel: s.Parallel,
|
||||
Client: s.Client,
|
||||
Log: s.Log,
|
||||
}
|
||||
return s.Log
|
||||
}
|
||||
|
||||
func (s *Source) client() *http.Client {
|
||||
if s.Client == nil {
|
||||
return &http.Client{Timeout: 2 * time.Minute}
|
||||
// url generates the GFS URL for one (date, runHour, _, step, levelSet).
|
||||
// member is unused for GFS.
|
||||
func (s *Source) url(date string, runHour, _, step int, ls wgfs.LevelSet) string {
|
||||
if ls == wgfs.LevelSetB {
|
||||
return s.Variant.GribURLB(date, runHour, step)
|
||||
}
|
||||
return s.Client
|
||||
return s.Variant.GribURL(date, runHour, step)
|
||||
}
|
||||
|
||||
func (s *Source) parallel() int {
|
||||
if s.Parallel <= 0 {
|
||||
return 8
|
||||
}
|
||||
return s.Parallel
|
||||
}
|
||||
|
||||
// LatestEpoch returns the most recent run NOAA has finished publishing,
|
||||
// determined by HEAD-ing the .idx for the final forecast hour. Walks back
|
||||
// up to 8 runs (48 hours) before giving up.
|
||||
// LatestEpoch returns the most recent run NOAA has finished publishing.
|
||||
func (s *Source) LatestEpoch(ctx context.Context) (time.Time, error) {
|
||||
now := time.Now().UTC()
|
||||
hour := now.Hour() - (now.Hour() % 6)
|
||||
current := time.Date(now.Year(), now.Month(), now.Day(), hour, 0, 0, 0, time.UTC)
|
||||
|
||||
client := s.Client
|
||||
if client == nil {
|
||||
client = &http.Client{Timeout: 2 * time.Minute}
|
||||
}
|
||||
log := s.Log
|
||||
if log == nil {
|
||||
log = zap.NewNop()
|
||||
}
|
||||
|
||||
for range 8 {
|
||||
date := current.Format("20060102")
|
||||
url := wgfs.GribURL(date, current.Hour(), wgfs.MaxHour) + ".idx"
|
||||
|
||||
url := s.Variant.GribURL(date, current.Hour(), s.Variant.MaxHour) + ".idx"
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodHead, url, nil)
|
||||
if err == nil {
|
||||
resp, err := s.client().Do(req)
|
||||
resp, err := client.Do(req)
|
||||
if err == nil {
|
||||
resp.Body.Close()
|
||||
if resp.StatusCode == http.StatusOK {
|
||||
s.log().Info("latest GFS run discovered",
|
||||
log.Info("latest run discovered",
|
||||
zap.String("variant", s.Variant.ID),
|
||||
zap.Time("run", current),
|
||||
zap.String("verified_url", url))
|
||||
return current, nil
|
||||
|
|
@ -88,343 +99,40 @@ func (s *Source) LatestEpoch(ctx context.Context) (time.Time, error) {
|
|||
}
|
||||
current = current.Add(-6 * time.Hour)
|
||||
}
|
||||
return time.Time{}, fmt.Errorf("no recent GFS run found (checked 8 runs)")
|
||||
return time.Time{}, fmt.Errorf("no recent %s run found (checked 8 runs)", s.Variant.ID)
|
||||
}
|
||||
|
||||
// Coverage returns the geographic and temporal extent of id.
|
||||
func (s *Source) Coverage(id datasets.DatasetID) datasets.Coverage {
|
||||
v := s.Variant
|
||||
cov := datasets.Coverage{
|
||||
Region: datasets.Region{MinLat: -90, MaxLat: 90, MinLng: 0, MaxLng: 360},
|
||||
StartTime: id.Epoch,
|
||||
EndTime: id.Epoch.Add(time.Duration(v.MaxHour) * time.Hour),
|
||||
}
|
||||
if r := id.Subset.Region; r != nil {
|
||||
cov.Region = *r
|
||||
}
|
||||
if h := id.Subset.HourRange; h != nil {
|
||||
cov.StartTime = id.Epoch.Add(time.Duration(h.MinHour) * time.Hour)
|
||||
cov.EndTime = id.Epoch.Add(time.Duration(h.MaxHour) * time.Hour)
|
||||
}
|
||||
return cov
|
||||
}
|
||||
|
||||
// Open loads a stored dataset as a WindField.
|
||||
func (s *Source) Open(_ context.Context, epoch time.Time, store datasets.Storage) (weather.WindField, error) {
|
||||
if !store.Exists(epoch) {
|
||||
return nil, fmt.Errorf("epoch %s not found", epoch.Format(time.RFC3339))
|
||||
func (s *Source) Open(_ context.Context, id datasets.DatasetID, store datasets.Storage) (weather.WindField, error) {
|
||||
if !store.Exists(id) {
|
||||
return nil, fmt.Errorf("dataset %s not found", id.Filename())
|
||||
}
|
||||
file, err := wgfs.Open(store.Path(epoch), epoch.UTC())
|
||||
file, err := wgfs.Open(store.Path(id), s.Variant, id.Epoch.UTC())
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return wgfs.NewWind(file), nil
|
||||
}
|
||||
|
||||
// neededVariables is the GRIB variable set we extract.
|
||||
var neededVariables = map[string]bool{"HGT": true, "UGRD": true, "VGRD": true}
|
||||
|
||||
// Download fetches the full dataset for epoch in parallel, resuming any
|
||||
// previously-completed work units. Honours ctx cancellation and prog
|
||||
// (which may be nil).
|
||||
func (s *Source) Download(ctx context.Context, epoch time.Time, store datasets.Storage, prog datasets.ProgressSink, throttle datasets.Throttle) error {
|
||||
if prog == nil {
|
||||
prog = noopSink{}
|
||||
}
|
||||
|
||||
handle, err := store.BeginWrite(epoch)
|
||||
if err != nil {
|
||||
return fmt.Errorf("begin write: %w", err)
|
||||
}
|
||||
manifest := handle.Manifest()
|
||||
|
||||
// Open or create the temp file. If a previous attempt left a partial
|
||||
// file of the right size, reuse it (resume); otherwise Create.
|
||||
file, err := openOrCreateCube(handle.Path())
|
||||
if err != nil {
|
||||
_ = handle.Abort()
|
||||
return err
|
||||
}
|
||||
|
||||
date := epoch.UTC().Format("20060102")
|
||||
runHour := epoch.UTC().Hour()
|
||||
steps := wgfs.Hours()
|
||||
totalUnits := len(steps) * 2
|
||||
|
||||
prog.SetTotal(totalUnits)
|
||||
// Pre-count already-done units so progress is accurate on resume.
|
||||
for _, u := range manifest.Units() {
|
||||
_ = u
|
||||
prog.StepComplete()
|
||||
}
|
||||
|
||||
start := time.Now()
|
||||
g, ctx := errgroup.WithContext(ctx)
|
||||
g.SetLimit(s.parallel())
|
||||
|
||||
// fileMu serialises concurrent BlitGribData calls because the underlying
|
||||
// mmap is shared and SetVal isn't atomic.
|
||||
var fileMu sync.Mutex
|
||||
|
||||
for _, step := range steps {
|
||||
hourIdx := wgfs.HourIndex(step)
|
||||
if hourIdx < 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
for _, ls := range []wgfs.LevelSet{wgfs.LevelSetA, wgfs.LevelSetB} {
|
||||
unit := unitKey(step, ls)
|
||||
if manifest.Has(unit) {
|
||||
continue
|
||||
}
|
||||
|
||||
g.Go(func() error {
|
||||
var url string
|
||||
switch ls {
|
||||
case wgfs.LevelSetA:
|
||||
url = wgfs.GribURL(date, runHour, step)
|
||||
case wgfs.LevelSetB:
|
||||
url = wgfs.GribURLB(date, runHour, step)
|
||||
}
|
||||
if err := s.downloadAndBlit(ctx, file, &fileMu, url, hourIdx, ls, prog, throttle); err != nil {
|
||||
return fmt.Errorf("step %d %s: %w", step, levelSetLabel(ls), err)
|
||||
}
|
||||
if err := manifest.Mark(unit); err != nil {
|
||||
return fmt.Errorf("mark unit: %w", err)
|
||||
}
|
||||
prog.StepComplete()
|
||||
return nil
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
if err := g.Wait(); err != nil {
|
||||
_ = file.Close()
|
||||
// Don't Abort on context cancellation — preserve progress for resume.
|
||||
if errors.Is(err, context.Canceled) {
|
||||
return err
|
||||
}
|
||||
// Other errors: abort if no progress was made; otherwise leave for resume.
|
||||
if len(manifest.Units()) == 0 {
|
||||
_ = handle.Abort()
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
if err := file.Flush(); err != nil {
|
||||
_ = file.Close()
|
||||
return fmt.Errorf("flush: %w", err)
|
||||
}
|
||||
if err := file.Close(); err != nil {
|
||||
return fmt.Errorf("close: %w", err)
|
||||
}
|
||||
if err := handle.Commit(); err != nil {
|
||||
return fmt.Errorf("commit: %w", err)
|
||||
}
|
||||
|
||||
s.log().Info("download complete",
|
||||
zap.Time("epoch", epoch),
|
||||
zap.Duration("elapsed", time.Since(start)))
|
||||
return nil
|
||||
// Download fetches the dataset for id. GFS ignores Subset.Members.
|
||||
func (s *Source) Download(ctx context.Context, id datasets.DatasetID, store datasets.Storage, prog datasets.ProgressSink, throttle datasets.Throttle) error {
|
||||
return s.downloader().Run(ctx, id, 0, store, prog, throttle)
|
||||
}
|
||||
|
||||
// openOrCreateCube returns a writable cube file at path, creating it if the
|
||||
// file does not exist or has the wrong size.
|
||||
func openOrCreateCube(path string) (*wgfs.File, error) {
|
||||
info, err := os.Stat(path)
|
||||
if err == nil && info.Size() == wgfs.DatasetSize {
|
||||
return wgfs.OpenWritable(path)
|
||||
}
|
||||
if err != nil && !errors.Is(err, os.ErrNotExist) {
|
||||
return nil, fmt.Errorf("stat cube: %w", err)
|
||||
}
|
||||
// Wrong-size or missing — truncate-create.
|
||||
return wgfs.Create(path)
|
||||
}
|
||||
|
||||
// downloadAndBlit fetches and decodes one (URL, level-set) chunk and writes
|
||||
// it into the dataset.
|
||||
func (s *Source) downloadAndBlit(
|
||||
ctx context.Context,
|
||||
file *wgfs.File,
|
||||
fileMu *sync.Mutex,
|
||||
baseURL string,
|
||||
hourIdx int,
|
||||
ls wgfs.LevelSet,
|
||||
prog datasets.ProgressSink,
|
||||
throttle datasets.Throttle,
|
||||
) error {
|
||||
idxBody, err := s.httpGet(ctx, baseURL+".idx", throttle, prog)
|
||||
if err != nil {
|
||||
return fmt.Errorf("idx: %w", err)
|
||||
}
|
||||
entries := ParseIdx(idxBody)
|
||||
filtered := FilterIdx(entries, neededVariables)
|
||||
|
||||
var relevant []IdxEntry
|
||||
for _, e := range filtered {
|
||||
set, ok := wgfs.PressureLevelSet(e.LevelMB)
|
||||
if ok && set == ls {
|
||||
relevant = append(relevant, e)
|
||||
}
|
||||
}
|
||||
if len(relevant) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
ranges := EntriesToRanges(relevant)
|
||||
tmp, err := os.CreateTemp("", "gfs-msg-*.tmp")
|
||||
if err != nil {
|
||||
return fmt.Errorf("temp: %w", err)
|
||||
}
|
||||
tmpPath := tmp.Name()
|
||||
defer os.Remove(tmpPath)
|
||||
|
||||
for _, r := range ranges {
|
||||
body, err := s.httpGetRange(ctx, baseURL, r.Start, r.End, throttle, prog)
|
||||
if err != nil {
|
||||
tmp.Close()
|
||||
return fmt.Errorf("range %d-%d: %w", r.Start, r.End, err)
|
||||
}
|
||||
if _, err := tmp.Write(body); err != nil {
|
||||
tmp.Close()
|
||||
return fmt.Errorf("write tmp: %w", err)
|
||||
}
|
||||
}
|
||||
if err := tmp.Close(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
f, err := os.Open(tmpPath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
messages, err := griblib.ReadMessages(f)
|
||||
f.Close()
|
||||
if err != nil {
|
||||
return fmt.Errorf("read grib: %w", err)
|
||||
}
|
||||
|
||||
for _, msg := range messages {
|
||||
if msg.Section4.ProductDefinitionTemplateNumber != 0 {
|
||||
continue
|
||||
}
|
||||
p := msg.Section4.ProductDefinitionTemplate
|
||||
varIdx := wgfs.VariableIndex(int(p.ParameterCategory), int(p.ParameterNumber))
|
||||
if varIdx < 0 {
|
||||
continue
|
||||
}
|
||||
if p.FirstSurface.Type != 100 { // isobaric only
|
||||
continue
|
||||
}
|
||||
pressureMB := int(math.Round(float64(p.FirstSurface.Value) / 100.0))
|
||||
levelIdx := wgfs.PressureIndex(pressureMB)
|
||||
if levelIdx < 0 {
|
||||
continue
|
||||
}
|
||||
data := msg.Data()
|
||||
fileMu.Lock()
|
||||
err := file.BlitGribData(hourIdx, levelIdx, varIdx, data)
|
||||
fileMu.Unlock()
|
||||
if err != nil {
|
||||
return fmt.Errorf("blit: %w", err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// httpGet downloads a URL body with 3 retries and optional throttling.
|
||||
func (s *Source) httpGet(ctx context.Context, url string, throttle datasets.Throttle, prog datasets.ProgressSink) ([]byte, error) {
|
||||
var lastErr error
|
||||
for attempt := range 3 {
|
||||
if attempt > 0 {
|
||||
select {
|
||||
case <-time.After(time.Duration(attempt*2) * time.Second):
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
resp, err := s.client().Do(req)
|
||||
if err != nil {
|
||||
lastErr = err
|
||||
continue
|
||||
}
|
||||
body, err := readThrottled(ctx, resp.Body, throttle, prog)
|
||||
resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
lastErr = fmt.Errorf("HTTP %d for %s", resp.StatusCode, url)
|
||||
continue
|
||||
}
|
||||
if err != nil {
|
||||
lastErr = err
|
||||
continue
|
||||
}
|
||||
return body, nil
|
||||
}
|
||||
return nil, fmt.Errorf("after 3 attempts: %w", lastErr)
|
||||
}
|
||||
|
||||
// httpGetRange downloads an inclusive byte range with 3 retries and throttling.
|
||||
func (s *Source) httpGetRange(ctx context.Context, url string, start, end int64, throttle datasets.Throttle, prog datasets.ProgressSink) ([]byte, error) {
|
||||
var lastErr error
|
||||
for attempt := range 3 {
|
||||
if attempt > 0 {
|
||||
select {
|
||||
case <-time.After(time.Duration(attempt*2) * time.Second):
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", start, end))
|
||||
resp, err := s.client().Do(req)
|
||||
if err != nil {
|
||||
lastErr = err
|
||||
continue
|
||||
}
|
||||
body, err := readThrottled(ctx, resp.Body, throttle, prog)
|
||||
resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusPartialContent && resp.StatusCode != http.StatusOK {
|
||||
lastErr = fmt.Errorf("HTTP %d for range %d-%d of %s", resp.StatusCode, start, end, url)
|
||||
continue
|
||||
}
|
||||
if err != nil {
|
||||
lastErr = err
|
||||
continue
|
||||
}
|
||||
return body, nil
|
||||
}
|
||||
return nil, fmt.Errorf("after 3 attempts: %w", lastErr)
|
||||
}
|
||||
|
||||
// readThrottled reads r into memory, consulting throttle (if non-nil) before
|
||||
// each chunk and reporting bytes to prog.
|
||||
func readThrottled(ctx context.Context, r io.Reader, throttle datasets.Throttle, prog datasets.ProgressSink) ([]byte, error) {
|
||||
buf := make([]byte, 0, 64*1024)
|
||||
chunk := make([]byte, 32*1024)
|
||||
for {
|
||||
if throttle != nil {
|
||||
if err := throttle.Wait(ctx, len(chunk)); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
n, err := r.Read(chunk)
|
||||
if n > 0 {
|
||||
buf = append(buf, chunk[:n]...)
|
||||
prog.Bytes(int64(n))
|
||||
}
|
||||
if errors.Is(err, io.EOF) {
|
||||
return buf, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func unitKey(step int, ls wgfs.LevelSet) string {
|
||||
return fmt.Sprintf("step%03d-%s", step, levelSetLabel(ls))
|
||||
}
|
||||
|
||||
func levelSetLabel(ls wgfs.LevelSet) string {
|
||||
if ls == wgfs.LevelSetB {
|
||||
return "B"
|
||||
}
|
||||
return "A"
|
||||
}
|
||||
|
||||
// noopSink discards progress events.
|
||||
type noopSink struct{}
|
||||
|
||||
func (noopSink) SetTotal(int) {}
|
||||
func (noopSink) StepComplete() {}
|
||||
func (noopSink) Bytes(int64) {}
|
||||
|
|
|
|||
369
internal/datasets/grib/downloader.go
Normal file
369
internal/datasets/grib/downloader.go
Normal file
|
|
@ -0,0 +1,369 @@
|
|||
package grib
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"math"
|
||||
"net/http"
|
||||
"os"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/nilsmagnus/grib/griblib"
|
||||
"go.uber.org/zap"
|
||||
"golang.org/x/sync/errgroup"
|
||||
|
||||
"predictor-refactored/internal/datasets"
|
||||
wgfs "predictor-refactored/internal/weather/gfs"
|
||||
)
|
||||
|
||||
// URLFunc returns the GRIB URL for one (date, runHour, member, step, levelSet).
|
||||
// Sources that don't have members (GFS) ignore the member argument.
|
||||
type URLFunc func(date string, runHour, member, step int, ls wgfs.LevelSet) string
|
||||
|
||||
// Downloader is the generic GRIB-cube downloader.
|
||||
//
|
||||
// A Source plugs in its variant, URL templating, and member-resolution
|
||||
// logic; the Downloader runs the parallel idx fetch, byte-range download,
|
||||
// GRIB decode, and blit loop with manifest-based resume.
|
||||
type Downloader struct {
|
||||
Variant *wgfs.Variant
|
||||
URLs URLFunc
|
||||
Parallel int
|
||||
Client *http.Client
|
||||
Log *zap.Logger
|
||||
}
|
||||
|
||||
func (d *Downloader) log() *zap.Logger {
|
||||
if d.Log == nil {
|
||||
return zap.NewNop()
|
||||
}
|
||||
return d.Log
|
||||
}
|
||||
|
||||
func (d *Downloader) client() *http.Client {
|
||||
if d.Client == nil {
|
||||
return &http.Client{Timeout: 2 * time.Minute}
|
||||
}
|
||||
return d.Client
|
||||
}
|
||||
|
||||
func (d *Downloader) parallel() int {
|
||||
if d.Parallel <= 0 {
|
||||
return 8
|
||||
}
|
||||
return d.Parallel
|
||||
}
|
||||
|
||||
// neededVariables is the GRIB variable set every source extracts.
|
||||
var neededVariables = map[string]bool{"HGT": true, "UGRD": true, "VGRD": true}
|
||||
|
||||
// Run downloads the dataset for id, member into store. The caller may
|
||||
// pass member=0 for non-ensemble sources.
|
||||
func (d *Downloader) Run(ctx context.Context, id datasets.DatasetID, member int, store datasets.Storage, prog datasets.ProgressSink, throttle datasets.Throttle) error {
|
||||
if prog == nil {
|
||||
prog = noopSink{}
|
||||
}
|
||||
|
||||
handle, err := store.BeginWrite(id)
|
||||
if err != nil {
|
||||
return fmt.Errorf("begin write: %w", err)
|
||||
}
|
||||
manifest := handle.Manifest()
|
||||
|
||||
file, err := openOrCreateCube(handle.Path(), d.Variant)
|
||||
if err != nil {
|
||||
_ = handle.Abort()
|
||||
return err
|
||||
}
|
||||
|
||||
epoch := id.Epoch.UTC()
|
||||
date := epoch.Format("20060102")
|
||||
runHour := epoch.Hour()
|
||||
|
||||
steps := d.Variant.Hours()
|
||||
if hr := id.Subset.HourRange; hr != nil {
|
||||
filtered := steps[:0]
|
||||
for _, step := range steps {
|
||||
if step >= hr.MinHour && step <= hr.MaxHour {
|
||||
filtered = append(filtered, step)
|
||||
}
|
||||
}
|
||||
steps = filtered
|
||||
}
|
||||
prog.SetTotal(len(steps) * 2)
|
||||
for range manifest.Units() {
|
||||
prog.StepComplete()
|
||||
}
|
||||
|
||||
start := time.Now()
|
||||
g, ctx := errgroup.WithContext(ctx)
|
||||
g.SetLimit(d.parallel())
|
||||
var fileMu sync.Mutex
|
||||
|
||||
for _, step := range steps {
|
||||
hourIdx := d.Variant.HourIndex(step)
|
||||
if hourIdx < 0 {
|
||||
continue
|
||||
}
|
||||
for _, ls := range []wgfs.LevelSet{wgfs.LevelSetA, wgfs.LevelSetB} {
|
||||
unit := unitKey(step, ls)
|
||||
if manifest.Has(unit) {
|
||||
continue
|
||||
}
|
||||
g.Go(func() error {
|
||||
url := d.URLs(date, runHour, member, step, ls)
|
||||
if err := d.downloadAndBlit(ctx, file, &fileMu, url, hourIdx, ls, prog, throttle); err != nil {
|
||||
return fmt.Errorf("step %d %s: %w", step, levelSetLabel(ls), err)
|
||||
}
|
||||
if err := manifest.Mark(unit); err != nil {
|
||||
return fmt.Errorf("mark unit: %w", err)
|
||||
}
|
||||
prog.StepComplete()
|
||||
return nil
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
if err := g.Wait(); err != nil {
|
||||
_ = file.Close()
|
||||
if errors.Is(err, context.Canceled) {
|
||||
return err
|
||||
}
|
||||
if len(manifest.Units()) == 0 {
|
||||
_ = handle.Abort()
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
if err := file.Flush(); err != nil {
|
||||
_ = file.Close()
|
||||
return fmt.Errorf("flush: %w", err)
|
||||
}
|
||||
if err := file.Close(); err != nil {
|
||||
return fmt.Errorf("close: %w", err)
|
||||
}
|
||||
if err := handle.Commit(); err != nil {
|
||||
return fmt.Errorf("commit: %w", err)
|
||||
}
|
||||
|
||||
d.log().Info("download complete",
|
||||
zap.String("variant", d.Variant.ID),
|
||||
zap.Time("epoch", epoch),
|
||||
zap.Duration("elapsed", time.Since(start)))
|
||||
return nil
|
||||
}
|
||||
|
||||
// openOrCreateCube opens an existing cube at path if it matches variant's
|
||||
// expected size, else truncate-creates a new one.
|
||||
func openOrCreateCube(path string, variant *wgfs.Variant) (*wgfs.File, error) {
|
||||
info, err := os.Stat(path)
|
||||
if err == nil && info.Size() == variant.DatasetSize() {
|
||||
return wgfs.OpenWritable(path, variant)
|
||||
}
|
||||
if err != nil && !errors.Is(err, os.ErrNotExist) {
|
||||
return nil, fmt.Errorf("stat cube: %w", err)
|
||||
}
|
||||
return wgfs.Create(path, variant)
|
||||
}
|
||||
|
||||
// downloadAndBlit fetches and decodes one (URL, level-set) chunk.
|
||||
func (d *Downloader) downloadAndBlit(
|
||||
ctx context.Context,
|
||||
file *wgfs.File,
|
||||
fileMu *sync.Mutex,
|
||||
baseURL string,
|
||||
hourIdx int,
|
||||
ls wgfs.LevelSet,
|
||||
prog datasets.ProgressSink,
|
||||
throttle datasets.Throttle,
|
||||
) error {
|
||||
idxBody, err := d.httpGet(ctx, baseURL+".idx", throttle, prog)
|
||||
if err != nil {
|
||||
return fmt.Errorf("idx: %w", err)
|
||||
}
|
||||
entries := ParseIdx(idxBody)
|
||||
filtered := FilterIdx(entries, neededVariables)
|
||||
|
||||
var relevant []IdxEntry
|
||||
for _, e := range filtered {
|
||||
set, ok := d.Variant.PressureLevelSet(e.LevelMB)
|
||||
if ok && set == ls {
|
||||
relevant = append(relevant, e)
|
||||
}
|
||||
}
|
||||
if len(relevant) == 0 {
|
||||
return nil
|
||||
}
|
||||
ranges := EntriesToRanges(relevant)
|
||||
|
||||
tmp, err := os.CreateTemp("", "grib-msg-*.tmp")
|
||||
if err != nil {
|
||||
return fmt.Errorf("temp: %w", err)
|
||||
}
|
||||
tmpPath := tmp.Name()
|
||||
defer os.Remove(tmpPath)
|
||||
|
||||
for _, r := range ranges {
|
||||
body, err := d.httpGetRange(ctx, baseURL, r.Start, r.End, throttle, prog)
|
||||
if err != nil {
|
||||
tmp.Close()
|
||||
return fmt.Errorf("range: %w", err)
|
||||
}
|
||||
if _, err := tmp.Write(body); err != nil {
|
||||
tmp.Close()
|
||||
return fmt.Errorf("write tmp: %w", err)
|
||||
}
|
||||
}
|
||||
if err := tmp.Close(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
f, err := os.Open(tmpPath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
messages, err := griblib.ReadMessages(f)
|
||||
f.Close()
|
||||
if err != nil {
|
||||
return fmt.Errorf("read grib: %w", err)
|
||||
}
|
||||
|
||||
for _, msg := range messages {
|
||||
if msg.Section4.ProductDefinitionTemplateNumber != 0 {
|
||||
continue
|
||||
}
|
||||
p := msg.Section4.ProductDefinitionTemplate
|
||||
varIdx := d.Variant.VariableIndex(int(p.ParameterCategory), int(p.ParameterNumber))
|
||||
if varIdx < 0 {
|
||||
continue
|
||||
}
|
||||
if p.FirstSurface.Type != 100 {
|
||||
continue
|
||||
}
|
||||
pressureMB := int(math.Round(float64(p.FirstSurface.Value) / 100.0))
|
||||
levelIdx := d.Variant.PressureIndex(pressureMB)
|
||||
if levelIdx < 0 {
|
||||
continue
|
||||
}
|
||||
data := msg.Data()
|
||||
fileMu.Lock()
|
||||
err := file.BlitGribData(hourIdx, levelIdx, varIdx, data)
|
||||
fileMu.Unlock()
|
||||
if err != nil {
|
||||
return fmt.Errorf("blit: %w", err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d *Downloader) httpGet(ctx context.Context, url string, throttle datasets.Throttle, prog datasets.ProgressSink) ([]byte, error) {
|
||||
var lastErr error
|
||||
for attempt := range 3 {
|
||||
if attempt > 0 {
|
||||
select {
|
||||
case <-time.After(time.Duration(attempt*2) * time.Second):
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
resp, err := d.client().Do(req)
|
||||
if err != nil {
|
||||
lastErr = err
|
||||
continue
|
||||
}
|
||||
body, err := readThrottled(ctx, resp.Body, throttle, prog)
|
||||
resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
lastErr = fmt.Errorf("HTTP %d for %s", resp.StatusCode, url)
|
||||
continue
|
||||
}
|
||||
if err != nil {
|
||||
lastErr = err
|
||||
continue
|
||||
}
|
||||
return body, nil
|
||||
}
|
||||
return nil, fmt.Errorf("after 3 attempts: %w", lastErr)
|
||||
}
|
||||
|
||||
func (d *Downloader) httpGetRange(ctx context.Context, url string, start, end int64, throttle datasets.Throttle, prog datasets.ProgressSink) ([]byte, error) {
|
||||
var lastErr error
|
||||
for attempt := range 3 {
|
||||
if attempt > 0 {
|
||||
select {
|
||||
case <-time.After(time.Duration(attempt*2) * time.Second):
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", start, end))
|
||||
resp, err := d.client().Do(req)
|
||||
if err != nil {
|
||||
lastErr = err
|
||||
continue
|
||||
}
|
||||
body, err := readThrottled(ctx, resp.Body, throttle, prog)
|
||||
resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusPartialContent && resp.StatusCode != http.StatusOK {
|
||||
lastErr = fmt.Errorf("HTTP %d for range", resp.StatusCode)
|
||||
continue
|
||||
}
|
||||
if err != nil {
|
||||
lastErr = err
|
||||
continue
|
||||
}
|
||||
return body, nil
|
||||
}
|
||||
return nil, fmt.Errorf("after 3 attempts: %w", lastErr)
|
||||
}
|
||||
|
||||
func readThrottled(ctx context.Context, r io.Reader, throttle datasets.Throttle, prog datasets.ProgressSink) ([]byte, error) {
|
||||
buf := make([]byte, 0, 64*1024)
|
||||
chunk := make([]byte, 32*1024)
|
||||
for {
|
||||
if throttle != nil {
|
||||
if err := throttle.Wait(ctx, len(chunk)); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
n, err := r.Read(chunk)
|
||||
if n > 0 {
|
||||
buf = append(buf, chunk[:n]...)
|
||||
prog.Bytes(int64(n))
|
||||
}
|
||||
if errors.Is(err, io.EOF) {
|
||||
return buf, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func unitKey(step int, ls wgfs.LevelSet) string {
|
||||
return fmt.Sprintf("step%03d-%s", step, levelSetLabel(ls))
|
||||
}
|
||||
|
||||
func levelSetLabel(ls wgfs.LevelSet) string {
|
||||
if ls == wgfs.LevelSetB {
|
||||
return "B"
|
||||
}
|
||||
return "A"
|
||||
}
|
||||
|
||||
type noopSink struct{}
|
||||
|
||||
func (noopSink) SetTotal(int) {}
|
||||
func (noopSink) StepComplete() {}
|
||||
func (noopSink) Bytes(int64) {}
|
||||
|
|
@ -1,4 +1,8 @@
|
|||
package gfs
|
||||
// Package grib contains the GRIB-cube download skeleton shared by every
|
||||
// NOAA source (GFS, GEFS, future families). It exposes the .idx parser,
|
||||
// HTTP helpers, and a parallel download loop; source-specific URL
|
||||
// templating is injected by the caller.
|
||||
package grib
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
|
@ -1,4 +1,4 @@
|
|||
package gfs
|
||||
package grib
|
||||
|
||||
import "testing"
|
||||
|
||||
|
|
@ -27,23 +27,22 @@ const (
|
|||
|
||||
// JobInfo is the externally-visible snapshot of a download job.
|
||||
type JobInfo struct {
|
||||
ID string
|
||||
Source string
|
||||
Epoch time.Time
|
||||
Status JobStatus
|
||||
StartedAt time.Time
|
||||
EndedAt *time.Time
|
||||
Err string
|
||||
Total int
|
||||
Done int
|
||||
Bytes int64
|
||||
ID string
|
||||
Source string
|
||||
Dataset DatasetID
|
||||
Status JobStatus
|
||||
StartedAt time.Time
|
||||
EndedAt *time.Time
|
||||
Err string
|
||||
Total int
|
||||
Done int
|
||||
Bytes int64
|
||||
}
|
||||
|
||||
// jobEntry is the Manager's mutable record for one job.
|
||||
type jobEntry struct {
|
||||
id string
|
||||
source string
|
||||
epoch time.Time
|
||||
dataset DatasetID
|
||||
startedAt time.Time
|
||||
cancel context.CancelFunc
|
||||
|
||||
|
|
@ -60,7 +59,7 @@ type jobEntry struct {
|
|||
func (e *jobEntry) snapshot() JobInfo {
|
||||
e.mu.Lock()
|
||||
info := JobInfo{
|
||||
ID: e.id, Source: e.source, Epoch: e.epoch,
|
||||
ID: e.id, Source: e.source, Dataset: e.dataset,
|
||||
StartedAt: e.startedAt, Status: e.status, Err: e.errStr,
|
||||
}
|
||||
if !e.endedAt.IsZero() {
|
||||
|
|
@ -74,14 +73,20 @@ func (e *jobEntry) snapshot() JobInfo {
|
|||
return info
|
||||
}
|
||||
|
||||
// jobProgress is the ProgressSink wired into a jobEntry.
|
||||
type jobProgress struct{ e *jobEntry }
|
||||
|
||||
func (p jobProgress) SetTotal(n int) { p.e.total.Store(int64(n)) }
|
||||
func (p jobProgress) StepComplete() { p.e.done.Add(1) }
|
||||
func (p jobProgress) Bytes(n int64) { p.e.bytes.Add(n) }
|
||||
|
||||
// Manager coordinates dataset downloads and exposes the active WindField.
|
||||
// loadedDataset bundles a loaded WindField with its identity and coverage.
|
||||
type loadedDataset struct {
|
||||
ID DatasetID
|
||||
Field weather.WindField
|
||||
Coverage Coverage
|
||||
}
|
||||
|
||||
// Manager coordinates dataset downloads and exposes the active WindFields.
|
||||
type Manager struct {
|
||||
src Source
|
||||
store Storage
|
||||
|
|
@ -89,18 +94,15 @@ type Manager struct {
|
|||
log *zap.Logger
|
||||
|
||||
activeMu sync.RWMutex
|
||||
active weather.WindField
|
||||
active []loadedDataset
|
||||
|
||||
jobsMu sync.RWMutex
|
||||
jobs map[string]*jobEntry
|
||||
|
||||
// inFlight maps an epoch's RFC3339 representation to its jobID, enforcing
|
||||
// single-flight per epoch.
|
||||
inFlight sync.Map
|
||||
inFlight sync.Map // key: dataset filename, value: jobID
|
||||
}
|
||||
|
||||
// New returns a Manager wiring source, store, and an optional throttle.
|
||||
// A nil log uses zap.NewNop().
|
||||
// New wires a Manager.
|
||||
func New(src Source, store Storage, throttle Throttle, log *zap.Logger) *Manager {
|
||||
if log == nil {
|
||||
log = zap.NewNop()
|
||||
|
|
@ -119,18 +121,65 @@ func New(src Source, store Storage, throttle Throttle, log *zap.Logger) *Manager
|
|||
// Source returns the underlying source ID.
|
||||
func (m *Manager) Source() string { return m.src.ID() }
|
||||
|
||||
// Active returns the currently-loaded WindField, or nil.
|
||||
// Active returns the currently-loaded global WindField (the dataset with
|
||||
// IsGlobal subset, most recently loaded). Returns nil if no global
|
||||
// dataset is loaded; in cluster setups with only regional subsets, callers
|
||||
// should use SelectFor.
|
||||
func (m *Manager) Active() weather.WindField {
|
||||
m.activeMu.RLock()
|
||||
defer m.activeMu.RUnlock()
|
||||
return m.active
|
||||
for _, d := range m.active {
|
||||
if d.ID.Subset.IsGlobal() {
|
||||
return d.Field
|
||||
}
|
||||
}
|
||||
if len(m.active) > 0 {
|
||||
return m.active[0].Field
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Ready reports whether a dataset is currently loaded.
|
||||
// Ready reports whether at least one dataset is loaded.
|
||||
func (m *Manager) Ready() bool { return m.Active() != nil }
|
||||
|
||||
// ListEpochs returns all stored dataset epochs, newest first.
|
||||
func (m *Manager) ListEpochs() ([]time.Time, error) { return m.store.List() }
|
||||
// SelectFor returns a loaded WindField whose coverage contains (t, lat, lng).
|
||||
// Returns nil when no loaded dataset covers the query.
|
||||
func (m *Manager) SelectFor(t time.Time, lat, lng float64) weather.WindField {
|
||||
m.activeMu.RLock()
|
||||
defer m.activeMu.RUnlock()
|
||||
for _, d := range m.active {
|
||||
if d.Coverage.Covers(t, lat, lng) {
|
||||
return d.Field
|
||||
}
|
||||
}
|
||||
// Fallback: any global dataset is permissive about region.
|
||||
for _, d := range m.active {
|
||||
if d.ID.Subset.IsGlobal() {
|
||||
return d.Field
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// LoadedDatasets returns snapshots of every currently-loaded dataset.
|
||||
func (m *Manager) LoadedDatasets() []LoadedDatasetInfo {
|
||||
m.activeMu.RLock()
|
||||
defer m.activeMu.RUnlock()
|
||||
out := make([]LoadedDatasetInfo, 0, len(m.active))
|
||||
for _, d := range m.active {
|
||||
out = append(out, LoadedDatasetInfo{ID: d.ID, Coverage: d.Coverage})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// LoadedDatasetInfo is a serializable snapshot of one active dataset.
|
||||
type LoadedDatasetInfo struct {
|
||||
ID DatasetID
|
||||
Coverage Coverage
|
||||
}
|
||||
|
||||
// ListEpochs returns all stored datasets, newest first.
|
||||
func (m *Manager) ListEpochs() ([]DatasetID, error) { return m.store.List() }
|
||||
|
||||
// ListJobs returns snapshots of every job recorded since startup.
|
||||
func (m *Manager) ListJobs() []JobInfo {
|
||||
|
|
@ -143,7 +192,7 @@ func (m *Manager) ListJobs() []JobInfo {
|
|||
return out
|
||||
}
|
||||
|
||||
// GetJob returns the snapshot for a job, or false if id is unknown.
|
||||
// GetJob returns the snapshot for a job.
|
||||
func (m *Manager) GetJob(id string) (JobInfo, bool) {
|
||||
m.jobsMu.RLock()
|
||||
e, ok := m.jobs[id]
|
||||
|
|
@ -154,8 +203,7 @@ func (m *Manager) GetJob(id string) (JobInfo, bool) {
|
|||
return e.snapshot(), true
|
||||
}
|
||||
|
||||
// CancelJob cancels a running job. Returns false if id is unknown or the
|
||||
// job is already terminal.
|
||||
// CancelJob cancels a running job.
|
||||
func (m *Manager) CancelJob(id string) bool {
|
||||
m.jobsMu.RLock()
|
||||
e, ok := m.jobs[id]
|
||||
|
|
@ -173,28 +221,31 @@ func (m *Manager) CancelJob(id string) bool {
|
|||
return true
|
||||
}
|
||||
|
||||
// RemoveEpoch deletes a stored dataset. If epoch is currently active, the
|
||||
// active field is cleared.
|
||||
func (m *Manager) RemoveEpoch(epoch time.Time) error {
|
||||
epoch = epoch.UTC()
|
||||
if active := m.Active(); active != nil && active.Epoch().Equal(epoch) {
|
||||
m.activeMu.Lock()
|
||||
m.active = nil
|
||||
m.activeMu.Unlock()
|
||||
// Remove deletes a stored dataset. If the dataset is currently loaded,
|
||||
// it is unloaded first.
|
||||
func (m *Manager) Remove(id DatasetID) error {
|
||||
m.activeMu.Lock()
|
||||
out := m.active[:0]
|
||||
var removed *loadedDataset
|
||||
for i := range m.active {
|
||||
d := m.active[i]
|
||||
if d.ID.Equals(id) {
|
||||
removed = &d
|
||||
continue
|
||||
}
|
||||
out = append(out, d)
|
||||
}
|
||||
return m.store.Remove(epoch)
|
||||
m.active = out
|
||||
m.activeMu.Unlock()
|
||||
if removed != nil {
|
||||
closeField(removed.Field, m.log)
|
||||
}
|
||||
return m.store.Remove(id)
|
||||
}
|
||||
|
||||
// Download starts (or resumes) a download job for epoch in the background.
|
||||
// Returns the JobID. If a job for the same epoch is already running, its
|
||||
// existing JobID is returned.
|
||||
//
|
||||
// If the dataset is already present on disk, a synthetic completed JobInfo
|
||||
// is recorded and its JobID returned.
|
||||
func (m *Manager) Download(epoch time.Time) string {
|
||||
epoch = epoch.UTC()
|
||||
key := epoch.Format(time.RFC3339)
|
||||
|
||||
// Download starts (or resumes) a download job for id in the background.
|
||||
func (m *Manager) Download(id DatasetID) string {
|
||||
key := id.Filename()
|
||||
if existing, ok := m.inFlight.Load(key); ok {
|
||||
return existing.(string)
|
||||
}
|
||||
|
|
@ -209,7 +260,7 @@ func (m *Manager) Download(epoch time.Time) string {
|
|||
e := &jobEntry{
|
||||
id: jobID,
|
||||
source: m.src.ID(),
|
||||
epoch: epoch,
|
||||
dataset: id,
|
||||
startedAt: now,
|
||||
status: JobPending,
|
||||
cancel: cancel,
|
||||
|
|
@ -218,8 +269,7 @@ func (m *Manager) Download(epoch time.Time) string {
|
|||
m.jobs[jobID] = e
|
||||
m.jobsMu.Unlock()
|
||||
|
||||
if m.store.Exists(epoch) {
|
||||
// Skip the download but still record the job for traceability.
|
||||
if m.store.Exists(id) {
|
||||
go m.completeShortCircuit(ctx, e)
|
||||
return jobID
|
||||
}
|
||||
|
|
@ -227,46 +277,54 @@ func (m *Manager) Download(epoch time.Time) string {
|
|||
return jobID
|
||||
}
|
||||
|
||||
// LoadEpoch swaps the active WindField to epoch's stored dataset.
|
||||
func (m *Manager) LoadEpoch(ctx context.Context, epoch time.Time) error {
|
||||
epoch = epoch.UTC()
|
||||
if !m.store.Exists(epoch) {
|
||||
return fmt.Errorf("epoch %s not present on disk", epoch.Format(time.RFC3339))
|
||||
// Load swaps in id's stored dataset, making it available to predictions.
|
||||
func (m *Manager) Load(ctx context.Context, id DatasetID) error {
|
||||
if !m.store.Exists(id) {
|
||||
return fmt.Errorf("dataset %s not present on disk", id.Filename())
|
||||
}
|
||||
field, err := m.src.Open(ctx, epoch, m.store)
|
||||
field, err := m.src.Open(ctx, id, m.store)
|
||||
if err != nil {
|
||||
return fmt.Errorf("open epoch: %w", err)
|
||||
return fmt.Errorf("open dataset: %w", err)
|
||||
}
|
||||
m.swapActive(field)
|
||||
cov := m.src.Coverage(id)
|
||||
m.activeMu.Lock()
|
||||
// Replace any previously-loaded dataset with the same ID.
|
||||
for i := range m.active {
|
||||
if m.active[i].ID.Equals(id) {
|
||||
closeField(m.active[i].Field, m.log)
|
||||
m.active[i] = loadedDataset{ID: id, Field: field, Coverage: cov}
|
||||
m.activeMu.Unlock()
|
||||
return nil
|
||||
}
|
||||
}
|
||||
m.active = append(m.active, loadedDataset{ID: id, Field: field, Coverage: cov})
|
||||
m.activeMu.Unlock()
|
||||
m.log.Info("loaded dataset",
|
||||
zap.Time("epoch", epoch),
|
||||
zap.String("filename", id.Filename()),
|
||||
zap.String("source", m.src.ID()))
|
||||
return nil
|
||||
}
|
||||
|
||||
// Refresh ensures the most recent upstream dataset is downloaded and active.
|
||||
//
|
||||
// If the freshest stored dataset is newer than retentionTTL old, no upstream
|
||||
// check is performed. Otherwise the source's LatestEpoch is consulted; if it
|
||||
// is newer than the active dataset, a download is started and on completion
|
||||
// the new dataset becomes active.
|
||||
// Refresh ensures the freshest global dataset is downloaded and active.
|
||||
//
|
||||
// Returns the JobID started, or empty string when nothing was scheduled.
|
||||
func (m *Manager) Refresh(ctx context.Context, freshnessTTL time.Duration) (string, error) {
|
||||
if active := m.Active(); active != nil && time.Since(active.Epoch()) < freshnessTTL {
|
||||
if a := m.activeGlobal(); a != nil && time.Since(a.ID.Epoch) < freshnessTTL {
|
||||
return "", nil
|
||||
}
|
||||
|
||||
// Try loading the freshest existing dataset before going to the network.
|
||||
if epochs, err := m.store.List(); err == nil {
|
||||
for _, e := range epochs {
|
||||
if time.Since(e) > freshnessTTL {
|
||||
if datasets, err := m.store.List(); err == nil {
|
||||
for _, id := range datasets {
|
||||
if !id.Subset.IsGlobal() {
|
||||
continue
|
||||
}
|
||||
if active := m.Active(); active != nil && active.Epoch().Equal(e) {
|
||||
if time.Since(id.Epoch) > freshnessTTL {
|
||||
continue
|
||||
}
|
||||
if a := m.activeGlobal(); a != nil && a.ID.Equals(id) {
|
||||
return "", nil
|
||||
}
|
||||
if err := m.LoadEpoch(ctx, e); err == nil {
|
||||
if err := m.Load(ctx, id); err == nil {
|
||||
return "", nil
|
||||
}
|
||||
}
|
||||
|
|
@ -276,37 +334,50 @@ func (m *Manager) Refresh(ctx context.Context, freshnessTTL time.Duration) (stri
|
|||
if err != nil {
|
||||
return "", fmt.Errorf("latest epoch: %w", err)
|
||||
}
|
||||
if active := m.Active(); active != nil && !latest.After(active.Epoch()) {
|
||||
id := DatasetID{Epoch: latest}
|
||||
if a := m.activeGlobal(); a != nil && !latest.After(a.ID.Epoch) {
|
||||
return "", nil
|
||||
}
|
||||
|
||||
jobID := m.Download(latest)
|
||||
|
||||
// Spawn a watcher that loads the dataset on successful completion.
|
||||
go func() {
|
||||
for {
|
||||
info, ok := m.GetJob(jobID)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
switch info.Status {
|
||||
case JobComplete:
|
||||
if err := m.LoadEpoch(context.Background(), latest); err != nil {
|
||||
m.log.Error("load after download", zap.Error(err))
|
||||
}
|
||||
return
|
||||
case JobFailed, JobCancelled:
|
||||
return
|
||||
}
|
||||
time.Sleep(2 * time.Second)
|
||||
}
|
||||
}()
|
||||
jobID := m.Download(id)
|
||||
go m.loadAfterCompletion(jobID, id)
|
||||
return jobID, nil
|
||||
}
|
||||
|
||||
// runDownload executes one Source.Download invocation and records its outcome.
|
||||
// activeGlobal returns the currently-loaded global dataset, if any.
|
||||
func (m *Manager) activeGlobal() *loadedDataset {
|
||||
m.activeMu.RLock()
|
||||
defer m.activeMu.RUnlock()
|
||||
for i := range m.active {
|
||||
if m.active[i].ID.Subset.IsGlobal() {
|
||||
d := m.active[i]
|
||||
return &d
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *Manager) loadAfterCompletion(jobID string, id DatasetID) {
|
||||
for {
|
||||
info, ok := m.GetJob(jobID)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
switch info.Status {
|
||||
case JobComplete:
|
||||
if err := m.Load(context.Background(), id); err != nil {
|
||||
m.log.Error("load after download", zap.Error(err))
|
||||
}
|
||||
return
|
||||
case JobFailed, JobCancelled:
|
||||
return
|
||||
}
|
||||
time.Sleep(2 * time.Second)
|
||||
}
|
||||
}
|
||||
|
||||
func (m *Manager) runDownload(ctx context.Context, e *jobEntry) {
|
||||
defer m.inFlight.Delete(e.epoch.Format(time.RFC3339))
|
||||
defer m.inFlight.Delete(e.dataset.Filename())
|
||||
|
||||
e.mu.Lock()
|
||||
e.status = JobRunning
|
||||
|
|
@ -314,9 +385,9 @@ func (m *Manager) runDownload(ctx context.Context, e *jobEntry) {
|
|||
|
||||
m.log.Info("download started",
|
||||
zap.String("job", e.id),
|
||||
zap.Time("epoch", e.epoch))
|
||||
zap.String("dataset", e.dataset.Filename()))
|
||||
|
||||
err := m.src.Download(ctx, e.epoch, m.store, jobProgress{e: e}, m.throttle)
|
||||
err := m.src.Download(ctx, e.dataset, m.store, jobProgress{e: e}, m.throttle)
|
||||
now := time.Now().UTC()
|
||||
|
||||
e.mu.Lock()
|
||||
|
|
@ -339,10 +410,9 @@ func (m *Manager) runDownload(ctx context.Context, e *jobEntry) {
|
|||
zap.NamedError("err", err))
|
||||
}
|
||||
|
||||
// completeShortCircuit records a job as complete without performing any work.
|
||||
func (m *Manager) completeShortCircuit(ctx context.Context, e *jobEntry) {
|
||||
_ = ctx
|
||||
defer m.inFlight.Delete(e.epoch.Format(time.RFC3339))
|
||||
defer m.inFlight.Delete(e.dataset.Filename())
|
||||
now := time.Now().UTC()
|
||||
e.mu.Lock()
|
||||
e.status = JobComplete
|
||||
|
|
@ -350,20 +420,6 @@ func (m *Manager) completeShortCircuit(ctx context.Context, e *jobEntry) {
|
|||
e.mu.Unlock()
|
||||
}
|
||||
|
||||
// swapActive replaces the active field and closes the previous one if it
|
||||
// implements io.Closer.
|
||||
func (m *Manager) swapActive(f weather.WindField) {
|
||||
m.activeMu.Lock()
|
||||
old := m.active
|
||||
m.active = f
|
||||
m.activeMu.Unlock()
|
||||
if c, ok := old.(interface{ Close() error }); ok && c != nil {
|
||||
if err := c.Close(); err != nil {
|
||||
m.log.Warn("close old dataset", zap.Error(err))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Close releases all resources, cancelling any in-flight jobs.
|
||||
func (m *Manager) Close() error {
|
||||
m.jobsMu.Lock()
|
||||
|
|
@ -373,11 +429,18 @@ func (m *Manager) Close() error {
|
|||
m.jobsMu.Unlock()
|
||||
|
||||
m.activeMu.Lock()
|
||||
active := m.active
|
||||
for _, d := range m.active {
|
||||
closeField(d.Field, m.log)
|
||||
}
|
||||
m.active = nil
|
||||
m.activeMu.Unlock()
|
||||
if c, ok := active.(interface{ Close() error }); ok && c != nil {
|
||||
return c.Close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func closeField(f weather.WindField, log *zap.Logger) {
|
||||
if c, ok := f.(interface{ Close() error }); ok && c != nil {
|
||||
if err := c.Close(); err != nil && log != nil {
|
||||
log.Warn("close dataset", zap.Error(err))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -14,15 +14,16 @@ import (
|
|||
//
|
||||
// Layout under Root:
|
||||
//
|
||||
// <epoch>.bin — committed dataset (binary cube)
|
||||
// <epoch>.bin.downloading — in-progress dataset
|
||||
// <epoch>.bin.manifest.json — manifest of completed work units
|
||||
// <filename>.bin — committed dataset
|
||||
// <filename>.bin.downloading — in-progress dataset
|
||||
// <filename>.bin.manifest.json — completed work units
|
||||
//
|
||||
// The .bin suffix exists to differentiate from sidecars in directory listings;
|
||||
// epoch is formatted as "20060102T150405Z" (UTC).
|
||||
// where <filename> is DatasetID.Filename() — typically
|
||||
// "20060102T150405Z" for the global subset or
|
||||
// "20060102T150405Z_r-10.10.-30.30_h0.72" for a subset.
|
||||
type LocalStore struct {
|
||||
Root string
|
||||
Source string // source ID, recorded for safety but currently advisory
|
||||
Source string
|
||||
Extension string // default ".bin"
|
||||
}
|
||||
|
||||
|
|
@ -37,8 +38,6 @@ func NewLocalStore(root, sourceID string) (*LocalStore, error) {
|
|||
// SourceID returns the source ID this store is configured for.
|
||||
func (s *LocalStore) SourceID() string { return s.Source }
|
||||
|
||||
const epochFormat = "20060102T150405Z"
|
||||
|
||||
func (s *LocalStore) ext() string {
|
||||
if s.Extension == "" {
|
||||
return ".bin"
|
||||
|
|
@ -46,32 +45,32 @@ func (s *LocalStore) ext() string {
|
|||
return s.Extension
|
||||
}
|
||||
|
||||
// Path returns the canonical path for an epoch's committed dataset file.
|
||||
func (s *LocalStore) Path(epoch time.Time) string {
|
||||
return filepath.Join(s.Root, epoch.UTC().Format(epochFormat)+s.ext())
|
||||
// Path returns the canonical path for id's committed dataset.
|
||||
func (s *LocalStore) Path(id DatasetID) string {
|
||||
return filepath.Join(s.Root, id.Filename()+s.ext())
|
||||
}
|
||||
|
||||
func (s *LocalStore) tempPath(epoch time.Time) string {
|
||||
return s.Path(epoch) + ".downloading"
|
||||
func (s *LocalStore) tempPath(id DatasetID) string {
|
||||
return s.Path(id) + ".downloading"
|
||||
}
|
||||
|
||||
func (s *LocalStore) manifestPath(epoch time.Time) string {
|
||||
return s.Path(epoch) + ".manifest.json"
|
||||
func (s *LocalStore) manifestPath(id DatasetID) string {
|
||||
return s.Path(id) + ".manifest.json"
|
||||
}
|
||||
|
||||
// Exists reports whether a committed dataset for epoch is present.
|
||||
func (s *LocalStore) Exists(epoch time.Time) bool {
|
||||
info, err := os.Stat(s.Path(epoch))
|
||||
// Exists reports whether a committed dataset for id is present.
|
||||
func (s *LocalStore) Exists(id DatasetID) bool {
|
||||
info, err := os.Stat(s.Path(id))
|
||||
return err == nil && !info.IsDir()
|
||||
}
|
||||
|
||||
// List returns all committed epochs, newest first.
|
||||
func (s *LocalStore) List() ([]time.Time, error) {
|
||||
// List returns all committed dataset IDs, newest first.
|
||||
func (s *LocalStore) List() ([]DatasetID, error) {
|
||||
entries, err := os.ReadDir(s.Root)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read store: %w", err)
|
||||
}
|
||||
var out []time.Time
|
||||
var out []DatasetID
|
||||
ext := s.ext()
|
||||
for _, e := range entries {
|
||||
if e.IsDir() {
|
||||
|
|
@ -82,24 +81,47 @@ func (s *LocalStore) List() ([]time.Time, error) {
|
|||
continue
|
||||
}
|
||||
stem := strings.TrimSuffix(name, ext)
|
||||
// skip in-progress files (their stem already has .bin.downloading...)
|
||||
// Skip in-progress files (their stem ends in .downloading or .manifest)
|
||||
if strings.Contains(stem, ".") {
|
||||
continue
|
||||
}
|
||||
t, err := time.Parse(epochFormat, stem)
|
||||
if err != nil {
|
||||
id, ok := parseFilename(stem)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
out = append(out, t.UTC())
|
||||
out = append(out, id)
|
||||
}
|
||||
sort.Slice(out, func(i, j int) bool { return out[i].After(out[j]) })
|
||||
sort.Slice(out, func(i, j int) bool {
|
||||
if !out[i].Epoch.Equal(out[j].Epoch) {
|
||||
return out[i].Epoch.After(out[j].Epoch)
|
||||
}
|
||||
return out[i].Subset.Key() < out[j].Subset.Key()
|
||||
})
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// Remove deletes the committed dataset and any sidecar files for epoch.
|
||||
func (s *LocalStore) Remove(epoch time.Time) error {
|
||||
// parseFilename inverts DatasetID.Filename(). The subset portion is not
|
||||
// fully reversible (Key encoding is one-way for floats), so List returns
|
||||
// IDs whose Subset is zero — the storage layer treats names as opaque
|
||||
// identifiers. Callers wanting structured subset metadata should keep an
|
||||
// out-of-band record.
|
||||
func parseFilename(stem string) (DatasetID, bool) {
|
||||
parts := strings.SplitN(stem, "_", 2)
|
||||
epoch, err := time.Parse("20060102T150405Z", parts[0])
|
||||
if err != nil {
|
||||
return DatasetID{}, false
|
||||
}
|
||||
id := DatasetID{Epoch: epoch.UTC()}
|
||||
// Subset key is opaque on disk; we don't reconstruct its parameters
|
||||
// here. Admin callers track subset specs separately when they need
|
||||
// the structured form.
|
||||
return id, true
|
||||
}
|
||||
|
||||
// Remove deletes the committed dataset and any sidecar files for id.
|
||||
func (s *LocalStore) Remove(id DatasetID) error {
|
||||
var errs []error
|
||||
for _, p := range []string{s.Path(epoch), s.tempPath(epoch), s.manifestPath(epoch)} {
|
||||
for _, p := range []string{s.Path(id), s.tempPath(id), s.manifestPath(id)} {
|
||||
if err := os.Remove(p); err != nil && !errors.Is(err, os.ErrNotExist) {
|
||||
errs = append(errs, err)
|
||||
}
|
||||
|
|
@ -110,55 +132,46 @@ func (s *LocalStore) Remove(epoch time.Time) error {
|
|||
return nil
|
||||
}
|
||||
|
||||
// BeginWrite opens or resumes a TempHandle for epoch.
|
||||
//
|
||||
// If a partial download is already present, its file and manifest are reused
|
||||
// so the new download picks up where the previous one stopped.
|
||||
func (s *LocalStore) BeginWrite(epoch time.Time) (TempHandle, error) {
|
||||
man, err := LoadManifest(s.manifestPath(epoch))
|
||||
// BeginWrite opens or resumes a TempHandle for id.
|
||||
func (s *LocalStore) BeginWrite(id DatasetID) (TempHandle, error) {
|
||||
man, err := LoadManifest(s.manifestPath(id))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &localHandle{
|
||||
store: s,
|
||||
epoch: epoch,
|
||||
manifest: man,
|
||||
}, nil
|
||||
return &localHandle{store: s, id: id, manifest: man}, nil
|
||||
}
|
||||
|
||||
type localHandle struct {
|
||||
store *LocalStore
|
||||
epoch time.Time
|
||||
id DatasetID
|
||||
manifest *Manifest
|
||||
closed bool
|
||||
}
|
||||
|
||||
func (h *localHandle) Path() string { return h.store.tempPath(h.epoch) }
|
||||
func (h *localHandle) Path() string { return h.store.tempPath(h.id) }
|
||||
func (h *localHandle) Manifest() *Manifest { return h.manifest }
|
||||
|
||||
// Commit promotes the temp file to its final path and removes the manifest.
|
||||
func (h *localHandle) Commit() error {
|
||||
if h.closed {
|
||||
return nil
|
||||
}
|
||||
h.closed = true
|
||||
if err := os.Rename(h.store.tempPath(h.epoch), h.store.Path(h.epoch)); err != nil {
|
||||
if err := os.Rename(h.store.tempPath(h.id), h.store.Path(h.id)); err != nil {
|
||||
return fmt.Errorf("commit rename: %w", err)
|
||||
}
|
||||
if err := os.Remove(h.store.manifestPath(h.epoch)); err != nil && !errors.Is(err, os.ErrNotExist) {
|
||||
if err := os.Remove(h.store.manifestPath(h.id)); err != nil && !errors.Is(err, os.ErrNotExist) {
|
||||
return fmt.Errorf("commit remove manifest: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Abort removes the in-progress file and manifest.
|
||||
func (h *localHandle) Abort() error {
|
||||
if h.closed {
|
||||
return nil
|
||||
}
|
||||
h.closed = true
|
||||
var firstErr error
|
||||
for _, p := range []string{h.store.tempPath(h.epoch), h.store.manifestPath(h.epoch)} {
|
||||
for _, p := range []string{h.store.tempPath(h.id), h.store.manifestPath(h.id)} {
|
||||
if err := os.Remove(p); err != nil && !errors.Is(err, os.ErrNotExist) && firstErr == nil {
|
||||
firstErr = err
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,7 +2,6 @@ package datasets
|
|||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
|
@ -14,8 +13,8 @@ func TestLocalStoreBeginWriteResume(t *testing.T) {
|
|||
t.Fatalf("NewLocalStore: %v", err)
|
||||
}
|
||||
|
||||
epoch := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)
|
||||
h, err := store.BeginWrite(epoch)
|
||||
id := DatasetID{Epoch: time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)}
|
||||
h, err := store.BeginWrite(id)
|
||||
if err != nil {
|
||||
t.Fatalf("BeginWrite: %v", err)
|
||||
}
|
||||
|
|
@ -27,7 +26,7 @@ func TestLocalStoreBeginWriteResume(t *testing.T) {
|
|||
}
|
||||
|
||||
// Re-open should see the previous manifest entry.
|
||||
h2, err := store.BeginWrite(epoch)
|
||||
h2, err := store.BeginWrite(id)
|
||||
if err != nil {
|
||||
t.Fatalf("BeginWrite resume: %v", err)
|
||||
}
|
||||
|
|
@ -35,48 +34,59 @@ func TestLocalStoreBeginWriteResume(t *testing.T) {
|
|||
t.Errorf("resumed manifest missing step000-A; units = %v", h2.Manifest().Units())
|
||||
}
|
||||
|
||||
// Commit promotes the temp file and removes the manifest.
|
||||
if err := h2.Commit(); err != nil {
|
||||
t.Fatalf("Commit: %v", err)
|
||||
}
|
||||
if !store.Exists(epoch) {
|
||||
if !store.Exists(id) {
|
||||
t.Errorf("Exists after commit returned false")
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(dir, store.manifestPath(epoch))); !os.IsNotExist(err) {
|
||||
if _, err := os.Stat(store.manifestPath(id)); !os.IsNotExist(err) {
|
||||
t.Errorf("manifest should be removed, got err=%v", err)
|
||||
}
|
||||
|
||||
// Listing finds the committed epoch.
|
||||
epochs, err := store.List()
|
||||
stored, err := store.List()
|
||||
if err != nil {
|
||||
t.Fatalf("List: %v", err)
|
||||
}
|
||||
if len(epochs) != 1 || !epochs[0].Equal(epoch) {
|
||||
t.Errorf("List = %v, want [%v]", epochs, epoch)
|
||||
if len(stored) != 1 || !stored[0].Epoch.Equal(id.Epoch) {
|
||||
t.Errorf("List = %v, want one item with epoch %v", stored, id.Epoch)
|
||||
}
|
||||
|
||||
// Remove cleans up.
|
||||
if err := store.Remove(epoch); err != nil {
|
||||
if err := store.Remove(id); err != nil {
|
||||
t.Fatalf("Remove: %v", err)
|
||||
}
|
||||
if store.Exists(epoch) {
|
||||
if store.Exists(id) {
|
||||
t.Errorf("Exists after remove returned true")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLocalStoreAbort(t *testing.T) {
|
||||
func TestLocalStoreSubsetPath(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
store, _ := NewLocalStore(dir, "gfs-test")
|
||||
epoch := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)
|
||||
|
||||
h, _ := store.BeginWrite(epoch)
|
||||
os.WriteFile(h.Path(), []byte("x"), 0o644)
|
||||
h.Manifest().Mark("step000-A")
|
||||
|
||||
if err := h.Abort(); err != nil {
|
||||
t.Fatalf("Abort: %v", err)
|
||||
regional := DatasetID{
|
||||
Epoch: epoch,
|
||||
Subset: SubsetSpec{
|
||||
Region: &Region{MinLat: -10, MaxLat: 10, MinLng: 0, MaxLng: 30},
|
||||
HourRange: &HourRange{MinHour: 0, MaxHour: 72},
|
||||
},
|
||||
}
|
||||
if _, err := os.Stat(h.Path()); !os.IsNotExist(err) {
|
||||
t.Errorf("temp file should be removed after abort, got %v", err)
|
||||
global := DatasetID{Epoch: epoch}
|
||||
if store.Path(global) == store.Path(regional) {
|
||||
t.Errorf("global and regional should have distinct paths")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSubsetSpecCoverage(t *testing.T) {
|
||||
r := Region{MinLat: -10, MaxLat: 10, MinLng: 350, MaxLng: 10} // wraps antimeridian
|
||||
s := SubsetSpec{Region: &r}
|
||||
if !s.IncludesLatLng(0, 0) {
|
||||
t.Errorf("(0,0) should be inside antimeridian region")
|
||||
}
|
||||
if !s.IncludesLatLng(0, 359) {
|
||||
t.Errorf("(0,359) should be inside antimeridian region")
|
||||
}
|
||||
if s.IncludesLatLng(0, 180) {
|
||||
t.Errorf("(0,180) should be outside antimeridian region")
|
||||
}
|
||||
}
|
||||
|
|
|
|||
156
internal/datasets/subset.go
Normal file
156
internal/datasets/subset.go
Normal file
|
|
@ -0,0 +1,156 @@
|
|||
package datasets
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"slices"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// SubsetSpec describes which portion of a dataset to download.
|
||||
//
|
||||
// A zero-value SubsetSpec means "the full dataset". The Region and
|
||||
// HourRange fields independently restrict what is fetched and stored.
|
||||
type SubsetSpec struct {
|
||||
// Region restricts the geographic extent. nil means global.
|
||||
Region *Region `json:"region,omitempty"`
|
||||
|
||||
// HourRange restricts the forecast horizon. nil means the source's
|
||||
// full horizon (e.g. 0..192h for GFS 0.5°).
|
||||
HourRange *HourRange `json:"hour_range,omitempty"`
|
||||
|
||||
// Members restricts ensemble members for sources that support them (GEFS).
|
||||
// nil means all available members.
|
||||
Members []int `json:"members,omitempty"`
|
||||
}
|
||||
|
||||
// Region is an axis-aligned geographic bounding box.
|
||||
//
|
||||
// Longitudes are in [0, 360); a box crossing the antimeridian has
|
||||
// MinLng > MaxLng.
|
||||
type Region struct {
|
||||
MinLat float64 `json:"min_lat"`
|
||||
MaxLat float64 `json:"max_lat"`
|
||||
MinLng float64 `json:"min_lng"`
|
||||
MaxLng float64 `json:"max_lng"`
|
||||
}
|
||||
|
||||
// HourRange is an inclusive forecast-hour range.
|
||||
type HourRange struct {
|
||||
MinHour int `json:"min_hour"`
|
||||
MaxHour int `json:"max_hour"`
|
||||
}
|
||||
|
||||
// IsGlobal reports whether the spec selects the entire dataset.
|
||||
func (s SubsetSpec) IsGlobal() bool {
|
||||
return s.Region == nil && s.HourRange == nil && len(s.Members) == 0
|
||||
}
|
||||
|
||||
// IncludesLatLng reports whether (lat, lng) lies inside the spec's Region,
|
||||
// or the spec has no Region.
|
||||
func (s SubsetSpec) IncludesLatLng(lat, lng float64) bool {
|
||||
if s.Region == nil {
|
||||
return true
|
||||
}
|
||||
r := s.Region
|
||||
if lat < r.MinLat || lat > r.MaxLat {
|
||||
return false
|
||||
}
|
||||
if r.MinLng <= r.MaxLng {
|
||||
return lng >= r.MinLng && lng <= r.MaxLng
|
||||
}
|
||||
// Wraps the antimeridian.
|
||||
return lng >= r.MinLng || lng <= r.MaxLng
|
||||
}
|
||||
|
||||
// IncludesHour reports whether the forecast hour is in range.
|
||||
func (s SubsetSpec) IncludesHour(h int) bool {
|
||||
if s.HourRange == nil {
|
||||
return true
|
||||
}
|
||||
return h >= s.HourRange.MinHour && h <= s.HourRange.MaxHour
|
||||
}
|
||||
|
||||
// IncludesMember reports whether the ensemble member is in range.
|
||||
func (s SubsetSpec) IncludesMember(m int) bool {
|
||||
if len(s.Members) == 0 {
|
||||
return true
|
||||
}
|
||||
return slices.Contains(s.Members, m)
|
||||
}
|
||||
|
||||
// Key returns a deterministic short identifier for the spec. The empty
|
||||
// string represents the global subset.
|
||||
func (s SubsetSpec) Key() string {
|
||||
if s.IsGlobal() {
|
||||
return ""
|
||||
}
|
||||
var b strings.Builder
|
||||
if s.Region != nil {
|
||||
fmt.Fprintf(&b, "r%g.%g.%g.%g", s.Region.MinLat, s.Region.MaxLat, s.Region.MinLng, s.Region.MaxLng)
|
||||
}
|
||||
if s.HourRange != nil {
|
||||
if b.Len() > 0 {
|
||||
b.WriteByte('_')
|
||||
}
|
||||
fmt.Fprintf(&b, "h%d.%d", s.HourRange.MinHour, s.HourRange.MaxHour)
|
||||
}
|
||||
if len(s.Members) > 0 {
|
||||
if b.Len() > 0 {
|
||||
b.WriteByte('_')
|
||||
}
|
||||
fmt.Fprintf(&b, "m")
|
||||
for i, m := range s.Members {
|
||||
if i > 0 {
|
||||
b.WriteByte('.')
|
||||
}
|
||||
fmt.Fprintf(&b, "%d", m)
|
||||
}
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// DatasetID identifies one storable dataset.
|
||||
type DatasetID struct {
|
||||
Epoch time.Time
|
||||
Subset SubsetSpec
|
||||
}
|
||||
|
||||
// Equals reports whether two DatasetIDs refer to the same dataset.
|
||||
// DatasetID is not comparable with == because SubsetSpec contains slices.
|
||||
func (id DatasetID) Equals(other DatasetID) bool {
|
||||
return id.Epoch.Equal(other.Epoch) && id.Subset.Key() == other.Subset.Key()
|
||||
}
|
||||
|
||||
// Filename returns the canonical filename stem for the dataset. The
|
||||
// extension is appended by the Storage implementation.
|
||||
func (id DatasetID) Filename() string {
|
||||
stem := id.Epoch.UTC().Format("20060102T150405Z")
|
||||
if k := id.Subset.Key(); k != "" {
|
||||
return stem + "_" + k
|
||||
}
|
||||
return stem
|
||||
}
|
||||
|
||||
// Coverage is the spatial and temporal extent of a loaded dataset, used by
|
||||
// the Manager to select which dataset can serve a given query.
|
||||
type Coverage struct {
|
||||
Region Region `json:"region"`
|
||||
StartTime time.Time `json:"start_time"`
|
||||
EndTime time.Time `json:"end_time"`
|
||||
}
|
||||
|
||||
// Covers reports whether (t, lat, lng) lies inside the coverage.
|
||||
func (c Coverage) Covers(t time.Time, lat, lng float64) bool {
|
||||
if t.Before(c.StartTime) || t.After(c.EndTime) {
|
||||
return false
|
||||
}
|
||||
r := c.Region
|
||||
if lat < r.MinLat || lat > r.MaxLat {
|
||||
return false
|
||||
}
|
||||
if r.MinLng <= r.MaxLng {
|
||||
return lng >= r.MinLng && lng <= r.MaxLng
|
||||
}
|
||||
return lng >= r.MinLng || lng <= r.MaxLng
|
||||
}
|
||||
|
|
@ -11,87 +11,75 @@ import (
|
|||
//
|
||||
// Implementations download dataset files in a transactional, resumable
|
||||
// manner and load them as weather.WindField. A Source must be safe for
|
||||
// concurrent use across multiple Manager calls.
|
||||
// concurrent use across many Manager calls.
|
||||
type Source interface {
|
||||
// ID is a stable identifier, e.g. "noaa-gfs-0p50".
|
||||
// ID is a stable identifier, e.g. "gfs-0p50-3h".
|
||||
ID() string
|
||||
|
||||
// LatestEpoch returns the most recent dataset epoch this source can provide.
|
||||
LatestEpoch(ctx context.Context) (time.Time, error)
|
||||
|
||||
// Download fetches the dataset for epoch into store. Sources must honour
|
||||
// any partial progress recorded in store's manifest and skip
|
||||
// already-completed work, so re-invocation after a crash resumes cleanly.
|
||||
// Download fetches the dataset identified by id into store. Sources
|
||||
// must honour any partial progress recorded in store's manifest and
|
||||
// skip already-completed work so re-invocation after a crash resumes
|
||||
// cleanly.
|
||||
//
|
||||
// prog receives progress events; nil is acceptable.
|
||||
// throttle, if non-nil, is consulted before each network read for
|
||||
// bandwidth limiting; nil means no throttling.
|
||||
Download(ctx context.Context, epoch time.Time, store Storage, prog ProgressSink, throttle Throttle) error
|
||||
Download(ctx context.Context, id DatasetID, store Storage, prog ProgressSink, throttle Throttle) error
|
||||
|
||||
// Open loads epoch's stored dataset and returns it as a WindField.
|
||||
Open(ctx context.Context, epoch time.Time, store Storage) (weather.WindField, error)
|
||||
// Open loads id's stored dataset and returns it as a WindField.
|
||||
Open(ctx context.Context, id DatasetID, store Storage) (weather.WindField, error)
|
||||
|
||||
// Coverage returns the geographical/temporal extent of a downloaded
|
||||
// dataset. Used by the Manager to decide which loaded dataset can
|
||||
// serve a given prediction query.
|
||||
Coverage(id DatasetID) Coverage
|
||||
}
|
||||
|
||||
// Storage abstracts the on-disk location of dataset files and their manifests.
|
||||
//
|
||||
// Atomicity: only datasets promoted via TempHandle.Commit appear in Exists or
|
||||
// List. Aborted or in-progress downloads are invisible to readers.
|
||||
// Atomicity: only datasets promoted via TempHandle.Commit appear in Exists
|
||||
// or List. Aborted or in-progress downloads are invisible to readers.
|
||||
type Storage interface {
|
||||
// SourceID identifies the data source these files belong to. Mixing
|
||||
// sources in one Storage is not supported.
|
||||
// SourceID identifies the data source these files belong to.
|
||||
SourceID() string
|
||||
|
||||
// Path returns the canonical local path for epoch's dataset. The path
|
||||
// is valid even when the dataset has not been written.
|
||||
Path(epoch time.Time) string
|
||||
// Path returns the canonical local path for id's dataset.
|
||||
Path(id DatasetID) string
|
||||
|
||||
// Exists reports whether a committed dataset for epoch is present.
|
||||
Exists(epoch time.Time) bool
|
||||
// Exists reports whether a committed dataset for id is present.
|
||||
Exists(id DatasetID) bool
|
||||
|
||||
// List returns all committed epochs available, newest first.
|
||||
List() ([]time.Time, error)
|
||||
// List returns all committed dataset IDs available, newest first.
|
||||
List() ([]DatasetID, error)
|
||||
|
||||
// Remove deletes the dataset and any sidecar manifest for epoch.
|
||||
Remove(epoch time.Time) error
|
||||
// Remove deletes the dataset and any sidecar manifest for id.
|
||||
Remove(id DatasetID) error
|
||||
|
||||
// BeginWrite opens (or resumes) a transactional handle for downloading
|
||||
// epoch's dataset. Callers must Commit or Abort the returned handle.
|
||||
BeginWrite(epoch time.Time) (TempHandle, error)
|
||||
// id's dataset.
|
||||
BeginWrite(id DatasetID) (TempHandle, error)
|
||||
}
|
||||
|
||||
// TempHandle is the storage state for one in-progress download.
|
||||
type TempHandle interface {
|
||||
// Path returns the path of the in-progress file. Sources write directly here.
|
||||
Path() string
|
||||
|
||||
// Manifest is the tracker of completed work units for resume support.
|
||||
Manifest() *Manifest
|
||||
|
||||
// Commit promotes the temp file to its canonical location and removes
|
||||
// the manifest. Subsequent calls are no-ops.
|
||||
Commit() error
|
||||
|
||||
// Abort discards the temp file and manifest. Subsequent calls are no-ops.
|
||||
Abort() error
|
||||
}
|
||||
|
||||
// ProgressSink receives progress events during a download.
|
||||
//
|
||||
// All methods are safe to call concurrently.
|
||||
type ProgressSink interface {
|
||||
// SetTotal sets the total number of work units this download expects.
|
||||
// May be called multiple times if discovery happens incrementally.
|
||||
SetTotal(n int)
|
||||
// StepComplete records one work unit as completed.
|
||||
StepComplete()
|
||||
// Bytes records n bytes received from the network.
|
||||
Bytes(n int64)
|
||||
}
|
||||
|
||||
// Throttle is an optional bandwidth limiter consulted by sources before
|
||||
// each network read.
|
||||
type Throttle interface {
|
||||
// Wait blocks until n bytes can be consumed from the budget,
|
||||
// or returns ctx's error if the context is cancelled while waiting.
|
||||
Wait(ctx context.Context, n int) error
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue