use our own dataloader and simplify batch replaceable requests, removing bugs.

sdk/addressable_loader.go

@@ -2,14 +2,11 @@ package sdk
 
 import (
 	"context"
-	"errors"
-	"fmt"
-	"strconv"
 	"sync"
 	"time"
 
-	"github.com/graph-gophers/dataloader/v7"
 	"github.com/nbd-wtf/go-nostr"
+	"github.com/nbd-wtf/go-nostr/sdk/dataloader"
 )
 
 // this is similar to replaceable_loader and reuses logic from that.
@@ -33,28 +30,21 @@ func (sys *System) initializeAddressableDataloaders() {
 
 func (sys *System) createAddressableDataloader(kind int) *dataloader.Loader[string, []*nostr.Event] {
 	return dataloader.NewBatchedLoader(
-		func(_ context.Context, pubkeys []string) []*dataloader.Result[[]*nostr.Event] {
-			return sys.batchLoadAddressableEvents(kind, pubkeys)
+		func(ctxs []context.Context, pubkeys []string) map[string]dataloader.Result[[]*nostr.Event] {
+			return sys.batchLoadAddressableEvents(ctxs, kind, pubkeys)
 		},
 		dataloader.WithBatchCapacity[string, []*nostr.Event](30),
-		dataloader.WithClearCacheOnBatch[string, []*nostr.Event](),
-		dataloader.WithCache(&dataloader.NoCache[string, []*nostr.Event]{}),
 		dataloader.WithWait[string, []*nostr.Event](time.Millisecond*350),
 	)
 }
 
 func (sys *System) batchLoadAddressableEvents(
+	ctxs []context.Context,
 	kind int,
 	pubkeys []string,
-) []*dataloader.Result[[]*nostr.Event] {
-	ctx, cancel := context.WithTimeoutCause(context.Background(), time.Second*6,
-		errors.New("batch addressable load took too long"),
-	)
-	defer cancel()
-
+) map[string]dataloader.Result[[]*nostr.Event] {
 	batchSize := len(pubkeys)
-	results := make([]*dataloader.Result[[]*nostr.Event], batchSize)
-	keyPositions := make(map[string]int) // { [pubkey]: slice_index }
+	results := make(map[string]dataloader.Result[[]*nostr.Event], batchSize)
 	relayFilter := make([]nostr.DirectedFilter, 0, max(3, batchSize*2))
 	relayFilterIndex := make(map[string]int, max(3, batchSize*2))
 
@@ -62,36 +52,16 @@ func (sys *System) batchLoadAddressableEvents(
 	wg.Add(len(pubkeys))
 	cm := sync.Mutex{}
 
+	aggregatedContext, aggregatedCancel := context.WithCancel(context.Background())
+	waiting := len(pubkeys)
+
 	for i, pubkey := range pubkeys {
+		ctx := ctxs[i]
+
 		// build batched queries for the external relays
-		keyPositions[pubkey] = i // this is to help us know where to save the result later
-
 		go func(i int, pubkey string) {
-			defer wg.Done()
-
-			// if we're attempting this query with a short key (last 8 characters), stop here
-			if len(pubkey) != 64 {
-				results[i] = &dataloader.Result[[]*nostr.Event]{
-					Error: fmt.Errorf("won't proceed to query relays with a shortened key (%d)", kind),
-				}
-				return
-			}
-
-			// save attempts here so we don't try the same failed query over and over
-			if doItNow := doThisNotMoreThanOnceAnHour("repl:" + strconv.Itoa(kind) + pubkey); !doItNow {
-				results[i] = &dataloader.Result[[]*nostr.Event]{
-					Error: fmt.Errorf("last attempt failed, waiting more to try again"),
-				}
-				return
-			}
-
 			// gather relays we'll use for this pubkey
-			relays := sys.determineRelaysToQuery(pubkey, kind)
-
-			// by default we will return an error (this will be overwritten when we find an event)
-			results[i] = &dataloader.Result[[]*nostr.Event]{
-				Error: fmt.Errorf("couldn't find a kind %d event anywhere %v", kind, relays),
-			}
+			relays := sys.determineRelaysToQuery(ctx, pubkey, kind)
 
 			cm.Lock()
 			for _, relay := range relays {
@@ -116,6 +86,13 @@ func (sys *System) batchLoadAddressableEvents(
 				relayFilter[idx] = dfilter
 			}
 			cm.Unlock()
+			wg.Done()
+
+			<-ctx.Done()
+			waiting--
+			if waiting == 0 {
+				aggregatedCancel()
+			}
 		}(i, pubkey)
 	}
 
@@ -123,7 +100,7 @@ func (sys *System) batchLoadAddressableEvents(
 	wg.Wait()
 
 	// query all relays with the prepared filters
-	multiSubs := sys.Pool.BatchedSubManyEose(ctx, relayFilter)
+	multiSubs := sys.Pool.BatchedSubManyEose(aggregatedContext, relayFilter)
 nextEvent:
 	for {
 		select {
@@ -132,13 +109,10 @@ nextEvent:
 				return results
 			}
 
-			// insert this event at the desired position
-			pos := keyPositions[ie.PubKey] // @unchecked: it must succeed because it must be a key we passed
-
-			events := results[pos].Data
+			events := results[ie.PubKey].Data
 			if events == nil {
 				// no events found, so just add this and end
-				results[pos] = &dataloader.Result[[]*nostr.Event]{Data: []*nostr.Event{ie.Event}}
+				results[ie.PubKey] = dataloader.Result[[]*nostr.Event]{Data: []*nostr.Event{ie.Event}}
 				continue nextEvent
 			}
 
@@ -158,10 +132,9 @@ nextEvent:
 				}
 			}
 
-			// there is no match, so add to the end
 			events = append(events, ie.Event)
-			results[pos].Data = events
-		case <-ctx.Done():
+			results[ie.PubKey] = dataloader.Result[[]*nostr.Event]{Data: events}
+		case <-aggregatedContext.Done():
 			return results
 		}
 	}