// Package fetch downloads a list of files concurrently, skipping any that are // already on disk. // // It is deliberately source-agnostic: it knows nothing about provinces, exam // years or spreadsheet formats. Each source in internal/sources produces a // []Item and this package moves the bytes. package fetch import ( "context" "fmt" "io" "net/http" "os" "path/filepath" "sync" "time" ) // Item is one file to download. type Item struct { // Name is the human label used in progress output, e.g. "An Giang". Name string URL string // Path is the absolute destination, including filename. Path string } // Status is the outcome of one item. type Status int const ( // StatusOK means the file was downloaded during this run. StatusOK Status = iota // StatusSkip means a non-empty file was already present. StatusSkip // StatusFail means the download did not complete; Result.Err says why. StatusFail ) // Result records what happened to one item. type Result struct { Item Item Status Status Size int64 Err error } // Options configures a run. The zero value is usable: Concurrency and Timeout // fall back to defaults. type Options struct { // Concurrency is the number of files downloaded at once. Concurrency int // Timeout bounds each individual request, headers and body together, so a // single stalled connection cannot hang the whole run. Timeout time.Duration // Headers are sent with every request. The 2017 CDN rejects requests // without a browser User-Agent and a matching Referer. Headers map[string]string // OnResult, if set, is called once per completed item. Calls are // serialised, so it does not need its own locking, but they arrive in // completion order rather than list order. OnResult func(done, total int, r Result) } const ( defaultConcurrency = 6 defaultTimeout = 2 * time.Minute ) // Run downloads every item, returning one Result each. It does not return an // error for a failed download — that is reported per item — only for a problem // that stops the run as a whole, such as an unusable output directory. // // Results come back in completion order, not input order. func Run(ctx context.Context, items []Item, opts Options) ([]Result, error) { if opts.Concurrency <= 0 { opts.Concurrency = defaultConcurrency } if opts.Timeout <= 0 { opts.Timeout = defaultTimeout } // Every item's parent directory must exist before any worker starts, so a // missing output directory fails once here rather than N times in parallel. for _, it := range items { if err := os.MkdirAll(filepath.Dir(it.Path), 0o755); err != nil { return nil, fmt.Errorf("cannot create output directory: %w", err) } } client := &http.Client{Timeout: opts.Timeout} var ( mu sync.Mutex results = make([]Result, 0, len(items)) next int ) work := make(chan Item) var wg sync.WaitGroup for range opts.Concurrency { wg.Add(1) go func() { defer wg.Done() for it := range work { r := download(ctx, client, it, opts.Headers) mu.Lock() results = append(results, r) next++ if opts.OnResult != nil { opts.OnResult(next, len(items), r) } mu.Unlock() } }() } for _, it := range items { select { case work <- it: case <-ctx.Done(): close(work) wg.Wait() return results, ctx.Err() } } close(work) wg.Wait() return results, nil } // download fetches one item, or reports that it was already present. // // The body is streamed to a .part file and renamed into place only once it is // complete. Without that, an interrupted run leaves a truncated file at the // final path — and because the skip check only tests for a non-empty file, // every later run would skip it and the corruption would persist silently. // // A .part left behind by an abrupt kill needs no cleanup: it does not satisfy // the skip check, os.Create truncates it, and the next run re-fetches the file. // The parser ignores it in the meantime, since it reads only .xls and .xlsx. func download(ctx context.Context, client *http.Client, it Item, headers map[string]string) Result { if st, err := os.Stat(it.Path); err == nil && st.Size() > 0 { return Result{Item: it, Status: StatusSkip, Size: st.Size()} } fail := func(err error) Result { return Result{Item: it, Status: StatusFail, Err: err} } req, err := http.NewRequestWithContext(ctx, http.MethodGet, it.URL, nil) if err != nil { return fail(err) } for k, v := range headers { req.Header.Set(k, v) } resp, err := client.Do(req) if err != nil { return fail(err) } defer resp.Body.Close() if resp.StatusCode != http.StatusOK { return fail(fmt.Errorf("HTTP %d", resp.StatusCode)) } part := it.Path + ".part" f, err := os.Create(part) if err != nil { return fail(err) } n, err := io.Copy(f, resp.Body) if cerr := f.Close(); err == nil { err = cerr } if err != nil { os.Remove(part) return fail(err) } if n == 0 { os.Remove(part) return fail(fmt.Errorf("empty response body")) } if err := os.Rename(part, it.Path); err != nil { os.Remove(part) return fail(err) } return Result{Item: it, Status: StatusOK, Size: n} } // Tally counts results by status. func Tally(results []Result) (ok, skip int, failed []Result) { for _, r := range results { switch r.Status { case StatusOK: ok++ case StatusSkip: skip++ case StatusFail: failed = append(failed, r) } } return ok, skip, failed }