acfunlive-remake/avatar_cache_test.go

package main

import (
	"errors"
	"fmt"
	"image"
	"image/png"
	"io"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"sync/atomic"
	"testing"
	"time"
)

// 换掉真的反查动作，返回调用次数
func stubLookupAvatarURL(t *testing.T, f func(uid int64) string) *int64 {
	t.Helper()
	var calls int64
	old := lookupAvatarURL
	lookupAvatarURL = func(uid int64) string {
		atomic.AddInt64(&calls, 1)
		return f(uid)
	}
	t.Cleanup(func() { lookupAvatarURL = old })
	return &calls
}

// 清空进程内的链接记忆，模拟「重启进程」，用完恢复
func resetAvatarURLMemo(t *testing.T) {
	t.Helper()
	avatarURLMemo.Lock()
	prevM, prevAt, prevDir, prevLoaded := avatarURLMemo.m, avatarURLMemo.at, avatarURLMemo.dir, avatarURLMemo.loaded
	avatarURLMemo.m, avatarURLMemo.at, avatarURLMemo.dir, avatarURLMemo.loaded = nil, nil, "", false
	avatarURLMemo.Unlock()
	t.Cleanup(func() {
		avatarURLMemo.Lock()
		avatarURLMemo.m, avatarURLMemo.at, avatarURLMemo.dir, avatarURLMemo.loaded = prevM, prevAt, prevDir, prevLoaded
		avatarURLMemo.Unlock()
	})
}

// 直接写落盘文件（绕过 append，方便造历史记录）
func writeAvatarURLCacheFile(t *testing.T, dir string, lines []string) string {
	t.Helper()
	path := filepath.Join(dir, avatarURLCacheFile)
	if err := os.MkdirAll(dir, 0755); err != nil {
		t.Fatal(err)
	}
	body := strings.Join(lines, "\n")
	if body != "" {
		body += "\n"
	}
	if err := os.WriteFile(path, []byte(body), 0644); err != nil {
		t.Fatal(err)
	}
	return path
}

// 查到的链接要落盘，重启进程后不再重复反查
func TestAvatarURLCacheSurvivesRestart(t *testing.T) {
	dir := t.TempDir()
	resetAvatarURLMemo(t)
	calls := stubLookupAvatarURL(t, func(uid int64) string {
		return fmt.Sprintf("https://example.invalid/%d.png", uid)
	})

	url1, fetched1 := avatarURLByUID(42, dir)
	if !fetched1 {
		t.Fatal("第一次应当真的去查")
	}
	if url1 != "https://example.invalid/42.png" {
		t.Fatalf("拿到的链接不对：%q", url1)
	}
	if atomic.LoadInt64(calls) != 1 {
		t.Fatalf("第一次应当查 1 次，实际 %d", atomic.LoadInt64(calls))
	}

	// 同进程第二次：内存记忆命中，不再查
	if _, fetched := avatarURLByUID(42, dir); fetched {
		t.Fatal("同进程内第二次不该再查")
	}

	// 模拟重启：清空内存，只留落盘文件
	resetAvatarURLMemo(t)
	url2, fetched2 := avatarURLByUID(42, dir)
	if fetched2 {
		t.Fatal("重启后应当直接用落盘记录，不再向 AcFun 查")
	}
	if url2 != url1 {
		t.Fatalf("落盘记录里的链接应当和当初查到的一样，实际 %q", url2)
	}
	if n := atomic.LoadInt64(calls); n != 1 {
		t.Fatalf("整个过程只该查 1 次，实际 %d 次", n)
	}
}

// 过期的记录要重新查（给主播换头像留出路），没到期的直接用
func TestAvatarURLCacheTTL(t *testing.T) {
	dir := t.TempDir()
	now := time.Now()
	writeAvatarURLCacheFile(t, dir, []string{
		fmt.Sprintf(`{"uid":1,"url":"https://old/1.png","at":%d}`, now.Add(-8*24*time.Hour).Unix()),
		fmt.Sprintf(`{"uid":2,"url":"https://fresh/2.png","at":%d}`, now.Add(-time.Hour).Unix()),
	})
	resetAvatarURLMemo(t)
	stubLookupAvatarURL(t, func(uid int64) string { return "https://rechecked/x.png" })

	if u, fetched := avatarURLByUID(2, dir); fetched || u != "https://fresh/2.png" {
		t.Fatalf("没到期的记录应当直接用，实际 (%q, fetched=%v)", u, fetched)
	}
	if u, fetched := avatarURLByUID(1, dir); !fetched || u != "https://rechecked/x.png" {
		t.Fatalf("过了有效期的记录应当重新查，实际 (%q, fetched=%v)", u, fetched)
	}
}

// 文件里有重复行、过期行和坏行时，加载后顺手压一遍，别让它一直长
func TestAvatarURLCacheCompactsFile(t *testing.T) {
	dir := t.TempDir()
	now := time.Now()
	path := writeAvatarURLCacheFile(t, dir, []string{
		fmt.Sprintf(`{"uid":7,"url":"https://a/7.png","at":%d}`, now.Add(-3*time.Minute).Unix()),
		fmt.Sprintf(`{"uid":7,"url":"https://b/7.png","at":%d}`, now.Add(-2*time.Minute).Unix()),
		fmt.Sprintf(`{"uid":7,"url":"https://c/7.png","at":%d}`, now.Add(-time.Minute).Unix()),
		fmt.Sprintf(`{"uid":8,"url":"https://old/8.png","at":%d}`, now.Add(-30*24*time.Hour).Unix()),
		`这不是 json`,
	})
	resetAvatarURLMemo(t)
	stubLookupAvatarURL(t, func(uid int64) string { return "" })

	if u, fetched := avatarURLByUID(7, dir); fetched || u != "https://c/7.png" {
		t.Fatalf("同一 uid 应当以后写的为准（c），实际 (%q, fetched=%v)", u, fetched)
	}

	data, err := os.ReadFile(path)
	if err != nil {
		t.Fatal(err)
	}
	lines := strings.Split(strings.TrimSpace(string(data)), "\n")
	if len(lines) != 1 {
		t.Fatalf("压缩后应当只剩 1 行有效记录，实际 %d 行：%v", len(lines), lines)
	}
	if !strings.Contains(lines[0], "https://c/7.png") {
		t.Fatalf("压缩后保留的应当是最新的那条：%s", lines[0])
	}
}

// 新查到的链接要追加到落盘文件里
func TestAvatarURLCacheAppendsNewEntries(t *testing.T) {
	dir := t.TempDir()
	resetAvatarURLMemo(t)
	stubLookupAvatarURL(t, func(uid int64) string { return fmt.Sprintf("https://x/%d.png", uid) })

	if _, fetched := avatarURLByUID(11, dir); !fetched {
		t.Fatal("第一次应当去查")
	}
	data, err := os.ReadFile(filepath.Join(dir, avatarURLCacheFile))
	if err != nil {
		t.Fatalf("落盘文件应当已生成：%v", err)
	}
	if !strings.Contains(string(data), `"uid":11`) || !strings.Contains(string(data), "https://x/11.png") {
		t.Fatalf("落盘内容不对：%s", data)
	}
}

// 头像并发获取：首次全部下载，第二次全部命中缓存，不会重复下载
func TestFetchAvatarsReusesCache(t *testing.T) {
	withJobFlags(t, true, false, false)
	dir := t.TempDir()

	var hits int64
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		atomic.AddInt64(&hits, 1)
		_ = png.Encode(w, image.NewRGBA(image.Rect(0, 0, 8, 8)))
	}))
	defer srv.Close()

	const n = 20
	uids := make([]int64, 0, n)
	urlByUID := make(map[int64]string, n)
	for i := 0; i < n; i++ {
		uid := int64(1000 + i)
		uids = append(uids, uid)
		urlByUID[uid] = fmt.Sprintf("%s/av%d.png", srv.URL, uid)
	}

	paths, stat, err := fetchAvatars(uids, urlByUID, dir, "")
	if err != nil {
		t.Fatalf("获取头像不该出错：%v", err)
	}
	if len(paths) != n {
		t.Fatalf("应当拿到 %d 个头像，实际 %d", n, len(paths))
	}
	if stat.imgDownloaded != n || stat.imgReused != 0 || stat.failed != 0 {
		t.Fatalf("首次应当全部下载：下载 %d 复用 %d 失败 %d", stat.imgDownloaded, stat.imgReused, stat.failed)
	}
	if got := atomic.LoadInt64(&hits); got != n {
		t.Fatalf("首次应当请求 %d 次，实际 %d", n, got)
	}

	// 第二次：全部命中缓存，一次网络都不该发
	paths2, stat2, err := fetchAvatars(uids, urlByUID, dir, "")
	if err != nil {
		t.Fatalf("第二次获取头像不该出错：%v", err)
	}
	if len(paths2) != n {
		t.Fatalf("第二次也应当拿到 %d 个头像，实际 %d", n, len(paths2))
	}
	if stat2.imgReused != n || stat2.imgDownloaded != 0 {
		t.Fatalf("第二次应当全部复用缓存：复用 %d 下载 %d", stat2.imgReused, stat2.imgDownloaded)
	}
	if got := atomic.LoadInt64(&hits); got != n {
		t.Fatalf("第二次不该再下载，累计请求数应当还是 %d，实际 %d", n, got)
	}
}

// 统计要能区分「链接复用」和「真的去查了」
func TestFetchAvatarsCountsLinkLookups(t *testing.T) {
	withJobFlags(t, true, false, false)
	dir := t.TempDir()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_ = png.Encode(w, image.NewRGBA(image.Rect(0, 0, 8, 8)))
	}))
	defer srv.Close()

	resetAvatarURLMemo(t)
	lookups := stubLookupAvatarURL(t, func(uid int64) string { return srv.URL + "/looked.png" })

	uids := []int64{1, 2, 3}
	urlByUID := map[int64]string{
		1: srv.URL + "/own.png", // 自带的链接
		2: "",
		3: "",
	}
	paths, stat, err := fetchAvatars(uids, urlByUID, dir, "")
	if err != nil {
		t.Fatalf("不该出错：%v", err)
	}
	if len(paths) != 3 {
		t.Fatalf("三个用户都该拿到头像，实际 %d", len(paths))
	}
	if stat.fetched != 2 {
		t.Fatalf("应当有 2 个走了反查，实际 %d", stat.fetched)
	}
	if stat.reusedLink != 0 {
		t.Fatalf("第一次没有可复用的链接，实际 %d", stat.reusedLink)
	}
	if got := atomic.LoadInt64(lookups); got != 2 {
		t.Fatalf("反查动作应当只发生 2 次，实际 %d", got)
	}

	// 再跑一次：这两个 uid 的链接已经记下了，不该再查
	_, stat2, err := fetchAvatars(uids, urlByUID, dir, "")
	if err != nil {
		t.Fatalf("第二次不该出错：%v", err)
	}
	if stat2.reusedLink != 2 || stat2.fetched != 0 {
		t.Fatalf("第二次应当复用链接：复用 %d 反查 %d", stat2.reusedLink, stat2.fetched)
	}
	if got := atomic.LoadInt64(lookups); got != 2 {
		t.Fatalf("第二次不该再查链接，累计应当还是 2 次，实际 %d", got)
	}
}

// 任务被停止时，获取头像要立刻收工并返回 errJobStopped
func TestFetchAvatarsStopsWhenAsked(t *testing.T) {
	withJobFlags(t, true, false, true)
	dir := t.TempDir()
	_, _, err := fetchAvatars([]int64{1}, map[int64]string{1: "http://127.0.0.1:1/x.png"}, dir, "")
	if !errors.Is(err, errJobStopped) {
		t.Fatalf("应当返回 errJobStopped，实际 %v", err)
	}
}

// 统计里的「新下载」要能识别出「这个文件本来就在」
func TestFetchAvatarsCountsReuseAfterRestart(t *testing.T) {
	withJobFlags(t, true, false, false)
	dir := t.TempDir()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_ = png.Encode(w, image.NewRGBA(image.Rect(0, 0, 8, 8)))
	}))
	defer srv.Close()

	uids := []int64{5001}
	urlByUID := map[int64]string{5001: srv.URL + "/a.png"}

	// 先下载一次，再把内存里的链接记忆清掉（模拟重启后再合并同一场）
	if _, stat, err := fetchAvatars(uids, urlByUID, dir, ""); err != nil || stat.imgDownloaded != 1 {
		t.Fatalf("首次应当下载 1 个：%v %+v", err, stat)
	}
	resetAvatarURLMemo(t)
	if _, stat, err := fetchAvatars(uids, urlByUID, dir, ""); err != nil || stat.imgReused != 1 || stat.imgDownloaded != 0 {
		t.Fatalf("重启后应当复用缓存：%v %+v", err, stat)
	}
}

// 落盘文件写坏/读不出时不能影响合并，只是这次多查几遍
func TestAvatarURLCacheBrokenFileIsIgnored(t *testing.T) {
	dir := t.TempDir()
	writeAvatarURLCacheFile(t, dir, []string{"{坏的一行", "还是坏的"})
	resetAvatarURLMemo(t)
	calls := stubLookupAvatarURL(t, func(uid int64) string { return "https://ok/1.png" })

	u, fetched := avatarURLByUID(1, dir)
	if !fetched || u != "https://ok/1.png" {
		t.Fatalf("坏文件应当忽略并照常反查，实际 (%q, fetched=%v)", u, fetched)
	}
	if atomic.LoadInt64(calls) != 1 {
		t.Fatalf("应当查 1 次，实际 %d", atomic.LoadInt64(calls))
	}
}

// 确认落盘用的是头像缓存目录里的固定文件名（别改到别处去）
func TestAvatarURLCacheFilePath(t *testing.T) {
	dir := t.TempDir()
	resetAvatarURLMemo(t)
	stubLookupAvatarURL(t, func(uid int64) string { return "https://x/1.png" })
	avatarURLByUID(1, dir)

	want := filepath.Join(dir, avatarURLCacheFile)
	if _, err := os.Stat(want); err != nil {
		t.Fatalf("落盘文件应当在 %s：%v", want, err)
	}
	// 顺带确认它能被当成普通 jsonl 读出来（一行一条）
	f, err := os.Open(want)
	if err != nil {
		t.Fatal(err)
	}
	defer f.Close()
	data, err := io.ReadAll(f)
	if err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(string(data), "https://x/1.png") {
		t.Fatalf("内容不对：%s", data)
	}
}

// 头像确实是并发获取的：串行做几百个用户要好几分钟，这条锁住「同时有多个请求在跑」
func TestFetchAvatarsRunsConcurrently(t *testing.T) {
	withJobFlags(t, true, false, false)
	dir := t.TempDir()

	var active, maxActive int64
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		cur := atomic.AddInt64(&active, 1)
		for {
			old := atomic.LoadInt64(&maxActive)
			if cur <= old || atomic.CompareAndSwapInt64(&maxActive, old, cur) {
				break
			}
		}
		time.Sleep(40 * time.Millisecond) // 让并发的重叠看得见
		atomic.AddInt64(&active, -1)
		_ = png.Encode(w, image.NewRGBA(image.Rect(0, 0, 8, 8)))
	}))
	defer srv.Close()

	const n = 12
	uids := make([]int64, 0, n)
	urlByUID := make(map[int64]string, n)
	for i := 0; i < n; i++ {
		uid := int64(2000 + i)
		uids = append(uids, uid)
		urlByUID[uid] = fmt.Sprintf("%s/c%d.png", srv.URL, uid)
	}

	if _, _, err := fetchAvatars(uids, urlByUID, dir, ""); err != nil {
		t.Fatalf("不该出错：%v", err)
	}
	if got := atomic.LoadInt64(&maxActive); got < 2 {
		t.Fatalf("头像应当是并发获取的，同时最多只有 %d 个请求在跑（串行会把几百人拖成几分钟）", got)
	}
}