eventstore: boltdb (copied from lmdb)

eventstore/boltdb/count.go

@@ -0,0 +1,228 @@
+package bolt
+
+import (
+	"bytes"
+	"encoding/binary"
+	"encoding/hex"
+	"slices"
+
+	"fiatjaf.com/nostr"
+	"fiatjaf.com/nostr/eventstore/codec/betterbinary"
+	"fiatjaf.com/nostr/nip45"
+	"fiatjaf.com/nostr/nip45/hyperloglog"
+	"go.etcd.io/bbolt"
+)
+
+func (b *BoltBackend) CountEvents(filter nostr.Filter) (uint32, error) {
+	var count uint32 = 0
+
+	queries, extraAuthors, extraKinds, extraTagKey, extraTagValues, since, err := b.prepareQueries(filter)
+	if err != nil {
+		return 0, err
+	}
+
+	err = b.DB.View(func(txn *bbolt.Tx) error {
+		rawBucket := txn.Bucket(rawEventStore)
+
+		// actually iterate
+		for _, q := range queries {
+			cursor := txn.Bucket(q.bucket).Cursor()
+
+			it := &iterator{cursor: cursor}
+			it.seek(q.startingPoint)
+
+			for {
+				// we already have a k and a v and an err from the cursor setup, so check and use these
+				if it.err != nil ||
+					!bytes.HasPrefix(it.key, q.prefix) {
+					// either iteration has errored or we reached the end of this prefix
+					break // stop this cursor and move to the next one
+				}
+
+				createdAt := binary.BigEndian.Uint32(it.key[len(it.key)-4:])
+				if createdAt < since {
+					break
+				}
+
+				if extraAuthors == nil && extraKinds == nil && extraTagValues == nil {
+					count++
+				} else {
+					// fetch actual event
+					bin := rawBucket.Get(it.currIdPtr)
+					if bin == nil {
+						it.next()
+						continue
+					}
+
+					// check it against pubkeys without decoding the entire thing
+					if !slices.Contains(extraAuthors, betterbinary.GetPubKey(bin)) {
+						it.next()
+						continue
+					}
+
+					// check it against kinds without decoding the entire thing
+					if !slices.Contains(extraKinds, betterbinary.GetKind(bin)) {
+						it.next()
+						continue
+					}
+
+					evt := &nostr.Event{}
+					if err := betterbinary.Unmarshal(bin, evt); err != nil {
+						it.next()
+						continue
+					}
+
+					// if there is still a tag to be checked, do it now
+					if !evt.Tags.ContainsAny(extraTagKey, extraTagValues) {
+						it.next()
+						continue
+					}
+
+					count++
+				}
+			}
+		}
+
+		return nil
+	})
+
+	return count, err
+}
+
+// CountEventsHLL is like CountEvents, but it will build a hyperloglog value while iterating through results,
+// following NIP-45
+func (b *BoltBackend) CountEventsHLL(filter nostr.Filter, offset int) (uint32, *hyperloglog.HyperLogLog, error) {
+	if useCache, _ := b.EnableHLLCacheFor(filter.Kinds[0]); useCache {
+		return b.countEventsHLLCached(filter)
+	}
+
+	var count uint32 = 0
+
+	// this is different than CountEvents because some of these extra checks are not applicable in HLL-valid filters
+	queries, _, extraKinds, extraTagKey, extraTagValues, since, err := b.prepareQueries(filter)
+	if err != nil {
+		return 0, nil, err
+	}
+
+	hll := hyperloglog.New(offset)
+
+	err = b.DB.View(func(txn *bbolt.Tx) error {
+		rawBucket := txn.Bucket(rawEventStore)
+
+		// actually iterate
+		for _, q := range queries {
+			cursor := txn.Bucket(q.bucket).Cursor()
+
+			it := &iterator{cursor: cursor}
+			it.seek(q.startingPoint)
+
+			for {
+				// we already have a k and a v and an err from the cursor setup, so check and use these
+				if it.err != nil ||
+					!bytes.HasPrefix(it.key, q.prefix) {
+					// either iteration has errored or we reached the end of this prefix
+					break // stop this cursor and move to the next one
+				}
+
+				createdAt := binary.BigEndian.Uint32(it.key[len(it.key)-4:])
+				if createdAt < since {
+					break
+				}
+
+				// fetch actual event (we need it regardless because we need the pubkey for the hll)
+				bin := rawBucket.Get(it.currIdPtr)
+				if bin == nil {
+					continue
+				}
+
+				if extraKinds == nil && extraTagValues == nil {
+					// nothing extra to check
+					count++
+					hll.AddBytes(betterbinary.GetPubKey(bin))
+				} else {
+					// check it against kinds without decoding the entire thing
+					if !slices.Contains(extraKinds, betterbinary.GetKind(bin)) {
+						it.next()
+						continue
+					}
+
+					evt := &nostr.Event{}
+					if err := betterbinary.Unmarshal(bin, evt); err != nil {
+						it.next()
+						continue
+					}
+
+					// if there is still a tag to be checked, do it now
+					if !evt.Tags.ContainsAny(extraTagKey, extraTagValues) {
+						it.next()
+						continue
+					}
+
+					count++
+					hll.Add(evt.PubKey)
+				}
+			}
+		}
+
+		return nil
+	})
+
+	return count, hll, err
+}
+
+// countEventsHLLCached will just return a cached value from disk (and presumably we don't even have the events required to compute this anymore).
+func (b *BoltBackend) countEventsHLLCached(filter nostr.Filter) (uint32, *hyperloglog.HyperLogLog, error) {
+	cacheKey := make([]byte, 2+8)
+	binary.BigEndian.PutUint16(cacheKey[0:2], uint16(filter.Kinds[0]))
+	switch filter.Kinds[0] {
+	case 3:
+		hex.Decode(cacheKey[2:2+8], []byte(filter.Tags["p"][0][0:8*2]))
+	case 7:
+		hex.Decode(cacheKey[2:2+8], []byte(filter.Tags["e"][0][0:8*2]))
+	case 1111:
+		hex.Decode(cacheKey[2:2+8], []byte(filter.Tags["E"][0][0:8*2]))
+	}
+
+	var count uint32
+	var hll *hyperloglog.HyperLogLog
+
+	err := b.DB.View(func(txn *bbolt.Tx) error {
+		val := txn.Bucket(hllCache).Get(cacheKey)
+		if val == nil {
+			return nil
+		}
+		hll = hyperloglog.NewWithRegisters(val, 0) // offset doesn't matter here
+		count = uint32(hll.Count())
+		return nil
+	})
+
+	return count, hll, err
+}
+
+func (b *BoltBackend) updateHyperLogLogCachedValues(txn *bbolt.Tx, evt nostr.Event) error {
+	cacheKey := make([]byte, 2+8)
+	binary.BigEndian.PutUint16(cacheKey[0:2], uint16(evt.Kind))
+	hllBucket := txn.Bucket(hllCache)
+
+	for ref, offset := range nip45.HyperLogLogEventPubkeyOffsetsAndReferencesForEvent(evt) {
+		// setup cache key (reusing buffer)
+		hex.Decode(cacheKey[2:2+8], []byte(ref[0:8*2]))
+
+		// fetch hll value from cache db
+		hll := hyperloglog.New(offset)
+		val := hllBucket.Get(cacheKey)
+		if val != nil {
+			hll.SetRegisters(val)
+		}
+
+		// add this event
+		hll.Add(evt.PubKey)
+
+		// save values back again
+		if err := hllBucket.Put(cacheKey, hll.GetRegisters()); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}