diff --git a/eventstore/mmm/fuzz_test.go b/eventstore/mmm/fuzz_test.go index a9c0203..c352e66 100644 --- a/eventstore/mmm/fuzz_test.go +++ b/eventstore/mmm/fuzz_test.go @@ -3,6 +3,7 @@ package mmm import ( "fmt" "iter" + "maps" "math/rand/v2" "os" "slices" @@ -101,13 +102,15 @@ func FuzzTest(f *testing.F) { require.Equal(t, count, len(storedByLayer[layer.name]), "layer %d ('%s')", i, layer.name) // call ComputeStats - stats, err := layer.ComputeStats() + stats, err := layer.ComputeStats(StatsOptions{}) require.NoError(t, err, "ComputeStats failed for layer %d ('%s')", i, layer.name) - require.NotNil(t, stats, "ComputeStats returned nil for layer %d ('%s')", i, layer.name) require.Equal(t, stats.Total, uint(count)) if count > 0 { require.GreaterOrEqual(t, len(stats.PerWeek), 1) - require.Len(t, stats.PerPubKeyPrefix, 1) + require.Len(t, stats.PerPubKey, 1) + for pk := range maps.Keys(stats.PerPubKey) { + require.Equal(t, pk, sk.Public()) + } } } diff --git a/eventstore/mmm/stats.go b/eventstore/mmm/stats.go index 7031730..f79bb80 100644 --- a/eventstore/mmm/stats.go +++ b/eventstore/mmm/stats.go @@ -29,8 +29,12 @@ type PubKeyStats struct { PerKindPerWeek map[nostr.Kind][]uint } -func (il *IndexingLayer) ComputeStats() (*EventStats, error) { - stats := &EventStats{ +type StatsOptions struct { + OnlyPubKey nostr.PubKey +} + +func (il *IndexingLayer) ComputeStats(opts StatsOptions) (EventStats, error) { + stats := EventStats{ Total: 0, PerWeek: make([]uint, 0, 24), PerPubKey: make(map[nostr.PubKey]PubKeyStats, 30), @@ -47,22 +51,40 @@ func (il *IndexingLayer) ComputeStats() (*EventStats, error) { var currentPubKeyPrefix []byte var currentPubKey nostr.PubKey - for { - key, val, err := cursor.Get(nil, nil, lmdb.Next) + // position cursor based on options + var initialKey []byte + if opts.OnlyPubKey != nostr.ZeroPK { + // position cursor at the start of this author's data + initialKey = make([]byte, 8+4+4) + copy(initialKey[0:8], opts.OnlyPubKey[0:8]) + } + + var key []byte + var val []byte + if initialKey == nil { + key, val, err = cursor.Get(nil, nil, lmdb.Next) if lmdb.IsNotFound(err) { - break + return nil } if err != nil { return err } - - if len(key) < 14 { - continue + } else { + key, val, err = cursor.Get(initialKey, nil, lmdb.SetRange) + if err != nil { + return err } + } + for { // parse key: [8 bytes pubkey][2 bytes kind][4 bytes timestamp] pubkeyPrefix := key[0:8] if !bytes.Equal(pubkeyPrefix, currentPubKeyPrefix) { + if opts.OnlyPubKey != nostr.ZeroPK && len(currentPubKeyPrefix) > 0 { + // stop scanning now as we're filtering for a specific pubkey + break + } + // load pubkey from event (otherwise will use the same from before) pos := positionFromBytes(val) currentPubKey = betterbinary.GetPubKey(il.mmmm.mmapf[pos.start : pos.start+uint64(pos.size)]) @@ -115,6 +137,14 @@ func (il *IndexingLayer) ComputeStats() (*EventStats, error) { Total: 1, } } + + key, val, err = cursor.Get(nil, nil, lmdb.Next) + if lmdb.IsNotFound(err) { + break + } + if err != nil { + return err + } } return nil diff --git a/eventstore/mmm/stats_test.go b/eventstore/mmm/stats_test.go new file mode 100644 index 0000000..6841d09 --- /dev/null +++ b/eventstore/mmm/stats_test.go @@ -0,0 +1,92 @@ +package mmm + +import ( + "os" + "testing" + + "fiatjaf.com/nostr" + "github.com/stretchr/testify/require" +) + +func TestComputeStats(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "mmm_stats_test") + require.NoError(t, err) + defer os.RemoveAll(tmpDir) + mmmm := &MultiMmapManager{ + Dir: tmpDir, + } + err = mmmm.Init() + require.NoError(t, err) + defer mmmm.Close() + il, err := mmmm.EnsureLayer("testlayer") + require.NoError(t, err) + + // generate 5 random keys + keys := make([]nostr.SecretKey, 5) + pubkeys := make([]nostr.PubKey, 5) + for i := 0; i < 5; i++ { + privkey := nostr.Generate() + keys[i] = privkey + pubkeys[i] = privkey.Public() + } + + // add 10 events from each key, alternating between kinds 1 and 11 + for i := 0; i < 5; i++ { + for j := 0; j < 10; j++ { + kind := nostr.Kind(1) + if j%2 == 1 { + kind = 11 + } + + evt := nostr.Event{ + PubKey: pubkeys[i], + CreatedAt: nostr.Now() - nostr.Timestamp(j)*3600, // j hours ago + Kind: kind, + Tags: nil, + Content: "test event", + } + err := evt.Sign(keys[i]) + require.NoError(t, err) + + // save event + err = il.SaveEvent(evt) + require.NoError(t, err) + } + } + + // test ComputeStats with no options + stats, err := il.ComputeStats(StatsOptions{}) + require.NoError(t, err) + + // verify total count + require.Equal(t, uint(50), stats.Total) + + // verify we have stats for all 5 pubkeys + require.Len(t, stats.PerPubKey, 5) + + // verify each pubkey has 10 events + for _, pubkey := range pubkeys { + pkStats, _ := stats.PerPubKey[pubkey] + require.Equal(t, uint(10), pkStats.Total) + } + + // verify we have stats for both kinds + require.Len(t, stats.PerKind, 2) + + // verify kind counts (should be 25 each for kinds 1 and 11) + kindStats1, exists := stats.PerKind[1] + require.True(t, exists, "missing stats for kind 1") + require.Equal(t, uint(25), kindStats1.Total, "expected 25 events for kind 1, got %d", kindStats1.Total) + + kindStats11, exists := stats.PerKind[11] + require.True(t, exists, "missing stats for kind 11") + require.Equal(t, uint(25), kindStats11.Total, "expected 25 events for kind 11, got %d", kindStats11.Total) + + // test ComputeStats with OnlyPubKey option + firstPubkey := pubkeys[0] + stats, err = il.ComputeStats(StatsOptions{OnlyPubKey: firstPubkey}) + require.NoError(t, err, "failed to compute stats with OnlyPubKey: %v", err) + + require.Equal(t, uint(10), stats.Total) + require.Len(t, stats.PerPubKey, 1) +}