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.
This commit is contained in:
MHSanaei
2026-11-01 08:33:19 +01:00
parent 7f64a67ba5
commit 2024d34928
2 changed files with 46 additions and 24 deletions
+32
View File
@@ -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)
+14 -24
View File
@@ -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)