diff --git a/khatru/dispatcher_test.go b/khatru/dispatcher_test.go index df2848c..79c404b 100644 --- a/khatru/dispatcher_test.go +++ b/khatru/dispatcher_test.go @@ -89,6 +89,7 @@ func FuzzDispatcherCandidates(f *testing.F) { } ssid := d.addSubscription(sub) + active[ssid] = sub activeSSIDs = append(activeSSIDs, ssid) } else { @@ -96,14 +97,17 @@ func FuzzDispatcherCandidates(f *testing.F) { ssid := activeSSIDs[idx] d.removeSubscription(ssid) + delete(active, ssid) activeSSIDs = append(activeSSIDs[:idx], activeSSIDs[idx+1:]...) } for range int(checks%7) + 1 { event := fuzzDispatcherEvent(&state) + expected := expectedDispatcherCandidates(active, event) actual := collectedDispatcherCandidates(&d, event) + require.Equalf(t, expected, actual, "seed=%d advance=%d event=%s active=%v", seed, advance, event.String(), active) } } @@ -203,7 +207,7 @@ func fuzzDispatcherTagMap(seed *fuzzState) nostr.TagMap { count := seed.next(3) if count == 0 { - return nostr.TagMap{} + return nil } tags := make(nostr.TagMap, count) diff --git a/khatru/listener.go b/khatru/listener.go index e6f6cfd..3a66a79 100644 --- a/khatru/listener.go +++ b/khatru/listener.go @@ -31,19 +31,21 @@ type subscription struct { } type dispatcher struct { - serial int - subscriptions *xsync.MapOf[int, subscription] - byAuthor *xsync.MapOf[nostr.PubKey, set.Set[int]] - byKind *xsync.MapOf[nostr.Kind, set.Set[int]] - fallback set.Set[int] + serial int + subscriptions *xsync.MapOf[int, subscription] + byAuthor *xsync.MapOf[nostr.PubKey, set.Set[int]] + byKind *xsync.MapOf[nostr.Kind, set.Set[int]] + fallbackTags set.Set[int] + fallbackNothing set.Set[int] } func newDispatcher() dispatcher { return dispatcher{ - subscriptions: xsync.NewMapOf[int, subscription](), - byAuthor: xsync.NewMapOf[nostr.PubKey, set.Set[int]](), - byKind: xsync.NewMapOf[nostr.Kind, set.Set[int]](), - fallback: set.NewSliceSet[int](), + subscriptions: xsync.NewMapOf[int, subscription](), + byAuthor: xsync.NewMapOf[nostr.PubKey, set.Set[int]](), + byKind: xsync.NewMapOf[nostr.Kind, set.Set[int]](), + fallbackTags: set.NewSliceSet[int](), + fallbackNothing: set.NewSliceSet[int](), } } @@ -57,72 +59,80 @@ func (d *dispatcher) addSubscription(sub subscription) int { if sub.filter.Authors != nil { indexed = true for _, author := range sub.filter.Authors { - s, ok := d.byAuthor.Load(author) - if !ok { - s = set.NewSliceSet[int]() - d.byAuthor.Store(author, s) - } - s.Add(ssid) + d.byAuthor.Compute(author, func(s set.Set[int], loaded bool) (set.Set[int], bool) { + if !loaded { + s = set.NewSliceSet[int]() + } + s.Add(ssid) + return s, false + }) } } if sub.filter.Kinds != nil { indexed = true for _, kind := range sub.filter.Kinds { - s, ok := d.byKind.Load(kind) - if !ok { - s = set.NewSliceSet[int]() - d.byKind.Store(kind, s) - } - s.Add(ssid) + d.byKind.Compute(kind, func(s set.Set[int], loaded bool) (set.Set[int], bool) { + if !loaded { + s = set.NewSliceSet[int]() + } + s.Add(ssid) + return s, false + }) } } if !indexed { - d.fallback.Add(ssid) + if sub.filter.Tags != nil { + d.fallbackTags.Add(ssid) + } else { + d.fallbackNothing.Add(ssid) + } } return ssid } func (d *dispatcher) removeSubscription(ssid int) { - sub, ok := d.subscriptions.LoadAndDelete(ssid) - if !ok { - return - } + d.subscriptions.Compute(ssid, func(sub subscription, loaded bool) (subscription, bool) { + indexed := false - indexed := false - if sub.filter.Authors != nil { - indexed = true - for _, author := range sub.filter.Authors { - s, ok := d.byAuthor.Load(author) - if !ok { - return - } - s.Remove(ssid) - if s.Len() == 0 { - d.byAuthor.Delete(author) + if sub.filter.Authors != nil { + indexed = true + for _, author := range sub.filter.Authors { + d.byAuthor.Compute(author, func(s set.Set[int], loaded bool) (set.Set[int], bool) { + if !loaded { + return s, true + } + s.Remove(ssid) + return s, s.Len() == 0 + }) } } - } - if sub.filter.Kinds != nil { - indexed = true - for _, kind := range sub.filter.Kinds { - s, ok := d.byKind.Load(kind) - if !ok { - return - } - s.Remove(ssid) - if s.Len() == 0 { - d.byKind.Delete(kind) + if sub.filter.Kinds != nil { + indexed = true + for _, kind := range sub.filter.Kinds { + d.byKind.Compute(kind, func(s set.Set[int], loaded bool) (set.Set[int], bool) { + if !loaded { + return s, true + } + s.Remove(ssid) + return s, s.Len() == 0 + }) } } - } - if !indexed { - d.fallback.Remove(ssid) - } + if !indexed { + if sub.filter.Tags != nil { + d.fallbackTags.Remove(ssid) + } else { + d.fallbackNothing.Remove(ssid) + } + } + + return sub, true + }) } func (d *dispatcher) candidates(event nostr.Event) iter.Seq[subscription] { @@ -188,10 +198,21 @@ func (d *dispatcher) candidates(event nostr.Event) iter.Seq[subscription] { } } - for _, ssid := range d.fallback.Slice() { - sub, _ := d.subscriptions.Load(ssid) + if len(event.Tags) > 0 { + for _, ssid := range d.fallbackTags.Slice() { + sub, _ := d.subscriptions.Load(ssid) - if filterMatchesTimestampConstraintsAndTags(sub.filter, event) { + if filterMatchesTimestampConstraintsAndTags(sub.filter, event) { + if !yield(sub) { + return + } + } + } + } + + for _, ssid := range d.fallbackNothing.Slice() { + sub, _ := d.subscriptions.Load(ssid) + if filterMatchesTimestampConstraints(sub.filter, event) { if !yield(sub) { return } @@ -201,7 +222,7 @@ func (d *dispatcher) candidates(event nostr.Event) iter.Seq[subscription] { } //go:inline -func filterMatchesTimestampConstraintsAndTags(filter nostr.Filter, event nostr.Event) bool { +func filterMatchesTimestampConstraints(filter nostr.Filter, event nostr.Event) bool { if filter.Since != 0 && event.CreatedAt < filter.Since { return false } @@ -210,6 +231,15 @@ func filterMatchesTimestampConstraintsAndTags(filter nostr.Filter, event nostr.E return false } + return true +} + +//go:inline +func filterMatchesTimestampConstraintsAndTags(filter nostr.Filter, event nostr.Event) bool { + if !filterMatchesTimestampConstraints(filter, event) { + return false + } + for f, v := range filter.Tags { if !event.Tags.ContainsAny(f, v) { return false diff --git a/khatru/testdata/fuzz/FuzzDispatcherCandidates/030fbe034592967a b/khatru/testdata/fuzz/FuzzDispatcherCandidates/030fbe034592967a new file mode 100644 index 0000000..95df224 --- /dev/null +++ b/khatru/testdata/fuzz/FuzzDispatcherCandidates/030fbe034592967a @@ -0,0 +1,5 @@ +go test fuzz v1 +int(-180) +int(92) +byte('{') +byte('\n') diff --git a/khatru/testdata/fuzz/FuzzDispatcherCandidates/0be355dc7fe62f98 b/khatru/testdata/fuzz/FuzzDispatcherCandidates/0be355dc7fe62f98 new file mode 100644 index 0000000..91823f0 --- /dev/null +++ b/khatru/testdata/fuzz/FuzzDispatcherCandidates/0be355dc7fe62f98 @@ -0,0 +1,5 @@ +go test fuzz v1 +int(140) +int(-52) +byte('"') +byte('h')