provider: batch queue

Author: guillaumemichel

dual/provider/provider.go

@@ -3,23 +3,27 @@ package provider
 import (
 	"context"
 	"errors"
-	"fmt"
 
 	"github.com/ipfs/go-cid"
+	logging "github.com/ipfs/go-log/v2"
 	dht "github.com/libp2p/go-libp2p-kad-dht"
 	"github.com/libp2p/go-libp2p-kad-dht/dual"
+	"github.com/libp2p/go-libp2p-kad-dht/internal/batch"
 	"github.com/libp2p/go-libp2p-kad-dht/provider"
 	"github.com/libp2p/go-libp2p-kad-dht/provider/datastore"
 	mh "github.com/multiformats/go-multihash"
 )
 
+var logger = logging.Logger(provider.LoggerName)
+
 // SweepingProvider manages provides and reprovides for both DHT swarms (LAN
 // and WAN) in the dual DHT setup.
 type SweepingProvider struct {
-	dht      *dual.DHT
-	LAN      *provider.SweepingProvider
-	WAN      *provider.SweepingProvider
-	keyStore datastore.KeyStore
+	dht           *dual.DHT
+	lan           *provider.SweepingProvider
+	wan           *provider.SweepingProvider
+	keyStore      datastore.KeyStore
+	keyStoreQueue *batch.Processor
 }
 
 // New creates a new SweepingProvider that manages provides and reprovides for
@@ -70,32 +74,44 @@ func New(d *dual.DHT, opts ...Option) (*SweepingProvider, error) {
 		}
 	}
 
-	return &SweepingProvider{
+	s := &SweepingProvider{
 		dht:      d,
-		LAN:      sweepingProviders[0],
-		WAN:      sweepingProviders[1],
+		lan:      sweepingProviders[0],
+		wan:      sweepingProviders[1],
 		keyStore: cfg.keyStore,
-	}, nil
+	}
+	s.keyStoreQueue = batch.NewProcessor(func(keys []mh.Multihash, force bool) {
+		newKeys, err := s.keyStore.Put(context.Background(), keys...)
+		if err != nil {
+			logger.Errorf("failed to store multihashes: %v", err)
+			return
+		}
+		if !force {
+			if len(newKeys) == 0 {
+				return
+			}
+			keys = newKeys
+		}
+
+		// Add to schedule and to provide queue of both systems
+		s.runOnBoth(func(p *provider.SweepingProvider) {
+			p.AddToScheduleAndProvide(newKeys...)
+		})
+	})
+
+	return s, nil
 }
 
 // runOnBoth runs the provided function on both the LAN and WAN providers in
 // parallel and waits for both to complete.
-func (s *SweepingProvider) runOnBoth(f func(*provider.SweepingProvider) error) error {
-	var errs [2]error
+func (s *SweepingProvider) runOnBoth(f func(*provider.SweepingProvider)) {
 	done := make(chan struct{})
 	go func() {
 		defer close(done)
-		err := f(s.LAN)
-		if err != nil {
-			errs[0] = fmt.Errorf("LAN provider: %w", err)
-		}
+		f(s.lan)
 	}()
-	err := f(s.WAN)
-	if err != nil {
-		errs[1] = fmt.Errorf("WAN provider: %w", err)
-	}
+	f(s.wan)
 	<-done
-	return errors.Join(errs[:]...)
 }
 
 // ProvideOnce sends provider records for the specified keys to both DHT swarms
@@ -109,9 +125,9 @@ func (s *SweepingProvider) runOnBoth(f func(*provider.SweepingProvider) error) e
 // (either never bootstrapped, or disconnected since more than `OfflineDelay`).
 // The schedule and provide queue depend on the network size, hence recent
 // network connectivity is essential.
-func (s *SweepingProvider) ProvideOnce(keys ...mh.Multihash) error {
-	return s.runOnBoth(func(p *provider.SweepingProvider) error {
-		return p.ProvideOnce(keys...)
+func (s *SweepingProvider) ProvideOnce(keys ...mh.Multihash) {
+	s.runOnBoth(func(p *provider.SweepingProvider) {
+		p.ProvideOnce(keys...)
 	})
 }
 
@@ -132,22 +148,8 @@ func (s *SweepingProvider) ProvideOnce(keys ...mh.Multihash) error {
 // (either never bootstrapped, or disconnected since more than `OfflineDelay`).
 // The schedule and provide queue depend on the network size, hence recent
 // network connectivity is essential.
-func (s *SweepingProvider) StartProviding(force bool, keys ...mh.Multihash) error {
-	ctx := context.Background()
-	newKeys, err := s.keyStore.Put(ctx, keys...)
-	if err != nil {
-		return fmt.Errorf("failed to store multihashes: %w", err)
-	}
-
-	s.runOnBoth(func(p *provider.SweepingProvider) error {
-		return p.AddToSchedule(newKeys...)
-	})
-
-	if !force {
-		keys = newKeys
-	}
-
-	return s.ProvideOnce(keys...)
+func (s *SweepingProvider) StartProviding(force bool, keys ...mh.Multihash) {
+	s.keyStoreQueue.Enqueue(force, keys...)
 }
 
 // StopProviding stops reproviding the given keys to both DHT swarms. The node
@@ -157,12 +159,12 @@ func (s *SweepingProvider) StartProviding(force bool, keys ...mh.Multihash) erro
 // Remove the `keys` from the schedule and return immediately. Valid records
 // can remain in the DHT swarms up to the provider record TTL after calling
 // `StopProviding`.
-func (s *SweepingProvider) StopProviding(keys ...mh.Multihash) error {
+func (s *SweepingProvider) StopProviding(keys ...mh.Multihash) {
+	// TODO: batch deletes
 	err := s.keyStore.Delete(context.Background(), keys...)
 	if err != nil {
-		return fmt.Errorf("failed to stop providing keys: %w", err)
+		logger.Errorf("failed to stop providing keys: %w", err)
 	}
-	return nil
 }
 
 // Clear clears the all the keys from the provide queues of both DHTs and
@@ -171,7 +173,7 @@ func (s *SweepingProvider) StopProviding(keys ...mh.Multihash) error {
 // The keys are not deleted from the keystore, so they will continue to be
 // reprovided as scheduled.
 func (s *SweepingProvider) Clear() int {
-	return s.LAN.Clear() + s.WAN.Clear()
+	return s.lan.Clear() + s.wan.Clear()
 }
 
 // RefreshSchedule scans the KeyStore for any keys that are not currently
@@ -186,10 +188,17 @@ func (s *SweepingProvider) Clear() int {
 // Offline (either never bootstrapped, or disconnected since more than
 // `OfflineDelay`). The schedule depends on the network size, hence recent
 // network connectivity is essential.
-func (s *SweepingProvider) RefreshSchedule() error {
-	return s.runOnBoth(func(p *provider.SweepingProvider) error {
-		return p.RefreshSchedule()
+func (s *SweepingProvider) RefreshSchedule() {
+	s.runOnBoth(func(p *provider.SweepingProvider) {
+		p.RefreshSchedule()
+	})
+}
+
+func (s *SweepingProvider) Close() error {
+	s.runOnBoth(func(p *provider.SweepingProvider) {
+		p.Close()
 	})
+	return nil
 }
 
 var (
@@ -200,9 +209,10 @@ var (
 // dhtProvider is the interface to ensure that SweepingProvider and
 // provider.SweepingProvider share the same interface.
 type dhtProvider interface {
-	StartProviding(force bool, keys ...mh.Multihash) error
-	StopProviding(keys ...mh.Multihash) error
-	ProvideOnce(keys ...mh.Multihash) error
+	StartProviding(force bool, keys ...mh.Multihash)
+	StopProviding(keys ...mh.Multihash)
+	ProvideOnce(keys ...mh.Multihash)
 	Clear() int
-	RefreshSchedule() error
+	RefreshSchedule()
+	Close() error
 }

internal/batch/processor.go

@@ -0,0 +1,132 @@
+package batch
+
+import (
+	"fmt"
+	"sync"
+
+	mh "github.com/multiformats/go-multihash"
+)
+
+type BatchHandler func(keys []mh.Multihash, special bool)
+
+const (
+	normalQueueID uint8 = iota
+	specialQueueID
+)
+
+type enqueueOp struct {
+	keys []mh.Multihash
+	id   uint8
+}
+
+type drainOp struct {
+	id uint8
+	ch chan []mh.Multihash
+}
+
+type Processor struct {
+	close     chan struct{}
+	closeOnce sync.Once
+	wg        sync.WaitGroup
+	enqueue   chan enqueueOp
+	drain     chan drainOp
+
+	signals [2]chan struct{}
+	bufs    [2][]mh.Multihash
+
+	f BatchHandler
+}
+
+func NewProcessor(f BatchHandler) *Processor {
+	q := &Processor{
+		close:   make(chan struct{}),
+		wg:      sync.WaitGroup{},
+		enqueue: make(chan enqueueOp),
+		drain:   make(chan drainOp),
+
+		signals: [2]chan struct{}{make(chan struct{}, 1), make(chan struct{}, 1)},
+		bufs:    [2][]mh.Multihash{make([]mh.Multihash, 0), make([]mh.Multihash, 0)},
+		f:       f,
+	}
+	q.wg.Add(2)
+	go q.in()
+	go q.out()
+	return q
+}
+
+func (q *Processor) in() {
+	defer q.wg.Done()
+
+	for {
+		select {
+		case <-q.close:
+			return
+		case op := <-q.enqueue:
+			q.bufs[op.id] = append(q.bufs[op.id], op.keys...)
+			// Signal there are keys in the queue.
+			select {
+			case q.signals[op.id] <- struct{}{}:
+			default:
+			}
+		case op := <-q.drain:
+			op.ch <- q.bufs[op.id]
+			q.bufs[op.id] = make([]mh.Multihash, 0)
+		}
+	}
+}
+
+func (q *Processor) out() {
+	defer q.wg.Done()
+
+	ch := make(chan []mh.Multihash)
+	handle := func(id uint8) {
+		select {
+		case <-q.close:
+			return
+		case q.drain <- drainOp{id: id, ch: ch}:
+		}
+
+		var keys []mh.Multihash
+		select {
+		case <-q.close:
+			return
+		case keys = <-ch:
+		}
+
+		// Run callback
+		q.f(keys, id == specialQueueID)
+	}
+
+	for {
+		select {
+		case <-q.close:
+			return
+		case <-q.signals[normalQueueID]:
+			handle(normalQueueID)
+		case <-q.signals[specialQueueID]:
+			handle(specialQueueID)
+		}
+	}
+}
+
+func (q *Processor) Close() error {
+	q.closeOnce.Do(func() {
+		close(q.enqueue)
+		close(q.close)
+	})
+	return nil
+}
+
+func (q *Processor) Enqueue(special bool, keys ...mh.Multihash) (err error) {
+	id := normalQueueID
+	if special {
+		id = specialQueueID
+	}
+	defer func() {
+		if r := recover(); r != nil {
+			err = fmt.Errorf("failed to enqueue item: %s", r)
+		}
+	}()
+	q.enqueue <- enqueueOp{keys: keys, id: id}
+	return
+}