feat: update collectModulePrefixRows to support prior snapshots and enhance CDN/domain prefix collection

Modified the collectModulePrefixRows function to accept an optional priorSnapshot parameter, allowing for more efficient data retrieval by skipping unnecessary CDN fetches. Refactored the logic for collecting prefix rows from AS, CDN, and domain sources to utilize dedicated functions, improving code organization and maintainability. Additionally, introduced caching for module prefix snapshots to optimize performance during refresh operations.

internal/pipeline/collect_parallel.go

@@ -0,0 +1,302 @@
+package pipeline
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"net/http"
+	"os"
+	"strings"
+	"sync"
+	"time"
+
+	"evobgp/internal/asnresolve"
+	"evobgp/internal/store"
+)
+
+func prefixRowsForSource(rows []store.PrefixRow, sourceKey string) []store.PrefixRow {
+	if len(rows) == 0 {
+		return nil
+	}
+	var out []store.PrefixRow
+	for _, row := range rows {
+		if row.Source == sourceKey {
+			out = append(out, row)
+		}
+	}
+	return out
+}
+
+func collectASPrefixRows(ctx context.Context, st store.Backend, hc *http.Client, tenantID string, mod *store.Module, list []*store.ASEntry) ([]store.PrefixRow, error) {
+	moduleID := mod.ID
+	legacy := strings.TrimSpace(os.Getenv("EVOBGP_ASN_RESOLVE")) == "0"
+	if legacy {
+		var rows []store.PrefixRow
+		for _, e := range list {
+			if !store.ValidASN(e.ASN) {
+				continue
+			}
+			comm := e.CommunityID
+			if comm == nil && mod.DefaultCommunityID != nil {
+				c := *mod.DefaultCommunityID
+				comm = &c
+			}
+			rows = append(rows, store.PrefixRow{Prefix: MaterializedASPrefixKey(e.ASN), CommunityID: comm, Source: "as_entry"})
+		}
+		return rows, nil
+	}
+
+	type entryResult struct {
+		rows   []store.PrefixRow
+		metaID string
+		asn    int64
+		holder string
+		count  int64
+		err    error
+	}
+
+	var valid []*store.ASEntry
+	for _, e := range list {
+		if e != nil && store.ValidASN(e.ASN) {
+			valid = append(valid, e)
+		}
+	}
+	sem := make(chan struct{}, collectConcurrency())
+	results := make([]entryResult, len(valid))
+	var wg sync.WaitGroup
+	for i, e := range valid {
+		wg.Add(1)
+		go func(idx int, entry *store.ASEntry) {
+			defer wg.Done()
+			sem <- struct{}{}
+			defer func() { <-sem }()
+
+			comm := entry.CommunityID
+			if comm == nil && mod.DefaultCommunityID != nil {
+				c := *mod.DefaultCommunityID
+				comm = &c
+			}
+			pfxs, err := asnresolve.AnnouncedPrefixes(ctx, hc, entry.ASN)
+			if err != nil {
+				results[idx] = entryResult{err: fmt.Errorf("resolve AS%d: %w", entry.ASN, err)}
+				return
+			}
+			holder := ""
+			if h, err := asnresolve.ASHolderName(ctx, hc, entry.ASN); err == nil {
+				holder = h
+			}
+			src := fmt.Sprintf("as:%d", entry.ASN)
+			var rows []store.PrefixRow
+			for _, pfx := range pfxs {
+				rows = append(rows, store.PrefixRow{Prefix: pfx.String(), CommunityID: comm, Source: src})
+			}
+			results[idx] = entryResult{
+				rows:   rows,
+				metaID: entry.ID,
+				asn:    entry.ASN,
+				holder: holder,
+				count:  int64(len(pfxs)),
+			}
+		}(i, e)
+	}
+	wg.Wait()
+
+	seenPfx := make(map[string]struct{})
+	var out []store.PrefixRow
+	now := time.Now().UTC()
+	for _, r := range results {
+		if r.err != nil {
+			return nil, r.err
+		}
+		if r.metaID != "" {
+			if err := st.UpdateASEntryResolveMeta(tenantID, moduleID, r.metaID, r.holder, r.count, now); err != nil {
+				return nil, fmt.Errorf("as entry meta AS%d: %w", r.asn, err)
+			}
+		}
+		for _, row := range r.rows {
+			k := row.Prefix
+			if _, ok := seenPfx[k]; ok {
+				continue
+			}
+			seenPfx[k] = struct{}{}
+			out = append(out, row)
+		}
+	}
+	return out, nil
+}
+
+func collectCDNPrefixRows(ctx context.Context, st store.Backend, hc *http.Client, tenantID string, mod *store.Module, sources []*store.CDNSource, priorSnapshot []store.PrefixRow) ([]store.PrefixRow, error) {
+	moduleID := mod.ID
+	now := time.Now().UTC()
+
+	var valid []*store.CDNSource
+	for _, s := range sources {
+		if s != nil {
+			valid = append(valid, s)
+		}
+	}
+	type srcResult struct {
+		rows []store.PrefixRow
+		err  error
+	}
+	results := make([]srcResult, len(valid))
+	sem := make(chan struct{}, collectConcurrency())
+	var wg sync.WaitGroup
+
+	for i, src := range valid {
+		wg.Add(1)
+		go func(idx int, src *store.CDNSource) {
+			defer wg.Done()
+			sem <- struct{}{}
+			defer func() { <-sem }()
+
+			sourceKey := "cdn:" + src.ID
+			if shouldSkipCDNSourceFetch(src, now) {
+				if cached := prefixRowsForSource(priorSnapshot, sourceKey); len(cached) > 0 {
+					results[idx] = srcResult{rows: cached}
+					return
+				}
+				if cached := latestCDNRowsBySource(st, tenantID)[sourceKey]; len(cached) > 0 {
+					results[idx] = srcResult{rows: cached}
+					return
+				}
+			}
+			u := strings.TrimSpace(src.URL)
+			if u == "" {
+				return
+			}
+			rows, err := fetchAndParseCDNSource(ctx, hc, st, tenantID, moduleID, mod, src, now)
+			if err != nil {
+				results[idx] = srcResult{err: err}
+				return
+			}
+			for i := range rows {
+				rows[i].Source = sourceKey
+			}
+			results[idx] = srcResult{rows: rows}
+		}(i, src)
+	}
+	wg.Wait()
+
+	var out []store.PrefixRow
+	for _, r := range results {
+		if r.err != nil {
+			return nil, r.err
+		}
+		out = append(out, r.rows...)
+	}
+	return out, nil
+}
+
+func fetchAndParseCDNSource(ctx context.Context, hc *http.Client, st store.Backend, tenantID, moduleID string, mod *store.Module, src *store.CDNSource, now time.Time) ([]store.PrefixRow, error) {
+	u := strings.TrimSpace(src.URL)
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, u, nil)
+	if err != nil {
+		return nil, err
+	}
+	resp, err := hc.Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("cdn fetch %s: %w", u, err)
+	}
+	if resp.StatusCode != http.StatusOK {
+		_, _ = io.Copy(io.Discard, resp.Body)
+		_ = resp.Body.Close()
+		return nil, fmt.Errorf("cdn url %s: %s", u, resp.Status)
+	}
+	body, err := io.ReadAll(io.LimitReader(resp.Body, 8<<20))
+	_ = resp.Body.Close()
+	if err != nil {
+		return nil, err
+	}
+	etag := strings.TrimSpace(resp.Header.Get("ETag"))
+	patch := &store.CDNSourcePatch{}
+	if etag != "" && etag != strings.TrimSpace(src.Etag) {
+		e := etag
+		patch.Etag = &e
+	}
+	refreshedAt := now
+	patch.LastRefreshedAt = &refreshedAt
+	_, _ = st.UpdateCDNSource(tenantID, moduleID, src.ID, patch)
+	pfxs, err := ExtractCIDRs(string(body), src.SourceKind, src.PrefixPath)
+	if err != nil {
+		return nil, fmt.Errorf("cdn parse %s: %w", u, err)
+	}
+	var rows []store.PrefixRow
+	for _, pfx := range pfxs {
+		comm := src.CommunityID
+		if comm == nil && mod.DefaultCommunityID != nil {
+			c := *mod.DefaultCommunityID
+			comm = &c
+		}
+		rows = append(rows, store.PrefixRow{Prefix: pfx.String(), CommunityID: comm})
+	}
+	return rows, nil
+}
+
+func collectDomainPrefixRows(ctx context.Context, hc *http.Client, mod *store.Module, profile *store.DohProfile, entries []*store.DomainEntry) ([]store.PrefixRow, error) {
+	var validDom []*store.DomainEntry
+	for _, e := range entries {
+		if e != nil {
+			validDom = append(validDom, e)
+		}
+	}
+	type domResult struct {
+		rows []store.PrefixRow
+		err  error
+	}
+	results := make([]domResult, len(validDom))
+	sem := make(chan struct{}, collectConcurrency())
+	var wg sync.WaitGroup
+
+	for i, e := range validDom {
+		wg.Add(1)
+		go func(idx int, entry *store.DomainEntry) {
+			defer wg.Done()
+			sem <- struct{}{}
+			defer func() { <-sem }()
+
+			comm := entry.CommunityID
+			if comm == nil && mod.DefaultCommunityID != nil {
+				c := *mod.DefaultCommunityID
+				comm = &c
+			}
+			addrs, err := resolveDomainIPs(ctx, hc, profile, entry.FQDN)
+			if err != nil {
+				results[idx] = domResult{err: fmt.Errorf("resolve domain %q: %w", entry.FQDN, err)}
+				return
+			}
+			src := "domain:" + strings.TrimSpace(entry.FQDN)
+			var rows []store.PrefixRow
+			for _, ip := range addrs {
+				cidr := ipToHostPrefix(ip)
+				if cidr == "" {
+					continue
+				}
+				rows = append(rows, store.PrefixRow{
+					Prefix:      cidr,
+					CommunityID: comm,
+					Source:      src,
+				})
+			}
+			results[idx] = domResult{rows: rows}
+		}(i, e)
+	}
+	wg.Wait()
+
+	seen := make(map[string]struct{})
+	var out []store.PrefixRow
+	for _, r := range results {
+		if r.err != nil {
+			return nil, r.err
+		}
+		for _, row := range r.rows {
+			key := row.Prefix + "|" + row.Source
+			if _, ok := seen[key]; ok {
+				continue
+			}
+			seen[key] = struct{}{}
+			out = append(out, row)
+		}
+	}
+	return out, nil
+}