EvoBGP/internal/pipeline/collect_parallel.go

package pipeline

import (
	"context"
	"evobgp/internal/logging"
	"fmt"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"

	"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 prefixCommunityKey(row store.PrefixRow) string {
	comm := ""
	if row.CommunityID != nil {
		comm = strings.TrimSpace(*row.CommunityID)
	}
	return row.Prefix + "\x00" + comm
}

func collectASPrefixRows(ctx context.Context, st store.Backend, hc *http.Client, tenantID string, mod *store.Module, list []*store.ASEntry, priorSnapshot []store.PrefixRow) ([]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, holder, err := resolveASNForEntry(ctx, st, hc, entry.ASN)
			if err != nil {
				if staleOnUpstreamError() {
					if staleRows, staleHolder, ok := staleASNPrefixes(st, priorSnapshot, entry.ASN); ok {
						logStaleUpstream("asn", fmt.Sprintf("AS%d: %v", entry.ASN, err))
						src := fmt.Sprintf("as:%d", entry.ASN)
						rows := append([]store.PrefixRow(nil), staleRows...)
						for i := range rows {
							rows[i].CommunityID = comm
							rows[i].Source = src
						}
						results[idx] = entryResult{
							rows:   rows,
							metaID: entry.ID,
							asn:    entry.ASN,
							holder: staleHolder,
							count:  int64(len(rows)),
						}
						return
					}
					if pfxs2, holder2, ok := asnCacheExpired(st, entry.ASN); ok {
						logStaleUpstream("asn", fmt.Sprintf("AS%d expired cache: %v", entry.ASN, err))
						src := fmt.Sprintf("as:%d", entry.ASN)
						var rows []store.PrefixRow
						for _, pfx := range pfxs2 {
							rows = append(rows, store.PrefixRow{Prefix: pfx.String(), CommunityID: comm, Source: src})
						}
						results[idx] = entryResult{
							rows:   rows,
							metaID: entry.ID,
							asn:    entry.ASN,
							holder: holder2,
							count:  int64(len(pfxs2)),
						}
						return
					}
				}
				results[idx] = entryResult{err: fmt.Errorf("resolve AS%d: %w", entry.ASN, err)}
				return
			}
			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
	var metaUpdates []store.ASEntryResolveMetaUpdate
	now := time.Now().UTC()
	for _, r := range results {
		if r.err != nil {
			return nil, r.err
		}
		if r.metaID != "" {
			metaUpdates = append(metaUpdates, store.ASEntryResolveMetaUpdate{
				EntryID:     r.metaID,
				ASNName:     r.holder,
				PrefixCount: r.count,
			})
		}
		for _, row := range r.rows {
			k := prefixCommunityKey(row)
			if _, ok := seenPfx[k]; ok {
				continue
			}
			seenPfx[k] = struct{}{}
			out = append(out, row)
		}
	}
	if len(metaUpdates) > 0 {
		if err := st.UpdateASEntryResolveMetaBatch(tenantID, moduleID, metaUpdates, now); err != nil {
			return nil, fmt.Errorf("as entry meta batch: %w", err)
		}
	}
	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 snap, ok, _ := st.GetModulePrefixSnapshot(tenantID, moduleID); ok && snap != nil {
					if cached := prefixRowsForSource(snap.Prefixes, sourceKey); len(cached) > 0 {
						results[idx] = srcResult{rows: cached}
						return
					}
				}
			}
			rows, err := fetchCDNSourceRows(ctx, st, hc, tenantID, moduleID, mod, src, priorSnapshot, now)
			if err != nil {
				if staleOnUpstreamError() {
					if cached, ok := staleCDNPrefixes(st, tenantID, moduleID, priorSnapshot, src.ID); ok {
						logStaleUpstream("cdn", fmt.Sprintf("source %s: %v", src.ID, err))
						results[idx] = srcResult{rows: cached}
						return
					}
				}
				results[idx] = srcResult{err: err}
				return
			}
			results[idx] = srcResult{rows: rows}
		}(i, src)
	}
	wg.Wait()

	var fetchedRows []store.PrefixRow
	var skipped int
	var fetchedIDs []string
	allIDs := make([]string, 0, len(valid))
	for _, src := range valid {
		allIDs = append(allIDs, src.ID)
	}
	for i, r := range results {
		if r.err != nil {
			if cdnPartialOK() {
				logging.Default().Info(fmt.Sprintf("pipeline: CDN partial skip source error: %v", r.err))
				skipped++
				continue
			}
			return nil, r.err
		}
		fetchedIDs = append(fetchedIDs, valid[i].ID)
		fetchedRows = append(fetchedRows, r.rows...)
	}
	if skipped > 0 && len(fetchedRows) == 0 && len(valid) > 0 {
		return nil, fmt.Errorf("cdn: all %d source(s) failed (partial ok)", len(valid))
	}
	out := fetchedRows
	if skipped > 0 {
		base := priorSnapshot
		if len(base) == 0 {
			if snap, ok, _ := st.GetModulePrefixSnapshot(tenantID, moduleID); ok && snap != nil {
				base = snap.Prefixes
			}
		}
		for _, row := range mergeSnapshotKeepSkippedCDN(base, fetchedIDs, allIDs) {
			if strings.HasPrefix(strings.TrimSpace(row.Source), "cdn:") {
				out = append(out, row)
			}
		}
	}
	if len(valid) > 0 {
		if err := mergeAllCDNSourcesIntoModuleSnapshot(st, tenantID, mod, priorSnapshot, fetchedRows, fetchedIDs, allIDs); err != nil {
			return nil, err
		}
	}
	return out, nil
}

func collectDomainPrefixRows(ctx context.Context, st store.Backend, hc *http.Client, mod *store.Module, profiles []*store.DohProfile, policy string, entries []*store.DomainEntry, priorSnapshot []store.PrefixRow) ([]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 := resolveDomainIPsCached(ctx, st, hc, profiles, policy, entry.FQDN)
			if err != nil {
				if staleOnUpstreamError() {
					if cached, ok := staleDomainPrefixes(priorSnapshot, entry.FQDN); ok {
						logStaleUpstream("domain", fmt.Sprintf("%q: %v", entry.FQDN, err))
						rows := append([]store.PrefixRow(nil), cached...)
						for i := range rows {
							if rows[i].CommunityID == nil {
								rows[i].CommunityID = comm
							}
						}
						results[idx] = domResult{rows: rows}
						return
					}
				}
				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
}