diff --git a/internal/sub/external_subscription.go b/internal/sub/external_subscription.go index b187c6284..7098e8db6 100644 --- a/internal/sub/external_subscription.go +++ b/internal/sub/external_subscription.go @@ -16,8 +16,9 @@ import ( // with a short timeout so a slow/dead provider can't stall a client's sub. const ( - subscriptionCacheTTL = 5 * time.Minute - subscriptionMaxBytes = 2 << 20 // 2 MiB + subscriptionCacheTTL = 5 * time.Minute + subscriptionMaxBytes = 2 << 20 // 2 MiB + subscriptionCacheCapacity = 256 ) var subscriptionHTTPClient = &http.Client{Timeout: 6 * time.Second} @@ -27,10 +28,19 @@ type subscriptionCacheEntry struct { fetchedAt time.Time } +type subscriptionFetch struct { + done chan struct{} + links []string +} + var subscriptionCache = struct { sync.Mutex - m map[string]subscriptionCacheEntry -}{m: make(map[string]subscriptionCacheEntry)} + m map[string]subscriptionCacheEntry + inflight map[string]*subscriptionFetch +}{ + m: make(map[string]subscriptionCacheEntry), + inflight: make(map[string]*subscriptionFetch), +} // fetchSubscriptionLinks returns the share links contained in a remote // subscription URL, using a short-lived cache. On any failure it returns the @@ -44,24 +54,59 @@ func fetchSubscriptionLinks(rawURL string) []string { subscriptionCache.Lock() cached, ok := subscriptionCache.m[rawURL] - subscriptionCache.Unlock() if ok && time.Since(cached.fetchedAt) < subscriptionCacheTTL { + subscriptionCache.Unlock() return cached.links } + if fetch, waiting := subscriptionCache.inflight[rawURL]; waiting { + subscriptionCache.Unlock() + <-fetch.done + return fetch.links + } + fetch := &subscriptionFetch{done: make(chan struct{})} + subscriptionCache.inflight[rawURL] = fetch + subscriptionCache.Unlock() + defer func() { + subscriptionCache.Lock() + close(fetch.done) + delete(subscriptionCache.inflight, rawURL) + subscriptionCache.Unlock() + }() links, err := doFetchSubscriptionLinks(rawURL) if err != nil { - // Serve stale on error rather than dropping the client's configs. if ok { - return cached.links + fetch.links = cached.links } - return nil + return fetch.links } subscriptionCache.Lock() subscriptionCache.m[rawURL] = subscriptionCacheEntry{links: links, fetchedAt: time.Now()} + trimSubscriptionCacheLocked(rawURL) subscriptionCache.Unlock() - return links + fetch.links = links + return fetch.links +} + +func trimSubscriptionCacheLocked(keep string) { + for len(subscriptionCache.m) > subscriptionCacheCapacity { + var oldestURL string + var oldest time.Time + for rawURL, entry := range subscriptionCache.m { + if rawURL == keep { + continue + } + if oldestURL == "" || entry.fetchedAt.Before(oldest) { + oldestURL = rawURL + oldest = entry.fetchedAt + } + } + if oldestURL == "" { + return + } + delete(subscriptionCache.m, oldestURL) + } } func doFetchSubscriptionLinks(rawURL string) ([]string, error) { diff --git a/internal/sub/external_subscription_test.go b/internal/sub/external_subscription_test.go index 0d1bca820..3602d5d7a 100644 --- a/internal/sub/external_subscription_test.go +++ b/internal/sub/external_subscription_test.go @@ -4,10 +4,150 @@ import ( "errors" "net/http" "net/http/httptest" + "strconv" "strings" + "sync" + "sync/atomic" "testing" + "time" ) +func resetSubscriptionCache(t *testing.T) { + t.Helper() + subscriptionCache.Lock() + previousEntries := subscriptionCache.m + previousInflight := subscriptionCache.inflight + subscriptionCache.m = make(map[string]subscriptionCacheEntry) + subscriptionCache.inflight = make(map[string]*subscriptionFetch) + subscriptionCache.Unlock() + t.Cleanup(func() { + subscriptionCache.Lock() + subscriptionCache.m = previousEntries + subscriptionCache.inflight = previousInflight + subscriptionCache.Unlock() + }) +} + +func TestFetchSubscriptionLinksSharesConcurrentRefresh(t *testing.T) { + resetSubscriptionCache(t) + var requests atomic.Int32 + release := make(chan struct{}) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requests.Add(1) + <-release + _, _ = w.Write([]byte("vless://uuid@example.com:443")) + })) + defer srv.Close() + + const callers = 16 + results := make(chan []string, callers) + var wg sync.WaitGroup + for range callers { + wg.Add(1) + go func() { + defer wg.Done() + results <- fetchSubscriptionLinks(srv.URL) + }() + } + + time.Sleep(100 * time.Millisecond) + close(release) + wg.Wait() + close(results) + + for links := range results { + if len(links) != 1 || links[0] != "vless://uuid@example.com:443" { + t.Fatalf("links = %#v", links) + } + } + if got := requests.Load(); got != 1 { + t.Fatalf("requests = %d, want 1", got) + } +} + +func TestFetchSubscriptionLinksBoundsCacheSize(t *testing.T) { + resetSubscriptionCache(t) + var requests atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requests.Add(1) + _, _ = w.Write([]byte("vless://uuid@example.com:443")) + })) + defer srv.Close() + + for i := range subscriptionCacheCapacity + 1 { + links := fetchSubscriptionLinks(srv.URL + "?id=" + strconv.Itoa(i)) + if len(links) != 1 { + t.Fatalf("links at %d = %#v", i, links) + } + } + + subscriptionCache.Lock() + entries := len(subscriptionCache.m) + subscriptionCache.Unlock() + if entries != subscriptionCacheCapacity { + t.Fatalf("cache entries = %d, want %d", entries, subscriptionCacheCapacity) + } + if got := requests.Load(); got != subscriptionCacheCapacity+1 { + t.Fatalf("requests = %d, want %d", got, subscriptionCacheCapacity+1) + } +} + +func TestFetchSubscriptionLinksSharesStaleResultAfterRefreshFailure(t *testing.T) { + resetSubscriptionCache(t) + stale := []string{"vless://stale@example.com:443"} + release := make(chan struct{}) + var staleRequests atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/stale" { + staleRequests.Add(1) + <-release + w.WriteHeader(http.StatusBadGateway) + return + } + _, _ = w.Write([]byte("vless://fresh@example.com:443")) + })) + defer srv.Close() + staleURL := srv.URL + "/stale" + + subscriptionCache.Lock() + subscriptionCache.m[staleURL] = subscriptionCacheEntry{ + links: stale, + fetchedAt: time.Now().Add(-subscriptionCacheTTL), + } + for i := range subscriptionCacheCapacity - 1 { + subscriptionCache.m["cached-"+strconv.Itoa(i)] = subscriptionCacheEntry{fetchedAt: time.Now()} + } + subscriptionCache.Unlock() + + const callers = 16 + results := make(chan []string, callers) + var wg sync.WaitGroup + for range callers { + wg.Add(1) + go func() { + defer wg.Done() + results <- fetchSubscriptionLinks(staleURL) + }() + } + + time.Sleep(100 * time.Millisecond) + if links := fetchSubscriptionLinks(srv.URL + "/fresh"); len(links) != 1 || links[0] != "vless://fresh@example.com:443" { + t.Fatalf("fresh links = %#v", links) + } + close(release) + wg.Wait() + close(results) + + for links := range results { + if len(links) != 1 || links[0] != stale[0] { + t.Fatalf("links = %#v, want %#v", links, stale) + } + } + if got := staleRequests.Load(); got != 1 { + t.Fatalf("requests = %d, want 1", got) + } +} + func TestDoFetchSubscriptionLinks_RejectsOversizedBody(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, _ = w.Write([]byte(strings.Repeat("a", subscriptionMaxBytes+1)))