diff --git a/internal/xray/geodata/geodata_test.go b/internal/xray/geodata/geodata_test.go index 00686a1b2..35ac99ac9 100644 --- a/internal/xray/geodata/geodata_test.go +++ b/internal/xray/geodata/geodata_test.go @@ -6,6 +6,7 @@ import ( "net/netip" "os" "path/filepath" + "slices" "strconv" "strings" "sync" @@ -407,6 +408,37 @@ func TestFileAboveSizeLimitIsRejected(t *testing.T) { } } +// Rebuilding the merged attribute list for every further record of a category +// made this sub-megabyte database take minutes to index instead of milliseconds. +func TestCategorySplitOverManyRecordsIsIndexedInOnePass(t *testing.T) { + dir := t.TempDir() + tagged := make([]*xraygeodata.Domain, 0, 20000) + for i := range 20000 { + tagged = append(tagged, domain(xraygeodata.Domain_Full, "a.example", strconv.Itoa(i))) + } + sites := []*xraygeodata.GeoSite{site("x", tagged...)} + for range 20000 { + sites = append(sites, site("x", domain(xraygeodata.Domain_Full, "a.example", "extra"))) + } + writeSiteDB(t, dir, "split.dat", sites...) + + started := time.Now() + page, err := NewStore(dir).Categories("split.dat", "", 0, 1) + if err != nil { + t.Fatalf("Categories() error = %v", err) + } + if elapsed := time.Since(started); elapsed > 10*time.Second { + t.Errorf("indexing took %s", elapsed) + } + got := page.Items[0] + if got.Entries != 40000 || len(got.Attributes) != 20001 { + t.Fatalf("category x: %d entries and %d attributes, want 40000 and 20001", got.Entries, len(got.Attributes)) + } + if !slices.IsSorted(got.Attributes) || !slices.Contains(got.Attributes, "extra") { + t.Errorf("attributes are not the sorted union of every record: %q ... %q", got.Attributes[0], got.Attributes[20000]) + } +} + func TestIndexCacheInvalidatedWhenFileChanges(t *testing.T) { dir := t.TempDir() sampleSiteDB(t, dir) diff --git a/internal/xray/geodata/reader.go b/internal/xray/geodata/reader.go index 88500303f..eaaa1bf4a 100644 --- a/internal/xray/geodata/reader.go +++ b/internal/xray/geodata/reader.go @@ -4,6 +4,7 @@ import ( "errors" "fmt" "io" + "maps" "net/netip" "os" "sort" @@ -53,6 +54,9 @@ type byteSpan struct { func scanIndex(data []byte, kind GeoKind) (*categoryScan, error) { scan := &categoryScan{kind: kind, spans: make(map[string][]byteSpan)} byCode := make(map[string]int) + // Parallel to scan.categories: a category split over many records adds each + // record's own attributes to its set and is sorted once, at the end. + var attributeSets []map[string]struct{} err := eachListEntry(data, func(entry []byte, span byteSpan) error { count := 0 @@ -88,20 +92,20 @@ func scanIndex(data []byte, kind GeoKind) (*categoryScan, error) { scan.spans[code] = append(scan.spans[code], span) if position, seen := byCode[code]; seen { scan.categories[position].Entries += count - scan.categories[position].Attributes = mergeAttributes(scan.categories[position].Attributes, attributes) + maps.Copy(attributeSets[position], attributes) return nil } byCode[code] = len(scan.categories) - scan.categories = append(scan.categories, GeoCategory{ - Code: code, - Entries: count, - Attributes: mergeAttributes(nil, attributes), - }) + scan.categories = append(scan.categories, GeoCategory{Code: code, Entries: count}) + attributeSets = append(attributeSets, attributes) return nil }) if err != nil { return nil, err } + for position, attributes := range attributeSets { + scan.categories[position].Attributes = sortedAttributes(attributes) + } sort.Slice(scan.categories, func(i, j int) bool { return scan.categories[i].Code < scan.categories[j].Code }) return scan, nil } @@ -461,25 +465,11 @@ func cidrBytes(payload []byte) ([]byte, bool, error) { return []byte(netip.PrefixFrom(addr, int(prefix)).String()), true, nil } -// mergeAttributes always returns a non-nil slice: the JSON contract declares -// attributes as an array, and a nil slice would marshal to null and break -// clients validating against it. -func mergeAttributes(existing []string, attributes map[string]struct{}) []string { - if len(attributes) == 0 { - if existing == nil { - return []string{} - } - return existing - } - merged := make(map[string]struct{}, len(existing)+len(attributes)) - for _, attribute := range existing { - merged[attribute] = struct{}{} - } +// sortedAttributes always returns a non-nil slice: the JSON contract declares +// attributes as an array, and a nil slice would marshal to null. +func sortedAttributes(attributes map[string]struct{}) []string { + out := make([]string, 0, len(attributes)) for attribute := range attributes { - merged[attribute] = struct{}{} - } - out := make([]string, 0, len(merged)) - for attribute := range merged { out = append(out, attribute) } sort.Strings(out)