From 2024d349287e5c32ad6c5da6482168517a459dc6 Mon Sep 17 00:00:00 2001 From: MHSanaei Date: Sun, 1 Nov 2026 08:33:19 +0100 Subject: [PATCH] fix(geodata): index a split category's attributes in one pass scanIndex merged each further record of an already-seen category with mergeAttributes, which rebuilt a map and re-sorted every attribute the category had so far. A category with A attributes split over N records cost A*N, about an hour and a half for a 5 MB file, while the scan lock was held and every geodata route waited on it. Each category now collects its attributes in one set during the walk, and the sorted lists are built once at the end. --- internal/xray/geodata/geodata_test.go | 32 ++++++++++++++++++++++ internal/xray/geodata/reader.go | 38 ++++++++++----------------- 2 files changed, 46 insertions(+), 24 deletions(-) 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)