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) } }