sdk: fix wot filter.

sdk/wot.go

@@ -2,13 +2,12 @@ package sdk
 
 import (
 	"context"
-	"maps"
-	"slices"
 	"strconv"
+	"sync"
+	"time"
 
 	"github.com/FastFilter/xorfilter"
 	"golang.org/x/sync/errgroup"
-	"sync"
 )
 
 func PubKeyToShid(pubkey string) uint64 {
@@ -16,59 +15,141 @@ func PubKeyToShid(pubkey string) uint64 {
 	return shid
 }
 
-func (sys *System) GetWoT(ctx context.Context, pubkey string) (map[uint64]struct{}, error) {
-	g, ctx := errgroup.WithContext(ctx)
-	g.SetLimit(30)
-
-	res := make(chan uint64, 100) // Add buffer to prevent blocking
-	result := make(map[uint64]struct{})
-	var resultMu sync.Mutex // Add mutex to protect map access
-
-	// Start consumer goroutine
-	done := make(chan struct{})
-	go func() {
-		defer close(done)
-		for shid := range res {
-			resultMu.Lock()
-			result[shid] = struct{}{}
-			resultMu.Unlock()
-		}
-	}()
-
-	// Process follow lists
-	for _, f := range sys.FetchFollowList(ctx, pubkey).Items {
-		f := f // Capture loop variable
-		g.Go(func() error {
-			for _, f2 := range sys.FetchFollowList(ctx, f.Pubkey).Items {
-				select {
-				case res <- PubKeyToShid(f2.Pubkey):
-				case <-ctx.Done():
-					return ctx.Err()
-				}
-			}
-			return nil
-		})
-	}
-
-	err := g.Wait()
-	close(res) // Close channel after all goroutines are done
-	<-done    // Wait for consumer to finish
-
-	return result, err
+type wotCall struct {
+	id          uint64 // basically the pubkey we're targeting here
+	mutex       sync.Mutex
+	resultbacks []chan WotXorFilter // all callers waiting for results
+	errorbacks  []chan error        // all callers waiting for errors
+	done        chan struct{}       // this is closed when this call is fully resolved and deleted
 }
 
-func (sys *System) GetWoTFilter(ctx context.Context, pubkey string) (WotXorFilter, error) {
-	m, err := sys.GetWoT(ctx, pubkey)
-	if err != nil {
-		return WotXorFilter{}, err
+const wotCallsSize = 8
+
+var (
+	wotCallsMutex   sync.Mutex
+	wotCallsInPlace [wotCallsSize]*wotCall
+)
+
+func (sys *System) LoadWoTFilter(ctx context.Context, pubkey string) (WotXorFilter, error) {
+	id := PubKeyToShid(pubkey)
+	pos := int(id % wotCallsSize)
+
+start:
+	wotCallsMutex.Lock()
+	wc := wotCallsInPlace[pos]
+	if wc == nil {
+		// we are the first to call at this position
+		wc = &wotCall{
+			id:          id,
+			resultbacks: make([]chan WotXorFilter, 0),
+			errorbacks:  make([]chan error, 0),
+			done:        make(chan struct{}),
+		}
+		wotCallsInPlace[pos] = wc
+		wotCallsMutex.Unlock()
+		goto actualcall
+	} else {
+		wotCallsMutex.Unlock()
 	}
 
-	xf, err := xorfilter.Populate(slices.Collect(maps.Keys(m)))
-	if err != nil {
-		return WotXorFilter{}, err
+	wc.mutex.Lock()
+	if wc.id == id {
+		// there is already a call for this exact pubkey ongoing, so we just wait
+		resch := make(chan WotXorFilter)
+		errch := make(chan error)
+		wc.resultbacks = append(wc.resultbacks, resch)
+		wc.errorbacks = append(wc.errorbacks, errch)
+		wc.mutex.Unlock()
+		select {
+		case res := <-resch:
+			return res, nil
+		case err := <-errch:
+			return WotXorFilter{}, err
+		}
+	} else {
+		wc.mutex.Unlock()
+		// there is already a call in this place, but it's for a different pubkey, so wait
+		<-wc.done
+		// when it's done restart
+		goto start
 	}
 
-	return WotXorFilter{*xf}, nil
+actualcall:
+	var res WotXorFilter
+	m, err := sys.loadWoT(ctx, pubkey)
+	if err != nil {
+		wc.mutex.Lock()
+		for _, ch := range wc.errorbacks {
+			ch <- err
+		}
+	} else {
+		res = makeWoTFilter(m)
+		wc.mutex.Lock()
+		for _, ch := range wc.resultbacks {
+			ch <- res
+		}
+	}
+
+	wotCallsMutex.Lock()
+	wotCallsInPlace[pos] = nil
+	wc.mutex.Unlock()
+	close(wc.done)
+	wotCallsMutex.Unlock()
+
+	return res, err
+}
+
+func (sys *System) loadWoT(ctx context.Context, pubkey string) (chan string, error) {
+	g, ctx := errgroup.WithContext(ctx)
+	g.SetLimit(45)
+
+	res := make(chan string)
+
+	// process follow lists
+	wg := sync.WaitGroup{}
+	wg.Add(1)
+
+	go func() {
+		for _, f := range sys.FetchFollowList(ctx, pubkey).Items {
+			wg.Add(1)
+
+			g.Go(func() error {
+				ctx, cancel := context.WithTimeout(ctx, time.Second*7)
+				defer cancel()
+
+				ff := sys.FetchFollowList(ctx, f.Pubkey).Items
+				for _, f2 := range ff {
+					res <- f2.Pubkey
+				}
+				wg.Done()
+				return nil
+			})
+		}
+
+		wg.Done()
+	}()
+
+	go func() {
+		wg.Wait()
+		close(res)
+	}()
+
+	return res, nil
+}
+
+func makeWoTFilter(m chan string) WotXorFilter {
+	shids := make([]uint64, 0, 60000)
+	shidMap := make(map[uint64]struct{}, 60000)
+	for pk := range m {
+		shid := PubKeyToShid(pk)
+		if _, alreadyAdded := shidMap[shid]; !alreadyAdded {
+			shidMap[shid] = struct{}{}
+			shids = append(shids, shid)
+		}
+	}
+
+	xf, _ := xorfilter.Populate(shids)
+	return WotXorFilter{*xf}
 }
 
 type WotXorFilter struct {