mirror of
https://github.com/MHSanaei/3x-ui.git
synced 2026-11-02 03:05:02 +03:00
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:
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user