lmdb/badger: turn getIndexKeysForEvent() into an iterator.

fiatjaf · Oct 7, 2024 · 6e79b94 · 6e79b94
1 parent 1f67814
commit 6e79b94
Show file tree

Hide file tree

Showing 6 changed files with 166 additions and 144 deletions.
diff --git a/badger/delete.go b/badger/delete.go
@@ -42,7 +42,7 @@ func (b *BadgerBackend) DeleteEvent(ctx context.Context, evt *nostr.Event) error
 		deletionHappened = true
 
 		// calculate all index keys we have for this event and delete them
-		for _, k := range b.getIndexKeysForEvent(evt, idx[1:]) {
+		for k := range b.getIndexKeysForEvent(evt, idx[1:]) {
 			if err := txn.Delete(k); err != nil {
 				return err
 			}

diff --git a/badger/helpers.go b/badger/helpers.go
@@ -4,6 +4,7 @@ import (
 	"cmp"
 	"encoding/binary"
 	"encoding/hex"
+	"iter"
 	"math"
 	"strconv"
 	"strings"
@@ -42,95 +43,104 @@ func getTagIndexPrefix(tagValue string) ([]byte, int) {
 	return k, offset
 }
 
-func (b BadgerBackend) getIndexKeysForEvent(evt *nostr.Event, idx []byte) [][]byte {
-	keys := make([][]byte, 0, 18)
-
-	// indexes
-	{
-		// ~ by id
-		idPrefix8, _ := hex.DecodeString(evt.ID[0 : 8*2])
-		k := make([]byte, 1+8+4)
-		k[0] = indexIdPrefix
-		copy(k[1:], idPrefix8)
-		copy(k[1+8:], idx)
-		keys = append(keys, k)
-	}
+func (b BadgerBackend) getIndexKeysForEvent(evt *nostr.Event, idx []byte) iter.Seq[[]byte] {
+	return func(yield func([]byte) bool) {
+		{
+			// ~ by id
+			idPrefix8, _ := hex.DecodeString(evt.ID[0 : 8*2])
+			k := make([]byte, 1+8+4)
+			k[0] = indexIdPrefix
+			copy(k[1:], idPrefix8)
+			copy(k[1+8:], idx)
+			if !yield(k) {
+				return
+			}
+		}
 
-	{
-		// ~ by pubkey+date
-		pubkeyPrefix8, _ := hex.DecodeString(evt.PubKey[0 : 8*2])
-		k := make([]byte, 1+8+4+4)
-		k[0] = indexPubkeyPrefix
-		copy(k[1:], pubkeyPrefix8)
-		binary.BigEndian.PutUint32(k[1+8:], uint32(evt.CreatedAt))
-		copy(k[1+8+4:], idx)
-		keys = append(keys, k)
-	}
+		{
+			// ~ by pubkey+date
+			pubkeyPrefix8, _ := hex.DecodeString(evt.PubKey[0 : 8*2])
+			k := make([]byte, 1+8+4+4)
+			k[0] = indexPubkeyPrefix
+			copy(k[1:], pubkeyPrefix8)
+			binary.BigEndian.PutUint32(k[1+8:], uint32(evt.CreatedAt))
+			copy(k[1+8+4:], idx)
+			if !yield(k) {
+				return
+			}
+		}
 
-	{
-		// ~ by kind+date
-		k := make([]byte, 1+2+4+4)
-		k[0] = indexKindPrefix
-		binary.BigEndian.PutUint16(k[1:], uint16(evt.Kind))
-		binary.BigEndian.PutUint32(k[1+2:], uint32(evt.CreatedAt))
-		copy(k[1+2+4:], idx)
-		keys = append(keys, k)
-	}
+		{
+			// ~ by kind+date
+			k := make([]byte, 1+2+4+4)
+			k[0] = indexKindPrefix
+			binary.BigEndian.PutUint16(k[1:], uint16(evt.Kind))
+			binary.BigEndian.PutUint32(k[1+2:], uint32(evt.CreatedAt))
+			copy(k[1+2+4:], idx)
+			if !yield(k) {
+				return
+			}
+		}
 
-	{
-		// ~ by pubkey+kind+date
-		pubkeyPrefix8, _ := hex.DecodeString(evt.PubKey[0 : 8*2])
-		k := make([]byte, 1+8+2+4+4)
-		k[0] = indexPubkeyKindPrefix
-		copy(k[1:], pubkeyPrefix8)
-		binary.BigEndian.PutUint16(k[1+8:], uint16(evt.Kind))
-		binary.BigEndian.PutUint32(k[1+8+2:], uint32(evt.CreatedAt))
-		copy(k[1+8+2+4:], idx)
-		keys = append(keys, k)
-	}
+		{
+			// ~ by pubkey+kind+date
+			pubkeyPrefix8, _ := hex.DecodeString(evt.PubKey[0 : 8*2])
+			k := make([]byte, 1+8+2+4+4)
+			k[0] = indexPubkeyKindPrefix
+			copy(k[1:], pubkeyPrefix8)
+			binary.BigEndian.PutUint16(k[1+8:], uint16(evt.Kind))
+			binary.BigEndian.PutUint32(k[1+8+2:], uint32(evt.CreatedAt))
+			copy(k[1+8+2+4:], idx)
+			if !yield(k) {
+				return
+			}
+		}
 
-	// ~ by tagvalue+date
-	customIndex := b.IndexLongerTag != nil
-	customSkip := b.SkipIndexingTag != nil
+		// ~ by tagvalue+date
+		customIndex := b.IndexLongerTag != nil
+		customSkip := b.SkipIndexingTag != nil
 
-	for i, tag := range evt.Tags {
-		if len(tag) < 2 || len(tag[0]) != 1 || len(tag[1]) == 0 || len(tag[1]) > 100 {
-			if !customIndex || !b.IndexLongerTag(evt, tag[0], tag[1]) {
-				// not indexable
-				continue
+		for i, tag := range evt.Tags {
+			if len(tag) < 2 || len(tag[0]) != 1 || len(tag[1]) == 0 || len(tag[1]) > 100 {
+				if !customIndex || !b.IndexLongerTag(evt, tag[0], tag[1]) {
+					// not indexable
+					continue
+				}
 			}
-		}
 
-		firstIndex := slices.IndexFunc(evt.Tags, func(t nostr.Tag) bool { return len(t) >= 2 && t[1] == tag[1] })
-		if firstIndex != i {
-			// duplicate
-			continue
-		}
+			firstIndex := slices.IndexFunc(evt.Tags, func(t nostr.Tag) bool { return len(t) >= 2 && t[1] == tag[1] })
+			if firstIndex != i {
+				// duplicate
+				continue
+			}
 
-		if customSkip && b.SkipIndexingTag(evt, tag[0], tag[1]) {
-			// purposefully skipped
-			continue
-		}
+			if customSkip && b.SkipIndexingTag(evt, tag[0], tag[1]) {
+				// purposefully skipped
+				continue
+			}
 
-		// get key prefix (with full length) and offset where to write the last parts
-		k, offset := getTagIndexPrefix(tag[1])
+			// get key prefix (with full length) and offset where to write the last parts
+			k, offset := getTagIndexPrefix(tag[1])
 
-		// write the last parts (created_at and idx)
-		binary.BigEndian.PutUint32(k[offset:], uint32(evt.CreatedAt))
-		copy(k[offset+4:], idx)
-		keys = append(keys, k)
-	}
+			// write the last parts (created_at and idx)
+			binary.BigEndian.PutUint32(k[offset:], uint32(evt.CreatedAt))
+			copy(k[offset+4:], idx)
+			if !yield(k) {
+				return
+			}
+		}
 
-	{
-		// ~ by date only
-		k := make([]byte, 1+4+4)
-		k[0] = indexCreatedAtPrefix
-		binary.BigEndian.PutUint32(k[1:], uint32(evt.CreatedAt))
-		copy(k[1+4:], idx)
-		keys = append(keys, k)
+		{
+			// ~ by date only
+			k := make([]byte, 1+4+4)
+			k[0] = indexCreatedAtPrefix
+			binary.BigEndian.PutUint32(k[1:], uint32(evt.CreatedAt))
+			copy(k[1+4:], idx)
+			if !yield(k) {
+				return
+			}
+		}
 	}
-
-	return keys
 }
 
 func getAddrTagElements(tagValue string) (kind uint16, pkb []byte, d string) {

diff --git a/badger/save.go b/badger/save.go
@@ -37,7 +37,7 @@ func (b *BadgerBackend) SaveEvent(ctx context.Context, evt *nostr.Event) error {
 			return err
 		}
 
-		for _, k := range b.getIndexKeysForEvent(evt, idx[1:]) {
+		for k := range b.getIndexKeysForEvent(evt, idx[1:]) {
 			if err := txn.Set(k, nil); err != nil {
 				return err
 			}

diff --git a/lmdb/delete.go b/lmdb/delete.go
@@ -21,7 +21,7 @@ func (b *LMDBBackend) DeleteEvent(ctx context.Context, evt *nostr.Event) error {
 		}
 
 		// calculate all index keys we have for this event and delete them
-		for _, k := range b.getIndexKeysForEvent(evt) {
+		for k := range b.getIndexKeysForEvent(evt) {
 			err := txn.Del(k.dbi, k.key, idx)
 			k.free()
 			if err != nil {

diff --git a/lmdb/helpers.go b/lmdb/helpers.go
@@ -3,6 +3,7 @@ package lmdb
 import (
 	"encoding/binary"
 	"encoding/hex"
+	"iter"
 	"strconv"
 	"strings"
 	"sync"
@@ -25,79 +26,90 @@ func (key key) free() {
 	indexKeyPool.Put(key.key)
 }
 
-func (b *LMDBBackend) getIndexKeysForEvent(evt *nostr.Event) []key {
-	keys := make([]key, 0, 18)
-
-	// indexes
-	{
-		// ~ by id
-		k := indexKeyPool.Get().([]byte)
-		hex.Decode(k[0:8], []byte(evt.ID[0:8*2]))
-		keys = append(keys, key{dbi: b.indexId, key: k[0:8]})
-	}
-
-	{
-		// ~ by pubkey+date
-		k := indexKeyPool.Get().([]byte)
-		hex.Decode(k[0:8], []byte(evt.PubKey[0:8*2]))
-		binary.BigEndian.PutUint32(k[8:8+4], uint32(evt.CreatedAt))
-		keys = append(keys, key{dbi: b.indexPubkey, key: k[0 : 8+4]})
-	}
-
-	{
-		// ~ by kind+date
-		k := indexKeyPool.Get().([]byte)
-		binary.BigEndian.PutUint16(k[0:2], uint16(evt.Kind))
-		binary.BigEndian.PutUint32(k[2:2+4], uint32(evt.CreatedAt))
-		keys = append(keys, key{dbi: b.indexKind, key: k[0 : 2+4]})
-	}
-
-	{
-		// ~ by pubkey+kind+date
-		k := indexKeyPool.Get().([]byte)
-		hex.Decode(k[0:8], []byte(evt.PubKey[0:8*2]))
-		binary.BigEndian.PutUint16(k[8:8+2], uint16(evt.Kind))
-		binary.BigEndian.PutUint32(k[8+2:8+2+4], uint32(evt.CreatedAt))
-		keys = append(keys, key{dbi: b.indexPubkeyKind, key: k[0 : 8+2+4]})
-	}
-
-	// ~ by tagvalue+date
-	// ~ by p-tag+kind+date
-	for i, tag := range evt.Tags {
-		if len(tag) < 2 || len(tag[0]) != 1 || len(tag[1]) == 0 || len(tag[1]) > 100 {
-			// not indexable
-			continue
+func (b *LMDBBackend) getIndexKeysForEvent(evt *nostr.Event) iter.Seq[key] {
+	return func(yield func(key) bool) {
+		{
+			// ~ by id
+			k := indexKeyPool.Get().([]byte)
+			hex.Decode(k[0:8], []byte(evt.ID[0:8*2]))
+			if !yield(key{dbi: b.indexId, key: k[0:8]}) {
+				return
+			}
 		}
-		firstIndex := slices.IndexFunc(evt.Tags, func(t nostr.Tag) bool { return len(t) >= 2 && t[1] == tag[1] })
-		if firstIndex != i {
-			// duplicate
-			continue
+
+		{
+			// ~ by pubkey+date
+			k := indexKeyPool.Get().([]byte)
+			hex.Decode(k[0:8], []byte(evt.PubKey[0:8*2]))
+			binary.BigEndian.PutUint32(k[8:8+4], uint32(evt.CreatedAt))
+			if !yield(key{dbi: b.indexPubkey, key: k[0 : 8+4]}) {
+				return
+			}
 		}
 
-		// get key prefix (with full length) and offset where to write the created_at
-		dbi, k, offset := b.getTagIndexPrefix(tag[1])
-		binary.BigEndian.PutUint32(k[offset:], uint32(evt.CreatedAt))
-		keys = append(keys, key{dbi: dbi, key: k})
+		{
+			// ~ by kind+date
+			k := indexKeyPool.Get().([]byte)
+			binary.BigEndian.PutUint16(k[0:2], uint16(evt.Kind))
+			binary.BigEndian.PutUint32(k[2:2+4], uint32(evt.CreatedAt))
+			if !yield(key{dbi: b.indexKind, key: k[0 : 2+4]}) {
+				return
+			}
+		}
 
-		// now the p-tag+kind+date
-		if dbi == b.indexTag32 && tag[0] == "p" {
+		{
+			// ~ by pubkey+kind+date
 			k := indexKeyPool.Get().([]byte)
-			hex.Decode(k[0:8], []byte(tag[1][0:8*2]))
+			hex.Decode(k[0:8], []byte(evt.PubKey[0:8*2]))
 			binary.BigEndian.PutUint16(k[8:8+2], uint16(evt.Kind))
 			binary.BigEndian.PutUint32(k[8+2:8+2+4], uint32(evt.CreatedAt))
-			dbi := b.indexPTagKind
-			keys = append(keys, key{dbi: dbi, key: k[0 : 8+2+4]})
+			if !yield(key{dbi: b.indexPubkeyKind, key: k[0 : 8+2+4]}) {
+				return
+			}
 		}
-	}
 
-	{
-		// ~ by date only
-		k := indexKeyPool.Get().([]byte)
-		binary.BigEndian.PutUint32(k[0:4], uint32(evt.CreatedAt))
-		keys = append(keys, key{dbi: b.indexCreatedAt, key: k[0:4]})
-	}
+		// ~ by tagvalue+date
+		// ~ by p-tag+kind+date
+		for i, tag := range evt.Tags {
+			if len(tag) < 2 || len(tag[0]) != 1 || len(tag[1]) == 0 || len(tag[1]) > 100 {
+				// not indexable
+				continue
+			}
+			firstIndex := slices.IndexFunc(evt.Tags, func(t nostr.Tag) bool { return len(t) >= 2 && t[1] == tag[1] })
+			if firstIndex != i {
+				// duplicate
+				continue
+			}
+
+			// get key prefix (with full length) and offset where to write the created_at
+			dbi, k, offset := b.getTagIndexPrefix(tag[1])
+			binary.BigEndian.PutUint32(k[offset:], uint32(evt.CreatedAt))
+			if !yield(key{dbi: dbi, key: k}) {
+				return
+			}
 
-	return keys
+			// now the p-tag+kind+date
+			if dbi == b.indexTag32 && tag[0] == "p" {
+				k := indexKeyPool.Get().([]byte)
+				hex.Decode(k[0:8], []byte(tag[1][0:8*2]))
+				binary.BigEndian.PutUint16(k[8:8+2], uint16(evt.Kind))
+				binary.BigEndian.PutUint32(k[8+2:8+2+4], uint32(evt.CreatedAt))
+				dbi := b.indexPTagKind
+				if !yield(key{dbi: dbi, key: k[0 : 8+2+4]}) {
+					return
+				}
+			}
+		}
+
+		{
+			// ~ by date only
+			k := indexKeyPool.Get().([]byte)
+			binary.BigEndian.PutUint32(k[0:4], uint32(evt.CreatedAt))
+			if !yield(key{dbi: b.indexCreatedAt, key: k[0:4]}) {
+				return
+			}
+		}
+	}
 }
 
 func (b *LMDBBackend) getTagIndexPrefix(tagValue string) (lmdb.DBI, []byte, int) {

diff --git a/lmdb/save.go b/lmdb/save.go
@@ -38,7 +38,7 @@ func (b *LMDBBackend) SaveEvent(ctx context.Context, evt *nostr.Event) error {
 			return err
 		}
 
-		for _, k := range b.getIndexKeysForEvent(evt) {
+		for k := range b.getIndexKeysForEvent(evt) {
 			err := txn.Put(k.dbi, k.key, idx, 0)
 			k.free()
 			if err != nil {