feat: data lifecycle rewrite, download clients, wanted list, and central catalog index
Build & publish Arch package / arch-package (push) Successful in 2m12s
Search index maintenance / maintain-index (push) Successful in 2h22m28s

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:
2026-08-06 17:12:01 -04:00
co-authored by Claude Sonnet 5
parent d0d86f85d5
commit e190fd75b9
165 changed files with 31088 additions and 5192 deletions
+87
View File
@@ -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)
}
+415
View File
@@ -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
}
+206
View File
@@ -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)
}
}
+389
View File
@@ -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)
}
}
+444
View File
@@ -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)
}
}
+1 -1
View File
@@ -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
+20
View File
@@ -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"
)
+1 -1
View File
@@ -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
+73
View File
@@ -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())
}
}
+3 -1
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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()
+562
View File
@@ -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 := &parallelReader{
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)
}
+372
View File
@@ -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")
}
}
+2
View File
@@ -1,3 +1,5 @@
//go:build indexbuild
package explore
import (
+679
View File
@@ -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")
+716
View File
@@ -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)
}
}
+339
View File
@@ -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() }
+257
View File
@@ -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
View File
@@ -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
}
+8 -4
View File
@@ -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)
}
+6
View File
@@ -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
}
}
}
+213
View File
@@ -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
View File
@@ -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(
+108
View File
@@ -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
}