mmm: better ComputeStats()

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())
+				}
 			}
 		}
 

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

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)
+}