feat(similarity): implement Worker.tickOnce with track + artist passes

Replace stub with full tickOnce: drains played tracks/artists via
ListPlayedTracksNeedingSimilarity / ListPlayedArtistsNeedingSimilarity,
calls LB SimilarRecordings/SimilarArtists, filters to local library via
bulk MBID lookup, enforces top-K=20 by score, and upserts similarity
rows. 429 aborts the entire tick; other errors log-and-skip. Ten
integration tests covering no-op, library filtering, top-K cap, 7-day
freshness cap, stale refresh, 429 abort, transient skip, and artist
pass.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-04-28 20:06:47 -04:00
parent 9adad094b0
commit 893adf04c9
2 changed files with 548 additions and 4 deletions
+150 -4
View File
@@ -8,11 +8,15 @@ package similarity
import (
"context"
"errors"
"log/slog"
"sort"
"time"
"github.com/jackc/pgx/v5/pgtype"
"github.com/jackc/pgx/v5/pgxpool"
"git.fabledsword.com/bvandeusen/minstrel/internal/db/dbq"
"git.fabledsword.com/bvandeusen/minstrel/internal/scrobble/listenbrainz"
)
@@ -58,9 +62,151 @@ func (w *Worker) Run(ctx context.Context) {
}
}
// tickOnce drains one batch of tracks and one batch of artists. Stub —
// implementation lands in Task 6 along with the integration tests that
// drive it.
func (w *Worker) tickOnce(_ context.Context) error {
// tickOnce drains one batch of tracks and one batch of artists. Per-row
// errors are logged and skipped (passive retry via timer). 429 aborts the
// entire tick.
func (w *Worker) tickOnce(ctx context.Context) error {
q := dbq.New(w.pool)
if err := w.tickTracks(ctx, q); err != nil {
return err
}
return w.tickArtists(ctx, q)
}
func (w *Worker) tickTracks(ctx context.Context, q *dbq.Queries) error {
rows, err := q.ListPlayedTracksNeedingSimilarity(ctx, w.batch)
if err != nil {
return err
}
for _, r := range rows {
if r.Mbid == nil {
continue // defensive — query already filters NULL
}
results, err := w.client.SimilarRecordings(ctx, *r.Mbid, 100)
if err != nil {
var ra *listenbrainz.RetryAfterError
if errors.As(err, &ra) {
w.logger.Warn("similarity: 429 — aborting tick", "retry_after", ra.Wait)
return nil
}
w.logger.Warn("similarity: similar-recordings failed", "track_id", r.ID, "err", err)
continue
}
w.upsertTrackSimilar(ctx, q, r.ID, results)
}
return nil
}
func (w *Worker) tickArtists(ctx context.Context, q *dbq.Queries) error {
rows, err := q.ListPlayedArtistsNeedingSimilarity(ctx, w.batch)
if err != nil {
return err
}
for _, r := range rows {
if r.Mbid == nil {
continue
}
results, err := w.client.SimilarArtists(ctx, *r.Mbid, 100)
if err != nil {
var ra *listenbrainz.RetryAfterError
if errors.As(err, &ra) {
w.logger.Warn("similarity: 429 — aborting tick", "retry_after", ra.Wait)
return nil
}
w.logger.Warn("similarity: similar-artists failed", "artist_id", r.ID, "err", err)
continue
}
w.upsertArtistSimilar(ctx, q, r.ID, results)
}
return nil
}
// upsertTrackSimilar filters returned MBIDs to those in our library, takes
// top-K by score, and upserts rows.
func (w *Worker) upsertTrackSimilar(ctx context.Context, q *dbq.Queries, trackAID pgtype.UUID, results []listenbrainz.SimilarRecording) {
if len(results) == 0 {
return
}
sort.Slice(results, func(i, j int) bool { return results[i].Score > results[j].Score })
mbids := make([]string, 0, len(results))
for _, r := range results {
mbids = append(mbids, r.MBID)
}
rows, err := q.GetTracksByMBIDs(ctx, mbids)
if err != nil {
w.logger.Warn("similarity: GetTracksByMBIDs", "err", err)
return
}
idByMBID := make(map[string]pgtype.UUID, len(rows))
for _, r := range rows {
if r.Mbid != nil {
idByMBID[*r.Mbid] = r.ID
}
}
taken := 0
for _, r := range results {
if taken >= w.topK {
break
}
localID, ok := idByMBID[r.MBID]
if !ok {
continue
}
if localID == trackAID {
continue // defensive — DB CHECK constraint also catches self-edges
}
if uerr := q.UpsertTrackSimilarity(ctx, dbq.UpsertTrackSimilarityParams{
TrackAID: trackAID, TrackBID: localID, Score: r.Score,
}); uerr != nil {
w.logger.Warn("similarity: UpsertTrackSimilarity", "err", uerr)
continue
}
taken++
}
}
func (w *Worker) upsertArtistSimilar(ctx context.Context, q *dbq.Queries, artistAID pgtype.UUID, results []listenbrainz.SimilarArtist) {
if len(results) == 0 {
return
}
sort.Slice(results, func(i, j int) bool { return results[i].Score > results[j].Score })
mbids := make([]string, 0, len(results))
for _, r := range results {
mbids = append(mbids, r.MBID)
}
rows, err := q.GetArtistsByMBIDs(ctx, mbids)
if err != nil {
w.logger.Warn("similarity: GetArtistsByMBIDs", "err", err)
return
}
idByMBID := make(map[string]pgtype.UUID, len(rows))
for _, r := range rows {
if r.Mbid != nil {
idByMBID[*r.Mbid] = r.ID
}
}
taken := 0
for _, r := range results {
if taken >= w.topK {
break
}
localID, ok := idByMBID[r.MBID]
if !ok {
continue
}
if localID == artistAID {
continue
}
if uerr := q.UpsertArtistSimilarity(ctx, dbq.UpsertArtistSimilarityParams{
ArtistAID: artistAID, ArtistBID: localID, Score: r.Score,
}); uerr != nil {
w.logger.Warn("similarity: UpsertArtistSimilarity", "err", uerr)
continue
}
taken++
}
}
@@ -0,0 +1,398 @@
package similarity
import (
"context"
"encoding/json"
"fmt"
"io"
"log/slog"
"net/http"
"net/http/httptest"
"os"
"strings"
"sync/atomic"
"testing"
"time"
"github.com/jackc/pgx/v5/pgtype"
"github.com/jackc/pgx/v5/pgxpool"
"git.fabledsword.com/bvandeusen/minstrel/internal/db"
"git.fabledsword.com/bvandeusen/minstrel/internal/db/dbq"
"git.fabledsword.com/bvandeusen/minstrel/internal/scrobble/listenbrainz"
)
func testPool(t *testing.T) (*pgxpool.Pool, *dbq.Queries) {
t.Helper()
if testing.Short() {
t.Skip("skipping in -short mode")
}
dsn := os.Getenv("MINSTREL_TEST_DATABASE_URL")
if dsn == "" {
t.Skip("MINSTREL_TEST_DATABASE_URL not set")
}
if err := db.Migrate(dsn, slog.New(slog.NewTextHandler(io.Discard, nil))); err != nil {
t.Fatalf("migrate: %v", err)
}
pool, err := pgxpool.New(context.Background(), dsn)
if err != nil {
t.Fatalf("pool: %v", err)
}
t.Cleanup(pool.Close)
if _, err := pool.Exec(context.Background(),
"TRUNCATE artist_similarity, track_similarity, scrobble_queue, play_events, skip_events, play_sessions, sessions, users, tracks, albums, artists RESTART IDENTITY CASCADE"); err != nil {
t.Fatalf("truncate: %v", err)
}
return pool, dbq.New(pool)
}
type fixture struct {
pool *pgxpool.Pool
q *dbq.Queries
user pgtype.UUID
artist dbq.Artist
album dbq.Album
}
func newFixture(t *testing.T) fixture {
t.Helper()
pool, q := testPool(t)
ctx := context.Background()
u, err := q.CreateUser(ctx, dbq.CreateUserParams{
Username: "tester", PasswordHash: "x", ApiToken: "x", IsAdmin: false,
})
if err != nil {
t.Fatalf("user: %v", err)
}
artistMbid := "aaaaaaaa-1111-1111-1111-111111111111"
a, err := q.UpsertArtist(ctx, dbq.UpsertArtistParams{
Name: "X", SortName: "X", Mbid: &artistMbid,
})
if err != nil {
t.Fatalf("artist: %v", err)
}
al, err := q.UpsertAlbum(ctx, dbq.UpsertAlbumParams{
Title: "X", SortTitle: "X", ArtistID: a.ID,
})
if err != nil {
t.Fatalf("album: %v", err)
}
return fixture{pool: pool, q: q, user: u.ID, artist: a, album: al}
}
func seedTrack(t *testing.T, f fixture, title string, mbid *string) dbq.Track {
t.Helper()
tr, err := f.q.UpsertTrack(context.Background(), dbq.UpsertTrackParams{
Title: title,
AlbumID: f.album.ID,
ArtistID: f.artist.ID,
FilePath: "/tmp/" + title + ".flac",
DurationMs: 200_000,
Mbid: mbid,
})
if err != nil {
t.Fatalf("track: %v", err)
}
return tr
}
func markPlayed(t *testing.T, f fixture, trackID pgtype.UUID) {
t.Helper()
var sessionID pgtype.UUID
if err := f.pool.QueryRow(context.Background(),
`INSERT INTO play_sessions (user_id, started_at, last_event_at, client_id)
VALUES ($1, now() - interval '5 minutes', now(), 'test') RETURNING id`,
f.user).Scan(&sessionID); err != nil {
t.Fatalf("session: %v", err)
}
if _, err := f.pool.Exec(context.Background(),
`INSERT INTO play_events (user_id, track_id, session_id, started_at, ended_at, duration_played_ms, completion_ratio, was_skipped)
VALUES ($1, $2, $3, now() - interval '1 minute', now(), 250000, 0.83, false)`,
f.user, trackID, sessionID); err != nil {
t.Fatalf("play_event: %v", err)
}
}
func newTestWorker(f fixture, lbBaseURL string) *Worker {
logger := slog.New(slog.NewTextHandler(io.Discard, nil))
return &Worker{
pool: f.pool,
client: &listenbrainz.Client{BaseURL: lbBaseURL, HTTP: http.DefaultClient},
logger: logger,
tick: 1 * time.Hour,
batch: 5,
topK: 20,
}
}
// stubLB returns an httptest server that responds to similar-recordings and
// similar-artists with the given JSON payloads.
func stubLB(recordingsBody, artistsBody string, status int) *httptest.Server {
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if status != 0 && status != http.StatusOK {
w.WriteHeader(status)
return
}
if strings.Contains(r.URL.Path, "/similar-recordings/") {
_, _ = w.Write([]byte(recordingsBody))
return
}
if strings.Contains(r.URL.Path, "/similar-artists/") {
_, _ = w.Write([]byte(artistsBody))
return
}
w.WriteHeader(http.StatusNotFound)
}))
}
func countTrackSim(t *testing.T, f fixture, a pgtype.UUID) int {
t.Helper()
var n int
if err := f.pool.QueryRow(context.Background(),
`SELECT count(*) FROM track_similarity WHERE track_a_id = $1`, a).Scan(&n); err != nil {
t.Fatalf("count: %v", err)
}
return n
}
func countArtistSim(t *testing.T, f fixture, a pgtype.UUID) int {
t.Helper()
var n int
if err := f.pool.QueryRow(context.Background(),
`SELECT count(*) FROM artist_similarity WHERE artist_a_id = $1`, a).Scan(&n); err != nil {
t.Fatalf("count: %v", err)
}
return n
}
func TestTickOnce_NoPlayedTracks_NoOp(t *testing.T) {
f := newFixture(t)
srv := stubLB(`[]`, `[]`, http.StatusOK)
defer srv.Close()
w := newTestWorker(f, srv.URL)
if err := w.tickOnce(context.Background()); err != nil {
t.Fatalf("tickOnce: %v", err)
}
var n int
_ = f.pool.QueryRow(context.Background(), `SELECT count(*) FROM track_similarity`).Scan(&n)
if n != 0 {
t.Errorf("track_similarity rows = %d, want 0", n)
}
_ = f.pool.QueryRow(context.Background(), `SELECT count(*) FROM artist_similarity`).Scan(&n)
if n != 0 {
t.Errorf("artist_similarity rows = %d, want 0", n)
}
}
func TestTickOnce_MapsLBResponseToLocalLibrary(t *testing.T) {
f := newFixture(t)
mbidA := "11111111-1111-1111-1111-111111111111"
mbidB := "22222222-2222-2222-2222-222222222222"
mbidC := "99999999-9999-9999-9999-999999999999" // NOT in library
trackA := seedTrack(t, f, "A", &mbidA)
_ = seedTrack(t, f, "B", &mbidB)
markPlayed(t, f, trackA.ID)
body := `[
{"recording_mbid": "` + mbidB + `", "score": 0.9},
{"recording_mbid": "` + mbidC + `", "score": 0.7}
]`
srv := stubLB(body, `[]`, http.StatusOK)
defer srv.Close()
w := newTestWorker(f, srv.URL)
if err := w.tickOnce(context.Background()); err != nil {
t.Fatalf("tickOnce: %v", err)
}
if got := countTrackSim(t, f, trackA.ID); got != 1 {
t.Errorf("track_similarity rows = %d, want 1 (only mbidB is in-library)", got)
}
}
func TestTickOnce_TopKEnforced(t *testing.T) {
f := newFixture(t)
mbidSeed := "11111111-1111-1111-1111-111111111111"
seed := seedTrack(t, f, "Seed", &mbidSeed)
markPlayed(t, f, seed.ID)
type lbRow struct {
MBID string `json:"recording_mbid"`
Score float64 `json:"score"`
}
rows := make([]lbRow, 0, 25)
for i := 0; i < 25; i++ {
mbid := fmt.Sprintf("20000000-0000-0000-0000-%012d", i+1)
_ = seedTrack(t, f, fmt.Sprintf("T%02d", i+1), &mbid)
rows = append(rows, lbRow{MBID: mbid, Score: 1.0 - 0.01*float64(i)})
}
body, _ := json.Marshal(rows)
srv := stubLB(string(body), `[]`, http.StatusOK)
defer srv.Close()
w := newTestWorker(f, srv.URL)
if err := w.tickOnce(context.Background()); err != nil {
t.Fatalf("tickOnce: %v", err)
}
if got := countTrackSim(t, f, seed.ID); got != 20 {
t.Errorf("got %d rows, want exactly 20 (all 25 in-library; top 20 by score)", got)
}
}
func TestTickOnce_RespectsSevenDayCap(t *testing.T) {
f := newFixture(t)
mbid := "11111111-1111-1111-1111-111111111111"
trackA := seedTrack(t, f, "A", &mbid)
markPlayed(t, f, trackA.ID)
otherMbid := "55555555-5555-5555-5555-555555555555"
other := seedTrack(t, f, "Other", &otherMbid)
if _, err := f.pool.Exec(context.Background(),
`INSERT INTO track_similarity (track_a_id, track_b_id, score, source, fetched_at)
VALUES ($1, $2, 0.5, 'listenbrainz', now())`, trackA.ID, other.ID); err != nil {
t.Fatalf("seed sim: %v", err)
}
beforeCount := countTrackSim(t, f, trackA.ID)
srv := stubLB(`[{"recording_mbid":"`+otherMbid+`","score":0.99}]`, `[]`, http.StatusOK)
defer srv.Close()
w := newTestWorker(f, srv.URL)
_ = w.tickOnce(context.Background())
if got := countTrackSim(t, f, trackA.ID); got != beforeCount {
t.Errorf("worker re-queried fresh track: before=%d after=%d", beforeCount, got)
}
}
func TestTickOnce_RefreshesStaleRow(t *testing.T) {
f := newFixture(t)
mbid := "11111111-1111-1111-1111-111111111111"
trackA := seedTrack(t, f, "A", &mbid)
markPlayed(t, f, trackA.ID)
otherMbid := "55555555-5555-5555-5555-555555555555"
other := seedTrack(t, f, "Other", &otherMbid)
if _, err := f.pool.Exec(context.Background(),
`INSERT INTO track_similarity (track_a_id, track_b_id, score, source, fetched_at)
VALUES ($1, $2, 0.5, 'listenbrainz', now() - interval '8 days')`, trackA.ID, other.ID); err != nil {
t.Fatalf("seed sim: %v", err)
}
srv := stubLB(`[{"recording_mbid":"`+otherMbid+`","score":0.95}]`, `[]`, http.StatusOK)
defer srv.Close()
w := newTestWorker(f, srv.URL)
if err := w.tickOnce(context.Background()); err != nil {
t.Fatalf("tickOnce: %v", err)
}
var score float64
var fetchedAt time.Time
_ = f.pool.QueryRow(context.Background(),
`SELECT score, fetched_at FROM track_similarity WHERE track_a_id = $1 AND track_b_id = $2 AND source = 'listenbrainz'`,
trackA.ID, other.ID).Scan(&score, &fetchedAt)
if score != 0.95 {
t.Errorf("score = %v, want 0.95 (refreshed)", score)
}
if time.Since(fetchedAt) > time.Minute {
t.Errorf("fetched_at not bumped: %v ago", time.Since(fetchedAt))
}
}
func TestTickOnce_429AbortsTick(t *testing.T) {
f := newFixture(t)
mbid := "11111111-1111-1111-1111-111111111111"
trackA := seedTrack(t, f, "A", &mbid)
markPlayed(t, f, trackA.ID)
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Retry-After", "60")
w.WriteHeader(http.StatusTooManyRequests)
}))
defer srv.Close()
w := newTestWorker(f, srv.URL)
_ = w.tickOnce(context.Background())
if got := countTrackSim(t, f, trackA.ID); got != 0 {
t.Errorf("rows on 429 = %d, want 0 (tick aborted before any inserts)", got)
}
}
func TestTickOnce_TransientErrorSkipsTrack(t *testing.T) {
f := newFixture(t)
mbidA := "11111111-1111-1111-1111-111111111111"
mbidB := "22222222-2222-2222-2222-222222222222"
otherMbid := "55555555-5555-5555-5555-555555555555"
trackA := seedTrack(t, f, "A", &mbidA)
trackB := seedTrack(t, f, "B", &mbidB)
_ = seedTrack(t, f, "Other", &otherMbid)
markPlayed(t, f, trackA.ID)
markPlayed(t, f, trackB.ID)
var seen atomic.Int32
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if !strings.Contains(r.URL.Path, "/similar-recordings/") {
_, _ = w.Write([]byte(`[]`))
return
}
n := seen.Add(1)
if n == 1 {
w.WriteHeader(http.StatusServiceUnavailable)
return
}
_, _ = w.Write([]byte(`[{"recording_mbid":"` + otherMbid + `","score":0.7}]`))
}))
defer srv.Close()
w := newTestWorker(f, srv.URL)
_ = w.tickOnce(context.Background())
a := countTrackSim(t, f, trackA.ID)
b := countTrackSim(t, f, trackB.ID)
if a+b != 1 {
t.Errorf("expected exactly one track to succeed: a=%d b=%d", a, b)
}
}
func TestTickOnce_FiltersInLibrary(t *testing.T) {
f := newFixture(t)
mbid := "11111111-1111-1111-1111-111111111111"
trackA := seedTrack(t, f, "A", &mbid)
markPlayed(t, f, trackA.ID)
body := `[
{"recording_mbid":"99999999-9999-9999-9999-999999999991","score":0.9},
{"recording_mbid":"99999999-9999-9999-9999-999999999992","score":0.8},
{"recording_mbid":"99999999-9999-9999-9999-999999999993","score":0.7}
]`
srv := stubLB(body, `[]`, http.StatusOK)
defer srv.Close()
w := newTestWorker(f, srv.URL)
if err := w.tickOnce(context.Background()); err != nil {
t.Fatalf("tickOnce: %v", err)
}
if got := countTrackSim(t, f, trackA.ID); got != 0 {
t.Errorf("filtered out-of-library: got %d, want 0", got)
}
}
func TestTickOnce_ArtistPassMirrors(t *testing.T) {
f := newFixture(t)
mbid := "11111111-1111-1111-1111-111111111111"
trackA := seedTrack(t, f, "A", &mbid)
markPlayed(t, f, trackA.ID)
bMbid := "bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb"
_, err := f.q.UpsertArtist(context.Background(), dbq.UpsertArtistParams{
Name: "Other", SortName: "Other", Mbid: &bMbid,
})
if err != nil {
t.Fatalf("artist: %v", err)
}
body := `[{"artist_mbid":"` + bMbid + `","score":0.85}]`
srv := stubLB(`[]`, body, http.StatusOK)
defer srv.Close()
w := newTestWorker(f, srv.URL)
if err := w.tickOnce(context.Background()); err != nil {
t.Fatalf("tickOnce: %v", err)
}
if got := countArtistSim(t, f, f.artist.ID); got != 1 {
t.Errorf("artist_similarity rows = %d, want 1", got)
}
}
func TestTickOnce_NoMBIDOnTrack_Skipped(t *testing.T) {
f := newFixture(t)
trackNoMbid := seedTrack(t, f, "NoMbid", nil)
markPlayed(t, f, trackNoMbid.ID)
srv := stubLB(`[{"recording_mbid":"11111111-1111-1111-1111-111111111111","score":0.9}]`, `[]`, http.StatusOK)
defer srv.Close()
w := newTestWorker(f, srv.URL)
if err := w.tickOnce(context.Background()); err != nil {
t.Fatalf("tickOnce: %v", err)
}
if got := countTrackSim(t, f, trackNoMbid.ID); got != 0 {
t.Errorf("no-MBID track produced rows: %d", got)
}
}