feat: data lifecycle rewrite, download clients, wanted list, and central catalog index
Ships the fresh-start schema cleanup: rebuilt explore catalog index pipeline (dump import, artifact fetch/build, incremental listen-count refresh), a new download subsystem (Lidarr/Prowlarr/qBittorrent/SABnzbd/ slskd/yt-dlp providers, staging, reconciliation, wanted list), and the supporting schema/query/store changes across backend and frontend. Also includes two smaller follow-ups: bump the central index's rebuild-after cadence from 90 to 180 days, and remove the Explore "library only" online/offline toggle entirely (frontend-only, no backend counterpart) rather than carry unused UI/state. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Y2Agd9af5hE7qzti2ackiS
This commit is contained in:
@@ -0,0 +1,87 @@
|
||||
package explore
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"yellowjacket/backend/jobs"
|
||||
)
|
||||
|
||||
// Fetching and merging the prebuilt catalog artifact.
|
||||
//
|
||||
// This is the whole catalog build on a user's machine. The import that
|
||||
// derives the catalog from the MetaBrainz dumps lives behind the
|
||||
// `indexbuild` tag and runs only in CI (see dumpbuild_stub.go).
|
||||
|
||||
// Artifact build stages, shown in the jobs panel.
|
||||
const (
|
||||
artifactStageDownload = iota
|
||||
artifactStageMerge
|
||||
)
|
||||
|
||||
var artifactStageNames = [...]string{
|
||||
"Downloading catalog",
|
||||
"Merging catalog",
|
||||
}
|
||||
|
||||
// tryCoreArtifact fetches and merges the prebuilt catalog. Every
|
||||
// failure path is non-fatal by design: the caller falls back, and a
|
||||
// fresh install with no network still gets its own library in Explore.
|
||||
func (si *SearchIndex) tryCoreArtifact(ctx context.Context) error {
|
||||
si.mu.Lock()
|
||||
si.buildStatus = IndexStatus{
|
||||
Building: true,
|
||||
Tiers: []TierStatus{
|
||||
{Name: artifactStageNames[artifactStageDownload], State: "pending"},
|
||||
{Name: artifactStageNames[artifactStageMerge], State: "pending"},
|
||||
},
|
||||
}
|
||||
si.mu.Unlock()
|
||||
si.refreshStatusCounts()
|
||||
|
||||
fetcher, err := newArtifactFetcher(si)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
si.setTierStatus(artifactStageNames[artifactStageDownload], "running", 0, 0)
|
||||
si.logIndexJob(jobs.LevelInfo, "Fetching prebuilt catalog")
|
||||
|
||||
path, err := fetcher.fetch(ctx)
|
||||
if err != nil {
|
||||
si.setTierStatus(artifactStageNames[artifactStageDownload], "error", 0, 0)
|
||||
|
||||
return err
|
||||
}
|
||||
|
||||
si.setTierStatus(artifactStageNames[artifactStageDownload], "complete", 0, 0)
|
||||
si.setTierStatus(artifactStageNames[artifactStageMerge], "running", 0, 0)
|
||||
|
||||
if err := si.importCoreArtifact(ctx, path); err != nil {
|
||||
si.setTierStatus(artifactStageNames[artifactStageMerge], "error", 0, 0)
|
||||
si.removeArtifactFile(path)
|
||||
|
||||
return err
|
||||
}
|
||||
|
||||
si.removeArtifactFile(path)
|
||||
si.setTierStatus(artifactStageNames[artifactStageMerge], "complete", 0, 0)
|
||||
|
||||
si.mu.Lock()
|
||||
si.buildStatus.Building = false
|
||||
si.mu.Unlock()
|
||||
|
||||
// Fold the user's own library into the freshly-merged catalog:
|
||||
// owned entities the artifact does not cover are inserted, and
|
||||
// covered ones are flagged in_library.
|
||||
si.PopulateLocalCrossReferences()
|
||||
si.refreshStatusCounts()
|
||||
si.scheduleChampionRebuild()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// artifactAlreadyMerged reports whether this index already carries a
|
||||
// merged artifact, so a restart does not re-download one.
|
||||
func (si *SearchIndex) artifactAlreadyMerged() bool {
|
||||
return si.hasMeta(coreArtifactVersionKey)
|
||||
}
|
||||
@@ -0,0 +1,415 @@
|
||||
package explore
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/klauspost/compress/zstd"
|
||||
|
||||
"yellowjacket/backend/jobs"
|
||||
"yellowjacket/backend/system"
|
||||
)
|
||||
|
||||
// Download of the prebuilt core index artifact.
|
||||
//
|
||||
// The artifact is published by .gitea/workflows/index-artifact.yml to a
|
||||
// Gitea generic package registry under a fixed "latest" version, so the
|
||||
// client needs no package-listing API (which requires a token) — only an
|
||||
// anonymous GET of a predictable URL.
|
||||
//
|
||||
// Everything here degrades to "no artifact" rather than failing: an
|
||||
// offline install, a 404 before the first artifact is published, or a
|
||||
// corrupt download all leave the caller free to fall back to the normal
|
||||
// build path. A missing artifact must never be the reason Explore is
|
||||
// broken.
|
||||
|
||||
const (
|
||||
// defaultCoreArtifactBaseURL is where the CI-built artifact lives.
|
||||
// Overridable via YJ_CORE_INDEX_URL for testing and for anyone
|
||||
// self-hosting their own index.
|
||||
defaultCoreArtifactBaseURL = "https://git.ljones.me/api/packages/yonlu/" +
|
||||
"generic/yellowjacket-core-index/latest/"
|
||||
|
||||
// coreArtifactFile is the compressed artifact's filename, and
|
||||
// coreArtifactChecksumFile its detached sha256.
|
||||
coreArtifactFile = "core-index.db.zst"
|
||||
coreArtifactChecksum = "core-index.db.zst.sha256"
|
||||
|
||||
// artifactURLEnv overrides the base URL.
|
||||
artifactURLEnv = "YJ_CORE_INDEX_URL"
|
||||
|
||||
// artifactDiscoverTimeout bounds the checksum fetch, which doubles as
|
||||
// the availability probe. Short: a first run should not sit for
|
||||
// minutes deciding whether an artifact exists.
|
||||
artifactDiscoverTimeout = 30 * time.Second
|
||||
|
||||
// artifactMinFreeBytes is the free disk needed to fetch and unpack.
|
||||
// The compressed artifact plus its expansion plus merge headroom —
|
||||
// two orders of magnitude below the full dump import's 6GB floor,
|
||||
// which is much of the point.
|
||||
artifactMinFreeBytes = 1 << 30
|
||||
|
||||
// artifactMaxRetries bounds resume attempts for the body download.
|
||||
artifactMaxRetries = 5
|
||||
)
|
||||
|
||||
// ErrArtifactUnavailable means no artifact could be fetched. It is an
|
||||
// expected outcome (offline, not yet published), not a failure.
|
||||
var ErrArtifactUnavailable = errors.New("core index artifact unavailable")
|
||||
|
||||
// coreArtifactBaseURL resolves the artifact location, honouring the
|
||||
// environment override.
|
||||
func coreArtifactBaseURL() string {
|
||||
if v := strings.TrimSpace(os.Getenv(artifactURLEnv)); v != "" {
|
||||
if !strings.HasSuffix(v, "/") {
|
||||
return v + "/"
|
||||
}
|
||||
|
||||
return v
|
||||
}
|
||||
|
||||
return defaultCoreArtifactBaseURL
|
||||
}
|
||||
|
||||
// artifactFetcher downloads and unpacks the core index artifact.
|
||||
type artifactFetcher struct {
|
||||
si *SearchIndex
|
||||
client *http.Client
|
||||
baseURL string
|
||||
stagingDir string
|
||||
}
|
||||
|
||||
func newArtifactFetcher(si *SearchIndex) (*artifactFetcher, error) {
|
||||
dataDir, err := system.GetUserDataDirPath()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("core artifact: data dir: %w", err)
|
||||
}
|
||||
|
||||
stagingDir := filepath.Join(dataDir, "explore-staging")
|
||||
if err := os.MkdirAll(stagingDir, 0o755); err != nil {
|
||||
return nil, fmt.Errorf("core artifact: staging dir: %w", err)
|
||||
}
|
||||
|
||||
return &artifactFetcher{
|
||||
si: si,
|
||||
// No client-level timeout: a 70MB body over a slow link can take
|
||||
// a while. Stalls are handled by resume rather than by killing
|
||||
// the whole download.
|
||||
client: &http.Client{},
|
||||
baseURL: coreArtifactBaseURL(),
|
||||
stagingDir: stagingDir,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// compressedPath and unpackedPath are the two staging files. Both live
|
||||
// beside the dump importer's own staging data and are removed on success.
|
||||
func (f *artifactFetcher) compressedPath() string {
|
||||
return filepath.Join(f.stagingDir, coreArtifactFile)
|
||||
}
|
||||
|
||||
func (f *artifactFetcher) unpackedPath() string {
|
||||
return filepath.Join(f.stagingDir, "core-index.db")
|
||||
}
|
||||
|
||||
// fetch downloads, verifies and decompresses the artifact, returning the
|
||||
// path to the ready-to-merge database.
|
||||
func (f *artifactFetcher) fetch(ctx context.Context) (string, error) {
|
||||
if err := checkFreeDisk(f.stagingDir, artifactMinFreeBytes); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
want, err := f.fetchChecksum(ctx)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
if err := f.download(ctx); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
got, err := fileSHA256(f.compressedPath())
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
if got != want {
|
||||
// A partial file that resumed against a newer published artifact
|
||||
// would fail here forever; discard it so the next attempt starts
|
||||
// clean rather than re-resuming into the same mismatch.
|
||||
_ = os.Remove(f.compressedPath())
|
||||
|
||||
return "", fmt.Errorf("%w: checksum mismatch (got %s, want %s)",
|
||||
ErrArtifactUnusable, got, want)
|
||||
}
|
||||
|
||||
if err := f.decompress(ctx); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
// The compressed copy is dead weight once unpacked.
|
||||
_ = os.Remove(f.compressedPath())
|
||||
|
||||
return f.unpackedPath(), nil
|
||||
}
|
||||
|
||||
// fetchChecksum retrieves the expected sha256. This doubles as the
|
||||
// availability probe: it is a few bytes, so a missing or unreachable
|
||||
// artifact is discovered without starting a large download.
|
||||
func (f *artifactFetcher) fetchChecksum(ctx context.Context) (string, error) {
|
||||
probeCtx, cancel := context.WithTimeout(ctx, artifactDiscoverTimeout)
|
||||
defer cancel()
|
||||
|
||||
req, err := http.NewRequestWithContext(
|
||||
probeCtx, http.MethodGet, f.baseURL+coreArtifactChecksum, nil)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("%w: checksum request: %w", ErrArtifactUnavailable, err)
|
||||
}
|
||||
|
||||
req.Header.Set("User-Agent", lbUserAgent)
|
||||
|
||||
resp, err := f.client.Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("%w: %w", ErrArtifactUnavailable, err)
|
||||
}
|
||||
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return "", fmt.Errorf("%w: HTTP %d fetching checksum",
|
||||
ErrArtifactUnavailable, resp.StatusCode)
|
||||
}
|
||||
|
||||
// The file is `sha256sum` output: "<hex> <filename>".
|
||||
body, err := io.ReadAll(io.LimitReader(resp.Body, 4096))
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("%w: read checksum: %w", ErrArtifactUnavailable, err)
|
||||
}
|
||||
|
||||
sum := strings.TrimSpace(string(body))
|
||||
if i := strings.IndexAny(sum, " \t"); i > 0 {
|
||||
sum = sum[:i]
|
||||
}
|
||||
|
||||
if len(sum) != sha256.Size*2 {
|
||||
return "", fmt.Errorf("%w: malformed checksum %q", ErrArtifactUnusable, sum)
|
||||
}
|
||||
|
||||
return strings.ToLower(sum), nil
|
||||
}
|
||||
|
||||
// download fetches the artifact body, resuming a partial file with a
|
||||
// Range request rather than restarting it.
|
||||
func (f *artifactFetcher) download(ctx context.Context) error {
|
||||
var lastErr error
|
||||
|
||||
for attempt := range artifactMaxRetries {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if attempt > 0 {
|
||||
delay := min(streamRetryBaseDelay<<(attempt-1), streamRetryMaxDelay)
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-time.After(delay):
|
||||
}
|
||||
}
|
||||
|
||||
err := f.downloadOnce(ctx)
|
||||
if err == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
lastErr = err
|
||||
|
||||
f.si.logger.Warn("core artifact: download attempt failed",
|
||||
"attempt", attempt+1, "error", err,
|
||||
)
|
||||
}
|
||||
|
||||
return fmt.Errorf("%w: after %d attempts: %w",
|
||||
ErrArtifactUnavailable, artifactMaxRetries, lastErr)
|
||||
}
|
||||
|
||||
func (f *artifactFetcher) downloadOnce(ctx context.Context) error {
|
||||
// A file left by a previous attempt is resumed from its length.
|
||||
var have int64
|
||||
|
||||
if fi, err := os.Stat(f.compressedPath()); err == nil {
|
||||
have = fi.Size()
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(
|
||||
ctx, http.MethodGet, f.baseURL+coreArtifactFile, nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("artifact request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("User-Agent", lbUserAgent)
|
||||
|
||||
if have > 0 {
|
||||
req.Header.Set("Range", "bytes="+strconv.FormatInt(have, 10)+"-")
|
||||
}
|
||||
|
||||
resp, err := f.client.Do(req)
|
||||
if err != nil {
|
||||
return fmt.Errorf("artifact fetch: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
switch resp.StatusCode {
|
||||
case http.StatusPartialContent:
|
||||
// Resuming: append.
|
||||
case http.StatusOK:
|
||||
// The server ignored the Range header (or there was nothing to
|
||||
// resume), so this is a whole-file body and the partial must go.
|
||||
have = 0
|
||||
default:
|
||||
return fmt.Errorf("%w: HTTP %d fetching artifact",
|
||||
ErrArtifactUnavailable, resp.StatusCode)
|
||||
}
|
||||
|
||||
flags := os.O_CREATE | os.O_WRONLY
|
||||
if have > 0 {
|
||||
flags |= os.O_APPEND
|
||||
} else {
|
||||
flags |= os.O_TRUNC
|
||||
}
|
||||
|
||||
file, err := os.OpenFile(f.compressedPath(), flags, 0o644)
|
||||
if err != nil {
|
||||
return fmt.Errorf("open artifact file: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = file.Close() }()
|
||||
|
||||
total := have + resp.ContentLength
|
||||
|
||||
if _, err := io.Copy(file, f.progressReader(resp.Body, have, total)); err != nil {
|
||||
return fmt.Errorf("artifact download: %w", err)
|
||||
}
|
||||
|
||||
return file.Close()
|
||||
}
|
||||
|
||||
// progressReader wraps the body so download progress reaches the jobs
|
||||
// panel, since this is the one visible wait on a fresh install.
|
||||
func (f *artifactFetcher) progressReader(r io.Reader, done, total int64) io.Reader {
|
||||
return &artifactProgress{
|
||||
inner: r,
|
||||
done: done,
|
||||
total: total,
|
||||
si: f.si,
|
||||
last: time.Now(),
|
||||
}
|
||||
}
|
||||
|
||||
type artifactProgress struct {
|
||||
inner io.Reader
|
||||
done, total int64
|
||||
si *SearchIndex
|
||||
last time.Time
|
||||
}
|
||||
|
||||
func (p *artifactProgress) Read(b []byte) (int, error) {
|
||||
n, err := p.inner.Read(b)
|
||||
p.done += int64(n)
|
||||
|
||||
if time.Since(p.last) >= time.Second {
|
||||
p.last = time.Now()
|
||||
|
||||
detail := formatGB(p.done)
|
||||
if p.total > 0 {
|
||||
detail = fmt.Sprintf("%.0f%% of %s",
|
||||
100*float64(p.done)/float64(p.total), formatGB(p.total))
|
||||
}
|
||||
|
||||
// Reported in KiB so a multi-hundred-MB artifact cannot overflow
|
||||
// the int progress fields on a 32-bit build.
|
||||
p.si.setTierDetail(
|
||||
artifactStageNames[artifactStageDownload], "running",
|
||||
int(p.done>>10), int(p.total>>10), detail,
|
||||
)
|
||||
}
|
||||
|
||||
if err != nil && !errors.Is(err, io.EOF) {
|
||||
return n, fmt.Errorf("artifact body read: %w", err)
|
||||
}
|
||||
|
||||
return n, err //nolint:wrapcheck // io.EOF must reach the caller unwrapped.
|
||||
}
|
||||
|
||||
// decompress expands the zstd artifact into the staging directory.
|
||||
func (f *artifactFetcher) decompress(ctx context.Context) error {
|
||||
src, err := os.Open(f.compressedPath())
|
||||
if err != nil {
|
||||
return fmt.Errorf("%w: open compressed artifact: %w", ErrArtifactUnusable, err)
|
||||
}
|
||||
|
||||
defer func() { _ = src.Close() }()
|
||||
|
||||
zr, err := zstd.NewReader(src)
|
||||
if err != nil {
|
||||
return fmt.Errorf("%w: zstd reader: %w", ErrArtifactUnusable, err)
|
||||
}
|
||||
|
||||
defer zr.Close()
|
||||
|
||||
dst, err := os.Create(f.unpackedPath())
|
||||
if err != nil {
|
||||
return fmt.Errorf("%w: create artifact db: %w", ErrArtifactUnusable, err)
|
||||
}
|
||||
|
||||
defer func() { _ = dst.Close() }()
|
||||
|
||||
f.si.logIndexJob(jobs.LevelInfo, "Unpacking prebuilt catalog")
|
||||
|
||||
if _, err := io.Copy(dst, zr.IOReadCloser()); err != nil {
|
||||
// A half-written database would be rejected by inspectArtifact,
|
||||
// but leaving it around means the next run re-reads the same
|
||||
// wreckage before rejecting it.
|
||||
_ = os.Remove(f.unpackedPath())
|
||||
|
||||
return fmt.Errorf("%w: decompress: %w", ErrArtifactUnusable, err)
|
||||
}
|
||||
|
||||
if err := ctx.Err(); err != nil {
|
||||
_ = os.Remove(f.unpackedPath())
|
||||
|
||||
return err
|
||||
}
|
||||
|
||||
return dst.Close()
|
||||
}
|
||||
|
||||
// fileSHA256 hashes a file. The whole file is hashed after the download
|
||||
// completes rather than incrementally, because a resumed download never
|
||||
// sees the bytes it skipped.
|
||||
func fileSHA256(path string) (string, error) {
|
||||
file, err := os.Open(path)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("%w: open for checksum: %w", ErrArtifactUnusable, err)
|
||||
}
|
||||
|
||||
defer func() { _ = file.Close() }()
|
||||
|
||||
h := sha256.New()
|
||||
if _, err := io.Copy(h, file); err != nil {
|
||||
return "", fmt.Errorf("%w: checksum read: %w", ErrArtifactUnusable, err)
|
||||
}
|
||||
|
||||
return hex.EncodeToString(h.Sum(nil)), nil
|
||||
}
|
||||
@@ -0,0 +1,206 @@
|
||||
package explore
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/klauspost/compress/zstd"
|
||||
|
||||
"yellowjacket/backend/database"
|
||||
)
|
||||
|
||||
// artifactServer serves a compressed artifact and its checksum the way
|
||||
// the Gitea generic package registry does.
|
||||
func artifactServer(t *testing.T, body []byte) *httptest.Server {
|
||||
t.Helper()
|
||||
|
||||
sum := sha256.Sum256(body)
|
||||
checksum := hex.EncodeToString(sum[:]) + " " + coreArtifactFile + "\n"
|
||||
|
||||
mux := http.NewServeMux()
|
||||
|
||||
mux.HandleFunc("/"+coreArtifactChecksum,
|
||||
func(w http.ResponseWriter, _ *http.Request) {
|
||||
_, _ = w.Write([]byte(checksum))
|
||||
})
|
||||
|
||||
// ServeContent gives the stub real Range support, so the resume path
|
||||
// is exercised against the same semantics as the package registry.
|
||||
mux.HandleFunc("/"+coreArtifactFile,
|
||||
func(w http.ResponseWriter, r *http.Request) {
|
||||
http.ServeContent(
|
||||
w, r, coreArtifactFile, time.Time{}, bytes.NewReader(body))
|
||||
})
|
||||
|
||||
srv := httptest.NewServer(mux)
|
||||
t.Cleanup(srv.Close)
|
||||
|
||||
return srv
|
||||
}
|
||||
|
||||
// compressArtifact zstd-compresses a file the way CI publishes it.
|
||||
func compressArtifact(t *testing.T, path string) []byte {
|
||||
t.Helper()
|
||||
|
||||
raw, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatalf("read artifact: %v", err)
|
||||
}
|
||||
|
||||
enc, err := zstd.NewWriter(nil)
|
||||
if err != nil {
|
||||
t.Fatalf("zstd writer: %v", err)
|
||||
}
|
||||
|
||||
defer func() { _ = enc.Close() }()
|
||||
|
||||
return enc.EncodeAll(raw, nil)
|
||||
}
|
||||
|
||||
// withArtifactEnv points the fetcher at a stub server and gives it an
|
||||
// isolated data directory.
|
||||
func withArtifactEnv(t *testing.T, baseURL string) {
|
||||
t.Helper()
|
||||
|
||||
t.Setenv(artifactURLEnv, baseURL)
|
||||
t.Setenv("YJ_HOME", t.TempDir())
|
||||
}
|
||||
|
||||
func TestFetchAndMergeArtifactEndToEnd(t *testing.T) {
|
||||
src := writeTestArtifact(t, validMeta(), []artifactRow{
|
||||
{"artist", artA, "Artist A", "Artist A", artA, 5000},
|
||||
{"release_group", rgA, "Album A", "Artist A", artA, 3000},
|
||||
{"recording", recA, "Song A", "Artist A", artA, 2000},
|
||||
})
|
||||
|
||||
srv := artifactServer(t, compressArtifact(t, src))
|
||||
withArtifactEnv(t, srv.URL+"/")
|
||||
|
||||
db := database.NewTestDB(t)
|
||||
si := NewSearchIndex(db, nil, nil, testLogger())
|
||||
|
||||
if err := si.tryCoreArtifact(context.Background()); err != nil {
|
||||
t.Fatalf("tryCoreArtifact: %v", err)
|
||||
}
|
||||
|
||||
var got int
|
||||
if err := db.QueryRowWriter(
|
||||
"SELECT COUNT(*) FROM explore_index",
|
||||
).Scan(&got); err != nil {
|
||||
t.Fatalf("count rows: %v", err)
|
||||
}
|
||||
|
||||
if got != 3 {
|
||||
t.Errorf("merged %d rows, want 3", got)
|
||||
}
|
||||
|
||||
if !si.artifactAlreadyMerged() {
|
||||
t.Error("artifact merge not recorded; a restart would re-download it")
|
||||
}
|
||||
|
||||
// Staging must not keep a few hundred MB around after success.
|
||||
staging := filepath.Join(os.Getenv("YJ_HOME"), "data", "explore-staging")
|
||||
for _, name := range []string{coreArtifactFile, "core-index.db"} {
|
||||
if _, err := os.Stat(filepath.Join(staging, name)); err == nil {
|
||||
t.Errorf("%s left behind in staging after import", name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestFetchArtifactRejectsBadChecksum(t *testing.T) {
|
||||
src := writeTestArtifact(t, validMeta(), []artifactRow{
|
||||
{"artist", artA, "Artist A", "Artist A", artA, 5000},
|
||||
})
|
||||
|
||||
body := compressArtifact(t, src)
|
||||
|
||||
// Serve a checksum for different content.
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("/"+coreArtifactChecksum,
|
||||
func(w http.ResponseWriter, _ *http.Request) {
|
||||
_, _ = w.Write([]byte(strings.Repeat("a", 64) + " " + coreArtifactFile))
|
||||
})
|
||||
mux.HandleFunc("/"+coreArtifactFile,
|
||||
func(w http.ResponseWriter, _ *http.Request) {
|
||||
_, _ = w.Write(body)
|
||||
})
|
||||
|
||||
srv := httptest.NewServer(mux)
|
||||
defer srv.Close()
|
||||
|
||||
withArtifactEnv(t, srv.URL+"/")
|
||||
|
||||
db := database.NewTestDB(t)
|
||||
si := NewSearchIndex(db, nil, nil, testLogger())
|
||||
|
||||
err := si.tryCoreArtifact(context.Background())
|
||||
if err == nil {
|
||||
t.Fatal("expected checksum rejection")
|
||||
}
|
||||
|
||||
if !strings.Contains(err.Error(), "checksum mismatch") {
|
||||
t.Errorf("error = %v, want a checksum mismatch", err)
|
||||
}
|
||||
|
||||
// A corrupt download must not leave the index claiming a catalog.
|
||||
if si.hasMeta(dumpImportDoneKey) || si.artifactAlreadyMerged() {
|
||||
t.Error("failed download still marked the catalog as imported")
|
||||
}
|
||||
}
|
||||
|
||||
// A missing artifact is an ordinary outcome — the app has to keep
|
||||
// working before the first one is ever published.
|
||||
func TestFetchArtifactMissingIsUnavailable(t *testing.T) {
|
||||
srv := httptest.NewServer(http.NotFoundHandler())
|
||||
defer srv.Close()
|
||||
|
||||
withArtifactEnv(t, srv.URL+"/")
|
||||
|
||||
db := database.NewTestDB(t)
|
||||
si := NewSearchIndex(db, nil, nil, testLogger())
|
||||
|
||||
err := si.tryCoreArtifact(context.Background())
|
||||
if !errors.Is(err, ErrArtifactUnavailable) {
|
||||
t.Errorf("error = %v, want ErrArtifactUnavailable", err)
|
||||
}
|
||||
}
|
||||
|
||||
// The download resumes rather than restarting, which is what makes a
|
||||
// large artifact survive a flaky connection.
|
||||
func TestFetchArtifactResumesPartialDownload(t *testing.T) {
|
||||
src := writeTestArtifact(t, validMeta(), []artifactRow{
|
||||
{"artist", artA, "Artist A", "Artist A", artA, 5000},
|
||||
})
|
||||
|
||||
body := compressArtifact(t, src)
|
||||
srv := artifactServer(t, body)
|
||||
withArtifactEnv(t, srv.URL+"/")
|
||||
|
||||
db := database.NewTestDB(t)
|
||||
si := NewSearchIndex(db, nil, nil, testLogger())
|
||||
|
||||
fetcher, err := newArtifactFetcher(si)
|
||||
if err != nil {
|
||||
t.Fatalf("newArtifactFetcher: %v", err)
|
||||
}
|
||||
|
||||
// Pre-seed a truncated download.
|
||||
half := len(body) / 2
|
||||
if err := os.WriteFile(fetcher.compressedPath(), body[:half], 0o644); err != nil {
|
||||
t.Fatalf("seed partial download: %v", err)
|
||||
}
|
||||
|
||||
if _, err := fetcher.fetch(context.Background()); err != nil {
|
||||
t.Fatalf("fetch after partial: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,389 @@
|
||||
package explore
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
_ "modernc.org/sqlite" // SQLite driver for reading the artifact file.
|
||||
|
||||
"yellowjacket/backend/jobs"
|
||||
)
|
||||
|
||||
// Import of the prebuilt "core" index artifact.
|
||||
//
|
||||
// The catalog half of the index is identical for every user, so deriving
|
||||
// it on each machine means every install streams ~89GB from
|
||||
// data.metabrainz.org — a server that caps a client at roughly 2MB/s, so
|
||||
// better than half a day of downloading to reach a result everyone else
|
||||
// already has. Instead CI runs that import once (cmd/indexbuild),
|
||||
// exports a subset (cmd/indexexport), and clients merge the resulting
|
||||
// artifact in seconds.
|
||||
//
|
||||
// The artifact is an ordinary SQLite database holding two tables:
|
||||
// explore_index with the global catalog columns only, and artifact_meta
|
||||
// describing what it is. It deliberately carries no FTS table and no
|
||||
// triggers — rows land in the client's own explore_index, whose AFTER
|
||||
// INSERT trigger populates the search index as a side effect.
|
||||
//
|
||||
// Merging goes through the same ON CONFLICT rules as every other index
|
||||
// write (upsertIndexConflictSQL), so an artifact can be applied over an
|
||||
// existing index without clobbering better data: non-empty values win
|
||||
// over empty, higher listen counts win over lower, and the personal
|
||||
// columns the artifact does not carry are left untouched.
|
||||
|
||||
const (
|
||||
// coreArtifactVersionKey records which artifact version was merged,
|
||||
// so a client can tell whether it already has one and skip re-import.
|
||||
coreArtifactVersionKey = "core_artifact_version"
|
||||
|
||||
// supportedArtifactVersion is the artifact schema this build knows how
|
||||
// to read. The exporter stamps it into artifact_meta; a mismatch is
|
||||
// refused rather than guessed at, because an artifact written against
|
||||
// a different explore_index schema would merge wrong columns.
|
||||
supportedArtifactVersion = "1"
|
||||
)
|
||||
|
||||
// artifactMergeBatch bounds how many rows are merged per transaction.
|
||||
// Large enough that per-transaction overhead disappears, small enough
|
||||
// that a cancelled import doesn't roll back minutes of work. A var so
|
||||
// tests can shrink it and still cross several batch boundaries.
|
||||
var artifactMergeBatch = 50_000
|
||||
|
||||
var (
|
||||
// ErrArtifactUnusable means the file is not a core index artifact this
|
||||
// build can merge. Callers treat it as "fall back to a normal build"
|
||||
// rather than as a fatal error.
|
||||
ErrArtifactUnusable = errors.New("core index artifact unusable")
|
||||
|
||||
// ErrArtifactVersion is a version mismatch between the artifact and
|
||||
// this build.
|
||||
ErrArtifactVersion = errors.New("core index artifact version mismatch")
|
||||
)
|
||||
|
||||
// artifactInfo is what the artifact declares about itself.
|
||||
type artifactInfo struct {
|
||||
version string
|
||||
|
||||
// builtAt is when the source index finished importing, not when the
|
||||
// artifact was exported.
|
||||
builtAt string
|
||||
|
||||
// listensSeries is the incremental listens dump the artifact's
|
||||
// popularity numbers are baselined on. Stamped into the client's
|
||||
// index so RefreshListenCounts resumes from the right point instead
|
||||
// of reapplying deltas already folded in.
|
||||
listensSeries string
|
||||
|
||||
rows int
|
||||
}
|
||||
|
||||
// artifactCatalogColumns are the columns an artifact carries. It is the
|
||||
// global catalog only: the personal columns (in_library, is_similar,
|
||||
// local_*) describe one person's library and are recomputed locally by
|
||||
// PopulateLocalCrossReferences.
|
||||
//
|
||||
// Kept in sync with cmd/indexexport's catalogColumns by
|
||||
// TestArtifactColumnsMatchExporter.
|
||||
const artifactCatalogColumns = `entity_type, mbid, title, artist_name, artist_mbid,
|
||||
aliases, popularity, listener_count, duration, caa_release_mbid,
|
||||
release_name, primary_type, secondary_types, release_date,
|
||||
artist_type, country, disambiguation, sort_name, discog_fetched`
|
||||
|
||||
// inspectArtifact opens the artifact read-only and reports what it
|
||||
// declares, without touching the live index. Validation happens here so
|
||||
// a bad download is rejected before anything is attached.
|
||||
func inspectArtifact(path string) (artifactInfo, error) {
|
||||
var info artifactInfo
|
||||
|
||||
db, err := sql.Open("sqlite", "file:"+path+"?mode=ro")
|
||||
if err != nil {
|
||||
return info, fmt.Errorf("%w: open: %w", ErrArtifactUnusable, err)
|
||||
}
|
||||
|
||||
defer func() { _ = db.Close() }()
|
||||
|
||||
meta := map[string]string{}
|
||||
|
||||
rows, err := db.Query("SELECT key, value FROM artifact_meta")
|
||||
if err != nil {
|
||||
return info, fmt.Errorf("%w: read artifact_meta: %w", ErrArtifactUnusable, err)
|
||||
}
|
||||
|
||||
defer func() { _ = rows.Close() }()
|
||||
|
||||
for rows.Next() {
|
||||
var k, v string
|
||||
if err := rows.Scan(&k, &v); err != nil {
|
||||
return info, fmt.Errorf("%w: scan artifact_meta: %w", ErrArtifactUnusable, err)
|
||||
}
|
||||
|
||||
meta[k] = v
|
||||
}
|
||||
|
||||
if err := rows.Err(); err != nil {
|
||||
return info, fmt.Errorf("%w: read artifact_meta: %w", ErrArtifactUnusable, err)
|
||||
}
|
||||
|
||||
info.version = meta["artifact_version"]
|
||||
info.builtAt = meta["built_at"]
|
||||
info.listensSeries = meta["listens_applied_series"]
|
||||
|
||||
if info.version != supportedArtifactVersion {
|
||||
return info, fmt.Errorf("%w: artifact is version %q, this build reads %q",
|
||||
ErrArtifactVersion, info.version, supportedArtifactVersion)
|
||||
}
|
||||
|
||||
// A structurally valid but empty artifact would merge cleanly and
|
||||
// leave Explore just as empty as before, while stamping the index as
|
||||
// imported. Refuse it.
|
||||
if err := db.QueryRow(
|
||||
"SELECT COUNT(*) FROM explore_index",
|
||||
).Scan(&info.rows); err != nil {
|
||||
return info, fmt.Errorf("%w: count rows: %w", ErrArtifactUnusable, err)
|
||||
}
|
||||
|
||||
if info.rows == 0 {
|
||||
return info, fmt.Errorf("%w: artifact contains no rows", ErrArtifactUnusable)
|
||||
}
|
||||
|
||||
return info, nil
|
||||
}
|
||||
|
||||
// importCoreArtifact merges a validated artifact at path into the live
|
||||
// index. It is idempotent — the merge is an upsert keyed by MBID, so a
|
||||
// re-run over an already-imported artifact is a no-op in effect.
|
||||
//
|
||||
// The caller keeps ownership of the file; nothing here deletes it.
|
||||
func (si *SearchIndex) importCoreArtifact(ctx context.Context, path string) error {
|
||||
info, err := inspectArtifact(path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
si.logger.Info("core artifact: merging",
|
||||
"rows", info.rows,
|
||||
"builtAt", info.builtAt,
|
||||
"listensSeries", info.listensSeries,
|
||||
)
|
||||
si.logIndexJob(jobs.LevelInfo, fmt.Sprintf(
|
||||
"Merging prebuilt catalog (%s rows, built %s)",
|
||||
formatCount(info.rows), info.builtAt,
|
||||
))
|
||||
|
||||
// ATTACH cannot run inside a transaction, and it binds to a single
|
||||
// connection — which is why every statement below goes through the
|
||||
// writer (SetMaxOpenConns(1)). Reads must not use db.QueryContext:
|
||||
// that routes to the separate read pool, where "core" does not exist.
|
||||
if _, err := si.db.ExecContext(`ATTACH DATABASE ? AS core`, path); err != nil {
|
||||
return fmt.Errorf("%w: attach: %w", ErrArtifactUnusable, err)
|
||||
}
|
||||
|
||||
defer func() {
|
||||
if _, err := si.db.ExecContext(`DETACH DATABASE core`); err != nil {
|
||||
si.logger.Warn("core artifact: detach failed", "error", err)
|
||||
}
|
||||
}()
|
||||
|
||||
// Per-row FTS maintenance across a million inserts costs far more
|
||||
// than the inserts themselves (~31 rows/s against ~4,700), so the
|
||||
// search index is rebuilt once at the end instead.
|
||||
ftsSuspended := true
|
||||
|
||||
if err := si.db.SuspendExploreIndexFTS(); err != nil {
|
||||
si.logger.Warn("core artifact: could not suspend FTS sync", "error", err)
|
||||
|
||||
ftsSuspended = false
|
||||
}
|
||||
|
||||
merged, mergeErr := si.mergeArtifactRows(ctx, info.rows)
|
||||
|
||||
if ftsSuspended {
|
||||
start := time.Now()
|
||||
|
||||
if err := si.db.ResumeExploreIndexFTS(); err != nil {
|
||||
// Leaving search unindexed is worse than a slow import: this
|
||||
// needs a rebuild to recover, so it is loud.
|
||||
si.logger.Error("core artifact: FTS rebuild failed — search index is stale",
|
||||
"error", err,
|
||||
)
|
||||
} else {
|
||||
si.logger.Info("core artifact: FTS index rebuilt",
|
||||
"elapsed", time.Since(start).Round(time.Millisecond),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
if mergeErr != nil {
|
||||
return mergeErr
|
||||
}
|
||||
|
||||
si.stampArtifactMeta(info)
|
||||
si.analyzeIndex()
|
||||
|
||||
si.logger.Info("core artifact: merge complete", "rows", merged)
|
||||
si.logIndexJob(jobs.LevelInfo, fmt.Sprintf(
|
||||
"Prebuilt catalog merged (%s rows)", formatCount(merged),
|
||||
))
|
||||
|
||||
si.MarkReadyIfPopulated()
|
||||
si.refreshStatusCounts()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// analyzeIndex refreshes the query planner's table statistics.
|
||||
//
|
||||
// It runs here rather than at schema creation because an empty database
|
||||
// has nothing to measure: the numbers that matter only exist once the
|
||||
// catalog has been merged. Without them the planner mis-estimates the
|
||||
// partial expression indexes on explore_index and falls back to scanning
|
||||
// a million rows for queries that should seek.
|
||||
func (si *SearchIndex) analyzeIndex() {
|
||||
start := time.Now()
|
||||
|
||||
if _, err := si.db.ExecContext("ANALYZE"); err != nil {
|
||||
// Only a performance loss, so it must not fail the import.
|
||||
si.logger.Warn("core artifact: ANALYZE failed", "error", err)
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
si.logger.Info("core artifact: query planner statistics refreshed",
|
||||
"elapsed", time.Since(start).Round(time.Millisecond),
|
||||
)
|
||||
}
|
||||
|
||||
// mergeArtifactRows copies the attached artifact into explore_index in
|
||||
// bounded batches, walking the artifact's MBID primary key so each batch
|
||||
// is an index range scan and a cancelled import leaves committed work
|
||||
// behind rather than rolling it all back.
|
||||
func (si *SearchIndex) mergeArtifactRows(ctx context.Context, total int) (int, error) {
|
||||
insertSQL := `
|
||||
INSERT INTO explore_index (` + artifactCatalogColumns + `)
|
||||
SELECT ` + artifactCatalogColumns + `
|
||||
FROM core.explore_index
|
||||
WHERE mbid > ?` + upsertIndexConflictSQL
|
||||
|
||||
// The final batch has no upper bound, so the range predicate is
|
||||
// appended only while one exists.
|
||||
insertRangeSQL := `
|
||||
INSERT INTO explore_index (` + artifactCatalogColumns + `)
|
||||
SELECT ` + artifactCatalogColumns + `
|
||||
FROM core.explore_index
|
||||
WHERE mbid > ? AND mbid <= ?` + upsertIndexConflictSQL
|
||||
|
||||
var (
|
||||
cursor string
|
||||
merged int
|
||||
)
|
||||
|
||||
for {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return merged, err
|
||||
}
|
||||
|
||||
upper, hasUpper, err := si.artifactBatchBound(cursor)
|
||||
if err != nil {
|
||||
return merged, err
|
||||
}
|
||||
|
||||
var res sql.Result
|
||||
|
||||
if hasUpper {
|
||||
res, err = si.db.ExecContext(insertRangeSQL, cursor, upper)
|
||||
} else {
|
||||
res, err = si.db.ExecContext(insertSQL, cursor)
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
return merged, fmt.Errorf("%w: merge batch: %w", ErrArtifactUnusable, err)
|
||||
}
|
||||
|
||||
n, err := res.RowsAffected()
|
||||
if err != nil {
|
||||
return merged, fmt.Errorf("%w: merge batch rows: %w", ErrArtifactUnusable, err)
|
||||
}
|
||||
|
||||
merged += int(n)
|
||||
|
||||
si.setTierDetail(
|
||||
artifactStageNames[artifactStageMerge], "running", merged, total,
|
||||
fmt.Sprintf("%s of %s rows", formatCount(merged), formatCount(total)),
|
||||
)
|
||||
|
||||
if !hasUpper {
|
||||
return merged, nil
|
||||
}
|
||||
|
||||
cursor = upper
|
||||
}
|
||||
}
|
||||
|
||||
// artifactBatchBound returns the MBID that ends the next batch, and
|
||||
// whether one exists — no bound means the remainder is the last batch.
|
||||
func (si *SearchIndex) artifactBatchBound(cursor string) (string, bool, error) {
|
||||
var bound string
|
||||
|
||||
err := si.db.QueryRowWriter(
|
||||
`SELECT mbid FROM core.explore_index
|
||||
WHERE mbid > ? ORDER BY mbid LIMIT 1 OFFSET ?`,
|
||||
cursor, artifactMergeBatch-1,
|
||||
).Scan(&bound)
|
||||
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return "", false, nil
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
return "", false, fmt.Errorf("%w: batch bound: %w", ErrArtifactUnusable, err)
|
||||
}
|
||||
|
||||
return bound, true, nil
|
||||
}
|
||||
|
||||
// stampArtifactMeta records what the merge established: the catalog half
|
||||
// is populated, and popularity is baselined on the artifact's listens
|
||||
// series so the incremental refresh resumes from there.
|
||||
func (si *SearchIndex) stampArtifactMeta(info artifactInfo) {
|
||||
si.setMeta(coreArtifactVersionKey, info.version)
|
||||
|
||||
// The catalog is present, so nothing should trigger a full dump
|
||||
// import on top of the artifact it was meant to replace.
|
||||
si.setMeta(dumpImportDoneKey, time.Now().UTC().Format(time.RFC3339))
|
||||
|
||||
// Without a baseline series RefreshListenCounts refuses to run at
|
||||
// all, so an artifact exported before that key existed leaves the
|
||||
// index permanently frozen at its shipped popularity. Better to say
|
||||
// so than to fail silently.
|
||||
if info.listensSeries == "" {
|
||||
si.logger.Warn(
|
||||
"core artifact: no listens series recorded — " +
|
||||
"popularity refresh will not run until the next full import",
|
||||
)
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
if _, err := strconv.Atoi(info.listensSeries); err != nil {
|
||||
si.logger.Warn("core artifact: unparseable listens series",
|
||||
"value", info.listensSeries,
|
||||
)
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
si.setMeta(listensAppliedSeriesKey, info.listensSeries)
|
||||
}
|
||||
|
||||
// removeArtifactFile deletes a merged artifact. Best-effort: a leftover
|
||||
// file costs disk, not correctness.
|
||||
func (si *SearchIndex) removeArtifactFile(path string) {
|
||||
if err := os.Remove(path); err != nil && !os.IsNotExist(err) {
|
||||
si.logger.Warn("core artifact: cleanup failed", "path", path, "error", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,444 @@
|
||||
package explore
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"yellowjacket/backend/database"
|
||||
)
|
||||
|
||||
// artifactRow is one catalog row written into a test artifact.
|
||||
type artifactRow struct {
|
||||
entityType string
|
||||
mbid string
|
||||
title string
|
||||
artistName string
|
||||
artistMBID string
|
||||
popularity int
|
||||
}
|
||||
|
||||
// writeTestArtifact builds an artifact file matching what cmd/indexexport
|
||||
// produces: catalog columns only, no FTS, no triggers.
|
||||
func writeTestArtifact(
|
||||
t *testing.T, meta map[string]string, rows []artifactRow,
|
||||
) string {
|
||||
t.Helper()
|
||||
|
||||
path := filepath.Join(t.TempDir(), "core-index.db")
|
||||
|
||||
db, err := sql.Open("sqlite", "file:"+path)
|
||||
if err != nil {
|
||||
t.Fatalf("open artifact: %v", err)
|
||||
}
|
||||
|
||||
defer func() { _ = db.Close() }()
|
||||
|
||||
for _, stmt := range []string{
|
||||
`CREATE TABLE explore_index (
|
||||
entity_type TEXT NOT NULL,
|
||||
mbid TEXT NOT NULL,
|
||||
title TEXT NOT NULL,
|
||||
artist_name TEXT NOT NULL,
|
||||
artist_mbid TEXT NOT NULL,
|
||||
aliases TEXT NOT NULL DEFAULT '',
|
||||
popularity INTEGER NOT NULL DEFAULT 0,
|
||||
listener_count INTEGER NOT NULL DEFAULT 0,
|
||||
duration INTEGER NOT NULL DEFAULT 0,
|
||||
caa_release_mbid TEXT NOT NULL DEFAULT '',
|
||||
release_name TEXT NOT NULL DEFAULT '',
|
||||
primary_type TEXT NOT NULL DEFAULT '',
|
||||
secondary_types TEXT NOT NULL DEFAULT '',
|
||||
release_date TEXT NOT NULL DEFAULT '',
|
||||
artist_type TEXT NOT NULL DEFAULT '',
|
||||
country TEXT NOT NULL DEFAULT '',
|
||||
disambiguation TEXT NOT NULL DEFAULT '',
|
||||
sort_name TEXT NOT NULL DEFAULT '',
|
||||
discog_fetched INTEGER NOT NULL DEFAULT 0,
|
||||
PRIMARY KEY (mbid)
|
||||
) WITHOUT ROWID`,
|
||||
`CREATE TABLE artifact_meta (
|
||||
key TEXT PRIMARY KEY,
|
||||
value TEXT NOT NULL
|
||||
)`,
|
||||
} {
|
||||
if _, err := db.Exec(stmt); err != nil {
|
||||
t.Fatalf("create artifact schema: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
for k, v := range meta {
|
||||
if _, err := db.Exec(
|
||||
`INSERT INTO artifact_meta (key, value) VALUES (?, ?)`, k, v,
|
||||
); err != nil {
|
||||
t.Fatalf("stamp artifact meta: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
for _, r := range rows {
|
||||
if _, err := db.Exec(`
|
||||
INSERT INTO explore_index
|
||||
(entity_type, mbid, title, artist_name, artist_mbid, popularity)
|
||||
VALUES (?, ?, ?, ?, ?, ?)`,
|
||||
r.entityType, r.mbid, r.title, r.artistName, r.artistMBID, r.popularity,
|
||||
); err != nil {
|
||||
t.Fatalf("insert artifact row: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
return path
|
||||
}
|
||||
|
||||
// validMeta is the artifact_meta a well-formed artifact carries.
|
||||
func validMeta() map[string]string {
|
||||
return map[string]string{
|
||||
"artifact_version": supportedArtifactVersion,
|
||||
"built_at": "2026-07-17T03:53:43Z",
|
||||
"listens_applied_series": "2593",
|
||||
"source_rows": "2052168",
|
||||
}
|
||||
}
|
||||
|
||||
func TestImportCoreArtifactMergesCatalog(t *testing.T) {
|
||||
db := database.NewTestDB(t)
|
||||
si := NewSearchIndex(db, nil, nil, testLogger())
|
||||
|
||||
path := writeTestArtifact(t, validMeta(), []artifactRow{
|
||||
{"artist", artA, "Artist A", "Artist A", artA, 5000},
|
||||
{"artist", artB, "Artist B", "Artist B", artB, 4000},
|
||||
{"release_group", rgA, "Album A", "Artist A", artA, 3000},
|
||||
{"recording", recA, "Song A", "Artist A", artA, 2000},
|
||||
})
|
||||
|
||||
if err := si.importCoreArtifact(context.Background(), path); err != nil {
|
||||
t.Fatalf("importCoreArtifact: %v", err)
|
||||
}
|
||||
|
||||
var got int
|
||||
if err := db.QueryRowWriter(
|
||||
"SELECT COUNT(*) FROM explore_index",
|
||||
).Scan(&got); err != nil {
|
||||
t.Fatalf("count rows: %v", err)
|
||||
}
|
||||
|
||||
if got != 4 {
|
||||
t.Errorf("merged %d rows, want 4", got)
|
||||
}
|
||||
|
||||
// The merge must leave the index searchable: the FTS rebuild that
|
||||
// closes the bulk-load window is the only thing populating it, since
|
||||
// the triggers were dropped for the duration.
|
||||
var hits int
|
||||
if err := db.QueryRowWriter(
|
||||
`SELECT COUNT(*) FROM explore_index_fts WHERE explore_index_fts MATCH ?`,
|
||||
"Song",
|
||||
).Scan(&hits); err != nil {
|
||||
t.Fatalf("query fts: %v", err)
|
||||
}
|
||||
|
||||
if hits == 0 {
|
||||
t.Error("FTS index is empty after merge; search would return nothing")
|
||||
}
|
||||
}
|
||||
|
||||
func TestImportCoreArtifactStampsMeta(t *testing.T) {
|
||||
db := database.NewTestDB(t)
|
||||
si := NewSearchIndex(db, nil, nil, testLogger())
|
||||
|
||||
path := writeTestArtifact(t, validMeta(), []artifactRow{
|
||||
{"artist", artA, "Artist A", "Artist A", artA, 5000},
|
||||
})
|
||||
|
||||
if err := si.importCoreArtifact(context.Background(), path); err != nil {
|
||||
t.Fatalf("importCoreArtifact: %v", err)
|
||||
}
|
||||
|
||||
if !si.hasMeta(dumpImportDoneKey) {
|
||||
t.Error("dump_import_done not stamped; a full dump import would run anyway")
|
||||
}
|
||||
|
||||
// Without this the incremental refresh refuses to run and the
|
||||
// shipped popularity never updates again.
|
||||
if series, ok := si.metaInt(listensAppliedSeriesKey); !ok || series != 2593 {
|
||||
t.Errorf("listens_applied_series = %d (ok=%v), want 2593", series, ok)
|
||||
}
|
||||
|
||||
if !si.hasMeta(coreArtifactVersionKey) {
|
||||
t.Error("core_artifact_version not stamped")
|
||||
}
|
||||
}
|
||||
|
||||
// A merge must never downgrade what the index already holds, because a
|
||||
// user's own library rows and any lazily-fetched detail predate it.
|
||||
func TestImportCoreArtifactPreservesLocalData(t *testing.T) {
|
||||
db := database.NewTestDB(t)
|
||||
si := NewSearchIndex(db, nil, nil, testLogger())
|
||||
|
||||
si.upsertBatch([]SearchIndexResult{{
|
||||
EntityType: "recording",
|
||||
MBID: recA,
|
||||
Title: "Song A",
|
||||
ArtistName: "Artist A",
|
||||
ArtistMBID: artA,
|
||||
Popularity: 9999,
|
||||
Duration: 210000,
|
||||
InLibrary: true,
|
||||
DiscogFetched: true,
|
||||
}})
|
||||
|
||||
path := writeTestArtifact(t, validMeta(), []artifactRow{
|
||||
// Lower popularity and no duration: both must lose.
|
||||
{"recording", recA, "Song A", "Artist A", artA, 10},
|
||||
})
|
||||
|
||||
if err := si.importCoreArtifact(context.Background(), path); err != nil {
|
||||
t.Fatalf("importCoreArtifact: %v", err)
|
||||
}
|
||||
|
||||
var popularity, duration, inLibrary, discogFetched int
|
||||
|
||||
if err := db.QueryRowWriter(`
|
||||
SELECT popularity, duration, in_library, discog_fetched
|
||||
FROM explore_index WHERE mbid = ?`, recA,
|
||||
).Scan(&popularity, &duration, &inLibrary, &discogFetched); err != nil {
|
||||
t.Fatalf("read merged row: %v", err)
|
||||
}
|
||||
|
||||
if popularity != 9999 {
|
||||
t.Errorf("popularity = %d, want 9999 (higher must win)", popularity)
|
||||
}
|
||||
|
||||
if duration != 210000 {
|
||||
t.Errorf("duration = %d, want 210000 (artifact carries none)", duration)
|
||||
}
|
||||
|
||||
if inLibrary != 1 {
|
||||
t.Error("in_library was cleared; the artifact must not touch personal columns")
|
||||
}
|
||||
|
||||
if discogFetched != 1 {
|
||||
t.Error("discog_fetched was cleared by the merge")
|
||||
}
|
||||
}
|
||||
|
||||
// The batch walk is the part most likely to drop or duplicate rows, so
|
||||
// it is exercised across many batch boundaries rather than one.
|
||||
func TestImportCoreArtifactBatchWalkCoversAllRows(t *testing.T) {
|
||||
db := database.NewTestDB(t)
|
||||
si := NewSearchIndex(db, nil, nil, testLogger())
|
||||
|
||||
// Shrink the batch rather than growing the fixture: what matters is
|
||||
// crossing several boundaries, including a short final one.
|
||||
original := artifactMergeBatch
|
||||
artifactMergeBatch = 100
|
||||
|
||||
t.Cleanup(func() { artifactMergeBatch = original })
|
||||
|
||||
const total = 337
|
||||
|
||||
rows := make([]artifactRow, 0, total)
|
||||
for i := range total {
|
||||
rows = append(rows, artifactRow{
|
||||
entityType: "recording",
|
||||
mbid: syntheticMBID(i),
|
||||
title: "Song",
|
||||
artistName: "Artist",
|
||||
artistMBID: artA,
|
||||
popularity: i,
|
||||
})
|
||||
}
|
||||
|
||||
path := writeTestArtifact(t, validMeta(), rows)
|
||||
|
||||
if err := si.importCoreArtifact(context.Background(), path); err != nil {
|
||||
t.Fatalf("importCoreArtifact: %v", err)
|
||||
}
|
||||
|
||||
var got int
|
||||
if err := db.QueryRowWriter(
|
||||
"SELECT COUNT(*) FROM explore_index",
|
||||
).Scan(&got); err != nil {
|
||||
t.Fatalf("count rows: %v", err)
|
||||
}
|
||||
|
||||
if got != total {
|
||||
t.Errorf("merged %d rows, want %d", got, total)
|
||||
}
|
||||
}
|
||||
|
||||
func TestImportCoreArtifactRejectsBadArtifacts(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
meta map[string]string
|
||||
rows []artifactRow
|
||||
want error
|
||||
}{
|
||||
{
|
||||
name: "version mismatch",
|
||||
meta: map[string]string{"artifact_version": "99"},
|
||||
rows: []artifactRow{{"artist", artA, "A", "A", artA, 1}},
|
||||
want: ErrArtifactVersion,
|
||||
},
|
||||
{
|
||||
name: "no version",
|
||||
meta: map[string]string{},
|
||||
rows: []artifactRow{{"artist", artA, "A", "A", artA, 1}},
|
||||
want: ErrArtifactVersion,
|
||||
},
|
||||
{
|
||||
name: "empty catalog",
|
||||
meta: validMeta(),
|
||||
rows: nil,
|
||||
want: ErrArtifactUnusable,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
db := database.NewTestDB(t)
|
||||
si := NewSearchIndex(db, nil, nil, testLogger())
|
||||
|
||||
path := writeTestArtifact(t, tt.meta, tt.rows)
|
||||
|
||||
err := si.importCoreArtifact(context.Background(), path)
|
||||
if err == nil {
|
||||
t.Fatal("expected rejection, got nil")
|
||||
}
|
||||
|
||||
if !strings.Contains(err.Error(), tt.want.Error()) {
|
||||
t.Errorf("error = %v, want it to wrap %v", err, tt.want)
|
||||
}
|
||||
|
||||
// A rejected artifact must not leave the index claiming it
|
||||
// has a catalog, or the real build would never run.
|
||||
if si.hasMeta(dumpImportDoneKey) {
|
||||
t.Error("rejected artifact still stamped dump_import_done")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestInspectArtifactRejectsNonArtifactFile(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "garbage.db")
|
||||
|
||||
if err := os.WriteFile(path, []byte("not a database"), 0o644); err != nil {
|
||||
t.Fatalf("write file: %v", err)
|
||||
}
|
||||
|
||||
if _, err := inspectArtifact(path); err == nil {
|
||||
t.Error("expected rejection of a non-artifact file")
|
||||
}
|
||||
}
|
||||
|
||||
// syntheticMBID produces distinct, well-formed MBIDs for bulk fixtures.
|
||||
func syntheticMBID(i int) string {
|
||||
const hex = "0123456789abcdef"
|
||||
|
||||
buf := []byte("00000000-0000-0000-0000-000000000000")
|
||||
|
||||
for pos := len(buf) - 1; pos >= 0 && i > 0; pos-- {
|
||||
if buf[pos] == '-' {
|
||||
continue
|
||||
}
|
||||
|
||||
buf[pos] = hex[i%16]
|
||||
i /= 16
|
||||
}
|
||||
|
||||
return string(buf)
|
||||
}
|
||||
|
||||
// resolveArtistName falls back to the MBID when no name is available.
|
||||
// That value must never reach the index: the upsert no longer defends
|
||||
// against it, so the writer is the only thing standing in the way.
|
||||
func TestAddFromCacheNeverStoresMBIDAsName(t *testing.T) {
|
||||
db := database.NewTestDB(t)
|
||||
si := NewSearchIndex(db, nil, nil, testLogger())
|
||||
|
||||
si.AddFromCache(artA, artA, []MBReleaseGroup{
|
||||
{MBID: rgA, Title: "Album A", PrimaryType: "Album"},
|
||||
})
|
||||
|
||||
var artistRows int
|
||||
if err := db.QueryRowWriter(
|
||||
`SELECT COUNT(*) FROM explore_index
|
||||
WHERE entity_type = 'artist' AND title = mbid`,
|
||||
).Scan(&artistRows); err != nil {
|
||||
t.Fatalf("count artist rows: %v", err)
|
||||
}
|
||||
|
||||
if artistRows != 0 {
|
||||
t.Errorf("%d artist rows stored the MBID as their title", artistRows)
|
||||
}
|
||||
|
||||
var artistName string
|
||||
if err := db.QueryRowWriter(
|
||||
`SELECT artist_name FROM explore_index WHERE mbid = ?`, rgA,
|
||||
).Scan(&artistName); err != nil {
|
||||
t.Fatalf("read release group: %v", err)
|
||||
}
|
||||
|
||||
if artistName == artA {
|
||||
t.Error("release group stored the artist MBID as its artist_name")
|
||||
}
|
||||
|
||||
// A real name arriving later must still win over the empty one.
|
||||
si.AddFromCache("Real Artist", artA, []MBReleaseGroup{
|
||||
{MBID: rgA, Title: "Album A", PrimaryType: "Album"},
|
||||
})
|
||||
|
||||
if err := db.QueryRowWriter(
|
||||
`SELECT artist_name FROM explore_index WHERE mbid = ?`, rgA,
|
||||
).Scan(&artistName); err != nil {
|
||||
t.Fatalf("re-read release group: %v", err)
|
||||
}
|
||||
|
||||
if artistName != "Real Artist" {
|
||||
t.Errorf("artist_name = %q, want it filled in once known", artistName)
|
||||
}
|
||||
}
|
||||
|
||||
// The importer names the artifact's columns independently of the
|
||||
// exporter that writes them. If the two lists drift, the merge either
|
||||
// fails outright or silently shifts values into the wrong columns, so
|
||||
// they are compared directly against cmd/indexexport's source.
|
||||
func TestArtifactColumnsMatchExporter(t *testing.T) {
|
||||
src, err := os.ReadFile("../../cmd/indexexport/main.go")
|
||||
if err != nil {
|
||||
t.Fatalf("read exporter: %v", err)
|
||||
}
|
||||
|
||||
const marker = "const catalogColumns = `"
|
||||
|
||||
i := strings.Index(string(src), marker)
|
||||
if i < 0 {
|
||||
t.Fatalf("catalogColumns not found in cmd/indexexport/main.go")
|
||||
}
|
||||
|
||||
rest := string(src)[i+len(marker):]
|
||||
|
||||
j := strings.Index(rest, "`")
|
||||
if j < 0 {
|
||||
t.Fatal("unterminated catalogColumns literal")
|
||||
}
|
||||
|
||||
normalise := func(s string) string {
|
||||
out := make([]string, 0, 32)
|
||||
for _, f := range strings.Split(s, ",") {
|
||||
out = append(out, strings.Join(strings.Fields(f), ""))
|
||||
}
|
||||
|
||||
return strings.Join(out, ",")
|
||||
}
|
||||
|
||||
exporter := normalise(rest[:j])
|
||||
importer := normalise(artifactCatalogColumns)
|
||||
|
||||
if exporter != importer {
|
||||
t.Errorf("column lists have drifted:\n exporter: %s\n importer: %s",
|
||||
exporter, importer)
|
||||
}
|
||||
}
|
||||
@@ -36,7 +36,7 @@ const (
|
||||
artistImageTimeout = 10 * time.Second
|
||||
artistImageCacheTTL = 365 * 24 * time.Hour // positive results: ~permanent
|
||||
artistImageMissCacheTTL = 30 * 24 * time.Hour // negative results: retry monthly
|
||||
artistImageBaseDir = "artist-images"
|
||||
artistImageBaseDir = ArtistImageDirName
|
||||
artistImageMaxBytes = 2 * 1024 * 1024
|
||||
artistImageMaxSize = 500 // max dimension for stored full-res images
|
||||
maxImagesPerArtist = 10
|
||||
|
||||
@@ -0,0 +1,20 @@
|
||||
package explore
|
||||
|
||||
// Asset directory names under the user data directory.
|
||||
//
|
||||
// These are exported so the maintenance janitor can sweep them without
|
||||
// importing the internals of the providers that write them. Each is
|
||||
// catalogued in backend/datamap as cache data: rebuildable from the
|
||||
// network, but expensive enough that eviction is by age rather than
|
||||
// tied to the lifetime of anything else.
|
||||
const (
|
||||
// ArtistImageDirName holds one subdirectory per artist MBID, each
|
||||
// containing fetched photos, a primary.jpg with its thumbnails, and
|
||||
// possibly a .miss marker recording that no artwork was found.
|
||||
ArtistImageDirName = "artist-images"
|
||||
|
||||
// CoverArtCacheDirName holds Cover Art Archive thumbnails fetched
|
||||
// while browsing Explore, named by release-group MBID. Nothing in
|
||||
// the database references these files.
|
||||
CoverArtCacheDirName = "cover-art-cache"
|
||||
)
|
||||
@@ -23,7 +23,7 @@ var ErrCoverArt = errors.New("cover art fetch failed")
|
||||
const (
|
||||
// thumbnailDir is the subdirectory under the user data dir
|
||||
// where cached cover art thumbnails are stored.
|
||||
thumbnailDir = "cover-art-cache"
|
||||
thumbnailDir = CoverArtCacheDirName
|
||||
|
||||
// thumbnailTimeout is the HTTP timeout for fetching a thumbnail.
|
||||
thumbnailTimeout = 10 * time.Second
|
||||
|
||||
@@ -0,0 +1,73 @@
|
||||
//go:build !indexbuild
|
||||
|
||||
package explore
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
|
||||
"yellowjacket/backend/jobs"
|
||||
)
|
||||
|
||||
// The app binary does not carry the full dump import.
|
||||
//
|
||||
// Building the catalog from source means streaming ~89GB of listens dump
|
||||
// from a server that caps a client near 2MB/s — better than half a day
|
||||
// of downloading to derive a catalog that is identical for everyone.
|
||||
// That work happens once in CI (`go build -tags indexbuild ./cmd/indexbuild`)
|
||||
// and reaches users as a prebuilt artifact instead.
|
||||
//
|
||||
// So on the client the build path is: merge the artifact, then keep
|
||||
// popularity current with the daily incremental dumps. Artists outside
|
||||
// the artifact's coverage still resolve lazily, exactly as before.
|
||||
|
||||
// runDumpBuild populates the catalog. In the app build that means the
|
||||
// prebuilt artifact and nothing else; the tagged build in
|
||||
// dumpimport.go runs the real import.
|
||||
func (si *SearchIndex) runDumpBuild(ctx context.Context) {
|
||||
si.MarkReadyIfPopulated()
|
||||
|
||||
if si.hasMeta(dumpImportDoneKey) {
|
||||
si.logger.Info("search index: catalog already populated, skipping")
|
||||
si.refreshStatusCounts()
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
if si.artifactAlreadyMerged() {
|
||||
si.refreshStatusCounts()
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
if err := si.tryCoreArtifact(ctx); err != nil {
|
||||
if ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
|
||||
si.logArtifactFallback(err)
|
||||
si.refreshStatusCounts()
|
||||
}
|
||||
}
|
||||
|
||||
// logArtifactFallback explains why the catalog is not there. An empty
|
||||
// Explore with nothing in the log is the worst version of this failure.
|
||||
func (si *SearchIndex) logArtifactFallback(err error) {
|
||||
switch {
|
||||
case errors.Is(err, ErrArtifactUnavailable):
|
||||
si.logger.Info("search index: no prebuilt catalog available", "error", err)
|
||||
si.logIndexJob(jobs.LevelWarn,
|
||||
"No prebuilt catalog available — Explore will cover your own "+
|
||||
"library only until one can be fetched.")
|
||||
|
||||
case errors.Is(err, ErrArtifactVersion):
|
||||
si.logger.Warn("search index: prebuilt catalog is for a different app version",
|
||||
"error", err)
|
||||
si.logIndexJob(jobs.LevelWarn,
|
||||
"The published catalog does not match this app version; skipping it.")
|
||||
|
||||
default:
|
||||
si.logger.Warn("search index: prebuilt catalog import failed", "error", err)
|
||||
si.logIndexJob(jobs.LevelWarn, "Prebuilt catalog import failed: "+err.Error())
|
||||
}
|
||||
}
|
||||
@@ -1,3 +1,5 @@
|
||||
//go:build indexbuild
|
||||
|
||||
package explore
|
||||
|
||||
import (
|
||||
@@ -440,7 +442,7 @@ func (a *artistTopRG) add(rg uuid16, listens uint32) {
|
||||
func (imp *dumpImporter) scanCanonicalDump(
|
||||
ctx context.Context, url string, ks *keptSets,
|
||||
) (*canonicalScan, error) {
|
||||
stream := newResumableReader(ctx, imp.httpClient, url, 0)
|
||||
stream := imp.openDumpStream(ctx, url, 0)
|
||||
|
||||
defer func() { _ = stream.Close() }()
|
||||
|
||||
|
||||
+306
-215
@@ -1,9 +1,10 @@
|
||||
//go:build indexbuild
|
||||
|
||||
package explore
|
||||
|
||||
import (
|
||||
"archive/tar"
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/binary"
|
||||
"encoding/json"
|
||||
@@ -13,8 +14,8 @@ import (
|
||||
"os"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/parquet-go/parquet-go"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Stage 1 of the dump import: stream the ListenBrainz spark listens
|
||||
@@ -27,9 +28,6 @@ import (
|
||||
|
||||
const (
|
||||
// countKindRecording etc. tag entries in the counts map/file.
|
||||
countKindRecording = byte(1)
|
||||
countKindRelease = byte(2)
|
||||
countKindArtist = byte(3)
|
||||
|
||||
// countsFlushEveryMembers controls checkpoint frequency. Each
|
||||
// flush rewrites counts.bin (~1GB by the end), so this trades
|
||||
@@ -37,105 +35,35 @@ const (
|
||||
// 19GB of stream progress).
|
||||
countsFlushEveryMembers = 150
|
||||
|
||||
// countsProgressEveryMembers controls progress log frequency.
|
||||
countsProgressEveryMembers = 50
|
||||
// countsUIRefreshInterval is how often the live download line is
|
||||
// pushed to the UI. A parquet member is ~128MB, so member
|
||||
// boundaries are minutes apart on a typical connection — sampling
|
||||
// the stream position instead keeps the stage visibly moving.
|
||||
countsUIRefreshInterval = 3 * time.Second
|
||||
|
||||
// countsLogInterval and countsJobLogInterval throttle the two log
|
||||
// surfaces: the app log gets a line every few minutes, the jobs
|
||||
// panel a coarser one. Checkpoints always log to both.
|
||||
countsLogInterval = 2 * time.Minute
|
||||
countsJobLogInterval = 15 * time.Minute
|
||||
|
||||
// countsStallAfter is how long the stream position may stand still
|
||||
// before progress is reported as stalled rather than as a rate.
|
||||
countsStallAfter = 45 * time.Second
|
||||
|
||||
// countsRateSmoothing is the EWMA weight given to the newest
|
||||
// throughput sample, trading responsiveness against jitter.
|
||||
countsRateSmoothing = 0.25
|
||||
|
||||
// parquetParseWorkers is the number of concurrent parquet
|
||||
// decoders. Bounded to limit RAM: each worker holds one
|
||||
// ~128MB member buffer.
|
||||
parquetParseWorkers = 3
|
||||
|
||||
// maxParquetMemberSize guards against unexpected dump format
|
||||
// changes blowing out RAM.
|
||||
maxParquetMemberSize = 1 << 30
|
||||
|
||||
// countsFileMagic identifies + versions the counts file format.
|
||||
countsFileMagic = "YJCNTS01"
|
||||
)
|
||||
|
||||
// ErrDumpFormat is returned when dump contents don't match the
|
||||
// expected format.
|
||||
var ErrDumpFormat = errors.New("unexpected dump format")
|
||||
|
||||
// mbidKey is a parsed UUID plus an entity-kind tag, used as the counts
|
||||
// map key. 17 bytes instead of a 36-byte string keeps the ~40M-entry
|
||||
// map around 2GB.
|
||||
type mbidKey [17]byte
|
||||
|
||||
func makeMBIDKey(kind byte, mbid string) (mbidKey, bool) {
|
||||
var k mbidKey
|
||||
|
||||
k[0] = kind
|
||||
|
||||
if !parseUUID(mbid, k[1:]) {
|
||||
return k, false
|
||||
}
|
||||
|
||||
return k, true
|
||||
}
|
||||
|
||||
// parseUUID parses a canonical 36-char UUID string into 16 bytes.
|
||||
// Returns false for anything malformed.
|
||||
func parseUUID(s string, out []byte) bool {
|
||||
if len(s) != 36 || s[8] != '-' || s[13] != '-' || s[18] != '-' || s[23] != '-' {
|
||||
return false
|
||||
}
|
||||
|
||||
j := 0
|
||||
|
||||
for i := 0; i < 36; i++ {
|
||||
if i == 8 || i == 13 || i == 18 || i == 23 {
|
||||
continue
|
||||
}
|
||||
|
||||
hi := hexNibble(s[i])
|
||||
i++
|
||||
|
||||
lo := hexNibble(s[i])
|
||||
if hi == 0xFF || lo == 0xFF {
|
||||
return false
|
||||
}
|
||||
|
||||
out[j] = hi<<4 | lo
|
||||
j++
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
func hexNibble(c byte) byte {
|
||||
switch {
|
||||
case c >= '0' && c <= '9':
|
||||
return c - '0'
|
||||
case c >= 'a' && c <= 'f':
|
||||
return c - 'a' + 10
|
||||
case c >= 'A' && c <= 'F':
|
||||
return c - 'A' + 10
|
||||
default:
|
||||
return 0xFF
|
||||
}
|
||||
}
|
||||
|
||||
func formatUUID(b []byte) string {
|
||||
const hexdigits = "0123456789abcdef"
|
||||
|
||||
out := make([]byte, 36)
|
||||
j := 0
|
||||
|
||||
for i := range 16 {
|
||||
if i == 4 || i == 6 || i == 8 || i == 10 {
|
||||
out[j] = '-'
|
||||
j++
|
||||
}
|
||||
|
||||
out[j] = hexdigits[b[i]>>4]
|
||||
out[j+1] = hexdigits[b[i]&0x0F]
|
||||
j += 2
|
||||
}
|
||||
|
||||
return string(out)
|
||||
}
|
||||
|
||||
// countsState is the checkpointed stage-1 state: the counts map plus
|
||||
// the stream position it corresponds to.
|
||||
type countsState struct {
|
||||
@@ -157,14 +85,6 @@ type countsState struct {
|
||||
counts map[mbidKey]uint32
|
||||
}
|
||||
|
||||
// sparkListenRow is the projection of the spark listens parquet schema
|
||||
// that the aggregator reads. All other columns are skipped.
|
||||
type sparkListenRow struct {
|
||||
RecordingMBID string `parquet:"recording_mbid,optional"`
|
||||
ReleaseMBID string `parquet:"release_mbid,optional"`
|
||||
ArtistMBIDs []string `parquet:"artist_credit_mbids,optional,list"`
|
||||
}
|
||||
|
||||
type countParseJob struct {
|
||||
idx int
|
||||
endOffset int64 // exact offset of the next member header
|
||||
@@ -180,15 +100,47 @@ type countParseResult struct {
|
||||
|
||||
// aggregateListenCounts runs stage 1 to completion (or ctx cancel),
|
||||
// checkpointing to the staging counts file as it goes.
|
||||
//
|
||||
// Column projection is tried first: it downloads only the three MBID
|
||||
// columns the aggregator reads, which is well under half the archive.
|
||||
// It needs a Range-serving origin, so a server that won't range falls
|
||||
// back to streaming the whole tar.
|
||||
func (imp *dumpImporter) aggregateListenCounts(ctx context.Context, st *countsState) error {
|
||||
if st.counts == nil {
|
||||
st.counts = make(map[mbidKey]uint32, 1<<20)
|
||||
}
|
||||
|
||||
stream := newResumableReader(ctx, imp.httpClient, st.SparkURL, st.Offset)
|
||||
if size, ok := projectionSupported(ctx, imp.httpClient, st.SparkURL); ok {
|
||||
err := imp.aggregateProjected(ctx, st, size)
|
||||
if !errors.Is(err, errProjectionUnsupported) {
|
||||
return err
|
||||
}
|
||||
|
||||
imp.logger.Warn("dump import: column projection unavailable, streaming whole dump",
|
||||
"error", err,
|
||||
)
|
||||
}
|
||||
|
||||
return imp.aggregateStreamed(ctx, st)
|
||||
}
|
||||
|
||||
// aggregateStreamed is the fallback stage-1 path: read the tar end to
|
||||
// end and parse every parquet member in full.
|
||||
func (imp *dumpImporter) aggregateStreamed(ctx context.Context, st *countsState) error {
|
||||
stream := imp.openDumpStream(ctx, st.SparkURL, st.Offset)
|
||||
|
||||
defer func() { _ = stream.Close() }()
|
||||
|
||||
// A live reporter samples the stream position on a timer; without
|
||||
// it the stage would sit unchanged for minutes at a time between
|
||||
// parquet members, which reads as "hung" rather than "downloading".
|
||||
var awaitingWorkers atomic.Bool
|
||||
|
||||
stopReporter := imp.startCountsReporter(ctx, stream, &awaitingWorkers)
|
||||
defer stopReporter()
|
||||
|
||||
progress := &countsLogger{imp: imp, stream: stream, started: time.Now()}
|
||||
|
||||
buffered := bufio.NewReaderSize(stream, 1<<20)
|
||||
tr := tar.NewReader(buffered)
|
||||
|
||||
@@ -196,7 +148,7 @@ func (imp *dumpImporter) aggregateListenCounts(ctx context.Context, st *countsSt
|
||||
// tar reader has consumed: bytes delivered by HTTP minus bytes
|
||||
// still sitting in the bufio buffer.
|
||||
consumedOffset := func() int64 {
|
||||
return stream.Offset - int64(buffered.Buffered())
|
||||
return stream.Pos() - int64(buffered.Buffered())
|
||||
}
|
||||
|
||||
jobs := make(chan countParseJob)
|
||||
@@ -231,64 +183,13 @@ func (imp *dumpImporter) aggregateListenCounts(ctx context.Context, st *countsSt
|
||||
// checkpoint is a contiguous prefix of the stream. It owns
|
||||
// st.counts, st.Offset, and st.MemberIdx until applierDone closes;
|
||||
// on error it keeps draining results so nothing deadlocks.
|
||||
var applyErr error
|
||||
applier := newCountsApplier(imp, st, progress)
|
||||
|
||||
go func() {
|
||||
defer close(applierDone)
|
||||
|
||||
pending := make(map[int]countParseResult)
|
||||
next := st.MemberIdx
|
||||
lastFlushed := st.MemberIdx
|
||||
|
||||
for res := range results {
|
||||
if applyErr != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
pending[res.idx] = res
|
||||
|
||||
for {
|
||||
r, ok := pending[next]
|
||||
if !ok {
|
||||
break
|
||||
}
|
||||
|
||||
delete(pending, next)
|
||||
|
||||
if r.err != nil {
|
||||
applyErr = r.err
|
||||
|
||||
break
|
||||
}
|
||||
|
||||
for k, v := range r.deltas {
|
||||
st.counts[k] += v
|
||||
}
|
||||
|
||||
next++
|
||||
st.MemberIdx = next
|
||||
st.Offset = r.endOffset
|
||||
|
||||
if next-lastFlushed >= countsFlushEveryMembers {
|
||||
if err := imp.writeCountsFile(st); err != nil {
|
||||
applyErr = err
|
||||
|
||||
break
|
||||
}
|
||||
|
||||
lastFlushed = next
|
||||
|
||||
imp.logCountsProgress(next, r.endOffset, stream.Size, len(st.counts))
|
||||
|
||||
if err := imp.checkDiskHeadroom(); err != nil {
|
||||
applyErr = err
|
||||
|
||||
break
|
||||
}
|
||||
} else if next%countsProgressEveryMembers == 0 {
|
||||
imp.logCountsProgress(next, r.endOffset, stream.Size, len(st.counts))
|
||||
}
|
||||
}
|
||||
applier.apply(res, nil)
|
||||
}
|
||||
}()
|
||||
|
||||
@@ -342,9 +243,14 @@ readLoop:
|
||||
// uses PAX extension headers (those start at its header offset).
|
||||
endOffset := consumedOffset() + tarPadding(hdr.Size)
|
||||
|
||||
awaitingWorkers.Store(true)
|
||||
|
||||
select {
|
||||
case jobs <- countParseJob{idx: memberIdx, endOffset: endOffset, buf: buf}:
|
||||
awaitingWorkers.Store(false)
|
||||
case <-ctx.Done():
|
||||
awaitingWorkers.Store(false)
|
||||
|
||||
readErr = ctx.Err()
|
||||
|
||||
break readLoop
|
||||
@@ -359,7 +265,7 @@ readLoop:
|
||||
<-applierDone
|
||||
|
||||
if readErr == nil {
|
||||
readErr = applyErr
|
||||
readErr = applier.err
|
||||
}
|
||||
|
||||
if readErr != nil {
|
||||
@@ -378,9 +284,16 @@ readLoop:
|
||||
|
||||
imp.logger.Info("dump import: listen counts complete",
|
||||
"members", st.MemberIdx,
|
||||
"gb", fmt.Sprintf("%.1f", float64(st.Offset)/(1<<30)),
|
||||
"entities", len(st.counts),
|
||||
"elapsed", time.Since(progress.started).Truncate(time.Second).String(),
|
||||
)
|
||||
|
||||
imp.logJob(fmt.Sprintf(
|
||||
"Listen counts complete — %s of listens read, %s entities ranked",
|
||||
formatGB(st.Offset), formatCount(len(st.counts)),
|
||||
))
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -392,55 +305,6 @@ func tarPadding(size int64) int64 {
|
||||
return (block - size%block) % block
|
||||
}
|
||||
|
||||
// parseListenParquet decodes one parquet member and returns the
|
||||
// per-entity listen-count deltas.
|
||||
func parseListenParquet(buf []byte) (map[mbidKey]uint32, error) {
|
||||
reader := parquet.NewGenericReader[sparkListenRow](bytes.NewReader(buf))
|
||||
|
||||
defer func() { _ = reader.Close() }()
|
||||
|
||||
deltas := make(map[mbidKey]uint32, 1<<18)
|
||||
rows := make([]sparkListenRow, 4096)
|
||||
|
||||
for {
|
||||
n, err := reader.Read(rows)
|
||||
|
||||
for _, row := range rows[:n] {
|
||||
key, ok := makeMBIDKey(countKindRecording, row.RecordingMBID)
|
||||
if !ok {
|
||||
// Unmapped listen — no usable recording MBID.
|
||||
continue
|
||||
}
|
||||
|
||||
deltas[key]++
|
||||
|
||||
if relKey, relOK := makeMBIDKey(countKindRelease, row.ReleaseMBID); relOK {
|
||||
deltas[relKey]++
|
||||
}
|
||||
|
||||
for _, artist := range row.ArtistMBIDs {
|
||||
if artKey, artOK := makeMBIDKey(countKindArtist, artist); artOK {
|
||||
deltas[artKey]++
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if errors.Is(err, io.EOF) {
|
||||
break
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("parquet read: %w", err)
|
||||
}
|
||||
|
||||
if n == 0 {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
return deltas, nil
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// counts.bin persistence
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -567,18 +431,245 @@ func (imp *dumpImporter) readCountsFile() (*countsState, error) {
|
||||
return st, nil
|
||||
}
|
||||
|
||||
func (imp *dumpImporter) logCountsProgress(members int, offset, size int64, entities int) {
|
||||
pct := float64(0)
|
||||
if size > 0 {
|
||||
pct = float64(offset) / float64(size) * 100
|
||||
// ---------------------------------------------------------------------------
|
||||
// progress reporting
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// startCountsReporter runs a goroutine that samples the listens stream
|
||||
// position every few seconds and publishes it as the stage's UI
|
||||
// progress. The returned function stops the reporter and waits for it
|
||||
// to exit, so no stale "running" update can land after the stage is
|
||||
// marked complete.
|
||||
func (imp *dumpImporter) startCountsReporter(
|
||||
ctx context.Context, stream dumpStream, backlog *atomic.Bool,
|
||||
) func() {
|
||||
stop := make(chan struct{})
|
||||
exited := make(chan struct{})
|
||||
|
||||
rep := &countsReporter{
|
||||
imp: imp,
|
||||
stream: stream,
|
||||
backlog: backlog,
|
||||
lastSample: time.Now(),
|
||||
lastOffset: stream.Fetched(),
|
||||
lastMoved: time.Now(),
|
||||
}
|
||||
|
||||
imp.logger.Info("dump import: listen counts progress",
|
||||
go func() {
|
||||
defer close(exited)
|
||||
|
||||
ticker := time.NewTicker(countsUIRefreshInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-stop:
|
||||
return
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case now := <-ticker.C:
|
||||
rep.tick(now)
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
return func() {
|
||||
close(stop)
|
||||
<-exited
|
||||
}
|
||||
}
|
||||
|
||||
// countsReporter turns stream position samples into a percentage and a
|
||||
// human-readable throughput line. Only its own goroutine touches it.
|
||||
type countsReporter struct {
|
||||
imp *dumpImporter
|
||||
stream dumpStream
|
||||
|
||||
// backlog is set while the reader is blocked handing a member to
|
||||
// the parquet workers. The stream stops moving then too, and
|
||||
// calling that a network stall would be wrong.
|
||||
backlog *atomic.Bool
|
||||
|
||||
lastSample time.Time
|
||||
lastOffset int64
|
||||
lastMoved time.Time
|
||||
rate float64 // EWMA bytes/sec
|
||||
}
|
||||
|
||||
func (rep *countsReporter) tick(now time.Time) {
|
||||
// Track the downloader, not the consumer: parallel lanes buffer a
|
||||
// chunk at a time, so delivery stands still for a minute at the
|
||||
// start of a stream while the network is in fact saturated. Watching
|
||||
// Pos here would report that as a stall.
|
||||
offset := rep.stream.Fetched()
|
||||
size := rep.stream.Total()
|
||||
|
||||
if elapsed := now.Sub(rep.lastSample).Seconds(); elapsed > 0 {
|
||||
sample := float64(offset-rep.lastOffset) / elapsed
|
||||
if rep.rate == 0 {
|
||||
rep.rate = sample
|
||||
} else {
|
||||
rep.rate = countsRateSmoothing*sample + (1-countsRateSmoothing)*rep.rate
|
||||
}
|
||||
}
|
||||
|
||||
if offset != rep.lastOffset {
|
||||
rep.lastMoved = now
|
||||
}
|
||||
|
||||
rep.lastSample = now
|
||||
rep.lastOffset = offset
|
||||
|
||||
// A stream that has stopped moving is either reconnecting or held
|
||||
// up by the parsers; either way, reporting a rate that is really
|
||||
// just an average of nothing would be misleading.
|
||||
var detail string
|
||||
|
||||
switch {
|
||||
case now.Sub(rep.lastMoved) <= countsStallAfter:
|
||||
rep.imp.countsRate.Store(uint64(max(rep.rate, 0)))
|
||||
|
||||
detail = formatStreamProgress(offset, size, rep.rate)
|
||||
case rep.backlog != nil && rep.backlog.Load():
|
||||
rep.imp.countsRate.Store(0)
|
||||
|
||||
detail = formatGB(offset) + " downloaded · parsing, download paused"
|
||||
default:
|
||||
rep.imp.countsRate.Store(0)
|
||||
|
||||
detail = formatGB(offset) + " downloaded · stalled, retrying…"
|
||||
}
|
||||
|
||||
rep.imp.setStageDetail(dumpStageCounts, streamPercent(offset, size), 100, detail)
|
||||
}
|
||||
|
||||
// countsLogger writes stage-1 progress to the app log and the jobs
|
||||
// panel on independent time-based schedules. Per-member lines would be
|
||||
// too sparse to reassure and too noisy to read; checkpoints, which are
|
||||
// the points a crash would resume from, always log to both.
|
||||
type countsLogger struct {
|
||||
imp *dumpImporter
|
||||
stream dumpStream
|
||||
|
||||
started time.Time
|
||||
lastLog time.Time
|
||||
lastJobLog time.Time
|
||||
}
|
||||
|
||||
// member reports an applied parquet member, logging only if enough time
|
||||
// has passed since the last line.
|
||||
func (l *countsLogger) member(members int, offset int64, entities int) {
|
||||
now := time.Now()
|
||||
|
||||
if now.Sub(l.lastLog) >= countsLogInterval {
|
||||
l.lastLog = now
|
||||
|
||||
l.logApp("dump import: listen counts progress", members, offset, entities)
|
||||
}
|
||||
|
||||
if now.Sub(l.lastJobLog) >= countsJobLogInterval {
|
||||
l.lastJobLog = now
|
||||
|
||||
l.logJob("Listen counts", members, offset, entities)
|
||||
}
|
||||
}
|
||||
|
||||
// checkpoint reports a counts.bin flush, which always logs — it is the
|
||||
// point an interrupted import would resume from.
|
||||
func (l *countsLogger) checkpoint(members int, offset int64, entities int) {
|
||||
now := time.Now()
|
||||
|
||||
l.lastLog = now
|
||||
l.lastJobLog = now
|
||||
|
||||
l.logApp("dump import: listen counts checkpoint", members, offset, entities)
|
||||
l.logJob("Listen counts checkpointed", members, offset, entities)
|
||||
}
|
||||
|
||||
func (l *countsLogger) logApp(msg string, members int, offset int64, entities int) {
|
||||
size := l.stream.Total()
|
||||
|
||||
l.imp.logger.Info(msg,
|
||||
"members", members,
|
||||
"gb", fmt.Sprintf("%.1f", float64(offset)/(1<<30)),
|
||||
"pct", fmt.Sprintf("%.1f", pct),
|
||||
"pct", streamPercent(offset, size),
|
||||
"rate", formatRate(l.imp.streamRate()),
|
||||
"eta", formatETA(offset, size, l.imp.streamRate()),
|
||||
"entities", entities,
|
||||
)
|
||||
}
|
||||
|
||||
imp.setStageProgress(dumpStageCounts, int(pct), 100)
|
||||
func (l *countsLogger) logJob(prefix string, members int, offset int64, entities int) {
|
||||
l.imp.logJob(fmt.Sprintf("%s: %s · %s members · %s entities",
|
||||
prefix,
|
||||
formatStreamProgress(offset, l.stream.Total(), l.imp.streamRate()),
|
||||
formatCount(members),
|
||||
formatCount(entities),
|
||||
))
|
||||
}
|
||||
|
||||
// streamRate returns the listens stream throughput most recently
|
||||
// measured by the reporter, in bytes/sec.
|
||||
func (imp *dumpImporter) streamRate() float64 {
|
||||
return float64(imp.countsRate.Load())
|
||||
}
|
||||
|
||||
// streamPercent is the whole-percent position in a stream of known
|
||||
// size; 0 when the size is not yet known.
|
||||
func streamPercent(offset, size int64) int {
|
||||
if size <= 0 {
|
||||
return 0
|
||||
}
|
||||
|
||||
return int(float64(offset) / float64(size) * 100)
|
||||
}
|
||||
|
||||
// formatStreamProgress renders "42.3 / 205.1 GB (20%) · 18 MB/s ·
|
||||
// ~3h20m left", degrading gracefully when the size or rate is unknown.
|
||||
func formatStreamProgress(offset, size int64, rate float64) string {
|
||||
parts := make([]string, 0, 3)
|
||||
|
||||
if size > 0 {
|
||||
parts = append(parts, fmt.Sprintf("%s / %s (%d%%)",
|
||||
formatGB(offset), formatGB(size), streamPercent(offset, size)))
|
||||
} else {
|
||||
parts = append(parts, formatGB(offset)+" downloaded")
|
||||
}
|
||||
|
||||
if rate > 0 {
|
||||
parts = append(parts, formatRate(rate))
|
||||
}
|
||||
|
||||
if eta := formatETA(offset, size, rate); eta != "" {
|
||||
parts = append(parts, "~"+eta+" left")
|
||||
}
|
||||
|
||||
return strings.Join(parts, " · ")
|
||||
}
|
||||
|
||||
func formatRate(bytesPerSec float64) string {
|
||||
if bytesPerSec <= 0 {
|
||||
return "—"
|
||||
}
|
||||
|
||||
return fmt.Sprintf("%.1f MB/s", bytesPerSec/(1<<20))
|
||||
}
|
||||
|
||||
// formatETA estimates remaining time at the current rate. Returns ""
|
||||
// when the total size or the rate is unknown.
|
||||
func formatETA(offset, size int64, rate float64) string {
|
||||
if size <= 0 || rate <= 0 || offset >= size {
|
||||
return ""
|
||||
}
|
||||
|
||||
remaining := time.Duration(float64(size-offset)/rate) * time.Second
|
||||
|
||||
switch {
|
||||
case remaining < time.Minute:
|
||||
return "<1m"
|
||||
case remaining < time.Hour:
|
||||
return fmt.Sprintf("%dm", int(remaining.Minutes()))
|
||||
default:
|
||||
return fmt.Sprintf("%dh%02dm", int(remaining.Hours()), int(remaining.Minutes())%60)
|
||||
}
|
||||
}
|
||||
|
||||
+113
-71
@@ -1,25 +1,33 @@
|
||||
//go:build indexbuild
|
||||
|
||||
package explore
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"os"
|
||||
"path"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"yellowjacket/backend/database"
|
||||
"yellowjacket/backend/jobs"
|
||||
"yellowjacket/backend/system"
|
||||
)
|
||||
|
||||
// Dump-based index population. Instead of crawling the ListenBrainz
|
||||
// API artist-by-artist, the index is built from two MetaBrainz dumps:
|
||||
//
|
||||
// 1. The spark listens dump (~170GB, streamed, never stored) yields
|
||||
// 1. The spark listens dump (~205GB on the server, of which only the
|
||||
// three MBID columns are downloaded — see dumpproject.go) yields
|
||||
// listen counts for every recording/release/artist MBID.
|
||||
// 2. The MusicBrainz canonical dump (~2GB, streamed) yields names and
|
||||
// MBIDs, filtered to entities above a popularity floor.
|
||||
@@ -30,14 +38,6 @@ import (
|
||||
// resumes from checkpoints after interruption.
|
||||
|
||||
const (
|
||||
// dumpImportDoneKey marks a completed import in explore_index_meta.
|
||||
dumpImportDoneKey = "dump_import_done"
|
||||
|
||||
// listensAppliedSeriesKey stores the dump series number whose listen
|
||||
// counts are folded into popularity (the high-water-mark for the
|
||||
// incremental refresh). Set to the full dump's series at import, then
|
||||
// advanced by each applied incremental.
|
||||
listensAppliedSeriesKey = "listens_applied_series"
|
||||
|
||||
// releaseToRGInsertBatch bounds how many rows are written per
|
||||
// transaction when persisting the release→release-group map.
|
||||
@@ -56,9 +56,6 @@ const (
|
||||
dumpStageAssembled = "assembled"
|
||||
)
|
||||
|
||||
// ErrDiskSpace is returned when free disk falls below the safety floor.
|
||||
var ErrDiskSpace = errors.New("insufficient free disk space")
|
||||
|
||||
// Dump import stages, mapped to status names shown in the UI.
|
||||
const (
|
||||
dumpStageCounts = iota
|
||||
@@ -115,6 +112,14 @@ type dumpImporter struct {
|
||||
// pendingArtists are kept artists whose names weren't derivable
|
||||
// from the canonical dump; the metadata patch pass resolves them.
|
||||
pendingArtists []string
|
||||
|
||||
// countsRate is the listens stream throughput in bytes/sec, written
|
||||
// by the stage-1 progress reporter and read by its loggers.
|
||||
countsRate atomic.Uint64
|
||||
|
||||
// ftsResumed guards the bulk-load window so the FTS rebuild runs
|
||||
// exactly once per import.
|
||||
ftsResumed bool
|
||||
}
|
||||
|
||||
func newDumpImporter(si *SearchIndex, lb *ListenBrainzClient) (*dumpImporter, error) {
|
||||
@@ -132,10 +137,10 @@ func newDumpImporter(si *SearchIndex, lb *ListenBrainzClient) (*dumpImporter, er
|
||||
si: si,
|
||||
lb: lb,
|
||||
logger: si.logger,
|
||||
// No client-level timeout: the listens stream runs for hours.
|
||||
// Discovery requests use per-request context timeouts, and
|
||||
// resumableReader recovers from stalled connections.
|
||||
httpClient: &http.Client{},
|
||||
// HTTP/1.1-only, no client-level timeout: the listens stream
|
||||
// runs for hours. Discovery requests use per-request context
|
||||
// timeouts, and the stream readers recover from stalls.
|
||||
httpClient: newDumpHTTPClient(),
|
||||
stagingDir: stagingDir,
|
||||
canonicalBaseURL: defaultCanonicalBaseURL,
|
||||
listensBaseURL: defaultListensBaseURL,
|
||||
@@ -192,6 +197,8 @@ func (imp *dumpImporter) run(ctx context.Context) error {
|
||||
}
|
||||
|
||||
imp.logger.Info("dump import: starting", "listensDump", sparkURL)
|
||||
imp.logJob("Streaming listens dump " + path.Base(sparkURL) +
|
||||
" — this stage reads tens of gigabytes and runs for hours")
|
||||
|
||||
counts = &countsState{SparkURL: sparkURL}
|
||||
} else if !counts.Done {
|
||||
@@ -200,6 +207,10 @@ func (imp *dumpImporter) run(ctx context.Context) error {
|
||||
"members", counts.MemberIdx,
|
||||
"entities", len(counts.counts),
|
||||
)
|
||||
imp.logJob(fmt.Sprintf(
|
||||
"Resuming listen counts from %s (%s members, %s entities so far)",
|
||||
formatGB(counts.Offset), formatCount(counts.MemberIdx), formatCount(len(counts.counts)),
|
||||
))
|
||||
}
|
||||
|
||||
// Record which dump series this import is baselined on, so the
|
||||
@@ -251,15 +262,15 @@ func (imp *dumpImporter) run(ctx context.Context) error {
|
||||
return err
|
||||
}
|
||||
|
||||
// One-time reset when migrating from the legacy API-crawled
|
||||
// index: its popularity values are on a different scale (the LB
|
||||
// API includes MLHD+ history) and would permanently outrank
|
||||
// dump-derived counts via the highest-wins upsert. Re-imports
|
||||
// (dump→dump) skip this — listen counts only grow.
|
||||
if !imp.si.hasMeta(dumpImportDoneKey) {
|
||||
if _, err := imp.si.db.ExecContext("DELETE FROM explore_index"); err == nil {
|
||||
imp.logger.Info("dump import: cleared legacy index for consistent popularity scale")
|
||||
}
|
||||
// Bulk-load window: the wipe below and the millions of upserts that
|
||||
// follow would otherwise each maintain the FTS5 index row by row,
|
||||
// which dominates the entire import (~31 rows/s vs ~4,700). The
|
||||
// index is rebuilt in one pass when the window closes.
|
||||
if err := imp.si.db.SuspendExploreIndexFTS(); err != nil {
|
||||
// Not fatal: the import still completes, just slowly.
|
||||
imp.logger.Warn("dump import: could not suspend FTS sync", "error", err)
|
||||
} else {
|
||||
defer imp.resumeFTS()
|
||||
}
|
||||
|
||||
if err := imp.assembleIndex(ctx, kept, scan); err != nil {
|
||||
@@ -272,6 +283,11 @@ func (imp *dumpImporter) run(ctx context.Context) error {
|
||||
// just built from.
|
||||
imp.persistReleaseToRG(ctx, scan.releaseToRG)
|
||||
|
||||
// Close the bulk-load window now rather than at return: the patch
|
||||
// passes below run at API rate, so per-row FTS upkeep costs nothing
|
||||
// there and keeps search current while they work.
|
||||
imp.resumeFTS()
|
||||
|
||||
imp.si.setTierStatus(dumpStageNames[dumpStageCatalog], "complete", 0, 0)
|
||||
|
||||
// Artists that need names from the metadata patch pass.
|
||||
@@ -299,16 +315,38 @@ func (imp *dumpImporter) run(ctx context.Context) error {
|
||||
return imp.finalize()
|
||||
}
|
||||
|
||||
// resumeFTS closes the bulk-load window, restoring the FTS sync
|
||||
// triggers and rebuilding the index. Idempotent, so it can run both at
|
||||
// its natural point in the pipeline and from a defer covering the
|
||||
// error and cancellation paths.
|
||||
func (imp *dumpImporter) resumeFTS() {
|
||||
if imp.ftsResumed {
|
||||
return
|
||||
}
|
||||
|
||||
imp.ftsResumed = true
|
||||
|
||||
start := time.Now()
|
||||
|
||||
if err := imp.si.db.ResumeExploreIndexFTS(); err != nil {
|
||||
// Leaving search unindexed is worse than a slow import, so this
|
||||
// is loud: it needs a rebuild to recover.
|
||||
imp.logger.Error("dump import: FTS rebuild failed — search index is stale",
|
||||
"error", err,
|
||||
)
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
imp.logger.Info("dump import: FTS index rebuilt",
|
||||
"elapsed", time.Since(start).Round(time.Millisecond),
|
||||
)
|
||||
}
|
||||
|
||||
// finalize records completion and removes all staging data.
|
||||
func (imp *dumpImporter) finalize() error {
|
||||
imp.si.setMeta(dumpImportDoneKey, time.Now().UTC().Format(time.RFC3339))
|
||||
|
||||
// Retire the legacy tier-crawl freshness keys.
|
||||
_, _ = imp.si.db.ExecContext(
|
||||
`DELETE FROM explore_index_meta
|
||||
WHERE key IN ('tier1_built', 'tier2_built', 'tier3_built', 'tier4_built')`,
|
||||
)
|
||||
|
||||
if err := os.RemoveAll(imp.stagingDir); err != nil {
|
||||
imp.logger.Warn("dump import: staging cleanup failed", "error", err)
|
||||
}
|
||||
@@ -326,25 +364,6 @@ func (imp *dumpImporter) finalize() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// dumpSeriesRe extracts the monotonic series number NNNN from a dump
|
||||
// URL or directory name (e.g. "listenbrainz-spark-dump-2593-…").
|
||||
var dumpSeriesRe = regexp.MustCompile(`listenbrainz-(?:spark-)?dump-(\d+)-`)
|
||||
|
||||
// parseDumpSeries pulls the series number out of a dump URL/name.
|
||||
func parseDumpSeries(url string) (int, bool) {
|
||||
m := dumpSeriesRe.FindStringSubmatch(url)
|
||||
if m == nil {
|
||||
return 0, false
|
||||
}
|
||||
|
||||
n, err := strconv.Atoi(m[1])
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
|
||||
return n, true
|
||||
}
|
||||
|
||||
// recordDumpSeries stores the baseline series number for this import.
|
||||
func (imp *dumpImporter) recordDumpSeries(sparkURL string) {
|
||||
series, ok := parseDumpSeries(sparkURL)
|
||||
@@ -374,7 +393,10 @@ func (imp *dumpImporter) persistReleaseToRG(ctx context.Context, m map[uuid16]rg
|
||||
written := 0
|
||||
pending := 0
|
||||
|
||||
tx, err := imp.si.db.BeginTx()
|
||||
// The statement is prepared per transaction rather than passed to
|
||||
// tx.Exec per row: re-parsing it for each of several million rows
|
||||
// costs an order of magnitude more than the insert itself.
|
||||
tx, stmt, err := beginReleaseToRGTx(imp.si.db)
|
||||
if err != nil {
|
||||
imp.logger.Warn("dump import: begin release_to_rg tx failed", "error", err)
|
||||
|
||||
@@ -383,13 +405,13 @@ func (imp *dumpImporter) persistReleaseToRG(ctx context.Context, m map[uuid16]rg
|
||||
|
||||
for rel, target := range m {
|
||||
if ctx.Err() != nil {
|
||||
_ = stmt.Close()
|
||||
_ = tx.Rollback()
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
if _, err := tx.Exec(
|
||||
"INSERT OR REPLACE INTO release_to_rg (release_mbid, rg_mbid) VALUES (?, ?)",
|
||||
if _, err := stmt.Exec(
|
||||
formatUUID(rel[:]), formatUUID(target.rg[:]),
|
||||
); err != nil {
|
||||
imp.logger.Warn("dump import: insert release_to_rg failed", "error", err)
|
||||
@@ -401,6 +423,8 @@ func (imp *dumpImporter) persistReleaseToRG(ctx context.Context, m map[uuid16]rg
|
||||
pending++
|
||||
|
||||
if pending >= releaseToRGInsertBatch {
|
||||
_ = stmt.Close()
|
||||
|
||||
if err := tx.Commit(); err != nil {
|
||||
imp.logger.Warn("dump import: commit release_to_rg batch failed", "error", err)
|
||||
|
||||
@@ -409,7 +433,7 @@ func (imp *dumpImporter) persistReleaseToRG(ctx context.Context, m map[uuid16]rg
|
||||
|
||||
pending = 0
|
||||
|
||||
tx, err = imp.si.db.BeginTx()
|
||||
tx, stmt, err = beginReleaseToRGTx(imp.si.db)
|
||||
if err != nil {
|
||||
imp.logger.Warn("dump import: begin release_to_rg tx failed", "error", err)
|
||||
|
||||
@@ -418,6 +442,8 @@ func (imp *dumpImporter) persistReleaseToRG(ctx context.Context, m map[uuid16]rg
|
||||
}
|
||||
}
|
||||
|
||||
_ = stmt.Close()
|
||||
|
||||
if err := tx.Commit(); err != nil {
|
||||
imp.logger.Warn("dump import: commit release_to_rg failed", "error", err)
|
||||
|
||||
@@ -427,6 +453,26 @@ func (imp *dumpImporter) persistReleaseToRG(ctx context.Context, m map[uuid16]rg
|
||||
imp.logger.Info("dump import: persisted release→release-group map", "rows", written)
|
||||
}
|
||||
|
||||
// beginReleaseToRGTx opens a batch transaction with the row insert
|
||||
// already prepared on it.
|
||||
func beginReleaseToRGTx(db *database.DB) (*sql.Tx, *sql.Stmt, error) {
|
||||
tx, err := db.BeginTx()
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("release_to_rg begin: %w", err)
|
||||
}
|
||||
|
||||
stmt, err := tx.Prepare(
|
||||
"INSERT OR REPLACE INTO release_to_rg (release_mbid, rg_mbid) VALUES (?, ?)",
|
||||
)
|
||||
if err != nil {
|
||||
_ = tx.Rollback()
|
||||
|
||||
return nil, nil, fmt.Errorf("release_to_rg prepare: %w", err)
|
||||
}
|
||||
|
||||
return tx, stmt, nil
|
||||
}
|
||||
|
||||
func (imp *dumpImporter) readState() (*dumpImportState, error) {
|
||||
state := &dumpImportState{}
|
||||
|
||||
@@ -470,28 +516,24 @@ func (imp *dumpImporter) setStageProgress(stage, completed, total int) {
|
||||
imp.si.setTierStatus(dumpStageNames[stage], "running", total, completed)
|
||||
}
|
||||
|
||||
// setStageDetail is setStageProgress plus a human-readable line for
|
||||
// stages where completed/total alone is uninformative.
|
||||
func (imp *dumpImporter) setStageDetail(stage, completed, total int, detail string) {
|
||||
imp.si.setTierDetail(dumpStageNames[stage], "running", total, completed, detail)
|
||||
}
|
||||
|
||||
// logJob appends a line to the index build's job log, which is what the
|
||||
// user sees in the jobs panel. A build with no registered job (tests,
|
||||
// headless imports) drops the line.
|
||||
func (imp *dumpImporter) logJob(message string) {
|
||||
imp.si.logIndexJob(jobs.LevelInfo, message)
|
||||
}
|
||||
|
||||
// checkDiskHeadroom aborts the import when free disk is critically low.
|
||||
func (imp *dumpImporter) checkDiskHeadroom() error {
|
||||
return checkFreeDisk(imp.stagingDir, imp.abortFreeBytes)
|
||||
}
|
||||
|
||||
// checkFreeDisk returns ErrDiskSpace when the volume holding path has
|
||||
// less than minBytes free. Unknown free space (unsupported platform)
|
||||
// passes.
|
||||
func checkFreeDisk(path string, minBytes uint64) error {
|
||||
free, ok := diskFreeBytes(path)
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
|
||||
if free < minBytes {
|
||||
return fmt.Errorf("%w: %d MB free, need %d MB",
|
||||
ErrDiskSpace, free>>20, minBytes>>20)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// SearchIndex integration
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
+232
-123
@@ -1,43 +1,26 @@
|
||||
//go:build indexbuild
|
||||
|
||||
package explore
|
||||
|
||||
import (
|
||||
"archive/tar"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/csv"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/klauspost/compress/zstd"
|
||||
"github.com/parquet-go/parquet-go"
|
||||
|
||||
"yellowjacket/backend/database"
|
||||
)
|
||||
|
||||
// Fixed MBIDs for fixtures.
|
||||
const (
|
||||
recA = "11111111-1111-1111-1111-111111111111"
|
||||
recB = "22222222-2222-2222-2222-222222222222"
|
||||
recC = "33333333-3333-3333-3333-333333333333"
|
||||
relA = "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"
|
||||
relB = "bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb"
|
||||
rgA = "cccccccc-cccc-cccc-cccc-cccccccccccc"
|
||||
rgB = "dddddddd-dddd-dddd-dddd-dddddddddddd"
|
||||
artA = "eeeeeeee-eeee-eeee-eeee-eeeeeeeeeeee"
|
||||
artB = "ffffffff-ffff-ffff-ffff-ffffffffffff"
|
||||
)
|
||||
|
||||
func testLogger() *slog.Logger {
|
||||
return slog.New(slog.DiscardHandler)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Unit tests: parsing helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -244,67 +227,6 @@ func TestArtistTopRGBoundedAndDeduped(t *testing.T) {
|
||||
// Fixture builders
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// sparkFixtureRow mimics the real spark listens schema: the aggregator
|
||||
// must project just recording/release/artist MBIDs out of it.
|
||||
type sparkFixtureRow struct {
|
||||
ListenedAt int64 `parquet:"listened_at"`
|
||||
UserID int64 `parquet:"user_id"`
|
||||
ArtistName string `parquet:"artist_name,optional"`
|
||||
RecordingMBID string `parquet:"recording_mbid,optional"`
|
||||
ReleaseMBID string `parquet:"release_mbid,optional"`
|
||||
ArtistMBIDs []string `parquet:"artist_credit_mbids,optional,list"`
|
||||
}
|
||||
|
||||
func makeParquet(t *testing.T, rows []sparkFixtureRow) []byte {
|
||||
t.Helper()
|
||||
|
||||
var buf bytes.Buffer
|
||||
|
||||
w := parquet.NewGenericWriter[sparkFixtureRow](&buf)
|
||||
|
||||
if _, err := w.Write(rows); err != nil {
|
||||
t.Fatalf("parquet write: %v", err)
|
||||
}
|
||||
|
||||
if err := w.Close(); err != nil {
|
||||
t.Fatalf("parquet close: %v", err)
|
||||
}
|
||||
|
||||
return buf.Bytes()
|
||||
}
|
||||
|
||||
func makeTar(t *testing.T, members map[string][]byte, order []string) []byte {
|
||||
t.Helper()
|
||||
|
||||
var buf bytes.Buffer
|
||||
|
||||
tw := tar.NewWriter(&buf)
|
||||
|
||||
for _, name := range order {
|
||||
data := members[name]
|
||||
hdr := &tar.Header{
|
||||
Name: name,
|
||||
Mode: 0o644,
|
||||
Size: int64(len(data)),
|
||||
Typeflag: tar.TypeReg,
|
||||
}
|
||||
|
||||
if err := tw.WriteHeader(hdr); err != nil {
|
||||
t.Fatalf("tar header: %v", err)
|
||||
}
|
||||
|
||||
if _, err := tw.Write(data); err != nil {
|
||||
t.Fatalf("tar write: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
if err := tw.Close(); err != nil {
|
||||
t.Fatalf("tar close: %v", err)
|
||||
}
|
||||
|
||||
return buf.Bytes()
|
||||
}
|
||||
|
||||
func zstdCompress(t *testing.T, data []byte) []byte {
|
||||
t.Helper()
|
||||
|
||||
@@ -339,23 +261,6 @@ func csvBytes(t *testing.T, rows [][]string) []byte {
|
||||
return buf.Bytes()
|
||||
}
|
||||
|
||||
// listensOf builds n identical listen rows for a recording.
|
||||
func listensOf(n int, recording, release string, artists []string) []sparkFixtureRow {
|
||||
rows := make([]sparkFixtureRow, n)
|
||||
for i := range rows {
|
||||
rows[i] = sparkFixtureRow{
|
||||
ListenedAt: 1700000000 + int64(i),
|
||||
UserID: int64(i),
|
||||
ArtistName: "Fixture Artist",
|
||||
RecordingMBID: recording,
|
||||
ReleaseMBID: release,
|
||||
ArtistMBIDs: artists,
|
||||
}
|
||||
}
|
||||
|
||||
return rows
|
||||
}
|
||||
|
||||
// canonicalDataCSV builds a canonical_musicbrainz_data.csv fixture.
|
||||
func canonicalDataCSV(t *testing.T) []byte {
|
||||
t.Helper()
|
||||
@@ -651,39 +556,37 @@ func TestDumpImportEndToEnd(t *testing.T) {
|
||||
|
||||
// A legacy API-crawled row with inflated popularity must be
|
||||
// cleared by the first dump import (scale consistency).
|
||||
legacyMBID := "99999999-9999-9999-9999-999999999999"
|
||||
|
||||
si.upsertBatch([]SearchIndexResult{{
|
||||
EntityType: "recording",
|
||||
MBID: legacyMBID,
|
||||
Title: "Legacy Row",
|
||||
ArtistName: "Old Crawl",
|
||||
ArtistMBID: artA,
|
||||
Popularity: 123_456_789,
|
||||
}})
|
||||
|
||||
if err := imp.run(context.Background()); err != nil {
|
||||
t.Fatalf("run: %v", err)
|
||||
}
|
||||
|
||||
legacyRows, err := db.QueryContext(
|
||||
"SELECT COUNT(*) FROM explore_index WHERE mbid = ?", legacyMBID,
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("legacy query: %v", err)
|
||||
}
|
||||
// The import bulk-loads with the FTS sync triggers suspended, so
|
||||
// the rebuild that closes that window is the only thing keeping
|
||||
// search usable: assembled rows must be findable afterwards.
|
||||
ftsCount := func(query string) int {
|
||||
t.Helper()
|
||||
|
||||
if legacyRows.Next() {
|
||||
var n int
|
||||
|
||||
_ = legacyRows.Scan(&n)
|
||||
|
||||
if n != 0 {
|
||||
t.Error("legacy API-crawled row survived the first dump import")
|
||||
rows, err := db.QueryContext(
|
||||
"SELECT COUNT(*) FROM explore_index_fts WHERE explore_index_fts MATCH ?", query,
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("fts query %q: %v", query, err)
|
||||
}
|
||||
|
||||
defer func() { _ = rows.Close() }()
|
||||
|
||||
n := 0
|
||||
|
||||
if rows.Next() {
|
||||
_ = rows.Scan(&n)
|
||||
}
|
||||
|
||||
return n
|
||||
}
|
||||
|
||||
_ = legacyRows.Close()
|
||||
if got := ftsCount("Song"); got == 0 {
|
||||
t.Error("FTS matches no assembled recordings; the rebuild did not run")
|
||||
}
|
||||
|
||||
// Index rows landed with dump-derived popularity.
|
||||
assertRow := func(mbid, entityType, title string, popularity int) {
|
||||
@@ -831,6 +734,212 @@ func TestDumpImportResumesAfterCancel(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Progress reporting
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestFormatStreamProgress(t *testing.T) {
|
||||
const gb = int64(1) << 30
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
offset int64
|
||||
size int64
|
||||
rate float64
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "size and rate known",
|
||||
offset: 40 * gb,
|
||||
size: 200 * gb,
|
||||
rate: 20 << 20,
|
||||
want: "40.0 GB / 200.0 GB (20%) · 20.0 MB/s · ~2h16m left",
|
||||
},
|
||||
{
|
||||
name: "size unknown before first response",
|
||||
offset: 2 * gb,
|
||||
rate: 10 << 20,
|
||||
want: "2.0 GB downloaded · 10.0 MB/s",
|
||||
},
|
||||
{
|
||||
name: "rate unknown on the first tick",
|
||||
offset: 10 * gb,
|
||||
size: 100 * gb,
|
||||
want: "10.0 GB / 100.0 GB (10%)",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
if got := formatStreamProgress(tt.offset, tt.size, tt.rate); got != tt.want {
|
||||
t.Errorf("formatStreamProgress() = %q, want %q", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestFormatETA(t *testing.T) {
|
||||
const gb = int64(1) << 30
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
offset int64
|
||||
size int64
|
||||
rate float64
|
||||
want string
|
||||
}{
|
||||
{name: "hours", offset: 0, size: 100 * gb, rate: 10 << 20, want: "2h50m"},
|
||||
{name: "minutes", offset: 0, size: gb, rate: 10 << 20, want: "1m"},
|
||||
{name: "seconds", offset: 0, size: 1 << 20, rate: 10 << 20, want: "<1m"},
|
||||
{name: "unknown size", offset: 0, size: -1, rate: 10 << 20, want: ""},
|
||||
{name: "stalled", offset: 0, size: 100 * gb, rate: 0, want: ""},
|
||||
{name: "past the end", offset: 2 * gb, size: gb, rate: 10 << 20, want: ""},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
if got := formatETA(tt.offset, tt.size, tt.rate); got != tt.want {
|
||||
t.Errorf("formatETA() = %q, want %q", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestFormatCount(t *testing.T) {
|
||||
tests := []struct {
|
||||
in int
|
||||
want string
|
||||
}{
|
||||
{0, "0"},
|
||||
{999, "999"},
|
||||
{1000, "1,000"},
|
||||
{12345, "12,345"},
|
||||
{1234567, "1,234,567"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
if got := formatCount(tt.in); got != tt.want {
|
||||
t.Errorf("formatCount(%d) = %q, want %q", tt.in, got, tt.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The reporter is what keeps the listens stage from looking frozen, so
|
||||
// it must publish a detail line for the stage while the stream runs.
|
||||
func TestCountsReporterPublishesDetail(t *testing.T) {
|
||||
db := database.NewTestDB(t)
|
||||
si := NewSearchIndex(db, nil, nil, testLogger())
|
||||
imp := &dumpImporter{si: si, logger: testLogger()}
|
||||
|
||||
stream := newResumableReader(context.Background(), nil, "", 0)
|
||||
stream.offset.Store(50 << 30)
|
||||
stream.size.Store(200 << 30)
|
||||
|
||||
rep := &countsReporter{
|
||||
imp: imp,
|
||||
stream: stream,
|
||||
lastSample: time.Now().Add(-time.Second),
|
||||
lastMoved: time.Now(),
|
||||
}
|
||||
|
||||
rep.tick(time.Now())
|
||||
|
||||
tier := findTier(t, si, dumpStageNames[dumpStageCounts])
|
||||
|
||||
if tier.Completed != 25 {
|
||||
t.Errorf("tier completed = %d, want 25", tier.Completed)
|
||||
}
|
||||
|
||||
if !strings.Contains(tier.Detail, "50.0 GB / 200.0 GB (25%)") {
|
||||
t.Errorf("tier detail = %q, want it to report GB progress", tier.Detail)
|
||||
}
|
||||
}
|
||||
|
||||
// A stream that stops moving is reported as stalled rather than as a
|
||||
// decaying transfer rate.
|
||||
func TestCountsReporterReportsStall(t *testing.T) {
|
||||
db := database.NewTestDB(t)
|
||||
si := NewSearchIndex(db, nil, nil, testLogger())
|
||||
imp := &dumpImporter{si: si, logger: testLogger()}
|
||||
|
||||
stream := newResumableReader(context.Background(), nil, "", 0)
|
||||
stream.offset.Store(50 << 30)
|
||||
stream.size.Store(200 << 30)
|
||||
|
||||
now := time.Now()
|
||||
rep := &countsReporter{
|
||||
imp: imp,
|
||||
stream: stream,
|
||||
lastSample: now.Add(-countsUIRefreshInterval),
|
||||
lastOffset: 50 << 30,
|
||||
lastMoved: now.Add(-2 * countsStallAfter),
|
||||
}
|
||||
|
||||
rep.tick(now)
|
||||
|
||||
tier := findTier(t, si, dumpStageNames[dumpStageCounts])
|
||||
|
||||
if !strings.Contains(tier.Detail, "stalled") {
|
||||
t.Errorf("tier detail = %q, want a stall notice", tier.Detail)
|
||||
}
|
||||
|
||||
if imp.streamRate() != 0 {
|
||||
t.Errorf("stalled rate = %v, want 0", imp.streamRate())
|
||||
}
|
||||
}
|
||||
|
||||
// A download paused by parser back-pressure is not a network stall and
|
||||
// must not be reported as one.
|
||||
func TestCountsReporterDistinguishesBacklog(t *testing.T) {
|
||||
db := database.NewTestDB(t)
|
||||
si := NewSearchIndex(db, nil, nil, testLogger())
|
||||
imp := &dumpImporter{si: si, logger: testLogger()}
|
||||
|
||||
stream := newResumableReader(context.Background(), nil, "", 0)
|
||||
stream.offset.Store(50 << 30)
|
||||
stream.size.Store(200 << 30)
|
||||
|
||||
var backlog atomic.Bool
|
||||
|
||||
backlog.Store(true)
|
||||
|
||||
now := time.Now()
|
||||
rep := &countsReporter{
|
||||
imp: imp,
|
||||
stream: stream,
|
||||
backlog: &backlog,
|
||||
lastSample: now.Add(-countsUIRefreshInterval),
|
||||
lastOffset: 50 << 30,
|
||||
lastMoved: now.Add(-2 * countsStallAfter),
|
||||
}
|
||||
|
||||
rep.tick(now)
|
||||
|
||||
tier := findTier(t, si, dumpStageNames[dumpStageCounts])
|
||||
|
||||
if strings.Contains(tier.Detail, "stalled") {
|
||||
t.Errorf("tier detail = %q, want parsing back-pressure, not a stall", tier.Detail)
|
||||
}
|
||||
|
||||
if !strings.Contains(tier.Detail, "parsing") {
|
||||
t.Errorf("tier detail = %q, want it to name the parsing pause", tier.Detail)
|
||||
}
|
||||
}
|
||||
|
||||
func findTier(t *testing.T, si *SearchIndex, name string) TierStatus {
|
||||
t.Helper()
|
||||
|
||||
for _, tier := range si.GetIndexStatus().Tiers {
|
||||
if tier.Name == name {
|
||||
return tier
|
||||
}
|
||||
}
|
||||
|
||||
t.Fatalf("tier %q not found", name)
|
||||
|
||||
return TierStatus{}
|
||||
}
|
||||
|
||||
func TestCheckFreeDisk(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
|
||||
|
||||
@@ -0,0 +1,562 @@
|
||||
//go:build indexbuild
|
||||
|
||||
package explore
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Parallel dump streaming. MetaBrainz shapes throughput per
|
||||
// connection, so a single sequential stream of the listens dump tops
|
||||
// out near 5 MB/s no matter how fast the link is — the 205GB stage-1
|
||||
// download alone runs for eleven hours that way. Several concurrent
|
||||
// Range requests lift that to tens of MB/s.
|
||||
//
|
||||
// parallelReader hides the concurrency behind the same sequential
|
||||
// io.Reader the tar decoders already consume: lanes fetch fixed-size
|
||||
// chunks ahead of the caller, and chunks are handed over strictly in
|
||||
// order. Pos() therefore keeps meaning "absolute offset of the next
|
||||
// byte to be delivered", which is what the stage-1 checkpoint records.
|
||||
//
|
||||
// For the listens dump this is now the fallback: dumpproject.go fetches
|
||||
// only the columns stage 1 reads, which is well under half the bytes.
|
||||
// This path still serves the canonical dump (a zstd stream that cannot
|
||||
// be range-projected) and any origin that refuses Range requests.
|
||||
//
|
||||
// A caution on the rates quoted below: measured again on 2026-07-28,
|
||||
// data.metabrainz.org served ~0.65 MB/s aggregate to one client and got
|
||||
// *slower* past eight connections — the shaping is per-IP, not per
|
||||
// connection, so raising dumpLanes buys throttling rather than speed.
|
||||
|
||||
const (
|
||||
// dumpLanes is the number of concurrent Range requests. One lane
|
||||
// gets ~5 MB/s because MetaBrainz shapes per connection; four reach
|
||||
// tens of MB/s. Deliberately conservative: these are a nonprofit's
|
||||
// servers, they answer sustained heavy use with 503s, and pushing
|
||||
// past this trades politeness for throughput that backoff eats
|
||||
// anyway.
|
||||
dumpLanes = 4
|
||||
|
||||
// dumpChunkSize is how much a lane fetches per request. Each
|
||||
// request pays a ramp-up cost, so small chunks squander the gain —
|
||||
// 32MB chunks measured roughly 40% slower than 128MB ones.
|
||||
dumpChunkSize = 128 << 20
|
||||
|
||||
// dumpWindowChunks bounds the chunks in flight or buffered awaiting
|
||||
// in-order delivery. Kept just above the lane count so lanes never
|
||||
// idle waiting for the consumer; costs dumpWindowChunks *
|
||||
// dumpChunkSize of buffer.
|
||||
dumpWindowChunks = dumpLanes + 2
|
||||
|
||||
// dumpProbeTimeout bounds the HEAD request that sizes a resource
|
||||
// before a parallel stream starts.
|
||||
dumpProbeTimeout = 30 * time.Second
|
||||
|
||||
// dumpMaxIdleConns keeps a pooled connection per lane so chunk
|
||||
// requests reuse TCP+TLS instead of reconnecting each time.
|
||||
dumpMaxIdleConns = dumpLanes * 2
|
||||
|
||||
// dumpRetryAfterCap bounds how long a server-supplied Retry-After is
|
||||
// honoured, so a bad header can't park a lane indefinitely.
|
||||
dumpRetryAfterCap = 60 * time.Second
|
||||
)
|
||||
|
||||
// dumpStream is the streaming surface the dump importers consume,
|
||||
// implemented by both parallelReader and resumableReader.
|
||||
type dumpStream interface {
|
||||
io.ReadCloser
|
||||
|
||||
// Pos is the absolute byte offset of the next byte to be delivered.
|
||||
// This is what an interrupted import checkpoints and resumes from.
|
||||
Pos() int64
|
||||
|
||||
// Fetched is the absolute byte offset the downloader has reached.
|
||||
// With prefetching lanes this runs ahead of Pos, and it — not Pos —
|
||||
// is what progress and stall reporting should watch: a reader
|
||||
// buffering a 128MB chunk is downloading, not stalled.
|
||||
Fetched() int64
|
||||
|
||||
// Total is the total resource size, or -1 while unknown.
|
||||
Total() int64
|
||||
}
|
||||
|
||||
// newDumpHTTPClient builds the client used for dump discovery and
|
||||
// streaming. HTTP/2 is disabled deliberately: it multiplexes every
|
||||
// lane onto a single TCP connection, which collapses parallel Range
|
||||
// requests back to one shaped stream (measured at 1-3 MB/s).
|
||||
func newDumpHTTPClient() *http.Client {
|
||||
transport := &http.Transport{
|
||||
// Explicitly HTTP/1.1: ALPN would otherwise negotiate h2 and
|
||||
// silently undo the parallelism below.
|
||||
TLSClientConfig: &tls.Config{
|
||||
NextProtos: []string{"http/1.1"},
|
||||
MinVersion: tls.VersionTLS12,
|
||||
},
|
||||
ForceAttemptHTTP2: false,
|
||||
TLSNextProto: map[string]func(string, *tls.Conn) http.RoundTripper{},
|
||||
MaxIdleConns: dumpMaxIdleConns,
|
||||
MaxIdleConnsPerHost: dumpMaxIdleConns,
|
||||
IdleConnTimeout: 90 * time.Second,
|
||||
}
|
||||
|
||||
// No client-level timeout: dump streams run for hours. Per-request
|
||||
// deadlines come from the caller's context, and chunk fetches retry
|
||||
// on their own.
|
||||
return &http.Client{Transport: transport}
|
||||
}
|
||||
|
||||
// probeDumpSize returns the resource size when the server advertises
|
||||
// one and supports Range requests, else (0, false).
|
||||
func probeDumpSize(ctx context.Context, client *http.Client, url string) (int64, bool) {
|
||||
reqCtx, cancel := context.WithTimeout(ctx, dumpProbeTimeout)
|
||||
defer cancel()
|
||||
|
||||
req, err := http.NewRequestWithContext(reqCtx, http.MethodHead, url, nil)
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
|
||||
req.Header.Set("User-Agent", lbUserAgent)
|
||||
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
if resp.StatusCode != http.StatusOK || resp.ContentLength <= 0 {
|
||||
return 0, false
|
||||
}
|
||||
|
||||
if resp.Header.Get("Accept-Ranges") != "bytes" {
|
||||
return 0, false
|
||||
}
|
||||
|
||||
return resp.ContentLength, true
|
||||
}
|
||||
|
||||
// chunkResult is one fetched chunk awaiting in-order delivery. bufp is
|
||||
// the pooled backing array, returned to the pool once drained.
|
||||
type chunkResult struct {
|
||||
bufp *[]byte
|
||||
n int
|
||||
err error
|
||||
}
|
||||
|
||||
// parallelReader streams [base, size) of an HTTP resource through
|
||||
// several concurrent Range requests, delivering bytes in order.
|
||||
type parallelReader struct {
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
client *http.Client
|
||||
url string
|
||||
logger *slog.Logger
|
||||
|
||||
base int64
|
||||
size int64
|
||||
totalChunks int64
|
||||
|
||||
// lanes, chunkSize and window are fields rather than constants so
|
||||
// tests can drive the reader with small chunks.
|
||||
lanes int
|
||||
chunkSize int64
|
||||
window int64
|
||||
|
||||
mu sync.Mutex
|
||||
cond *sync.Cond
|
||||
ready map[int64]*chunkResult
|
||||
nextDL int64 // next chunk index to hand to a lane
|
||||
nextOut int64 // next chunk index to deliver
|
||||
closed bool
|
||||
|
||||
delivered atomic.Int64
|
||||
|
||||
// fetched counts bytes pulled down by the lanes, including chunks
|
||||
// still buffered ahead of the caller.
|
||||
fetched atomic.Int64
|
||||
|
||||
// cur is the undelivered tail of the chunk being drained.
|
||||
cur []byte
|
||||
curBufp *[]byte
|
||||
|
||||
pool sync.Pool
|
||||
wg sync.WaitGroup
|
||||
}
|
||||
|
||||
// newParallelReader starts the lanes and returns a reader positioned at
|
||||
// offset. Returns nil when the resource can't be range-streamed, so
|
||||
// callers fall back to a sequential reader.
|
||||
func newParallelReader(
|
||||
ctx context.Context, client *http.Client, logger *slog.Logger, url string, offset int64,
|
||||
) *parallelReader {
|
||||
return newParallelReaderLogged(
|
||||
ctx, client, logger, url, offset, dumpLanes, dumpChunkSize, dumpWindowChunks,
|
||||
)
|
||||
}
|
||||
|
||||
// newParallelReaderWith is newParallelReader with the lane geometry
|
||||
// spelled out, so tests can exercise ordering and resume with chunks
|
||||
// small enough to be practical.
|
||||
func newParallelReaderWith(
|
||||
ctx context.Context, client *http.Client, url string, offset int64,
|
||||
lanes int, chunkSize, window int64,
|
||||
) *parallelReader {
|
||||
return newParallelReaderLogged(
|
||||
ctx, client, slog.New(slog.DiscardHandler), url, offset, lanes, chunkSize, window,
|
||||
)
|
||||
}
|
||||
|
||||
// newParallelReaderLogged is newParallelReaderWith with a logger, so a
|
||||
// throttled or retrying stream is diagnosable from the app log.
|
||||
func newParallelReaderLogged(
|
||||
ctx context.Context, client *http.Client, logger *slog.Logger,
|
||||
url string, offset int64, lanes int, chunkSize, window int64,
|
||||
) *parallelReader {
|
||||
size, ok := probeDumpSize(ctx, client, url)
|
||||
if !ok || offset >= size {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Below a couple of chunks there's nothing to parallelise.
|
||||
if size-offset < 2*chunkSize {
|
||||
return nil
|
||||
}
|
||||
|
||||
if logger == nil {
|
||||
logger = slog.New(slog.DiscardHandler)
|
||||
}
|
||||
|
||||
streamCtx, cancel := context.WithCancel(ctx)
|
||||
|
||||
p := ¶llelReader{
|
||||
ctx: streamCtx,
|
||||
cancel: cancel,
|
||||
client: client,
|
||||
logger: logger,
|
||||
url: url,
|
||||
base: offset,
|
||||
size: size,
|
||||
lanes: lanes,
|
||||
chunkSize: chunkSize,
|
||||
window: window,
|
||||
ready: make(map[int64]*chunkResult, window),
|
||||
pool: sync.Pool{New: func() any {
|
||||
b := make([]byte, chunkSize)
|
||||
|
||||
return &b
|
||||
}},
|
||||
}
|
||||
|
||||
remaining := size - offset
|
||||
p.totalChunks = (remaining + chunkSize - 1) / chunkSize
|
||||
p.cond = sync.NewCond(&p.mu)
|
||||
|
||||
for range lanes {
|
||||
p.wg.Add(1)
|
||||
|
||||
go p.lane()
|
||||
}
|
||||
|
||||
// A cancelled context must wake anyone blocked on the condition.
|
||||
go func() {
|
||||
<-streamCtx.Done()
|
||||
|
||||
p.mu.Lock()
|
||||
p.cond.Broadcast()
|
||||
p.mu.Unlock()
|
||||
}()
|
||||
|
||||
return p
|
||||
}
|
||||
|
||||
// Pos returns the absolute offset of the next byte to be delivered.
|
||||
func (p *parallelReader) Pos() int64 {
|
||||
return p.base + p.delivered.Load()
|
||||
}
|
||||
|
||||
// Fetched returns the absolute offset the lanes have downloaded to.
|
||||
func (p *parallelReader) Fetched() int64 {
|
||||
return p.base + p.fetched.Load()
|
||||
}
|
||||
|
||||
// Total returns the total resource size.
|
||||
func (p *parallelReader) Total() int64 {
|
||||
return p.size
|
||||
}
|
||||
|
||||
// lane fetches chunks until the window is exhausted or the stream ends.
|
||||
func (p *parallelReader) lane() {
|
||||
defer p.wg.Done()
|
||||
|
||||
for {
|
||||
p.mu.Lock()
|
||||
|
||||
for {
|
||||
if p.closed || p.ctx.Err() != nil || p.nextDL >= p.totalChunks {
|
||||
p.mu.Unlock()
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
// Stay inside the delivery window so buffered chunks can't
|
||||
// outrun the consumer.
|
||||
if p.nextDL < p.nextOut+p.window {
|
||||
break
|
||||
}
|
||||
|
||||
p.cond.Wait()
|
||||
}
|
||||
|
||||
idx := p.nextDL
|
||||
p.nextDL++
|
||||
|
||||
p.mu.Unlock()
|
||||
|
||||
bufp, n, err := p.fetchChunk(idx)
|
||||
if err == nil {
|
||||
p.fetched.Add(int64(n))
|
||||
}
|
||||
|
||||
p.mu.Lock()
|
||||
p.ready[idx] = &chunkResult{bufp: bufp, n: n, err: err}
|
||||
p.cond.Broadcast()
|
||||
p.mu.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
// fetchChunk retrieves one chunk, retrying transient failures. The
|
||||
// range is fully specified, so a retry simply re-requests it.
|
||||
func (p *parallelReader) fetchChunk(idx int64) (*[]byte, int, error) {
|
||||
lo := p.base + idx*p.chunkSize
|
||||
|
||||
hi := lo + p.chunkSize - 1
|
||||
if hi >= p.size {
|
||||
hi = p.size - 1
|
||||
}
|
||||
|
||||
want := int(hi - lo + 1)
|
||||
|
||||
bufp, _ := p.pool.Get().(*[]byte)
|
||||
|
||||
var (
|
||||
lastErr error
|
||||
wait time.Duration
|
||||
)
|
||||
|
||||
for attempt := 0; attempt <= maxStreamRetries; attempt++ {
|
||||
if attempt > 0 {
|
||||
delay := min(streamRetryBaseDelay<<(attempt-1), streamRetryMaxDelay)
|
||||
if wait > 0 {
|
||||
delay = wait
|
||||
}
|
||||
|
||||
select {
|
||||
case <-p.ctx.Done():
|
||||
p.pool.Put(bufp)
|
||||
|
||||
return nil, 0, p.ctx.Err()
|
||||
case <-time.After(delay):
|
||||
}
|
||||
}
|
||||
|
||||
n, retryAfter, err := p.fetchOnce(lo, hi, (*bufp)[:want])
|
||||
if err == nil && n == want {
|
||||
return bufp, n, nil
|
||||
}
|
||||
|
||||
if err == nil {
|
||||
err = io.ErrUnexpectedEOF
|
||||
}
|
||||
|
||||
lastErr = err
|
||||
wait = retryAfter
|
||||
|
||||
p.logger.Warn("dump import: chunk fetch failed, retrying",
|
||||
"chunk", idx,
|
||||
"attempt", attempt+1,
|
||||
"retryAfter", retryAfter,
|
||||
"error", err,
|
||||
)
|
||||
|
||||
if p.ctx.Err() != nil {
|
||||
p.pool.Put(bufp)
|
||||
|
||||
return nil, 0, p.ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
p.pool.Put(bufp)
|
||||
|
||||
return nil, 0, fmt.Errorf("%w: %s chunk %d after %d retries: %w",
|
||||
ErrDumpStream, p.url, idx, maxStreamRetries, lastErr)
|
||||
}
|
||||
|
||||
// fetchOnce performs a single Range request into buf. The second
|
||||
// return value is the server's requested Retry-After delay, if any.
|
||||
func (p *parallelReader) fetchOnce(lo, hi int64, buf []byte) (int, time.Duration, error) {
|
||||
req, err := http.NewRequestWithContext(p.ctx, http.MethodGet, p.url, nil)
|
||||
if err != nil {
|
||||
return 0, 0, fmt.Errorf("dump chunk request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("User-Agent", lbUserAgent)
|
||||
req.Header.Set("Range", "bytes="+strconv.FormatInt(lo, 10)+"-"+strconv.FormatInt(hi, 10))
|
||||
|
||||
resp, err := p.client.Do(req)
|
||||
if err != nil {
|
||||
return 0, 0, fmt.Errorf("dump chunk fetch: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
if resp.StatusCode != http.StatusPartialContent {
|
||||
// A loaded dump server answers with 503 (or 429) rather than
|
||||
// queueing; that is a "come back shortly", not a failure.
|
||||
return 0, parseRetryAfter(resp.Header.Get("Retry-After")),
|
||||
fmt.Errorf("%w: HTTP %d from %s", ErrDumpStream, resp.StatusCode, p.url)
|
||||
}
|
||||
|
||||
n, err := io.ReadFull(resp.Body, buf)
|
||||
if err != nil {
|
||||
return n, 0, fmt.Errorf("dump chunk read: %w", err)
|
||||
}
|
||||
|
||||
return n, 0, nil
|
||||
}
|
||||
|
||||
// parseRetryAfter reads a Retry-After header given in seconds, clamped
|
||||
// to dumpRetryAfterCap. Returns 0 when absent or unparseable.
|
||||
func parseRetryAfter(v string) time.Duration {
|
||||
if v == "" {
|
||||
return 0
|
||||
}
|
||||
|
||||
secs, err := strconv.Atoi(v)
|
||||
if err != nil || secs <= 0 {
|
||||
return 0
|
||||
}
|
||||
|
||||
return min(time.Duration(secs)*time.Second, dumpRetryAfterCap)
|
||||
}
|
||||
|
||||
func (p *parallelReader) Read(b []byte) (int, error) {
|
||||
for len(p.cur) == 0 {
|
||||
if err := p.ctx.Err(); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
res, err := p.nextChunk()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
p.curBufp = res.bufp
|
||||
p.cur = (*res.bufp)[:res.n]
|
||||
}
|
||||
|
||||
n := copy(b, p.cur)
|
||||
p.cur = p.cur[n:]
|
||||
p.delivered.Add(int64(n))
|
||||
|
||||
if len(p.cur) == 0 && p.curBufp != nil {
|
||||
p.pool.Put(p.curBufp)
|
||||
p.curBufp = nil
|
||||
}
|
||||
|
||||
return n, nil
|
||||
}
|
||||
|
||||
// nextChunk blocks until the next in-order chunk is available.
|
||||
func (p *parallelReader) nextChunk() (*chunkResult, error) {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
|
||||
for {
|
||||
if p.closed {
|
||||
return nil, io.ErrClosedPipe
|
||||
}
|
||||
|
||||
if err := p.ctx.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if p.nextOut >= p.totalChunks {
|
||||
return nil, io.EOF
|
||||
}
|
||||
|
||||
res, ok := p.ready[p.nextOut]
|
||||
if ok {
|
||||
delete(p.ready, p.nextOut)
|
||||
|
||||
p.nextOut++
|
||||
|
||||
// A freed window slot may unblock a waiting lane.
|
||||
p.cond.Broadcast()
|
||||
|
||||
if res.err != nil {
|
||||
return nil, res.err
|
||||
}
|
||||
|
||||
return res, nil
|
||||
}
|
||||
|
||||
p.cond.Wait()
|
||||
}
|
||||
}
|
||||
|
||||
// Close stops the lanes and releases buffered chunks.
|
||||
func (p *parallelReader) Close() error {
|
||||
p.mu.Lock()
|
||||
|
||||
if p.closed {
|
||||
p.mu.Unlock()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
p.closed = true
|
||||
p.cond.Broadcast()
|
||||
p.mu.Unlock()
|
||||
|
||||
p.cancel()
|
||||
p.wg.Wait()
|
||||
|
||||
p.mu.Lock()
|
||||
clear(p.ready)
|
||||
p.mu.Unlock()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// openDumpStream returns the best available stream for a dump URL,
|
||||
// preferring parallel Range lanes and falling back to a single
|
||||
// resumable connection when the server won't serve ranges.
|
||||
func (imp *dumpImporter) openDumpStream(
|
||||
ctx context.Context, url string, offset int64,
|
||||
) dumpStream {
|
||||
if p := newParallelReader(ctx, imp.httpClient, imp.logger, url, offset); p != nil {
|
||||
imp.logger.Info("dump import: streaming in parallel",
|
||||
"lanes", dumpLanes,
|
||||
"chunkMB", dumpChunkSize>>20,
|
||||
"url", url,
|
||||
)
|
||||
|
||||
return p
|
||||
}
|
||||
|
||||
imp.logger.Info("dump import: parallel streaming unavailable, using single stream",
|
||||
"url", url,
|
||||
)
|
||||
|
||||
return newResumableReader(ctx, imp.httpClient, url, offset)
|
||||
}
|
||||
@@ -0,0 +1,372 @@
|
||||
//go:build indexbuild
|
||||
|
||||
package explore
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"io"
|
||||
"math/rand"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strconv"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// serveBlob returns a Range-capable server for a fixed payload, plus a
|
||||
// counter of the GET requests it served.
|
||||
func serveBlob(t *testing.T, payload []byte) (*httptest.Server, *atomic.Int64) {
|
||||
t.Helper()
|
||||
|
||||
var gets atomic.Int64
|
||||
|
||||
modTime := time.Now()
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(
|
||||
func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method == http.MethodGet {
|
||||
gets.Add(1)
|
||||
}
|
||||
|
||||
http.ServeContent(w, r, "blob.bin", modTime, bytes.NewReader(payload))
|
||||
},
|
||||
))
|
||||
t.Cleanup(srv.Close)
|
||||
|
||||
return srv, &gets
|
||||
}
|
||||
|
||||
func randomPayload(n int) []byte {
|
||||
buf := make([]byte, n)
|
||||
|
||||
rng := rand.New(rand.NewSource(1)) //nolint:gosec // deterministic fixture
|
||||
_, _ = rng.Read(buf)
|
||||
|
||||
return buf
|
||||
}
|
||||
|
||||
// The whole point of the reader is that concurrency stays invisible:
|
||||
// bytes must come out in the same order a single stream would produce.
|
||||
func TestParallelReaderDeliversBytesInOrder(t *testing.T) {
|
||||
payload := randomPayload(200_000)
|
||||
srv, gets := serveBlob(t, payload)
|
||||
|
||||
p := newParallelReaderWith(
|
||||
context.Background(), srv.Client(), srv.URL, 0, 4, 8_192, 6,
|
||||
)
|
||||
if p == nil {
|
||||
t.Fatal("newParallelReaderWith returned nil, want a parallel reader")
|
||||
}
|
||||
|
||||
defer func() { _ = p.Close() }()
|
||||
|
||||
got, err := io.ReadAll(p)
|
||||
if err != nil {
|
||||
t.Fatalf("read: %v", err)
|
||||
}
|
||||
|
||||
if !bytes.Equal(got, payload) {
|
||||
t.Fatalf("payload mismatch: got %d bytes, want %d", len(got), len(payload))
|
||||
}
|
||||
|
||||
// Confirm it really did fan out rather than quietly falling back.
|
||||
if n := gets.Load(); n < 2 {
|
||||
t.Errorf("served %d GETs, want one per chunk", n)
|
||||
}
|
||||
}
|
||||
|
||||
// Pos is what the stage-1 checkpoint records, so it must track bytes
|
||||
// handed to the caller — not bytes fetched by the lanes running ahead.
|
||||
func TestParallelReaderPosTracksDeliveredBytes(t *testing.T) {
|
||||
payload := randomPayload(100_000)
|
||||
srv, _ := serveBlob(t, payload)
|
||||
|
||||
p := newParallelReaderWith(
|
||||
context.Background(), srv.Client(), srv.URL, 0, 4, 4_096, 8,
|
||||
)
|
||||
if p == nil {
|
||||
t.Fatal("newParallelReaderWith returned nil")
|
||||
}
|
||||
|
||||
defer func() { _ = p.Close() }()
|
||||
|
||||
if got := p.Total(); got != int64(len(payload)) {
|
||||
t.Errorf("Total() = %d, want %d", got, len(payload))
|
||||
}
|
||||
|
||||
buf := make([]byte, 1_000)
|
||||
|
||||
read, err := io.ReadFull(p, buf)
|
||||
if err != nil {
|
||||
t.Fatalf("read: %v", err)
|
||||
}
|
||||
|
||||
if got := p.Pos(); got != int64(read) {
|
||||
t.Errorf("Pos() = %d after reading %d bytes, want %d", got, read, read)
|
||||
}
|
||||
|
||||
// Let the lanes race ahead, then confirm Pos still reflects delivery.
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
if got := p.Pos(); got != int64(read) {
|
||||
t.Errorf("Pos() = %d after lanes prefetched, want %d", got, read)
|
||||
}
|
||||
}
|
||||
|
||||
// Progress reporting watches Fetched rather than Pos, because a reader
|
||||
// buffering a chunk ahead of the caller is downloading, not stalled.
|
||||
// Fetched must therefore outrun Pos while lanes prefetch.
|
||||
func TestParallelReaderFetchedOutrunsPos(t *testing.T) {
|
||||
payload := randomPayload(200_000)
|
||||
srv, _ := serveBlob(t, payload)
|
||||
|
||||
p := newParallelReaderWith(
|
||||
context.Background(), srv.Client(), srv.URL, 0, 4, 8_192, 8,
|
||||
)
|
||||
if p == nil {
|
||||
t.Fatal("newParallelReaderWith returned nil")
|
||||
}
|
||||
|
||||
defer func() { _ = p.Close() }()
|
||||
|
||||
// Read a single byte, then let the lanes fill the window.
|
||||
buf := make([]byte, 1)
|
||||
if _, err := io.ReadFull(p, buf); err != nil {
|
||||
t.Fatalf("read: %v", err)
|
||||
}
|
||||
|
||||
deadline := time.Now().Add(5 * time.Second)
|
||||
for time.Now().Before(deadline) && p.Fetched() <= p.Pos() {
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
|
||||
if p.Fetched() <= p.Pos() {
|
||||
t.Errorf("Fetched() = %d, Pos() = %d; want Fetched ahead while prefetching",
|
||||
p.Fetched(), p.Pos())
|
||||
}
|
||||
}
|
||||
|
||||
// Resuming an interrupted import constructs a reader at the
|
||||
// checkpointed offset; it must yield exactly the remaining tail.
|
||||
func TestParallelReaderResumesFromOffset(t *testing.T) {
|
||||
payload := randomPayload(120_000)
|
||||
srv, _ := serveBlob(t, payload)
|
||||
|
||||
const offset = 37_000
|
||||
|
||||
p := newParallelReaderWith(
|
||||
context.Background(), srv.Client(), srv.URL, offset, 3, 8_192, 5,
|
||||
)
|
||||
if p == nil {
|
||||
t.Fatal("newParallelReaderWith returned nil")
|
||||
}
|
||||
|
||||
defer func() { _ = p.Close() }()
|
||||
|
||||
if got := p.Pos(); got != offset {
|
||||
t.Errorf("Pos() = %d before reading, want %d", got, offset)
|
||||
}
|
||||
|
||||
got, err := io.ReadAll(p)
|
||||
if err != nil {
|
||||
t.Fatalf("read: %v", err)
|
||||
}
|
||||
|
||||
if !bytes.Equal(got, payload[offset:]) {
|
||||
t.Fatalf("resumed payload mismatch: got %d bytes, want %d",
|
||||
len(got), len(payload)-offset)
|
||||
}
|
||||
}
|
||||
|
||||
// A lane that hits a transient failure must retry its range rather than
|
||||
// tear down the whole multi-hour stream.
|
||||
func TestParallelReaderRetriesFailedChunk(t *testing.T) {
|
||||
payload := randomPayload(60_000)
|
||||
|
||||
var attempts atomic.Int64
|
||||
|
||||
modTime := time.Now()
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(
|
||||
func(w http.ResponseWriter, r *http.Request) {
|
||||
// Fail the third GET once; every other request succeeds.
|
||||
if r.Method == http.MethodGet && attempts.Add(1) == 3 {
|
||||
hj, ok := w.(http.Hijacker)
|
||||
if ok {
|
||||
conn, _, err := hj.Hijack()
|
||||
if err == nil {
|
||||
_ = conn.Close()
|
||||
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
http.ServeContent(w, r, "blob.bin", modTime, bytes.NewReader(payload))
|
||||
},
|
||||
))
|
||||
t.Cleanup(srv.Close)
|
||||
|
||||
p := newParallelReaderWith(
|
||||
context.Background(), srv.Client(), srv.URL, 0, 2, 8_192, 4,
|
||||
)
|
||||
if p == nil {
|
||||
t.Fatal("newParallelReaderWith returned nil")
|
||||
}
|
||||
|
||||
defer func() { _ = p.Close() }()
|
||||
|
||||
got, err := io.ReadAll(p)
|
||||
if err != nil {
|
||||
t.Fatalf("read after transient failure: %v", err)
|
||||
}
|
||||
|
||||
if !bytes.Equal(got, payload) {
|
||||
t.Fatalf("payload mismatch after retry: got %d bytes, want %d",
|
||||
len(got), len(payload))
|
||||
}
|
||||
}
|
||||
|
||||
// A loaded dump server answers with 503 rather than queueing. That is
|
||||
// "come back shortly", not a failure, so the lane must retry and the
|
||||
// stream must still complete — this is what stalled a real import.
|
||||
func TestParallelReaderRecoversFrom503(t *testing.T) {
|
||||
payload := randomPayload(60_000)
|
||||
|
||||
var gets atomic.Int64
|
||||
|
||||
modTime := time.Now()
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(
|
||||
func(w http.ResponseWriter, r *http.Request) {
|
||||
// Refuse the first two range GETs the way a busy
|
||||
// MetaBrainz mirror does.
|
||||
if r.Method == http.MethodGet && gets.Add(1) <= 2 {
|
||||
w.Header().Set("Retry-After", "1")
|
||||
w.WriteHeader(http.StatusServiceUnavailable)
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
http.ServeContent(w, r, "blob.bin", modTime, bytes.NewReader(payload))
|
||||
},
|
||||
))
|
||||
t.Cleanup(srv.Close)
|
||||
|
||||
p := newParallelReaderWith(
|
||||
context.Background(), srv.Client(), srv.URL, 0, 2, 8_192, 4,
|
||||
)
|
||||
if p == nil {
|
||||
t.Fatal("newParallelReaderWith returned nil")
|
||||
}
|
||||
|
||||
defer func() { _ = p.Close() }()
|
||||
|
||||
got, err := io.ReadAll(p)
|
||||
if err != nil {
|
||||
t.Fatalf("read after 503s: %v", err)
|
||||
}
|
||||
|
||||
if !bytes.Equal(got, payload) {
|
||||
t.Fatalf("payload mismatch after 503s: got %d bytes, want %d",
|
||||
len(got), len(payload))
|
||||
}
|
||||
}
|
||||
|
||||
// Retry-After is honoured but clamped, so a hostile or buggy header
|
||||
// can't park a download lane for hours.
|
||||
func TestParseRetryAfter(t *testing.T) {
|
||||
tests := []struct {
|
||||
header string
|
||||
want time.Duration
|
||||
}{
|
||||
{"", 0},
|
||||
{"5", 5 * time.Second},
|
||||
{"0", 0},
|
||||
{"-3", 0},
|
||||
{"not-a-number", 0},
|
||||
{"Wed, 21 Oct 2026 07:28:00 GMT", 0}, // HTTP-date form: ignored
|
||||
{"99999", dumpRetryAfterCap},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
if got := parseRetryAfter(tt.header); got != tt.want {
|
||||
t.Errorf("parseRetryAfter(%q) = %v, want %v", tt.header, got, tt.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Servers that won't serve ranges must fall back to the sequential
|
||||
// reader instead of failing the import.
|
||||
func TestParallelReaderDeclinesWithoutRangeSupport(t *testing.T) {
|
||||
payload := randomPayload(100_000)
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(
|
||||
func(w http.ResponseWriter, _ *http.Request) {
|
||||
// No Accept-Ranges header: a plain, non-seekable response.
|
||||
w.Header().Set("Content-Length", strconv.Itoa(len(payload)))
|
||||
_, _ = w.Write(payload)
|
||||
},
|
||||
))
|
||||
t.Cleanup(srv.Close)
|
||||
|
||||
if p := newParallelReaderWith(
|
||||
context.Background(), srv.Client(), srv.URL, 0, 4, 8_192, 6,
|
||||
); p != nil {
|
||||
_ = p.Close()
|
||||
|
||||
t.Fatal("got a parallel reader for a server without range support, want nil")
|
||||
}
|
||||
}
|
||||
|
||||
// Payloads too small to split aren't worth the fan-out.
|
||||
func TestParallelReaderDeclinesTinyPayload(t *testing.T) {
|
||||
payload := randomPayload(1_000)
|
||||
srv, _ := serveBlob(t, payload)
|
||||
|
||||
if p := newParallelReaderWith(
|
||||
context.Background(), srv.Client(), srv.URL, 0, 4, 8_192, 6,
|
||||
); p != nil {
|
||||
_ = p.Close()
|
||||
|
||||
t.Fatal("got a parallel reader for a sub-chunk payload, want nil")
|
||||
}
|
||||
}
|
||||
|
||||
// Cancelling the import must stop the lanes promptly rather than let
|
||||
// them keep pulling gigabytes in the background.
|
||||
func TestParallelReaderStopsOnCancel(t *testing.T) {
|
||||
payload := randomPayload(400_000)
|
||||
srv, _ := serveBlob(t, payload)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
|
||||
p := newParallelReaderWith(ctx, srv.Client(), srv.URL, 0, 4, 8_192, 6)
|
||||
if p == nil {
|
||||
t.Fatal("newParallelReaderWith returned nil")
|
||||
}
|
||||
|
||||
buf := make([]byte, 100)
|
||||
if _, err := io.ReadFull(p, buf); err != nil {
|
||||
t.Fatalf("initial read: %v", err)
|
||||
}
|
||||
|
||||
cancel()
|
||||
|
||||
// Close waits for every lane, so returning at all proves they exited.
|
||||
done := make(chan struct{})
|
||||
|
||||
go func() {
|
||||
_ = p.Close()
|
||||
|
||||
close(done)
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(10 * time.Second):
|
||||
t.Fatal("Close did not return after cancel; lanes are still running")
|
||||
}
|
||||
}
|
||||
@@ -1,3 +1,5 @@
|
||||
//go:build indexbuild
|
||||
|
||||
package explore
|
||||
|
||||
import (
|
||||
|
||||
@@ -0,0 +1,679 @@
|
||||
//go:build indexbuild
|
||||
|
||||
package explore
|
||||
|
||||
import (
|
||||
"archive/tar"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/binary"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/parquet-go/parquet-go"
|
||||
"github.com/parquet-go/parquet-go/format"
|
||||
)
|
||||
|
||||
// Column-projected fetching of the spark listens dump.
|
||||
//
|
||||
// Streaming the tar end to end pulls all 205GB even though the counts
|
||||
// aggregator reads three columns. Parquet is columnar and the dump
|
||||
// server serves Range requests, so the unread columns never have to
|
||||
// cross the wire: for each member we fetch the footer, look up the byte
|
||||
// ranges of the wanted column chunks, and download only those.
|
||||
//
|
||||
// Measured against listenbrainz-spark-dump-2593 (2026-07-28), the three
|
||||
// projected columns are 43.4% of the row-group bytes:
|
||||
//
|
||||
// recording_msid 26.9% (not read)
|
||||
// recording_mbid 24.1% ← wanted
|
||||
// recording_name 15.2% (not read)
|
||||
// release_mbid 13.5% ← wanted
|
||||
// release_name 6.0% (not read)
|
||||
// artist_credit_mbids.list.element 5.8% ← wanted
|
||||
// artist_name / artist_credit_id / user_id / created / listened_at
|
||||
//
|
||||
// The decoder is untouched: the fetched ranges are placed at their real
|
||||
// offsets in a member-sized buffer, so parseListenParquet still reads
|
||||
// an ordinary parquet file. It never touches the unfilled bytes,
|
||||
// because it only projects those three columns.
|
||||
|
||||
const (
|
||||
// projectedColumnPaths are the parquet leaf columns sparkListenRow
|
||||
// projects. Kept in sync with that struct by
|
||||
// TestProjectedColumnsMatchSchema.
|
||||
projectedRecordingMBID = "recording_mbid"
|
||||
projectedReleaseMBID = "release_mbid"
|
||||
projectedArtistMBIDs = "artist_credit_mbids.list.element"
|
||||
|
||||
// rangeGapCoalesce merges two wanted byte ranges separated by less
|
||||
// than this much unwanted data. Below a round-trip's worth of
|
||||
// bytes, downloading the gap is cheaper than a second request.
|
||||
rangeGapCoalesce = 2 << 20
|
||||
|
||||
// projectFetchLanes is how many Range requests one member's column
|
||||
// fetch issues concurrently. Total in-flight requests are this
|
||||
// times projectMembersInFlight, kept at the dumpLanes budget the
|
||||
// dump server tolerates.
|
||||
projectFetchLanes = 2
|
||||
|
||||
// tarHeaderSize is the size of a tar header block.
|
||||
tarHeaderSize = 512
|
||||
|
||||
// walkAheadMembers bounds how far the tar header walk runs ahead of
|
||||
// the fetchers. The walk is one small request per member and is
|
||||
// latency-bound, so it needs a long leash to stay off the critical
|
||||
// path.
|
||||
walkAheadMembers = 64
|
||||
)
|
||||
|
||||
// projectedColumns is the set of leaf column paths to download.
|
||||
var projectedColumns = []string{
|
||||
projectedRecordingMBID,
|
||||
projectedReleaseMBID,
|
||||
projectedArtistMBIDs,
|
||||
}
|
||||
|
||||
// tarMember is one regular file inside the dump tar.
|
||||
type tarMember struct {
|
||||
// headerOffset is the absolute offset of the member's tar header.
|
||||
// This is what the stage-1 checkpoint records: resuming means
|
||||
// restarting the walk from here.
|
||||
headerOffset int64
|
||||
|
||||
// dataOffset is where the member's contents begin, and size how
|
||||
// many bytes they occupy.
|
||||
dataOffset int64
|
||||
size int64
|
||||
|
||||
name string
|
||||
|
||||
// typeflag is the tar entry type. Selecting members by name alone
|
||||
// is not enough: an extension header can carry the name of the
|
||||
// member it describes, and reading one as data yields PAX records
|
||||
// where parquet is expected.
|
||||
typeflag byte
|
||||
}
|
||||
|
||||
// nextHeaderOffset is the absolute offset of the following tar header.
|
||||
func (m tarMember) nextHeaderOffset() int64 {
|
||||
return m.dataOffset + m.size + tarPadding(m.size)
|
||||
}
|
||||
|
||||
// byteRange is a half-open [lo, hi) span of a resource.
|
||||
type byteRange struct {
|
||||
lo, hi int64
|
||||
}
|
||||
|
||||
func (r byteRange) len() int64 { return r.hi - r.lo }
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Range fetching
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// rangeFetcher performs retrying HTTP Range reads against one URL.
|
||||
type rangeFetcher struct {
|
||||
ctx context.Context
|
||||
client *http.Client
|
||||
url string
|
||||
|
||||
// footerProbe is how much of a member's tail to fetch when reading
|
||||
// its parquet footer; zero means defaultFooterProbe. The real
|
||||
// dump's footers measure well under that, and an undersized guess
|
||||
// costs one extra request rather than failing.
|
||||
footerProbe int64
|
||||
}
|
||||
|
||||
// defaultFooterProbe is the tail size used when a fetcher does not
|
||||
// override it.
|
||||
const defaultFooterProbe = 64 << 10
|
||||
|
||||
func (f *rangeFetcher) probeBytes() int64 {
|
||||
if f.footerProbe > 0 {
|
||||
return f.footerProbe
|
||||
}
|
||||
|
||||
return defaultFooterProbe
|
||||
}
|
||||
|
||||
// fetch reads [lo, hi) into buf (which must be exactly that long),
|
||||
// retrying transient failures and honouring Retry-After.
|
||||
func (f *rangeFetcher) fetch(r byteRange, buf []byte) error {
|
||||
var (
|
||||
lastErr error
|
||||
wait time.Duration
|
||||
)
|
||||
|
||||
for attempt := 0; attempt <= maxStreamRetries; attempt++ {
|
||||
if attempt > 0 {
|
||||
delay := min(streamRetryBaseDelay<<(attempt-1), streamRetryMaxDelay)
|
||||
if wait > 0 {
|
||||
delay = wait
|
||||
}
|
||||
|
||||
select {
|
||||
case <-f.ctx.Done():
|
||||
return f.ctx.Err()
|
||||
case <-time.After(delay):
|
||||
}
|
||||
}
|
||||
|
||||
retryAfter, err := f.fetchOnce(r, buf)
|
||||
if err == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
lastErr = err
|
||||
wait = retryAfter
|
||||
|
||||
if f.ctx.Err() != nil {
|
||||
return f.ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
return fmt.Errorf("%w: %s bytes %d-%d after %d retries: %w",
|
||||
ErrDumpStream, f.url, r.lo, r.hi-1, maxStreamRetries, lastErr)
|
||||
}
|
||||
|
||||
func (f *rangeFetcher) fetchOnce(r byteRange, buf []byte) (time.Duration, error) {
|
||||
req, err := http.NewRequestWithContext(f.ctx, http.MethodGet, f.url, nil)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("dump range request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("User-Agent", lbUserAgent)
|
||||
req.Header.Set("Range", "bytes="+strconv.FormatInt(r.lo, 10)+
|
||||
"-"+strconv.FormatInt(r.hi-1, 10))
|
||||
|
||||
resp, err := f.client.Do(req)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("dump range fetch: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
if resp.StatusCode != http.StatusPartialContent {
|
||||
// A loaded dump server answers 503/429 rather than queueing.
|
||||
return parseRetryAfter(resp.Header.Get("Retry-After")),
|
||||
fmt.Errorf("%w: HTTP %d from %s", ErrDumpStream, resp.StatusCode, f.url)
|
||||
}
|
||||
|
||||
if _, err := io.ReadFull(resp.Body, buf); err != nil {
|
||||
return 0, fmt.Errorf("dump range read: %w", err)
|
||||
}
|
||||
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tar member walking
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// walkTarMembers reads tar headers by Range request, emitting every
|
||||
// regular member from startOffset onward. Only the 512-byte headers are
|
||||
// downloaded; member contents are skipped by arithmetic rather than by
|
||||
// pulling them over the wire, which is the whole point.
|
||||
//
|
||||
// PAX extension headers (which the dump uses for long names) are
|
||||
// themselves regular members and are walked like any other; their
|
||||
// contents are not needed because the aggregator selects members by
|
||||
// suffix and the extended name only ever restates the short one.
|
||||
func walkTarMembers(
|
||||
ctx context.Context, f *rangeFetcher, startOffset, total int64, out chan<- tarMember,
|
||||
) error {
|
||||
defer close(out)
|
||||
|
||||
hdr := make([]byte, tarHeaderSize)
|
||||
offset := startOffset
|
||||
|
||||
for offset+tarHeaderSize <= total {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if err := f.fetch(byteRange{offset, offset + tarHeaderSize}, hdr); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
m, ok, err := parseTarHeader(hdr, offset)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if !ok {
|
||||
// Two zero blocks mark end of archive; one is enough to stop.
|
||||
return nil
|
||||
}
|
||||
|
||||
// A GNU long-name entry would leave the following member's name
|
||||
// truncated in its ustar field, and this walk never reads member
|
||||
// bodies to recover it. The dump uses PAX, so this is a format
|
||||
// change rather than something to paper over.
|
||||
if m.typeflag == tar.TypeGNULongName || m.typeflag == tar.TypeGNULongLink {
|
||||
return fmt.Errorf(
|
||||
"%w: GNU long-name entry at %d is not supported", ErrDumpFormat, offset,
|
||||
)
|
||||
}
|
||||
|
||||
select {
|
||||
case out <- m:
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
|
||||
offset = m.nextHeaderOffset()
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// parseTarHeader decodes one 512-byte header block. Returns ok=false
|
||||
// at the end-of-archive marker.
|
||||
func parseTarHeader(hdr []byte, offset int64) (tarMember, bool, error) {
|
||||
if isZeroBlock(hdr) {
|
||||
return tarMember{}, false, nil
|
||||
}
|
||||
|
||||
// The fields are decoded by hand rather than with archive/tar. A
|
||||
// tar.Reader given a lone header block works for plain ustar entries
|
||||
// but fails on the PAX extension headers the real dump writes before
|
||||
// every member: Next() reads a PAX record's body to merge its
|
||||
// attributes, and here that body is not in the block. Those headers
|
||||
// only restate the name the ustar fields already carry, so decoding
|
||||
// the fixed fields is both sufficient and immune to that.
|
||||
if err := verifyTarChecksum(hdr); err != nil {
|
||||
return tarMember{}, false, fmt.Errorf(
|
||||
"%w: tar header at %d: %w", ErrDumpFormat, offset, err,
|
||||
)
|
||||
}
|
||||
|
||||
size, err := parseTarSize(hdr[124:136])
|
||||
if err != nil {
|
||||
return tarMember{}, false, fmt.Errorf(
|
||||
"%w: tar header at %d: %w", ErrDumpFormat, offset, err,
|
||||
)
|
||||
}
|
||||
|
||||
if size < 0 || offset+tarHeaderSize+size < 0 {
|
||||
return tarMember{}, false, fmt.Errorf(
|
||||
"%w: tar header at %d declares size %d", ErrDumpFormat, offset, size,
|
||||
)
|
||||
}
|
||||
|
||||
return tarMember{
|
||||
headerOffset: offset,
|
||||
dataOffset: offset + tarHeaderSize,
|
||||
size: size,
|
||||
name: tarName(hdr),
|
||||
typeflag: hdr[156],
|
||||
}, true, nil
|
||||
}
|
||||
|
||||
// tarName joins the ustar prefix and name fields. Long names split
|
||||
// across the two are rejoined; names carried only in a PAX record are
|
||||
// not needed, because member selection is by suffix and the dump's
|
||||
// ustar name field always holds the full path.
|
||||
func tarName(hdr []byte) string {
|
||||
name := trimTarField(hdr[0:100])
|
||||
|
||||
if string(hdr[257:262]) != "ustar" {
|
||||
return name
|
||||
}
|
||||
|
||||
if prefix := trimTarField(hdr[345:500]); prefix != "" {
|
||||
return prefix + "/" + name
|
||||
}
|
||||
|
||||
return name
|
||||
}
|
||||
|
||||
func trimTarField(b []byte) string {
|
||||
if i := bytes.IndexByte(b, 0); i >= 0 {
|
||||
b = b[:i]
|
||||
}
|
||||
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// parseTarSize decodes a tar size field, which is octal ASCII in the
|
||||
// common case and big-endian base-256 (high bit set) for sizes that do
|
||||
// not fit — GNU's encoding for files above 8GB.
|
||||
func parseTarSize(field []byte) (int64, error) {
|
||||
if len(field) > 0 && field[0]&0x80 != 0 {
|
||||
var n int64
|
||||
|
||||
// The high bit is a flag, not part of the magnitude.
|
||||
for i, c := range field {
|
||||
if i == 0 {
|
||||
c &= 0x7F
|
||||
}
|
||||
|
||||
n = n<<8 | int64(c)
|
||||
}
|
||||
|
||||
return n, nil
|
||||
}
|
||||
|
||||
trimmed := strings.Trim(string(field), " \x00")
|
||||
if trimmed == "" {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
n, err := strconv.ParseInt(trimmed, 8, 64)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("bad size field %q: %w", field, err)
|
||||
}
|
||||
|
||||
return n, nil
|
||||
}
|
||||
|
||||
// verifyTarChecksum checks the header's own checksum. The walk seeks to
|
||||
// computed offsets rather than reading forward, so this is what catches
|
||||
// a desync before it is mistaken for a member.
|
||||
func verifyTarChecksum(hdr []byte) error {
|
||||
stored, err := parseTarSize(hdr[148:156])
|
||||
if err != nil {
|
||||
return fmt.Errorf("bad checksum field: %w", err)
|
||||
}
|
||||
|
||||
var signed, unsigned int64
|
||||
|
||||
for i, c := range hdr {
|
||||
// The checksum field itself is treated as spaces.
|
||||
if i >= 148 && i < 156 {
|
||||
c = ' '
|
||||
}
|
||||
|
||||
unsigned += int64(c)
|
||||
signed += int64(int8(c))
|
||||
}
|
||||
|
||||
if stored != unsigned && stored != signed {
|
||||
return fmt.Errorf(
|
||||
"%w: checksum %d does not match %d", ErrDumpFormat, stored, unsigned,
|
||||
)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func isZeroBlock(b []byte) bool {
|
||||
for _, c := range b {
|
||||
if c != 0 {
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Column projection
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// projectedMemberRanges returns the byte ranges of a member that must be
|
||||
// downloaded to decode projectedColumns: the parquet header magic, the
|
||||
// wanted column chunks, and the footer. Offsets are member-relative.
|
||||
func projectedMemberRanges(meta *format.FileMetaData, size int64, footerLen int64) []byteRange {
|
||||
ranges := []byteRange{
|
||||
// Leading "PAR1" magic — parquet readers verify it.
|
||||
{0, 4},
|
||||
}
|
||||
|
||||
for i := range meta.RowGroups {
|
||||
for j := range meta.RowGroups[i].Columns {
|
||||
col := &meta.RowGroups[i].Columns[j]
|
||||
|
||||
path := strings.Join(col.MetaData.PathInSchema, ".")
|
||||
if !slices.Contains(projectedColumns, path) {
|
||||
continue
|
||||
}
|
||||
|
||||
lo := col.MetaData.DataPageOffset
|
||||
if col.MetaData.DictionaryPageOffset != 0 {
|
||||
lo = col.MetaData.DictionaryPageOffset
|
||||
}
|
||||
|
||||
hi := lo + col.MetaData.TotalCompressedSize
|
||||
if lo < 0 || hi > size || hi <= lo {
|
||||
continue
|
||||
}
|
||||
|
||||
ranges = append(ranges, byteRange{lo, hi})
|
||||
}
|
||||
}
|
||||
|
||||
// The footer was already fetched to get here, but including it keeps
|
||||
// the buffer self-describing for the decoder.
|
||||
ranges = append(ranges, byteRange{size - footerLen, size})
|
||||
|
||||
return coalesceRanges(ranges)
|
||||
}
|
||||
|
||||
// coalesceRanges sorts and merges overlapping or near-adjacent ranges so
|
||||
// each becomes one HTTP request.
|
||||
func coalesceRanges(ranges []byteRange) []byteRange {
|
||||
if len(ranges) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
slices.SortFunc(ranges, func(a, b byteRange) int {
|
||||
return int(a.lo - b.lo)
|
||||
})
|
||||
|
||||
merged := ranges[:1]
|
||||
|
||||
for _, r := range ranges[1:] {
|
||||
last := &merged[len(merged)-1]
|
||||
if r.lo <= last.hi+rangeGapCoalesce {
|
||||
last.hi = max(last.hi, r.hi)
|
||||
|
||||
continue
|
||||
}
|
||||
|
||||
merged = append(merged, r)
|
||||
}
|
||||
|
||||
return merged
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Member fetching
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// fetchProjectedMember downloads only the projected columns of one
|
||||
// parquet member and returns a member-sized buffer with those bytes at
|
||||
// their real offsets. The returned count is how many bytes actually
|
||||
// crossed the wire.
|
||||
func fetchProjectedMember(
|
||||
ctx context.Context, f *rangeFetcher, m tarMember, buf []byte,
|
||||
) (int64, error) {
|
||||
if int64(len(buf)) != m.size {
|
||||
return 0, fmt.Errorf("%w: buffer %d for member of %d bytes",
|
||||
ErrDumpFormat, len(buf), m.size)
|
||||
}
|
||||
|
||||
// Unfilled gaps must not carry data from a previous member: a stale
|
||||
// page header there could be read as valid parquet.
|
||||
clear(buf)
|
||||
|
||||
footerLen := min(f.probeBytes(), m.size)
|
||||
|
||||
fetchTail := func(n int64) error {
|
||||
return f.fetch(
|
||||
byteRange{m.dataOffset + m.size - n, m.dataOffset + m.size},
|
||||
buf[m.size-n:],
|
||||
)
|
||||
}
|
||||
|
||||
if err := fetchTail(footerLen); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
fetched := footerLen
|
||||
|
||||
// A member smaller than the probe arrived whole; there is nothing
|
||||
// left to project out of it.
|
||||
if footerLen == m.size {
|
||||
return fetched, nil
|
||||
}
|
||||
|
||||
// A footer larger than the probe leaves the metadata truncated; the
|
||||
// trailing length field says how much is really needed.
|
||||
if need := declaredFooterLen(buf, m.size); need > footerLen {
|
||||
footerLen = min(need, m.size)
|
||||
|
||||
if err := fetchTail(footerLen); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
fetched += footerLen
|
||||
}
|
||||
|
||||
meta, err := readFooterMetadata(buf, m.size, footerLen)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
ranges := projectedMemberRanges(meta, m.size, footerLen)
|
||||
|
||||
got, err := fetchRangesConcurrent(ctx, f, m.dataOffset, ranges, buf)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
return fetched + got, nil
|
||||
}
|
||||
|
||||
// declaredFooterLen reads the footer length a parquet file advertises in
|
||||
// its last eight bytes, plus those eight bytes. Returns 0 when the tail
|
||||
// in buf is too short to hold the field.
|
||||
func declaredFooterLen(buf []byte, size int64) int64 {
|
||||
const trailer = 8
|
||||
|
||||
if size < trailer {
|
||||
return 0
|
||||
}
|
||||
|
||||
return int64(binary.LittleEndian.Uint32(buf[size-trailer:size-4])) + trailer
|
||||
}
|
||||
|
||||
// readFooterMetadata parses the parquet footer sitting in the tail of
|
||||
// buf. parquet.OpenFile is given a reader over the tail alone, offset
|
||||
// so that footer-relative seeks land correctly.
|
||||
func readFooterMetadata(buf []byte, size, footerLen int64) (*format.FileMetaData, error) {
|
||||
// A parquet file ends with the 4-byte footer length followed by
|
||||
// "PAR1"; the metadata precedes it.
|
||||
if footerLen < 8 {
|
||||
return nil, fmt.Errorf("%w: member too small for a parquet footer", ErrDumpFormat)
|
||||
}
|
||||
|
||||
// Only the tail of buf holds real bytes at this point, which is all
|
||||
// footer parsing needs. SkipMagicBytes suppresses the leading-magic
|
||||
// check: those four bytes are fetched with the column ranges and are
|
||||
// verified by the decoder when the assembled buffer is parsed.
|
||||
file, err := parquet.OpenFile(bytes.NewReader(buf), size,
|
||||
parquet.SkipMagicBytes(true),
|
||||
parquet.SkipPageIndex(true),
|
||||
parquet.SkipBloomFilters(true),
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("%w: parquet footer: %w", ErrDumpFormat, err)
|
||||
}
|
||||
|
||||
return file.Metadata(), nil
|
||||
}
|
||||
|
||||
// fetchRangesConcurrent downloads ranges (member-relative) into buf,
|
||||
// using a few lanes so a member's chunks aren't fetched one round trip
|
||||
// at a time.
|
||||
func fetchRangesConcurrent(
|
||||
ctx context.Context, f *rangeFetcher, base int64, ranges []byteRange, buf []byte,
|
||||
) (int64, error) {
|
||||
type job struct{ r byteRange }
|
||||
|
||||
jobs := make(chan job)
|
||||
errs := make(chan error, projectFetchLanes)
|
||||
|
||||
var fetched atomic.Int64
|
||||
|
||||
fetchCtx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
|
||||
laneFetcher := &rangeFetcher{ctx: fetchCtx, client: f.client, url: f.url}
|
||||
|
||||
for range projectFetchLanes {
|
||||
go func() {
|
||||
var firstErr error
|
||||
|
||||
for j := range jobs {
|
||||
if firstErr != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
dst := buf[j.r.lo:j.r.hi]
|
||||
abs := byteRange{base + j.r.lo, base + j.r.hi}
|
||||
|
||||
if err := laneFetcher.fetch(abs, dst); err != nil {
|
||||
firstErr = err
|
||||
|
||||
continue
|
||||
}
|
||||
|
||||
fetched.Add(j.r.len())
|
||||
}
|
||||
|
||||
errs <- firstErr
|
||||
}()
|
||||
}
|
||||
|
||||
for _, r := range ranges {
|
||||
select {
|
||||
case jobs <- job{r}:
|
||||
case <-fetchCtx.Done():
|
||||
close(jobs)
|
||||
|
||||
for range projectFetchLanes {
|
||||
<-errs
|
||||
}
|
||||
|
||||
return 0, fetchCtx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
close(jobs)
|
||||
|
||||
var firstErr error
|
||||
|
||||
for range projectFetchLanes {
|
||||
if err := <-errs; err != nil && firstErr == nil {
|
||||
firstErr = err
|
||||
}
|
||||
}
|
||||
|
||||
if firstErr != nil {
|
||||
return 0, firstErr
|
||||
}
|
||||
|
||||
return fetched.Load(), nil
|
||||
}
|
||||
|
||||
// projectionSupported reports whether the dump server will serve the
|
||||
// Range requests column projection depends on.
|
||||
func projectionSupported(ctx context.Context, client *http.Client, url string) (int64, bool) {
|
||||
size, ok := probeDumpSize(ctx, client, url)
|
||||
if !ok || size <= 0 {
|
||||
return 0, false
|
||||
}
|
||||
|
||||
return size, true
|
||||
}
|
||||
|
||||
var errProjectionUnsupported = errors.New("column projection unavailable")
|
||||
@@ -0,0 +1,716 @@
|
||||
//go:build indexbuild
|
||||
|
||||
package explore
|
||||
|
||||
import (
|
||||
"archive/tar"
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"yellowjacket/backend/database"
|
||||
)
|
||||
|
||||
// serveRangeBlob serves payload with Range support, counting the bytes
|
||||
// actually delivered so a test can assert how much crossed the wire.
|
||||
func serveRangeBlob(t *testing.T, payload []byte) (*httptest.Server, *atomic.Int64) {
|
||||
t.Helper()
|
||||
|
||||
var served atomic.Int64
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(
|
||||
func(w http.ResponseWriter, r *http.Request) {
|
||||
rec := &countingResponseWriter{ResponseWriter: w, n: &served}
|
||||
http.ServeContent(rec, r, "dump.tar", time.Unix(0, 0), bytes.NewReader(payload))
|
||||
},
|
||||
))
|
||||
|
||||
t.Cleanup(srv.Close)
|
||||
|
||||
return srv, &served
|
||||
}
|
||||
|
||||
type countingResponseWriter struct {
|
||||
http.ResponseWriter
|
||||
|
||||
n *atomic.Int64
|
||||
}
|
||||
|
||||
func (w *countingResponseWriter) Write(b []byte) (int, error) {
|
||||
n, err := w.ResponseWriter.Write(b)
|
||||
w.n.Add(int64(n))
|
||||
|
||||
return n, err
|
||||
}
|
||||
|
||||
// bigSparkTar builds a tar whose single parquet member has enough rows
|
||||
// that the unprojected columns dominate its size.
|
||||
func bigSparkTar(t *testing.T) []byte {
|
||||
t.Helper()
|
||||
|
||||
rows := make([]sparkFixtureRow, 0, 20_000)
|
||||
|
||||
for i := range 20_000 {
|
||||
rows = append(rows, sparkFixtureRow{
|
||||
ListenedAt: int64(i),
|
||||
UserID: int64(i % 977),
|
||||
// A high-cardinality column that is not projected: it is what
|
||||
// projection must avoid downloading.
|
||||
ArtistName: strings.Repeat("padding-", 12) + string(rune('a'+i%26)) + itoa(i),
|
||||
RecordingMBID: recA,
|
||||
ReleaseMBID: relA,
|
||||
ArtistMBIDs: []string{artA},
|
||||
})
|
||||
}
|
||||
|
||||
name := "listenbrainz-spark-dump-1-20260101-000003-full/1.parquet"
|
||||
|
||||
return makeTar(t, map[string][]byte{name: makeParquet(t, rows)}, []string{name})
|
||||
}
|
||||
|
||||
func itoa(i int) string {
|
||||
if i == 0 {
|
||||
return "0"
|
||||
}
|
||||
|
||||
var b []byte
|
||||
|
||||
for i > 0 {
|
||||
b = append([]byte{byte('0' + i%10)}, b...)
|
||||
i /= 10
|
||||
}
|
||||
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func TestWalkTarMembersFindsMembersWithoutReadingThem(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
payload := fixtureSparkTar(t)
|
||||
srv, served := serveRangeBlob(t, payload)
|
||||
|
||||
f := &rangeFetcher{ctx: t.Context(), client: srv.Client(), url: srv.URL}
|
||||
out := make(chan tarMember, 8)
|
||||
|
||||
if err := walkTarMembers(t.Context(), f, 0, int64(len(payload)), out); err != nil {
|
||||
t.Fatalf("walk: %v", err)
|
||||
}
|
||||
|
||||
var members []tarMember
|
||||
for m := range out {
|
||||
members = append(members, m)
|
||||
}
|
||||
|
||||
if len(members) != 2 {
|
||||
t.Fatalf("got %d members, want 2", len(members))
|
||||
}
|
||||
|
||||
for i, m := range members {
|
||||
if !strings.HasSuffix(m.name, ".parquet") {
|
||||
t.Errorf("member %d: name %q", i, m.name)
|
||||
}
|
||||
|
||||
if m.size <= 0 {
|
||||
t.Errorf("member %d: size %d", i, m.size)
|
||||
}
|
||||
|
||||
// The member's declared bytes must match what the tar really holds.
|
||||
want := payload[m.dataOffset : m.dataOffset+m.size]
|
||||
if !bytes.HasPrefix(want, []byte("PAR1")) {
|
||||
t.Errorf("member %d: data offset %d does not start a parquet file", i, m.dataOffset)
|
||||
}
|
||||
}
|
||||
|
||||
// The walk must read headers only — not the multi-KB member bodies.
|
||||
if n := served.Load(); n > int64(4*tarHeaderSize) {
|
||||
t.Errorf("walk downloaded %d bytes, want only tar headers", n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFetchProjectedMemberMatchesFullParse(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
payload := bigSparkTar(t)
|
||||
srv, served := serveRangeBlob(t, payload)
|
||||
|
||||
f := &rangeFetcher{ctx: t.Context(), client: srv.Client(), url: srv.URL}
|
||||
out := make(chan tarMember, 4)
|
||||
|
||||
if err := walkTarMembers(t.Context(), f, 0, int64(len(payload)), out); err != nil {
|
||||
t.Fatalf("walk: %v", err)
|
||||
}
|
||||
|
||||
var member tarMember
|
||||
|
||||
for m := range out {
|
||||
if strings.HasSuffix(m.name, ".parquet") {
|
||||
member = m
|
||||
}
|
||||
}
|
||||
|
||||
if member.size == 0 {
|
||||
t.Fatal("no parquet member found")
|
||||
}
|
||||
|
||||
// Ground truth: parse the member as the sequential path would.
|
||||
full := payload[member.dataOffset : member.dataOffset+member.size]
|
||||
|
||||
wantDeltas, err := parseListenParquet(full)
|
||||
if err != nil {
|
||||
t.Fatalf("full parse: %v", err)
|
||||
}
|
||||
|
||||
served.Store(0)
|
||||
|
||||
buf := make([]byte, member.size)
|
||||
|
||||
fetched, err := fetchProjectedMember(t.Context(), f, member, buf)
|
||||
if err != nil {
|
||||
t.Fatalf("projected fetch: %v", err)
|
||||
}
|
||||
|
||||
gotDeltas, err := parseListenParquet(buf)
|
||||
if err != nil {
|
||||
t.Fatalf("projected parse: %v", err)
|
||||
}
|
||||
|
||||
if len(gotDeltas) != len(wantDeltas) {
|
||||
t.Fatalf("projected parse produced %d entities, want %d", len(gotDeltas), len(wantDeltas))
|
||||
}
|
||||
|
||||
for k, want := range wantDeltas {
|
||||
if got := gotDeltas[k]; got != want {
|
||||
t.Errorf("entity %x: count %d, want %d", k, got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// The point of the exercise: materially fewer bytes than the member.
|
||||
if fetched >= member.size {
|
||||
t.Errorf("projected fetch pulled %d bytes of a %d byte member", fetched, member.size)
|
||||
}
|
||||
|
||||
t.Logf("projected fetch: %d of %d bytes (%.1f%%)",
|
||||
fetched, member.size, 100*float64(fetched)/float64(member.size))
|
||||
}
|
||||
|
||||
func TestProjectedMemberRangesSkipsUnwantedColumns(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
payload := bigSparkTar(t)
|
||||
srv, _ := serveRangeBlob(t, payload)
|
||||
|
||||
f := &rangeFetcher{ctx: t.Context(), client: srv.Client(), url: srv.URL}
|
||||
out := make(chan tarMember, 4)
|
||||
|
||||
if err := walkTarMembers(t.Context(), f, 0, int64(len(payload)), out); err != nil {
|
||||
t.Fatalf("walk: %v", err)
|
||||
}
|
||||
|
||||
var member tarMember
|
||||
|
||||
for m := range out {
|
||||
if strings.HasSuffix(m.name, ".parquet") {
|
||||
member = m
|
||||
}
|
||||
}
|
||||
|
||||
buf := make([]byte, member.size)
|
||||
footerLen := min(f.probeBytes(), member.size)
|
||||
tailStart := member.dataOffset + member.size - footerLen
|
||||
copy(buf[member.size-footerLen:], payload[tailStart:member.dataOffset+member.size])
|
||||
|
||||
meta, err := readFooterMetadata(buf, member.size, footerLen)
|
||||
if err != nil {
|
||||
t.Fatalf("footer: %v", err)
|
||||
}
|
||||
|
||||
ranges := projectedMemberRanges(meta, member.size, footerLen)
|
||||
if len(ranges) == 0 {
|
||||
t.Fatal("no ranges computed")
|
||||
}
|
||||
|
||||
var total int64
|
||||
|
||||
for _, r := range ranges {
|
||||
if r.lo < 0 || r.hi > member.size || r.hi <= r.lo {
|
||||
t.Fatalf("range %v out of bounds for member size %d", r, member.size)
|
||||
}
|
||||
|
||||
total += r.len()
|
||||
}
|
||||
|
||||
if total >= member.size {
|
||||
t.Errorf("projected ranges cover %d of %d bytes", total, member.size)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCoalesceRangesMergesNeighbours(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
in []byteRange
|
||||
want []byteRange
|
||||
}{
|
||||
{
|
||||
name: "adjacent merge",
|
||||
in: []byteRange{{0, 100}, {100, 200}},
|
||||
want: []byteRange{{0, 200}},
|
||||
},
|
||||
{
|
||||
name: "small gap merges",
|
||||
in: []byteRange{{0, 100}, {100 + rangeGapCoalesce - 1, 500}},
|
||||
want: []byteRange{{0, 500}},
|
||||
},
|
||||
{
|
||||
name: "large gap stays split",
|
||||
in: []byteRange{{0, 100}, {100 + rangeGapCoalesce + 1, 500}},
|
||||
want: []byteRange{{0, 100}, {100 + rangeGapCoalesce + 1, 500}},
|
||||
},
|
||||
{
|
||||
name: "unsorted input",
|
||||
in: []byteRange{{10 * rangeGapCoalesce, 10*rangeGapCoalesce + 100}, {0, 100}},
|
||||
want: []byteRange{{0, 100}, {10 * rangeGapCoalesce, 10*rangeGapCoalesce + 100}},
|
||||
},
|
||||
{
|
||||
name: "overlap",
|
||||
in: []byteRange{{0, 300}, {100, 200}},
|
||||
want: []byteRange{{0, 300}},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
got := coalesceRanges(tt.in)
|
||||
if len(got) != len(tt.want) {
|
||||
t.Fatalf("got %v, want %v", got, tt.want)
|
||||
}
|
||||
|
||||
for i := range got {
|
||||
if got[i] != tt.want[i] {
|
||||
t.Fatalf("got %v, want %v", got, tt.want)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestFetchProjectedMemberRejectsWrongBuffer(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
f := &rangeFetcher{
|
||||
ctx: context.Background(),
|
||||
client: http.DefaultClient,
|
||||
url: "http://example.invalid",
|
||||
}
|
||||
|
||||
if _, err := fetchProjectedMember(
|
||||
t.Context(), f, tarMember{size: 100}, make([]byte, 50),
|
||||
); err == nil {
|
||||
t.Fatal("expected an error for a mis-sized buffer")
|
||||
}
|
||||
}
|
||||
|
||||
// serveDumpNoRanges serves the spark dump without advertising Range
|
||||
// support, which is what forces the streamed fallback.
|
||||
func serveDumpNoRanges(t *testing.T, sparkTar []byte) *httptest.Server {
|
||||
t.Helper()
|
||||
|
||||
const (
|
||||
listensDir = "listenbrainz-dump-1-20260101-000003-full"
|
||||
sparkFile = "listenbrainz-spark-dump-1-20260101-000003-full.tar"
|
||||
)
|
||||
|
||||
mux := http.NewServeMux()
|
||||
|
||||
mux.HandleFunc("/listens/", func(w http.ResponseWriter, r *http.Request) {
|
||||
switch strings.TrimPrefix(r.URL.Path, "/listens/") {
|
||||
case "":
|
||||
_, _ = fmt.Fprintf(w, `<a href="%s/">%s/</a>`, listensDir, listensDir)
|
||||
case listensDir + "/":
|
||||
_, _ = fmt.Fprintf(w, `<a href="%s">%s</a>`, sparkFile, sparkFile)
|
||||
case listensDir + "/" + sparkFile:
|
||||
// No Accept-Ranges, and Range headers are ignored.
|
||||
w.Header().Set("Content-Length", strconv.Itoa(len(sparkTar)))
|
||||
_, _ = w.Write(sparkTar)
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
})
|
||||
|
||||
srv := httptest.NewServer(mux)
|
||||
t.Cleanup(srv.Close)
|
||||
|
||||
return srv
|
||||
}
|
||||
|
||||
// The projected and streamed paths must agree exactly: projection is an
|
||||
// optimisation, not a different answer.
|
||||
func TestProjectedAndStreamedCountsAgree(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
sparkTar := bigSparkTar(t)
|
||||
|
||||
counts := func(srv *httptest.Server) *countsState {
|
||||
t.Helper()
|
||||
|
||||
db := database.NewTestDB(t)
|
||||
si := NewSearchIndex(db, nil, nil, testLogger())
|
||||
imp := testImporter(t, si, srv)
|
||||
|
||||
url := srv.URL + "/listens/listenbrainz-dump-1-20260101-000003-full/" +
|
||||
"listenbrainz-spark-dump-1-20260101-000003-full.tar"
|
||||
|
||||
st := &countsState{SparkURL: url}
|
||||
if err := imp.aggregateListenCounts(t.Context(), st); err != nil {
|
||||
t.Fatalf("aggregate: %v", err)
|
||||
}
|
||||
|
||||
if !st.Done {
|
||||
t.Fatal("aggregation did not complete")
|
||||
}
|
||||
|
||||
return st
|
||||
}
|
||||
|
||||
ranged := serveDumps(t, sparkTar, nil)
|
||||
plain := serveDumpNoRanges(t, sparkTar)
|
||||
|
||||
projected := counts(ranged)
|
||||
streamed := counts(plain)
|
||||
|
||||
if len(projected.counts) == 0 {
|
||||
t.Fatal("projected run produced no counts")
|
||||
}
|
||||
|
||||
if len(projected.counts) != len(streamed.counts) {
|
||||
t.Fatalf("projected %d entities, streamed %d",
|
||||
len(projected.counts), len(streamed.counts))
|
||||
}
|
||||
|
||||
for k, want := range streamed.counts {
|
||||
if got := projected.counts[k]; got != want {
|
||||
t.Errorf("entity %x: projected %d, streamed %d", k, got, want)
|
||||
}
|
||||
}
|
||||
|
||||
if projected.Offset != streamed.Offset {
|
||||
t.Errorf("checkpoint offset: projected %d, streamed %d",
|
||||
projected.Offset, streamed.Offset)
|
||||
}
|
||||
|
||||
if projected.MemberIdx != streamed.MemberIdx {
|
||||
t.Errorf("member index: projected %d, streamed %d",
|
||||
projected.MemberIdx, streamed.MemberIdx)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeclaredFooterLenReadsTrailer(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
payload := bigSparkTar(t)
|
||||
srv, _ := serveRangeBlob(t, payload)
|
||||
|
||||
f := &rangeFetcher{ctx: t.Context(), client: srv.Client(), url: srv.URL}
|
||||
out := make(chan tarMember, 4)
|
||||
|
||||
if err := walkTarMembers(t.Context(), f, 0, int64(len(payload)), out); err != nil {
|
||||
t.Fatalf("walk: %v", err)
|
||||
}
|
||||
|
||||
var m tarMember
|
||||
|
||||
for got := range out {
|
||||
if strings.HasSuffix(got.name, ".parquet") {
|
||||
m = got
|
||||
}
|
||||
}
|
||||
|
||||
member := payload[m.dataOffset : m.dataOffset+m.size]
|
||||
|
||||
got := declaredFooterLen(member, m.size)
|
||||
if got <= 8 || got > m.size {
|
||||
t.Fatalf("declared footer length %d for a %d byte member", got, m.size)
|
||||
}
|
||||
|
||||
// A tail of exactly that length must be enough to parse the footer.
|
||||
buf := make([]byte, m.size)
|
||||
copy(buf[m.size-got:], member[m.size-got:])
|
||||
|
||||
if _, err := readFooterMetadata(buf, m.size, got); err != nil {
|
||||
t.Fatalf("footer parse with declared length: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// A footer bigger than the initial probe must trigger a second, larger
|
||||
// tail fetch rather than failing.
|
||||
func TestFetchProjectedMemberRefetchesOversizedFooter(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
payload := bigSparkTar(t)
|
||||
srv, _ := serveRangeBlob(t, payload)
|
||||
|
||||
f := &rangeFetcher{ctx: t.Context(), client: srv.Client(), url: srv.URL}
|
||||
out := make(chan tarMember, 4)
|
||||
|
||||
if err := walkTarMembers(t.Context(), f, 0, int64(len(payload)), out); err != nil {
|
||||
t.Fatalf("walk: %v", err)
|
||||
}
|
||||
|
||||
var m tarMember
|
||||
|
||||
for got := range out {
|
||||
if strings.HasSuffix(got.name, ".parquet") {
|
||||
m = got
|
||||
}
|
||||
}
|
||||
|
||||
member := payload[m.dataOffset : m.dataOffset+m.size]
|
||||
|
||||
want, err := parseListenParquet(member)
|
||||
if err != nil {
|
||||
t.Fatalf("full parse: %v", err)
|
||||
}
|
||||
|
||||
// Force the probe to land short of the real footer.
|
||||
footer := declaredFooterLen(member, m.size)
|
||||
|
||||
short := &rangeFetcher{
|
||||
ctx: t.Context(),
|
||||
client: srv.Client(),
|
||||
url: srv.URL,
|
||||
footerProbe: footer / 2,
|
||||
}
|
||||
|
||||
buf := make([]byte, m.size)
|
||||
|
||||
if _, err := fetchProjectedMember(t.Context(), short, m, buf); err != nil {
|
||||
t.Fatalf("projected fetch with short probe: %v", err)
|
||||
}
|
||||
|
||||
got, err := parseListenParquet(buf)
|
||||
if err != nil {
|
||||
t.Fatalf("projected parse: %v", err)
|
||||
}
|
||||
|
||||
if len(got) != len(want) {
|
||||
t.Fatalf("got %d entities, want %d", len(got), len(want))
|
||||
}
|
||||
}
|
||||
|
||||
// makePaxTar writes members preceded by PAX extension headers, which is
|
||||
// how the real ListenBrainz dump is written. A lone header block from
|
||||
// such an archive cannot be decoded with archive/tar, so this is the
|
||||
// layout the walker must handle directly.
|
||||
func makePaxTar(t *testing.T, members map[string][]byte, order []string) []byte {
|
||||
t.Helper()
|
||||
|
||||
var buf bytes.Buffer
|
||||
|
||||
tw := tar.NewWriter(&buf)
|
||||
|
||||
for _, name := range order {
|
||||
data := members[name]
|
||||
hdr := &tar.Header{
|
||||
Name: name,
|
||||
Mode: 0o644,
|
||||
Size: int64(len(data)),
|
||||
Typeflag: tar.TypeReg,
|
||||
// Sub-second precision cannot be expressed in ustar, so the
|
||||
// writer emits a PAX extension header ahead of the member.
|
||||
ModTime: time.Unix(1700000000, 123456789),
|
||||
Format: tar.FormatPAX,
|
||||
}
|
||||
|
||||
if err := tw.WriteHeader(hdr); err != nil {
|
||||
t.Fatalf("tar header: %v", err)
|
||||
}
|
||||
|
||||
if _, err := tw.Write(data); err != nil {
|
||||
t.Fatalf("tar write: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
if err := tw.Close(); err != nil {
|
||||
t.Fatalf("tar close: %v", err)
|
||||
}
|
||||
|
||||
return buf.Bytes()
|
||||
}
|
||||
|
||||
func TestWalkTarMembersHandlesPaxHeaders(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
rows := listensOf(4, recA, relA, []string{artA})
|
||||
name := "listenbrainz-spark-dump-1-20260101-000003-full/1.parquet"
|
||||
member := makeParquet(t, rows)
|
||||
|
||||
payload := makePaxTar(t, map[string][]byte{name: member}, []string{name})
|
||||
|
||||
// Guard the premise: the fixture really does contain a PAX header.
|
||||
if !bytes.Contains(payload[:4096], []byte("PaxHeader")) {
|
||||
t.Fatal("fixture has no PAX extension header")
|
||||
}
|
||||
|
||||
srv, _ := serveRangeBlob(t, payload)
|
||||
|
||||
f := &rangeFetcher{ctx: t.Context(), client: srv.Client(), url: srv.URL}
|
||||
out := make(chan tarMember, 16)
|
||||
|
||||
if err := walkTarMembers(t.Context(), f, 0, int64(len(payload)), out); err != nil {
|
||||
t.Fatalf("walk: %v", err)
|
||||
}
|
||||
|
||||
var found tarMember
|
||||
|
||||
for m := range out {
|
||||
if isProjectableMember(m) {
|
||||
found = m
|
||||
}
|
||||
}
|
||||
|
||||
if found.size != int64(len(member)) {
|
||||
t.Fatalf("member size %d, want %d", found.size, len(member))
|
||||
}
|
||||
|
||||
if got := payload[found.dataOffset : found.dataOffset+4]; string(got) != "PAR1" {
|
||||
t.Fatalf("data offset %d does not start a parquet file (%q)", found.dataOffset, got)
|
||||
}
|
||||
}
|
||||
|
||||
// End to end through the aggregator, against a PAX archive.
|
||||
func TestAggregateProjectedWithPaxHeaders(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
prefix := "listenbrainz-spark-dump-1-20260101-000003-full/"
|
||||
m1 := prefix + "1.parquet"
|
||||
m2 := prefix + "2.parquet"
|
||||
|
||||
payload := makePaxTar(t,
|
||||
map[string][]byte{
|
||||
m1: makeParquet(t, listensOf(12, recA, relA, []string{artA})),
|
||||
m2: makeParquet(t, listensOf(11, recB, relA, []string{artA})),
|
||||
},
|
||||
[]string{m1, m2},
|
||||
)
|
||||
|
||||
srv := serveDumps(t, payload, nil)
|
||||
|
||||
db := database.NewTestDB(t)
|
||||
si := NewSearchIndex(db, nil, nil, testLogger())
|
||||
imp := testImporter(t, si, srv)
|
||||
|
||||
url := srv.URL + "/listens/listenbrainz-dump-1-20260101-000003-full/" +
|
||||
"listenbrainz-spark-dump-1-20260101-000003-full.tar"
|
||||
|
||||
st := &countsState{SparkURL: url}
|
||||
if err := imp.aggregateListenCounts(t.Context(), st); err != nil {
|
||||
t.Fatalf("aggregate: %v", err)
|
||||
}
|
||||
|
||||
if !st.Done {
|
||||
t.Fatal("aggregation did not complete")
|
||||
}
|
||||
|
||||
assert := func(kind byte, mbid string, want uint32) {
|
||||
t.Helper()
|
||||
|
||||
key, ok := makeMBIDKey(kind, mbid)
|
||||
if !ok {
|
||||
t.Fatalf("bad fixture mbid %s", mbid)
|
||||
}
|
||||
|
||||
if got := st.counts[key]; got != want {
|
||||
t.Errorf("count(kind=%d, %s) = %d, want %d", kind, mbid, got, want)
|
||||
}
|
||||
}
|
||||
|
||||
assert(countKindRecording, recA, 12)
|
||||
assert(countKindRecording, recB, 11)
|
||||
assert(countKindRelease, relA, 23)
|
||||
assert(countKindArtist, artA, 23)
|
||||
}
|
||||
|
||||
func TestParseTarHeaderRejectsCorruptBlock(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
payload := fixtureSparkTar(t)
|
||||
|
||||
valid := make([]byte, tarHeaderSize)
|
||||
copy(valid, payload[:tarHeaderSize])
|
||||
|
||||
if _, ok, err := parseTarHeader(valid, 0); err != nil || !ok {
|
||||
t.Fatalf("valid header rejected: ok=%v err=%v", ok, err)
|
||||
}
|
||||
|
||||
// A desynced walk lands mid-member; the checksum must catch it.
|
||||
corrupt := make([]byte, tarHeaderSize)
|
||||
copy(corrupt, valid)
|
||||
corrupt[10] ^= 0xFF
|
||||
|
||||
if _, _, err := parseTarHeader(corrupt, 0); err == nil {
|
||||
t.Fatal("corrupt header accepted")
|
||||
}
|
||||
}
|
||||
|
||||
// A member smaller than the footer probe arrives in the probe request;
|
||||
// it must not then be fetched a second time.
|
||||
func TestFetchProjectedMemberSkipsRefetchForTinyMembers(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
rows := listensOf(2, recA, relA, []string{artA})
|
||||
name := "listenbrainz-spark-dump-1-20260101-000003-full/1.parquet"
|
||||
member := makeParquet(t, rows)
|
||||
|
||||
if int64(len(member)) >= defaultFooterProbe {
|
||||
t.Skipf("fixture member is %d bytes, not smaller than the probe", len(member))
|
||||
}
|
||||
|
||||
payload := makeTar(t, map[string][]byte{name: member}, []string{name})
|
||||
srv, served := serveRangeBlob(t, payload)
|
||||
|
||||
f := &rangeFetcher{ctx: t.Context(), client: srv.Client(), url: srv.URL}
|
||||
out := make(chan tarMember, 8)
|
||||
|
||||
if err := walkTarMembers(t.Context(), f, 0, int64(len(payload)), out); err != nil {
|
||||
t.Fatalf("walk: %v", err)
|
||||
}
|
||||
|
||||
var m tarMember
|
||||
|
||||
for got := range out {
|
||||
if isProjectableMember(got) {
|
||||
m = got
|
||||
}
|
||||
}
|
||||
|
||||
served.Store(0)
|
||||
|
||||
buf := make([]byte, m.size)
|
||||
|
||||
fetched, err := fetchProjectedMember(t.Context(), f, m, buf)
|
||||
if err != nil {
|
||||
t.Fatalf("projected fetch: %v", err)
|
||||
}
|
||||
|
||||
if fetched != m.size {
|
||||
t.Errorf("fetched %d bytes for a %d byte member", fetched, m.size)
|
||||
}
|
||||
|
||||
if n := served.Load(); n > m.size {
|
||||
t.Errorf("server delivered %d bytes for a %d byte member", n, m.size)
|
||||
}
|
||||
|
||||
if _, err := parseListenParquet(buf); err != nil {
|
||||
t.Fatalf("parse: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,339 @@
|
||||
//go:build indexbuild
|
||||
|
||||
package explore
|
||||
|
||||
import (
|
||||
"archive/tar"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Stage 1 driven by column projection. The tar is never streamed: a
|
||||
// walker reads member headers by Range request, and workers download and
|
||||
// parse each parquet member's projected columns independently. Results
|
||||
// are applied in member order, so the checkpoint stays a contiguous
|
||||
// prefix of the archive exactly as it is on the streamed path.
|
||||
|
||||
const (
|
||||
// projectMemberWorkers is how many members are fetched and parsed
|
||||
// concurrently. Each holds one member-sized buffer, and each issues
|
||||
// projectFetchLanes concurrent Range requests, so the product is the
|
||||
// in-flight request count the dump server sees.
|
||||
projectMemberWorkers = 3
|
||||
|
||||
// minParquetMemberSize is the smallest member that can hold a
|
||||
// parquet footer ("PAR1" + length + "PAR1"). Anything shorter is
|
||||
// not a parquet file whatever its name says.
|
||||
minParquetMemberSize = 12
|
||||
)
|
||||
|
||||
// indexedMember is a parquet member with its position in the aggregation
|
||||
// order, which is what the applier reassembles results by.
|
||||
type indexedMember struct {
|
||||
idx int
|
||||
m tarMember
|
||||
}
|
||||
|
||||
// aggregateProjected runs stage 1 by downloading only the projected
|
||||
// columns of each parquet member. Returns errProjectionUnsupported
|
||||
// (wrapped) when the dump's layout defeats projection, so the caller can
|
||||
// fall back before any counts are applied.
|
||||
func (imp *dumpImporter) aggregateProjected(
|
||||
ctx context.Context, st *countsState, total int64,
|
||||
) error {
|
||||
imp.logger.Info("dump import: streaming listen counts with column projection",
|
||||
"columns", strings.Join(projectedColumns, ","),
|
||||
"workers", projectMemberWorkers,
|
||||
"lanesPerWorker", projectFetchLanes,
|
||||
"resumeOffset", st.Offset,
|
||||
)
|
||||
|
||||
runCtx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
|
||||
fetcher := &rangeFetcher{ctx: runCtx, client: imp.httpClient, url: st.SparkURL}
|
||||
|
||||
prog := &projectedProgress{total: total}
|
||||
prog.position.Store(st.Offset)
|
||||
|
||||
stopReporter := imp.startCountsReporter(runCtx, prog, nil)
|
||||
defer stopReporter()
|
||||
|
||||
progress := &countsLogger{imp: imp, stream: prog, started: time.Now()}
|
||||
|
||||
// The walker runs ahead of the workers: each member costs it one
|
||||
// small request, so given a leash it never becomes the bottleneck.
|
||||
rawMembers := make(chan tarMember, walkAheadMembers)
|
||||
walkErr := make(chan error, 1)
|
||||
|
||||
go func() {
|
||||
walkErr <- walkTarMembers(runCtx, fetcher, st.Offset, total, rawMembers)
|
||||
}()
|
||||
|
||||
// Number the parquet members the aggregation actually consumes,
|
||||
// continuing from the checkpoint so indices stay stable on resume.
|
||||
work := make(chan indexedMember)
|
||||
|
||||
go func() {
|
||||
defer close(work)
|
||||
|
||||
idx := st.MemberIdx
|
||||
|
||||
for m := range rawMembers {
|
||||
if !isProjectableMember(m) {
|
||||
continue
|
||||
}
|
||||
|
||||
select {
|
||||
case work <- indexedMember{idx: idx, m: m}:
|
||||
idx++
|
||||
case <-runCtx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
results := make(chan countParseResult, projectMemberWorkers)
|
||||
|
||||
var workerWG sync.WaitGroup
|
||||
|
||||
for range projectMemberWorkers {
|
||||
workerWG.Add(1)
|
||||
|
||||
go func() {
|
||||
defer workerWG.Done()
|
||||
|
||||
buf := []byte(nil)
|
||||
|
||||
for job := range work {
|
||||
if cap(buf) < int(job.m.size) {
|
||||
buf = make([]byte, job.m.size)
|
||||
}
|
||||
|
||||
buf = buf[:job.m.size]
|
||||
|
||||
fetched, err := fetchProjectedMember(runCtx, fetcher, job.m, buf)
|
||||
if err == nil {
|
||||
prog.addFetched(fetched)
|
||||
}
|
||||
|
||||
var deltas map[mbidKey]uint32
|
||||
|
||||
if err == nil {
|
||||
deltas, err = parseListenParquet(buf)
|
||||
}
|
||||
|
||||
results <- countParseResult{
|
||||
idx: job.idx,
|
||||
endOffset: job.m.nextHeaderOffset(),
|
||||
deltas: deltas,
|
||||
err: err,
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
applier := newCountsApplier(imp, st, progress)
|
||||
applierDone := make(chan struct{})
|
||||
|
||||
go func() {
|
||||
defer close(applierDone)
|
||||
|
||||
for res := range results {
|
||||
applier.apply(res, prog)
|
||||
}
|
||||
}()
|
||||
|
||||
workerWG.Wait()
|
||||
close(results)
|
||||
<-applierDone
|
||||
|
||||
// Drain the walker so its error (if any) is observed and its
|
||||
// goroutine cannot outlive this call.
|
||||
cancel()
|
||||
|
||||
for range rawMembers { //nolint:revive // draining
|
||||
}
|
||||
|
||||
err := applier.err
|
||||
if err == nil {
|
||||
err = walkFailure(ctx, <-walkErr)
|
||||
} else {
|
||||
<-walkErr
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
// Best-effort checkpoint so even a cancelled run resumes where
|
||||
// it left off.
|
||||
_ = imp.writeCountsFile(st)
|
||||
|
||||
return err
|
||||
}
|
||||
|
||||
st.Done = true
|
||||
|
||||
if err := imp.writeCountsFile(st); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
imp.logger.Info("dump import: listen counts complete",
|
||||
"members", st.MemberIdx,
|
||||
"gb", fmt.Sprintf("%.1f", float64(st.Offset)/(1<<30)),
|
||||
"downloadedGB", fmt.Sprintf("%.1f", float64(prog.Downloaded())/(1<<30)),
|
||||
"entities", len(st.counts),
|
||||
"elapsed", time.Since(progress.started).Truncate(time.Second).String(),
|
||||
)
|
||||
|
||||
imp.logJob(fmt.Sprintf(
|
||||
"Listen counts complete — %s of listens read (%s downloaded), %s entities ranked",
|
||||
formatGB(st.Offset), formatGB(prog.Downloaded()), formatCount(len(st.counts)),
|
||||
))
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// walkFailure reports a walker error worth surfacing. A walk cancelled
|
||||
// because the workers finished first is not a failure.
|
||||
func walkFailure(ctx context.Context, err error) error {
|
||||
if err == nil || (errors.Is(err, context.Canceled) && ctx.Err() == nil) {
|
||||
return nil
|
||||
}
|
||||
|
||||
return err
|
||||
}
|
||||
|
||||
// isProjectableMember reports whether a tar member is a parquet file the
|
||||
// aggregator should consume. This must match the streamed path's member
|
||||
// selection exactly, or the two paths would produce different counts.
|
||||
func isProjectableMember(m tarMember) bool {
|
||||
if m.typeflag != tar.TypeReg && m.typeflag != 0 {
|
||||
return false
|
||||
}
|
||||
|
||||
return strings.HasSuffix(m.name, ".parquet") && m.size >= minParquetMemberSize
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Applier
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// countsApplier merges per-member deltas into the counts map in member
|
||||
// order, checkpointing every countsFlushEveryMembers members. It owns
|
||||
// st.counts, st.Offset and st.MemberIdx for the duration of a stage.
|
||||
type countsApplier struct {
|
||||
imp *dumpImporter
|
||||
st *countsState
|
||||
progress *countsLogger
|
||||
|
||||
pending map[int]countParseResult
|
||||
next int
|
||||
lastFlushed int
|
||||
|
||||
err error
|
||||
}
|
||||
|
||||
func newCountsApplier(
|
||||
imp *dumpImporter, st *countsState, progress *countsLogger,
|
||||
) *countsApplier {
|
||||
return &countsApplier{
|
||||
imp: imp,
|
||||
st: st,
|
||||
progress: progress,
|
||||
pending: make(map[int]countParseResult),
|
||||
next: st.MemberIdx,
|
||||
lastFlushed: st.MemberIdx,
|
||||
}
|
||||
}
|
||||
|
||||
// apply buffers a result and folds in every member that is now
|
||||
// contiguous with the checkpoint. pos, when non-nil, is advanced to the
|
||||
// archive offset the checkpoint has reached.
|
||||
func (a *countsApplier) apply(res countParseResult, pos *projectedProgress) {
|
||||
if a.err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
a.pending[res.idx] = res
|
||||
|
||||
for {
|
||||
r, ok := a.pending[a.next]
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
delete(a.pending, a.next)
|
||||
|
||||
if r.err != nil {
|
||||
a.err = r.err
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
for k, v := range r.deltas {
|
||||
a.st.counts[k] += v
|
||||
}
|
||||
|
||||
a.next++
|
||||
a.st.MemberIdx = a.next
|
||||
a.st.Offset = r.endOffset
|
||||
|
||||
if pos != nil {
|
||||
pos.position.Store(r.endOffset)
|
||||
}
|
||||
|
||||
if a.next-a.lastFlushed >= countsFlushEveryMembers {
|
||||
if err := a.imp.writeCountsFile(a.st); err != nil {
|
||||
a.err = err
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
a.lastFlushed = a.next
|
||||
|
||||
a.progress.checkpoint(a.next, r.endOffset, len(a.st.counts))
|
||||
|
||||
if err := a.imp.checkDiskHeadroom(); err != nil {
|
||||
a.err = err
|
||||
|
||||
return
|
||||
}
|
||||
} else {
|
||||
a.progress.member(a.next, r.endOffset, len(a.st.counts))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Progress
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// projectedProgress presents the projected import to the stage-1
|
||||
// reporter through the same interface a sequential stream uses.
|
||||
//
|
||||
// Position — not bytes downloaded — is what is reported as stream
|
||||
// progress: with projection those diverge (under half the archive is
|
||||
// downloaded), and it is position that gives a percentage and an ETA the
|
||||
// user can act on. Bytes actually downloaded are tracked separately and
|
||||
// logged at the end.
|
||||
type projectedProgress struct {
|
||||
total int64
|
||||
|
||||
position atomic.Int64
|
||||
downloaded atomic.Int64
|
||||
}
|
||||
|
||||
func (p *projectedProgress) Read([]byte) (int, error) { return 0, errProjectionUnsupported }
|
||||
func (p *projectedProgress) Close() error { return nil }
|
||||
func (p *projectedProgress) Pos() int64 { return p.position.Load() }
|
||||
func (p *projectedProgress) Fetched() int64 { return p.position.Load() }
|
||||
func (p *projectedProgress) Total() int64 { return p.total }
|
||||
|
||||
func (p *projectedProgress) addFetched(n int64) { p.downloaded.Add(n) }
|
||||
|
||||
// Downloaded is how many bytes actually crossed the wire.
|
||||
func (p *projectedProgress) Downloaded() int64 { return p.downloaded.Load() }
|
||||
@@ -0,0 +1,257 @@
|
||||
package explore
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/parquet-go/parquet-go"
|
||||
)
|
||||
|
||||
// Plumbing shared between the client and the CI-only index builder.
|
||||
//
|
||||
// The full dump import (dumpimport.go and friends) is behind the
|
||||
// `indexbuild` build tag so it is not linked into the app: a user's
|
||||
// machine never streams the ~89GB listens dump, it merges the prebuilt
|
||||
// artifact instead. What stays here is what the client genuinely still
|
||||
// needs — the daily incremental refresh (dumpincremental.go) and the
|
||||
// artifact download (artifactfetch.go) — plus the small helpers both
|
||||
// sides share.
|
||||
|
||||
// mbidKey is a parsed UUID plus an entity-kind tag, used as the counts
|
||||
// map key. 17 bytes instead of a 36-byte string keeps the ~40M-entry
|
||||
// map around 2GB.
|
||||
type mbidKey [17]byte
|
||||
|
||||
func makeMBIDKey(kind byte, mbid string) (mbidKey, bool) {
|
||||
var k mbidKey
|
||||
|
||||
k[0] = kind
|
||||
|
||||
if !parseUUID(mbid, k[1:]) {
|
||||
return k, false
|
||||
}
|
||||
|
||||
return k, true
|
||||
}
|
||||
|
||||
func formatUUID(b []byte) string {
|
||||
const hexdigits = "0123456789abcdef"
|
||||
|
||||
out := make([]byte, 36)
|
||||
j := 0
|
||||
|
||||
for i := range 16 {
|
||||
if i == 4 || i == 6 || i == 8 || i == 10 {
|
||||
out[j] = '-'
|
||||
j++
|
||||
}
|
||||
|
||||
out[j] = hexdigits[b[i]>>4]
|
||||
out[j+1] = hexdigits[b[i]&0x0F]
|
||||
j += 2
|
||||
}
|
||||
|
||||
return string(out)
|
||||
}
|
||||
|
||||
func formatGB(n int64) string {
|
||||
return fmt.Sprintf("%.1f GB", float64(n)/(1<<30))
|
||||
}
|
||||
|
||||
// formatCount renders a number with thousands separators, matching how
|
||||
// the frontend shows counts.
|
||||
func formatCount(n int) string {
|
||||
s := strconv.Itoa(n)
|
||||
if len(s) <= 3 {
|
||||
return s
|
||||
}
|
||||
|
||||
var b strings.Builder
|
||||
|
||||
lead := len(s) % 3
|
||||
if lead > 0 {
|
||||
b.WriteString(s[:lead])
|
||||
}
|
||||
|
||||
for i := lead; i < len(s); i += 3 {
|
||||
if b.Len() > 0 {
|
||||
b.WriteByte(',')
|
||||
}
|
||||
|
||||
b.WriteString(s[i : i+3])
|
||||
}
|
||||
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// parseListenParquet decodes one parquet member and returns the
|
||||
// per-entity listen-count deltas.
|
||||
func parseListenParquet(buf []byte) (map[mbidKey]uint32, error) {
|
||||
reader := parquet.NewGenericReader[sparkListenRow](bytes.NewReader(buf))
|
||||
|
||||
defer func() { _ = reader.Close() }()
|
||||
|
||||
deltas := make(map[mbidKey]uint32, 1<<18)
|
||||
rows := make([]sparkListenRow, 4096)
|
||||
|
||||
for {
|
||||
n, err := reader.Read(rows)
|
||||
|
||||
for _, row := range rows[:n] {
|
||||
key, ok := makeMBIDKey(countKindRecording, row.RecordingMBID)
|
||||
if !ok {
|
||||
// Unmapped listen — no usable recording MBID.
|
||||
continue
|
||||
}
|
||||
|
||||
deltas[key]++
|
||||
|
||||
if relKey, relOK := makeMBIDKey(countKindRelease, row.ReleaseMBID); relOK {
|
||||
deltas[relKey]++
|
||||
}
|
||||
|
||||
for _, artist := range row.ArtistMBIDs {
|
||||
if artKey, artOK := makeMBIDKey(countKindArtist, artist); artOK {
|
||||
deltas[artKey]++
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if errors.Is(err, io.EOF) {
|
||||
break
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("parquet read: %w", err)
|
||||
}
|
||||
|
||||
if n == 0 {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
return deltas, nil
|
||||
}
|
||||
|
||||
// sparkListenRow is the projection of the spark listens parquet schema
|
||||
// that the aggregator reads. All other columns are skipped.
|
||||
type sparkListenRow struct {
|
||||
RecordingMBID string `parquet:"recording_mbid,optional"`
|
||||
ReleaseMBID string `parquet:"release_mbid,optional"`
|
||||
ArtistMBIDs []string `parquet:"artist_credit_mbids,optional,list"`
|
||||
}
|
||||
|
||||
// checkFreeDisk returns ErrDiskSpace when the volume holding path has
|
||||
// less than minBytes free. Unknown free space (unsupported platform)
|
||||
// passes.
|
||||
func checkFreeDisk(path string, minBytes uint64) error {
|
||||
free, ok := diskFreeBytes(path)
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
|
||||
if free < minBytes {
|
||||
return fmt.Errorf("%w: %d MB free, need %d MB",
|
||||
ErrDiskSpace, free>>20, minBytes>>20)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// dumpSeriesRe extracts the monotonic series number NNNN from a dump
|
||||
// URL or directory name (e.g. "listenbrainz-spark-dump-2593-…").
|
||||
var dumpSeriesRe = regexp.MustCompile(`listenbrainz-(?:spark-)?dump-(\d+)-`)
|
||||
|
||||
// parseDumpSeries pulls the series number out of a dump URL/name.
|
||||
func parseDumpSeries(url string) (int, bool) {
|
||||
m := dumpSeriesRe.FindStringSubmatch(url)
|
||||
if m == nil {
|
||||
return 0, false
|
||||
}
|
||||
|
||||
n, err := strconv.Atoi(m[1])
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
|
||||
return n, true
|
||||
}
|
||||
|
||||
// Meta keys describing the catalog's provenance. They are read by the
|
||||
// client (the incremental refresh and the artifact import) and written
|
||||
// by whichever path populated the index.
|
||||
const (
|
||||
// dumpImportDoneKey marks a populated catalog in explore_index_meta.
|
||||
dumpImportDoneKey = "dump_import_done"
|
||||
|
||||
// listensAppliedSeriesKey stores the listens dump series the
|
||||
// popularity numbers are folded up to — the high-water-mark the
|
||||
// incremental refresh resumes from.
|
||||
listensAppliedSeriesKey = "listens_applied_series"
|
||||
)
|
||||
|
||||
// Entity kinds, used as the first byte of an mbidKey so one map can hold
|
||||
// counts for all three entity types.
|
||||
const (
|
||||
countKindRecording = byte(1)
|
||||
countKindRelease = byte(2)
|
||||
countKindArtist = byte(3)
|
||||
)
|
||||
|
||||
// maxParquetMemberSize caps how large a single parquet member may be
|
||||
// before it is treated as a malformed dump rather than buffered whole.
|
||||
const maxParquetMemberSize = 1 << 30
|
||||
|
||||
// ErrDiskSpace is returned when free disk falls below the safety floor.
|
||||
var ErrDiskSpace = errors.New("insufficient free disk space")
|
||||
|
||||
// parseUUID parses a canonical 36-char UUID string into 16 bytes.
|
||||
// Returns false for anything malformed.
|
||||
func parseUUID(s string, out []byte) bool {
|
||||
if len(s) != 36 || s[8] != '-' || s[13] != '-' || s[18] != '-' || s[23] != '-' {
|
||||
return false
|
||||
}
|
||||
|
||||
j := 0
|
||||
|
||||
for i := 0; i < 36; i++ {
|
||||
if i == 8 || i == 13 || i == 18 || i == 23 {
|
||||
continue
|
||||
}
|
||||
|
||||
hi := hexNibble(s[i])
|
||||
i++
|
||||
|
||||
lo := hexNibble(s[i])
|
||||
if hi == 0xFF || lo == 0xFF {
|
||||
return false
|
||||
}
|
||||
|
||||
out[j] = hi<<4 | lo
|
||||
j++
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
func hexNibble(c byte) byte {
|
||||
switch {
|
||||
case c >= '0' && c <= '9':
|
||||
return c - '0'
|
||||
case c >= 'a' && c <= 'f':
|
||||
return c - 'a' + 10
|
||||
case c >= 'A' && c <= 'F':
|
||||
return c - 'A' + 10
|
||||
default:
|
||||
return 0xFF
|
||||
}
|
||||
}
|
||||
|
||||
// ErrDumpFormat is returned when dump contents don't match the
|
||||
// expected format.
|
||||
var ErrDumpFormat = errors.New("unexpected dump format")
|
||||
+102
-126
@@ -7,8 +7,8 @@ import (
|
||||
"io"
|
||||
"net/http"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strconv"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
|
||||
@@ -25,9 +25,14 @@ const (
|
||||
maxStreamRetries = 8
|
||||
|
||||
// streamRetryBaseDelay is the initial reconnect backoff; it
|
||||
// doubles per consecutive failure.
|
||||
// doubles per consecutive failure, up to streamRetryMaxDelay.
|
||||
streamRetryBaseDelay = 2 * time.Second
|
||||
|
||||
// streamRetryMaxDelay caps the backoff. Uncapped doubling reaches
|
||||
// four minutes by the last attempt, which is a long time to leave a
|
||||
// download lane idle over a transient 503 from a busy dump server.
|
||||
streamRetryMaxDelay = 20 * time.Second
|
||||
|
||||
// dumpDiscoveryTimeout bounds the small directory-listing
|
||||
// requests (not the multi-hour stream requests).
|
||||
dumpDiscoveryTimeout = 30 * time.Second
|
||||
@@ -42,103 +47,6 @@ var ErrDumpStream = errors.New("dump stream failed")
|
||||
|
||||
var hrefRe = regexp.MustCompile(`href="([^"?/][^"?]*)"`)
|
||||
|
||||
// listHrefs fetches an Apache-style index page and returns the href
|
||||
// values (directory entries end with a trailing slash).
|
||||
func listHrefs(ctx context.Context, client *http.Client, url string) ([]string, error) {
|
||||
reqCtx, cancel := context.WithTimeout(ctx, dumpDiscoveryTimeout)
|
||||
defer cancel()
|
||||
|
||||
req, err := http.NewRequestWithContext(reqCtx, http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("dump listing request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("User-Agent", lbUserAgent)
|
||||
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("dump listing fetch: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf(
|
||||
"%w: listing %s returned HTTP %d", ErrDumpDiscovery, url, resp.StatusCode,
|
||||
)
|
||||
}
|
||||
|
||||
body, err := io.ReadAll(io.LimitReader(resp.Body, 4<<20))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("dump listing read: %w", err)
|
||||
}
|
||||
|
||||
var hrefs []string
|
||||
|
||||
for _, m := range hrefRe.FindAllStringSubmatch(string(body), -1) {
|
||||
hrefs = append(hrefs, m[1])
|
||||
}
|
||||
|
||||
return hrefs, nil
|
||||
}
|
||||
|
||||
// discoverDumpFile walks a dump base directory, finds subdirectories
|
||||
// matching dirRe (newest first, lexicographically — MetaBrainz dump
|
||||
// directory names embed sortable timestamps), and returns the full URL
|
||||
// of the first file inside matching fileRe. Directories that don't
|
||||
// contain a matching file (e.g. partial uploads) are skipped.
|
||||
func discoverDumpFile(
|
||||
ctx context.Context,
|
||||
client *http.Client,
|
||||
baseURL string,
|
||||
dirRe, fileRe *regexp.Regexp,
|
||||
) (string, error) {
|
||||
hrefs, err := listHrefs(ctx, client, baseURL)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
var dirs []string
|
||||
|
||||
for _, h := range hrefs {
|
||||
trimmed := trimTrailingSlash(h)
|
||||
if dirRe.MatchString(trimmed) {
|
||||
dirs = append(dirs, trimmed)
|
||||
}
|
||||
}
|
||||
|
||||
if len(dirs) == 0 {
|
||||
return "", fmt.Errorf("%w: no dump directories under %s", ErrDumpDiscovery, baseURL)
|
||||
}
|
||||
|
||||
sort.Sort(sort.Reverse(sort.StringSlice(dirs)))
|
||||
|
||||
for _, dir := range dirs {
|
||||
dirURL := baseURL + dir + "/"
|
||||
|
||||
files, err := listHrefs(ctx, client, dirURL)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
for _, f := range files {
|
||||
if fileRe.MatchString(f) {
|
||||
return dirURL + f, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return "", fmt.Errorf("%w: no matching dump file under %s", ErrDumpDiscovery, baseURL)
|
||||
}
|
||||
|
||||
func trimTrailingSlash(s string) string {
|
||||
if len(s) > 0 && s[len(s)-1] == '/' {
|
||||
return s[:len(s)-1]
|
||||
}
|
||||
|
||||
return s
|
||||
}
|
||||
|
||||
// resumableReader is an io.Reader over an HTTP resource that survives
|
||||
// connection failures by reconnecting with a Range request at the
|
||||
// current offset. Offset is the absolute position of the next byte to
|
||||
@@ -149,27 +57,31 @@ type resumableReader struct {
|
||||
client *http.Client
|
||||
url string
|
||||
|
||||
// Offset is the absolute byte position of the next read.
|
||||
Offset int64
|
||||
|
||||
// Size is the total resource size, learned from the first
|
||||
// response. -1 until known.
|
||||
Size int64
|
||||
// offset is the absolute byte position of the next read, and size
|
||||
// the total resource size (-1 until the first response reveals it).
|
||||
// Both are atomic so a progress reporter on another goroutine can
|
||||
// sample them while the stream is being read.
|
||||
offset atomic.Int64
|
||||
size atomic.Int64
|
||||
|
||||
body io.ReadCloser
|
||||
retries int
|
||||
}
|
||||
|
||||
func newResumableReader(
|
||||
ctx context.Context, client *http.Client, url string, offset int64,
|
||||
) *resumableReader {
|
||||
return &resumableReader{
|
||||
ctx: ctx,
|
||||
client: client,
|
||||
url: url,
|
||||
Offset: offset,
|
||||
Size: -1,
|
||||
}
|
||||
// Pos returns the absolute byte position of the next read.
|
||||
func (r *resumableReader) Pos() int64 {
|
||||
return r.offset.Load()
|
||||
}
|
||||
|
||||
// Fetched matches Pos: a single sequential connection reads no further
|
||||
// ahead than it delivers.
|
||||
func (r *resumableReader) Fetched() int64 {
|
||||
return r.offset.Load()
|
||||
}
|
||||
|
||||
// Total returns the total resource size, or -1 while unknown.
|
||||
func (r *resumableReader) Total() int64 {
|
||||
return r.size.Load()
|
||||
}
|
||||
|
||||
func (r *resumableReader) Read(p []byte) (int, error) {
|
||||
@@ -185,7 +97,7 @@ func (r *resumableReader) Read(p []byte) (int, error) {
|
||||
}
|
||||
|
||||
n, err := r.body.Read(p)
|
||||
r.Offset += int64(n)
|
||||
offset := r.offset.Add(int64(n))
|
||||
|
||||
if n > 0 {
|
||||
r.retries = 0
|
||||
@@ -197,7 +109,7 @@ func (r *resumableReader) Read(p []byte) (int, error) {
|
||||
case errors.Is(err, io.EOF):
|
||||
// A server that closes early looks like EOF; only
|
||||
// trust it when we've seen the advertised size.
|
||||
if r.Size >= 0 && r.Offset < r.Size {
|
||||
if size := r.size.Load(); size >= 0 && offset < size {
|
||||
r.closeBody()
|
||||
|
||||
if retryErr := r.backoff(err); retryErr != nil {
|
||||
@@ -236,7 +148,7 @@ func (r *resumableReader) backoff(cause error) error {
|
||||
)
|
||||
}
|
||||
|
||||
delay := streamRetryBaseDelay << (r.retries - 1)
|
||||
delay := min(streamRetryBaseDelay<<(r.retries-1), streamRetryMaxDelay)
|
||||
|
||||
select {
|
||||
case <-r.ctx.Done():
|
||||
@@ -254,8 +166,9 @@ func (r *resumableReader) connect() error {
|
||||
|
||||
req.Header.Set("User-Agent", lbUserAgent)
|
||||
|
||||
if r.Offset > 0 {
|
||||
req.Header.Set("Range", "bytes="+strconv.FormatInt(r.Offset, 10)+"-")
|
||||
offset := r.offset.Load()
|
||||
if offset > 0 {
|
||||
req.Header.Set("Range", "bytes="+strconv.FormatInt(offset, 10)+"-")
|
||||
}
|
||||
|
||||
resp, err := r.client.Do(req)
|
||||
@@ -265,22 +178,22 @@ func (r *resumableReader) connect() error {
|
||||
|
||||
switch resp.StatusCode {
|
||||
case http.StatusPartialContent:
|
||||
if r.Size < 0 {
|
||||
r.Size = parseContentRangeTotal(resp.Header.Get("Content-Range"))
|
||||
if r.size.Load() < 0 {
|
||||
r.size.Store(parseContentRangeTotal(resp.Header.Get("Content-Range")))
|
||||
}
|
||||
|
||||
r.body = resp.Body
|
||||
|
||||
return nil
|
||||
case http.StatusOK:
|
||||
if r.Size < 0 && resp.ContentLength > 0 {
|
||||
r.Size = resp.ContentLength
|
||||
if r.size.Load() < 0 && resp.ContentLength > 0 {
|
||||
r.size.Store(resp.ContentLength)
|
||||
}
|
||||
|
||||
// Server ignored the Range header: discard the prefix so
|
||||
// the caller still reads from the requested offset.
|
||||
if r.Offset > 0 {
|
||||
if _, err := io.CopyN(io.Discard, resp.Body, r.Offset); err != nil {
|
||||
if offset > 0 {
|
||||
if _, err := io.CopyN(io.Discard, resp.Body, offset); err != nil {
|
||||
_ = resp.Body.Close()
|
||||
|
||||
return r.backoff(err)
|
||||
@@ -327,3 +240,66 @@ func parseContentRangeTotal(v string) int64 {
|
||||
|
||||
return -1
|
||||
}
|
||||
|
||||
func newResumableReader(
|
||||
ctx context.Context, client *http.Client, url string, offset int64,
|
||||
) *resumableReader {
|
||||
r := &resumableReader{
|
||||
ctx: ctx,
|
||||
client: client,
|
||||
url: url,
|
||||
}
|
||||
|
||||
r.offset.Store(offset)
|
||||
r.size.Store(-1)
|
||||
|
||||
return r
|
||||
}
|
||||
|
||||
// listHrefs fetches an Apache-style index page and returns the href
|
||||
// values (directory entries end with a trailing slash).
|
||||
func listHrefs(ctx context.Context, client *http.Client, url string) ([]string, error) {
|
||||
reqCtx, cancel := context.WithTimeout(ctx, dumpDiscoveryTimeout)
|
||||
defer cancel()
|
||||
|
||||
req, err := http.NewRequestWithContext(reqCtx, http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("dump listing request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("User-Agent", lbUserAgent)
|
||||
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("dump listing fetch: %w", err)
|
||||
}
|
||||
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf(
|
||||
"%w: listing %s returned HTTP %d", ErrDumpDiscovery, url, resp.StatusCode,
|
||||
)
|
||||
}
|
||||
|
||||
body, err := io.ReadAll(io.LimitReader(resp.Body, 4<<20))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("dump listing read: %w", err)
|
||||
}
|
||||
|
||||
var hrefs []string
|
||||
|
||||
for _, m := range hrefRe.FindAllStringSubmatch(string(body), -1) {
|
||||
hrefs = append(hrefs, m[1])
|
||||
}
|
||||
|
||||
return hrefs, nil
|
||||
}
|
||||
|
||||
func trimTrailingSlash(s string) string {
|
||||
if len(s) > 0 && s[len(s)-1] == '/' {
|
||||
return s[:len(s)-1]
|
||||
}
|
||||
|
||||
return s
|
||||
}
|
||||
|
||||
@@ -149,6 +149,12 @@ func (e *Service) StopIndexBuild() {
|
||||
e.index.StopBuild()
|
||||
}
|
||||
|
||||
// CoreCatalogImported reports whether a prebuilt catalog artifact has
|
||||
// been merged into this index.
|
||||
func (e *Service) CoreCatalogImported() bool {
|
||||
return e.index.artifactAlreadyMerged()
|
||||
}
|
||||
|
||||
// SetJobRegistry wires the background job registry into the search
|
||||
// index so its build reports progress and controls to the frontend.
|
||||
func (e *Service) SetJobRegistry(reg *jobs.Registry) {
|
||||
@@ -565,8 +571,7 @@ func (e *Service) resolveArtistName(artistMBID string, rgs []MBReleaseGroup) str
|
||||
// Check the index for a previously-indexed artist row.
|
||||
if indexed := e.index.LookupArtistByMBID(
|
||||
artistMBID,
|
||||
); indexed != nil && indexed.Title != "" &&
|
||||
indexed.Title != artistMBID {
|
||||
); indexed != nil && indexed.Title != "" {
|
||||
return indexed.Title
|
||||
}
|
||||
|
||||
@@ -1048,8 +1053,7 @@ func (e *Service) GetArtistImageURL(artistMBID string) string {
|
||||
|
||||
// GetArtistImageCached returns a base64 data URL for the artist's
|
||||
// photo ONLY if it's already on disk — no MB/Wikidata resolution
|
||||
// or Wikimedia fetch. Safe to call from library-only mode.
|
||||
// Returns "" if not cached.
|
||||
// or Wikimedia fetch. Returns "" if not cached.
|
||||
func (e *Service) GetArtistImageCached(artistMBID string) string {
|
||||
return e.artistImg.GetCachedImage(artistMBID)
|
||||
}
|
||||
|
||||
@@ -153,6 +153,12 @@ func (si *SearchIndex) applyStagesToJob(h *jobs.Handle, status IndexStatus) {
|
||||
phase = t.Name
|
||||
current = int64(t.Completed)
|
||||
total = int64(t.Total)
|
||||
|
||||
// The listens stream's raw completed/total is a bare
|
||||
// percentage; the detail line is what makes it legible.
|
||||
if t.Detail != "" {
|
||||
phase = t.Name + " — " + t.Detail
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,213 @@
|
||||
//go:build indexbuild
|
||||
|
||||
package explore
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"regexp"
|
||||
"sort"
|
||||
"sync"
|
||||
|
||||
"yellowjacket/backend/jobs"
|
||||
)
|
||||
|
||||
// Helpers used only while building the catalog from the MetaBrainz
|
||||
// dumps: the similar-artist patch pass and the per-stage error
|
||||
// reporting that goes with it. They live behind the `indexbuild` tag
|
||||
// with the importer they serve, so the app binary carries neither.
|
||||
|
||||
const (
|
||||
// indexSimilarPerArtist is how many similar artists to store per
|
||||
// library artist in similar_artist_map.
|
||||
indexSimilarPerArtist = 20
|
||||
|
||||
// similarArtistsBatchSize is the number of seed MBIDs processed in
|
||||
// one logging "batch". The labs multi-seed POST form is broken, so
|
||||
// one GET per seed is issued (concurrency bounded by indexerRate);
|
||||
// batching here just keeps progress log output bounded.
|
||||
similarArtistsBatchSize = 50
|
||||
)
|
||||
|
||||
// setTierError marks a build stage as errored.
|
||||
func (si *SearchIndex) setTierError(name, errMsg string) {
|
||||
si.mu.Lock()
|
||||
|
||||
for i := range si.buildStatus.Tiers {
|
||||
if si.buildStatus.Tiers[i].Name == name {
|
||||
si.buildStatus.Tiers[i].State = "error"
|
||||
si.buildStatus.Tiers[i].Error = errMsg
|
||||
si.mu.Unlock()
|
||||
|
||||
si.logIndexJob(jobs.LevelError, name+": "+errMsg)
|
||||
si.emitStatus()
|
||||
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
si.mu.Unlock()
|
||||
}
|
||||
|
||||
// fetchSimilarArtistsBatch queries the labs similar-artists endpoint
|
||||
// for multiple seed MBIDs. Despite the name, this actually fans
|
||||
// out one request per seed: the labs API's multi-seed mode is
|
||||
// broken (results for different seeds get mis-labeled, and some
|
||||
// seeds return zero), so batching with multiple artist_mbids is
|
||||
// not viable. Concurrency is bounded by indexerRate to respect
|
||||
// the labs rate limit; each call goes through the provided LB
|
||||
// client's rate limiter and cache.
|
||||
func (si *SearchIndex) fetchSimilarArtistsBatch(
|
||||
ctx context.Context, lb *ListenBrainzClient, seedMBIDs []string,
|
||||
) map[string][]lbSimilarArtistWire {
|
||||
if len(seedMBIDs) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
var (
|
||||
mu sync.Mutex
|
||||
grouped = make(map[string][]lbSimilarArtistWire, len(seedMBIDs))
|
||||
wg sync.WaitGroup
|
||||
)
|
||||
|
||||
sem := make(chan struct{}, indexerRate)
|
||||
|
||||
for _, seedMBID := range seedMBIDs {
|
||||
if ctx.Err() != nil {
|
||||
break
|
||||
}
|
||||
|
||||
sem <- struct{}{}
|
||||
|
||||
wg.Add(1)
|
||||
|
||||
go func(seed string) {
|
||||
defer func() {
|
||||
<-sem
|
||||
wg.Done()
|
||||
}()
|
||||
|
||||
// Use the LB client's per-seed GET form — goes through
|
||||
// the shared rate limiter and cache. The multi-seed
|
||||
// POST form is not viable (see function comment).
|
||||
similar, err := lb.SimilarArtists(ctx, seed)
|
||||
if err != nil || len(similar) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
// Convert to the internal wire type used by the caller
|
||||
// and trim to indexSimilarPerArtist.
|
||||
if len(similar) > indexSimilarPerArtist {
|
||||
similar = similar[:indexSimilarPerArtist]
|
||||
}
|
||||
|
||||
results := make([]lbSimilarArtistWire, len(similar))
|
||||
for i, s := range similar {
|
||||
results[i] = lbSimilarArtistWire{
|
||||
ArtistMBID: s.ArtistMBID,
|
||||
Name: s.Name,
|
||||
Score: int(s.Score),
|
||||
ReferenceMBID: seed,
|
||||
}
|
||||
}
|
||||
|
||||
mu.Lock()
|
||||
grouped[seed] = results
|
||||
mu.Unlock()
|
||||
}(seedMBID)
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
|
||||
return grouped
|
||||
}
|
||||
|
||||
// chunkStrings splits a slice into chunks of at most size n.
|
||||
func chunkStrings(s []string, n int) [][]string {
|
||||
var chunks [][]string
|
||||
|
||||
for i := 0; i < len(s); i += n {
|
||||
end := i + n
|
||||
if end > len(s) {
|
||||
end = len(s)
|
||||
}
|
||||
|
||||
chunks = append(chunks, s[i:end])
|
||||
}
|
||||
|
||||
return chunks
|
||||
}
|
||||
|
||||
// getLibraryArtistMBIDs returns MBIDs for all library artists that have one.
|
||||
// Used when Tier 3 was skipped but Tier 4 needs the library MBID list.
|
||||
func (si *SearchIndex) getLibraryArtistMBIDs() []string {
|
||||
rows, err := si.db.QueryContext(
|
||||
"SELECT DISTINCT mbid FROM artists WHERE mbid IS NOT NULL AND mbid != ''",
|
||||
)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
defer func() { _ = rows.Close() }()
|
||||
|
||||
var mbids []string
|
||||
|
||||
for rows.Next() {
|
||||
var mbid string
|
||||
if err := rows.Scan(&mbid); err == nil {
|
||||
mbids = append(mbids, mbid)
|
||||
}
|
||||
}
|
||||
|
||||
return mbids
|
||||
}
|
||||
|
||||
// discoverDumpFile walks a dump base directory, finds subdirectories
|
||||
// matching dirRe (newest first, lexicographically — MetaBrainz dump
|
||||
// directory names embed sortable timestamps), and returns the full URL
|
||||
// of the first file inside matching fileRe. Directories that don't
|
||||
// contain a matching file (e.g. partial uploads) are skipped.
|
||||
func discoverDumpFile(
|
||||
ctx context.Context,
|
||||
client *http.Client,
|
||||
baseURL string,
|
||||
dirRe, fileRe *regexp.Regexp,
|
||||
) (string, error) {
|
||||
hrefs, err := listHrefs(ctx, client, baseURL)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
var dirs []string
|
||||
|
||||
for _, h := range hrefs {
|
||||
trimmed := trimTrailingSlash(h)
|
||||
if dirRe.MatchString(trimmed) {
|
||||
dirs = append(dirs, trimmed)
|
||||
}
|
||||
}
|
||||
|
||||
if len(dirs) == 0 {
|
||||
return "", fmt.Errorf("%w: no dump directories under %s", ErrDumpDiscovery, baseURL)
|
||||
}
|
||||
|
||||
sort.Sort(sort.Reverse(sort.StringSlice(dirs)))
|
||||
|
||||
for _, dir := range dirs {
|
||||
dirURL := baseURL + dir + "/"
|
||||
|
||||
files, err := listHrefs(ctx, client, dirURL)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
for _, f := range files {
|
||||
if fileRe.MatchString(f) {
|
||||
return dirURL + f, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return "", fmt.Errorf("%w: no matching dump file under %s", ErrDumpDiscovery, baseURL)
|
||||
}
|
||||
+111
-229
@@ -94,23 +94,12 @@ const (
|
||||
// across launches instead of running for the better part of an hour.
|
||||
discogBackfillMaxPerRun = 2000
|
||||
|
||||
// indexSimilarPerArtist is how many similar artists to store
|
||||
// per library artist in similar_artist_map.
|
||||
indexSimilarPerArtist = 20
|
||||
|
||||
// labsBaseURL is the base URL for the ListenBrainz labs API.
|
||||
labsBaseURL = "https://labs.api.listenbrainz.org"
|
||||
|
||||
// labsSimilarAlgorithm is the algorithm parameter for the
|
||||
// similar-artists endpoint.
|
||||
labsSimilarAlgorithm = "session_based_days_7500_session_300_contribution_5_threshold_10_limit_100_filter_True_skip_30"
|
||||
|
||||
// similarArtistsBatchSize is the number of seed MBIDs processed
|
||||
// in one logging "batch" during Tier 4. The labs multi-seed
|
||||
// POST form is broken, so we actually issue one GET per seed
|
||||
// (concurrency bounded by indexerRate); batching here just
|
||||
// keeps progress log output bounded.
|
||||
similarArtistsBatchSize = 50
|
||||
)
|
||||
|
||||
// SearchIndexResult is a single hit from the local popularity index.
|
||||
@@ -160,16 +149,8 @@ type SearchIndexResult struct {
|
||||
LocalArtistID int64 `json:"localArtistId,omitempty"`
|
||||
LocalReleaseGroupID int64 `json:"localReleaseGroupId,omitempty"`
|
||||
LocalRecordingID int64 `json:"localRecordingId,omitempty"`
|
||||
|
||||
// Schema version for staleness detection.
|
||||
SchemaVersion int `json:"-"`
|
||||
}
|
||||
|
||||
// currentSchemaVersion is bumped when we add new fields that should
|
||||
// trigger re-indexing of existing rows. The build logic checks each
|
||||
// artist's rows against this version and re-fetches if stale.
|
||||
const currentSchemaVersion = 1
|
||||
|
||||
// lbSitewideArtist is the response shape from the LB sitewide
|
||||
// top-artists endpoint.
|
||||
type lbSitewideArtist struct {
|
||||
@@ -239,6 +220,11 @@ type TierStatus struct {
|
||||
Total int `json:"total"`
|
||||
Completed int `json:"completed"`
|
||||
Error string `json:"error,omitempty"`
|
||||
|
||||
// Detail is a human-readable progress line for stages whose raw
|
||||
// completed/total numbers say little on their own — the listens
|
||||
// stream reports "42.3 / 205.1 GB · 18 MB/s · ~3h20m left" here.
|
||||
Detail string `json:"detail,omitempty"`
|
||||
}
|
||||
|
||||
// IndexStatus is the full index build status, exposed to the frontend.
|
||||
@@ -456,7 +442,7 @@ func (si *SearchIndex) BackfillLibraryDiscographies(ctx context.Context) {
|
||||
// itself. Used to seed the discography fetch's artist entry.
|
||||
func (si *SearchIndex) artistDisplayName(mbid string) string {
|
||||
for _, q := range []string{
|
||||
"SELECT title FROM explore_index WHERE entity_type = 'artist' AND mbid = ? AND title != '' AND title != mbid LIMIT 1",
|
||||
"SELECT title FROM explore_index WHERE entity_type = 'artist' AND mbid = ? AND title != '' LIMIT 1",
|
||||
"SELECT name FROM artists WHERE mbid = ? AND name != '' LIMIT 1",
|
||||
} {
|
||||
rows, err := si.db.QueryContext(q, mbid)
|
||||
@@ -651,6 +637,11 @@ func (si *SearchIndex) refreshStatusCounts() {
|
||||
|
||||
// setTierStatus updates the build status for a named tier.
|
||||
func (si *SearchIndex) setTierStatus(name, state string, total, completed int) {
|
||||
si.setTierDetail(name, state, total, completed, "")
|
||||
}
|
||||
|
||||
// setTierDetail is setTierStatus plus a human-readable progress line.
|
||||
func (si *SearchIndex) setTierDetail(name, state string, total, completed int, detail string) {
|
||||
si.mu.Lock()
|
||||
|
||||
for i := range si.buildStatus.Tiers {
|
||||
@@ -659,6 +650,7 @@ func (si *SearchIndex) setTierStatus(name, state string, total, completed int) {
|
||||
si.buildStatus.Tiers[i].State = state
|
||||
si.buildStatus.Tiers[i].Total = total
|
||||
si.buildStatus.Tiers[i].Completed = completed
|
||||
si.buildStatus.Tiers[i].Detail = detail
|
||||
si.mu.Unlock()
|
||||
|
||||
if transitioned {
|
||||
@@ -676,32 +668,13 @@ func (si *SearchIndex) setTierStatus(name, state string, total, completed int) {
|
||||
State: state,
|
||||
Total: total,
|
||||
Completed: completed,
|
||||
Detail: detail,
|
||||
})
|
||||
|
||||
si.mu.Unlock()
|
||||
si.emitStatus()
|
||||
}
|
||||
|
||||
// setTierError marks a build stage as errored.
|
||||
func (si *SearchIndex) setTierError(name, errMsg string) {
|
||||
si.mu.Lock()
|
||||
|
||||
for i := range si.buildStatus.Tiers {
|
||||
if si.buildStatus.Tiers[i].Name == name {
|
||||
si.buildStatus.Tiers[i].State = "error"
|
||||
si.buildStatus.Tiers[i].Error = errMsg
|
||||
si.mu.Unlock()
|
||||
|
||||
si.logIndexJob(jobs.LevelError, name+": "+errMsg)
|
||||
si.emitStatus()
|
||||
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
si.mu.Unlock()
|
||||
}
|
||||
|
||||
// emitStatus pushes the current index status to the frontend via Wails event.
|
||||
func (si *SearchIndex) emitStatus() {
|
||||
if si.runtimeCtx == nil {
|
||||
@@ -1086,25 +1059,36 @@ func (si *SearchIndex) TopReleaseGroupsByArtist(artistMBID string, limit int) []
|
||||
return results
|
||||
}
|
||||
|
||||
// AddFromCache inserts entries from a cached discography browse
|
||||
// into the search index (Tier 5: organic growth). Called when a
|
||||
// user views an artist page and the discography is fetched.
|
||||
// AddFromCache inserts entries from a cached discography browse into
|
||||
// the search index. Called when a user views an artist page and the
|
||||
// discography is fetched — organic growth beyond the shipped catalog.
|
||||
func (si *SearchIndex) AddFromCache(artistName, artistMBID string, rgs []MBReleaseGroup) {
|
||||
if len(rgs) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
// resolveArtistName falls back to the MBID when it cannot find a
|
||||
// name, which is fine for a one-off render but must never be
|
||||
// persisted: an MBID stored as a title is unsearchable and shows up
|
||||
// as a UUID in the UI. Writing nothing lets the upsert's
|
||||
// "non-empty wins" rule keep whatever real name arrives later.
|
||||
if artistName == artistMBID {
|
||||
artistName = ""
|
||||
}
|
||||
|
||||
entries := make([]SearchIndexResult, 0, len(rgs)+1)
|
||||
|
||||
// Add the artist itself.
|
||||
entries = append(entries, SearchIndexResult{
|
||||
EntityType: "artist",
|
||||
MBID: artistMBID,
|
||||
Title: artistName,
|
||||
ArtistName: artistName,
|
||||
ArtistMBID: artistMBID,
|
||||
Popularity: 0, // Unknown from this path.
|
||||
})
|
||||
// Add the artist itself, unless there is no name to add.
|
||||
if artistName != "" {
|
||||
entries = append(entries, SearchIndexResult{
|
||||
EntityType: "artist",
|
||||
MBID: artistMBID,
|
||||
Title: artistName,
|
||||
ArtistName: artistName,
|
||||
ArtistMBID: artistMBID,
|
||||
Popularity: 0, // Unknown from this path.
|
||||
})
|
||||
}
|
||||
|
||||
for _, rg := range rgs {
|
||||
entries = append(entries, SearchIndexResult{
|
||||
@@ -1862,79 +1846,6 @@ type lbSimilarArtistWire struct {
|
||||
ReferenceMBID string `json:"reference_mbid"` // which seed artist this result belongs to
|
||||
}
|
||||
|
||||
// fetchSimilarArtistsBatch queries the labs similar-artists endpoint
|
||||
// for multiple seed MBIDs. Despite the name, this actually fans
|
||||
// out one request per seed: the labs API's multi-seed mode is
|
||||
// broken (results for different seeds get mis-labeled, and some
|
||||
// seeds return zero), so batching with multiple artist_mbids is
|
||||
// not viable. Concurrency is bounded by indexerRate to respect
|
||||
// the labs rate limit; each call goes through the provided LB
|
||||
// client's rate limiter and cache.
|
||||
func (si *SearchIndex) fetchSimilarArtistsBatch(
|
||||
ctx context.Context, lb *ListenBrainzClient, seedMBIDs []string,
|
||||
) map[string][]lbSimilarArtistWire {
|
||||
if len(seedMBIDs) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
var (
|
||||
mu sync.Mutex
|
||||
grouped = make(map[string][]lbSimilarArtistWire, len(seedMBIDs))
|
||||
wg sync.WaitGroup
|
||||
)
|
||||
|
||||
sem := make(chan struct{}, indexerRate)
|
||||
|
||||
for _, seedMBID := range seedMBIDs {
|
||||
if ctx.Err() != nil {
|
||||
break
|
||||
}
|
||||
|
||||
sem <- struct{}{}
|
||||
|
||||
wg.Add(1)
|
||||
|
||||
go func(seed string) {
|
||||
defer func() {
|
||||
<-sem
|
||||
wg.Done()
|
||||
}()
|
||||
|
||||
// Use the LB client's per-seed GET form — goes through
|
||||
// the shared rate limiter and cache. The multi-seed
|
||||
// POST form is not viable (see function comment).
|
||||
similar, err := lb.SimilarArtists(ctx, seed)
|
||||
if err != nil || len(similar) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
// Convert to the internal wire type used by the caller
|
||||
// and trim to indexSimilarPerArtist.
|
||||
if len(similar) > indexSimilarPerArtist {
|
||||
similar = similar[:indexSimilarPerArtist]
|
||||
}
|
||||
|
||||
results := make([]lbSimilarArtistWire, len(similar))
|
||||
for i, s := range similar {
|
||||
results[i] = lbSimilarArtistWire{
|
||||
ArtistMBID: s.ArtistMBID,
|
||||
Name: s.Name,
|
||||
Score: int(s.Score),
|
||||
ReferenceMBID: seed,
|
||||
}
|
||||
}
|
||||
|
||||
mu.Lock()
|
||||
grouped[seed] = results
|
||||
mu.Unlock()
|
||||
}(seedMBID)
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
|
||||
return grouped
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Shared: index artist discographies
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -2155,22 +2066,6 @@ func (si *SearchIndex) fetchTopRecordings(
|
||||
return results
|
||||
}
|
||||
|
||||
// chunkStrings splits a slice into chunks of at most size n.
|
||||
func chunkStrings(s []string, n int) [][]string {
|
||||
var chunks [][]string
|
||||
|
||||
for i := 0; i < len(s); i += n {
|
||||
end := i + n
|
||||
if end > len(s) {
|
||||
end = len(s)
|
||||
}
|
||||
|
||||
chunks = append(chunks, s[i:end])
|
||||
}
|
||||
|
||||
return chunks
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Database writes
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -2186,6 +2081,69 @@ func chunkStrings(s []string, n int) [][]string {
|
||||
// empty values, and numeric fields use "highest wins" for popularity/
|
||||
// listener_count/duration so older richer data survives refreshes.
|
||||
|
||||
// upsertIndexSQL is the single index write statement. It is kept as
|
||||
// a const so assembly can prepare it once per transaction instead of
|
||||
// re-parsing this large upsert for every row.
|
||||
const upsertIndexSQL = `
|
||||
INSERT INTO explore_index (
|
||||
entity_type, mbid, title, artist_name, artist_mbid, aliases,
|
||||
popularity, listener_count,
|
||||
duration, caa_release_mbid, release_name,
|
||||
primary_type, secondary_types, release_date,
|
||||
artist_type, country, disambiguation, sort_name,
|
||||
in_library, is_similar,
|
||||
local_artist_id, local_release_group_id, local_recording_id,
|
||||
discog_fetched
|
||||
) VALUES (
|
||||
?, ?, ?, ?, ?, ?,
|
||||
?, ?,
|
||||
?, ?, ?,
|
||||
?, ?, ?,
|
||||
?, ?, ?, ?,
|
||||
?, ?,
|
||||
NULLIF(?, 0), NULLIF(?, 0), NULLIF(?, 0),
|
||||
?
|
||||
)` + upsertIndexConflictSQL
|
||||
|
||||
// upsertIndexConflictSQL is the merge half of every index write, split
|
||||
// out so the bulk artifact import (which inserts by SELECT rather than
|
||||
// by parameter list) resolves conflicts identically instead of carrying
|
||||
// a second, drifting copy of these rules.
|
||||
const upsertIndexConflictSQL = `
|
||||
ON CONFLICT(mbid) DO UPDATE SET
|
||||
-- Title and artist info: never clobber a good value with an
|
||||
-- empty one. Writers are responsible for not offering an MBID
|
||||
-- as a name; AddFromCache is the path that used to.
|
||||
title = CASE WHEN excluded.title != '' THEN excluded.title ELSE title END,
|
||||
artist_name = CASE WHEN excluded.artist_name != '' THEN excluded.artist_name ELSE artist_name END,
|
||||
artist_mbid = CASE WHEN excluded.artist_mbid != '' THEN excluded.artist_mbid ELSE artist_mbid END,
|
||||
aliases = CASE WHEN excluded.aliases != '' THEN excluded.aliases ELSE aliases END,
|
||||
|
||||
-- Highest wins for popularity + listener_count (refreshes can go up).
|
||||
popularity = CASE WHEN excluded.popularity > popularity THEN excluded.popularity ELSE popularity END,
|
||||
listener_count = CASE WHEN excluded.listener_count > listener_count THEN excluded.listener_count ELSE listener_count END,
|
||||
|
||||
-- Non-empty wins for all other optional fields (never clobber with empty).
|
||||
duration = CASE WHEN excluded.duration > 0 THEN excluded.duration ELSE duration END,
|
||||
caa_release_mbid = CASE WHEN excluded.caa_release_mbid != '' THEN excluded.caa_release_mbid ELSE caa_release_mbid END,
|
||||
release_name = CASE WHEN excluded.release_name != '' THEN excluded.release_name ELSE release_name END,
|
||||
primary_type = CASE WHEN excluded.primary_type != '' THEN excluded.primary_type ELSE primary_type END,
|
||||
secondary_types = CASE WHEN excluded.secondary_types != '' THEN excluded.secondary_types ELSE secondary_types END,
|
||||
release_date = CASE WHEN excluded.release_date != '' THEN excluded.release_date ELSE release_date END,
|
||||
artist_type = CASE WHEN excluded.artist_type != '' THEN excluded.artist_type ELSE artist_type END,
|
||||
country = CASE WHEN excluded.country != '' THEN excluded.country ELSE country END,
|
||||
disambiguation = CASE WHEN excluded.disambiguation != '' THEN excluded.disambiguation ELSE disambiguation END,
|
||||
sort_name = CASE WHEN excluded.sort_name != '' THEN excluded.sort_name ELSE sort_name END,
|
||||
|
||||
-- Flags and cross-references: non-null wins.
|
||||
in_library = MAX(in_library, excluded.in_library),
|
||||
is_similar = MAX(is_similar, excluded.is_similar),
|
||||
discog_fetched = MAX(discog_fetched, excluded.discog_fetched),
|
||||
local_artist_id = COALESCE(excluded.local_artist_id, local_artist_id),
|
||||
local_release_group_id = COALESCE(excluded.local_release_group_id, local_release_group_id),
|
||||
local_recording_id = COALESCE(excluded.local_recording_id, local_recording_id)
|
||||
`
|
||||
|
||||
// upsertBatch writes a batch of SearchIndexResult entries to the index
|
||||
// inside a single transaction. This is the ONE function that all
|
||||
// writes go through. All fields are handled — callers don't need to
|
||||
@@ -2202,6 +2160,17 @@ func (si *SearchIndex) upsertBatch(entries []SearchIndexResult) {
|
||||
return
|
||||
}
|
||||
|
||||
stmt, err := tx.Prepare(upsertIndexSQL)
|
||||
if err != nil {
|
||||
si.logger.Warn("search index: prepare upsert error", "error", err)
|
||||
|
||||
_ = tx.Rollback()
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
defer func() { _ = stmt.Close() }()
|
||||
|
||||
for _, e := range entries {
|
||||
if e.MBID == "" {
|
||||
continue // skip entries without MBIDs — can't be looked up
|
||||
@@ -2222,69 +2191,7 @@ func (si *SearchIndex) upsertBatch(entries []SearchIndexResult) {
|
||||
discogFetched = 1
|
||||
}
|
||||
|
||||
if _, err := tx.Exec(`
|
||||
INSERT INTO explore_index (
|
||||
entity_type, mbid, title, artist_name, artist_mbid, aliases,
|
||||
popularity, listener_count,
|
||||
duration, caa_release_mbid, release_name,
|
||||
primary_type, secondary_types, release_date,
|
||||
artist_type, country, disambiguation, sort_name,
|
||||
in_library, is_similar,
|
||||
local_artist_id, local_release_group_id, local_recording_id,
|
||||
discog_fetched,
|
||||
schema_version
|
||||
) VALUES (
|
||||
?, ?, ?, ?, ?, ?,
|
||||
?, ?,
|
||||
?, ?, ?,
|
||||
?, ?, ?,
|
||||
?, ?, ?, ?,
|
||||
?, ?,
|
||||
NULLIF(?, 0), NULLIF(?, 0), NULLIF(?, 0),
|
||||
?,
|
||||
?
|
||||
)
|
||||
ON CONFLICT(mbid) DO UPDATE SET
|
||||
-- Title and artist info: don't clobber a good value with
|
||||
-- an empty string or with the MBID itself (which can sneak
|
||||
-- in via fallback paths in AddFromCache).
|
||||
title = CASE
|
||||
WHEN excluded.title != '' AND excluded.title != excluded.mbid THEN excluded.title
|
||||
ELSE title
|
||||
END,
|
||||
artist_name = CASE
|
||||
WHEN excluded.artist_name != '' AND excluded.artist_name != excluded.artist_mbid THEN excluded.artist_name
|
||||
ELSE artist_name
|
||||
END,
|
||||
artist_mbid = CASE WHEN excluded.artist_mbid != '' THEN excluded.artist_mbid ELSE artist_mbid END,
|
||||
aliases = CASE WHEN excluded.aliases != '' THEN excluded.aliases ELSE aliases END,
|
||||
|
||||
-- Highest wins for popularity + listener_count (refreshes can go up).
|
||||
popularity = CASE WHEN excluded.popularity > popularity THEN excluded.popularity ELSE popularity END,
|
||||
listener_count = CASE WHEN excluded.listener_count > listener_count THEN excluded.listener_count ELSE listener_count END,
|
||||
|
||||
-- Non-empty wins for all other optional fields (never clobber with empty).
|
||||
duration = CASE WHEN excluded.duration > 0 THEN excluded.duration ELSE duration END,
|
||||
caa_release_mbid = CASE WHEN excluded.caa_release_mbid != '' THEN excluded.caa_release_mbid ELSE caa_release_mbid END,
|
||||
release_name = CASE WHEN excluded.release_name != '' THEN excluded.release_name ELSE release_name END,
|
||||
primary_type = CASE WHEN excluded.primary_type != '' THEN excluded.primary_type ELSE primary_type END,
|
||||
secondary_types = CASE WHEN excluded.secondary_types != '' THEN excluded.secondary_types ELSE secondary_types END,
|
||||
release_date = CASE WHEN excluded.release_date != '' THEN excluded.release_date ELSE release_date END,
|
||||
artist_type = CASE WHEN excluded.artist_type != '' THEN excluded.artist_type ELSE artist_type END,
|
||||
country = CASE WHEN excluded.country != '' THEN excluded.country ELSE country END,
|
||||
disambiguation = CASE WHEN excluded.disambiguation != '' THEN excluded.disambiguation ELSE disambiguation END,
|
||||
sort_name = CASE WHEN excluded.sort_name != '' THEN excluded.sort_name ELSE sort_name END,
|
||||
|
||||
-- Flags and cross-references: non-null wins.
|
||||
in_library = MAX(in_library, excluded.in_library),
|
||||
is_similar = MAX(is_similar, excluded.is_similar),
|
||||
discog_fetched = MAX(discog_fetched, excluded.discog_fetched),
|
||||
local_artist_id = COALESCE(excluded.local_artist_id, local_artist_id),
|
||||
local_release_group_id = COALESCE(excluded.local_release_group_id, local_release_group_id),
|
||||
local_recording_id = COALESCE(excluded.local_recording_id, local_recording_id),
|
||||
|
||||
schema_version = MAX(schema_version, excluded.schema_version)
|
||||
`,
|
||||
if _, err := stmt.Exec(
|
||||
e.EntityType, e.MBID, e.Title, e.ArtistName, e.ArtistMBID, e.Aliases,
|
||||
e.Popularity, e.ListenerCount,
|
||||
e.Duration, e.CAAReleaseMBID, e.ReleaseName,
|
||||
@@ -2293,7 +2200,6 @@ func (si *SearchIndex) upsertBatch(entries []SearchIndexResult) {
|
||||
inLib, isSim,
|
||||
e.LocalArtistID, e.LocalReleaseGroupID, e.LocalRecordingID,
|
||||
discogFetched,
|
||||
currentSchemaVersion,
|
||||
); err != nil {
|
||||
si.logger.Warn("search index: upsert error",
|
||||
"mbid", e.MBID,
|
||||
@@ -2311,30 +2217,6 @@ func (si *SearchIndex) upsertBatch(entries []SearchIndexResult) {
|
||||
// Helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// getLibraryArtistMBIDs returns MBIDs for all library artists that have one.
|
||||
// Used when Tier 3 was skipped but Tier 4 needs the library MBID list.
|
||||
func (si *SearchIndex) getLibraryArtistMBIDs() []string {
|
||||
rows, err := si.db.QueryContext(
|
||||
"SELECT DISTINCT mbid FROM artists WHERE mbid IS NOT NULL AND mbid != ''",
|
||||
)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
defer func() { _ = rows.Close() }()
|
||||
|
||||
var mbids []string
|
||||
|
||||
for rows.Next() {
|
||||
var mbid string
|
||||
if err := rows.Scan(&mbid); err == nil {
|
||||
mbids = append(mbids, mbid)
|
||||
}
|
||||
}
|
||||
|
||||
return mbids
|
||||
}
|
||||
|
||||
// hasMeta reports whether a key exists in explore_index_meta.
|
||||
func (si *SearchIndex) hasMeta(key string) bool {
|
||||
rows, err := si.db.QueryContext(
|
||||
|
||||
@@ -0,0 +1,108 @@
|
||||
package explore
|
||||
|
||||
import (
|
||||
"archive/tar"
|
||||
"bytes"
|
||||
"log/slog"
|
||||
"testing"
|
||||
|
||||
"github.com/parquet-go/parquet-go"
|
||||
)
|
||||
|
||||
// Fixtures shared by the client tests and the tagged index-builder
|
||||
// tests. They live in an untagged file so both builds compile.
|
||||
|
||||
// Fixed MBIDs for fixtures.
|
||||
const (
|
||||
recA = "11111111-1111-1111-1111-111111111111"
|
||||
recB = "22222222-2222-2222-2222-222222222222"
|
||||
recC = "33333333-3333-3333-3333-333333333333"
|
||||
relA = "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"
|
||||
relB = "bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb"
|
||||
rgA = "cccccccc-cccc-cccc-cccc-cccccccccccc"
|
||||
rgB = "dddddddd-dddd-dddd-dddd-dddddddddddd"
|
||||
artA = "eeeeeeee-eeee-eeee-eeee-eeeeeeeeeeee"
|
||||
artB = "ffffffff-ffff-ffff-ffff-ffffffffffff"
|
||||
)
|
||||
|
||||
func testLogger() *slog.Logger {
|
||||
return slog.New(slog.DiscardHandler)
|
||||
}
|
||||
|
||||
func makeTar(t *testing.T, members map[string][]byte, order []string) []byte {
|
||||
t.Helper()
|
||||
|
||||
var buf bytes.Buffer
|
||||
|
||||
tw := tar.NewWriter(&buf)
|
||||
|
||||
for _, name := range order {
|
||||
data := members[name]
|
||||
hdr := &tar.Header{
|
||||
Name: name,
|
||||
Mode: 0o644,
|
||||
Size: int64(len(data)),
|
||||
Typeflag: tar.TypeReg,
|
||||
}
|
||||
|
||||
if err := tw.WriteHeader(hdr); err != nil {
|
||||
t.Fatalf("tar header: %v", err)
|
||||
}
|
||||
|
||||
if _, err := tw.Write(data); err != nil {
|
||||
t.Fatalf("tar write: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
if err := tw.Close(); err != nil {
|
||||
t.Fatalf("tar close: %v", err)
|
||||
}
|
||||
|
||||
return buf.Bytes()
|
||||
}
|
||||
|
||||
func makeParquet(t *testing.T, rows []sparkFixtureRow) []byte {
|
||||
t.Helper()
|
||||
|
||||
var buf bytes.Buffer
|
||||
|
||||
w := parquet.NewGenericWriter[sparkFixtureRow](&buf)
|
||||
|
||||
if _, err := w.Write(rows); err != nil {
|
||||
t.Fatalf("parquet write: %v", err)
|
||||
}
|
||||
|
||||
if err := w.Close(); err != nil {
|
||||
t.Fatalf("parquet close: %v", err)
|
||||
}
|
||||
|
||||
return buf.Bytes()
|
||||
}
|
||||
|
||||
// sparkFixtureRow mimics the real spark listens schema: the aggregator
|
||||
// must project just recording/release/artist MBIDs out of it.
|
||||
type sparkFixtureRow struct {
|
||||
ListenedAt int64 `parquet:"listened_at"`
|
||||
UserID int64 `parquet:"user_id"`
|
||||
ArtistName string `parquet:"artist_name,optional"`
|
||||
RecordingMBID string `parquet:"recording_mbid,optional"`
|
||||
ReleaseMBID string `parquet:"release_mbid,optional"`
|
||||
ArtistMBIDs []string `parquet:"artist_credit_mbids,optional,list"`
|
||||
}
|
||||
|
||||
// listensOf builds n identical listen rows for a recording.
|
||||
func listensOf(n int, recording, release string, artists []string) []sparkFixtureRow {
|
||||
rows := make([]sparkFixtureRow, n)
|
||||
for i := range rows {
|
||||
rows[i] = sparkFixtureRow{
|
||||
ListenedAt: 1700000000 + int64(i),
|
||||
UserID: int64(i),
|
||||
ArtistName: "Fixture Artist",
|
||||
RecordingMBID: recording,
|
||||
ReleaseMBID: release,
|
||||
ArtistMBIDs: artists,
|
||||
}
|
||||
}
|
||||
|
||||
return rows
|
||||
}
|
||||
Reference in New Issue
Block a user