fix: simplify implementation

Author: sbackend123

pkg/hive/gossip_buffer.go

@@ -5,7 +5,8 @@
 package hive
 
 import (
-	"math/rand/v2"
+	"maps"
+	"slices"
 	"sync"
 	"time"
 
@@ -14,7 +15,6 @@ import (
 
 const (
 	defaultGossipCoalesceInterval = time.Second
-	defaultGossipCoalesceJitter   = 100 * time.Millisecond
 	// coalesceThreshold: gossips with fewer peers are buffered; larger
 	// (already-batched) messages are dispatched immediately.
 	coalesceThreshold = 2
@@ -24,98 +24,78 @@ const (
 // flushed as one batched message.
 type gossipBuffer struct {
 	mu       sync.Mutex
-	pending  map[string]*pendingGossip // addressee bytestring -> buffered peers
+	pending  map[string]map[string]swarm.Address // addressee key -> peer key -> peer
 	interval time.Duration
-	jitter   time.Duration
 	maxBatch int
 }
 
-type pendingGossip struct {
+type gossipBatch struct {
 	addressee swarm.Address
-	peers     map[string]swarm.Address // peer bytestring -> address (set semantics)
-	deadline  time.Time
+	peers     []swarm.Address
 }
 
 func newGossipBuffer(interval time.Duration, maxBatch int) *gossipBuffer {
 	if interval == 0 {
 		interval = defaultGossipCoalesceInterval
 	}
 	return &gossipBuffer{
-		pending:  make(map[string]*pendingGossip),
+		pending:  make(map[string]map[string]swarm.Address),
 		interval: interval,
-		jitter:   defaultGossipCoalesceJitter,
 		maxBatch: maxBatch,
 	}
 }
 
-// add buffers peers for the addressee. If the buffer reaches maxBatch it is
+// stagePeers buffers peers for the addressee. If the buffer reaches maxBatch it is
 // removed and returned so the caller can flush it immediately.
-func (b *gossipBuffer) add(now time.Time, addressee swarm.Address, peers ...swarm.Address) *pendingGossip {
+func (b *gossipBuffer) stagePeers(addressee swarm.Address, peers ...swarm.Address) (flushPeers []swarm.Address, flush bool) {
 	b.mu.Lock()
 	defer b.mu.Unlock()
 
 	key := addressee.ByteString()
-	e, ok := b.pending[key]
+	peerSet, ok := b.pending[key]
 	if !ok {
-		var jitter time.Duration
-		if b.jitter > 0 {
-			jitter = time.Duration(rand.Int64N(int64(b.jitter)))
-		}
-		e = &pendingGossip{
-			addressee: addressee,
-			peers:     make(map[string]swarm.Address),
-			deadline:  now.Add(b.interval + jitter),
-		}
-		b.pending[key] = e
+		peerSet = make(map[string]swarm.Address)
+		b.pending[key] = peerSet
 	}
 	for _, p := range peers {
-		e.peers[p.ByteString()] = p
+		peerSet[p.ByteString()] = p
 	}
 
-	if len(e.peers) >= b.maxBatch {
+	if len(peerSet) >= b.maxBatch {
 		delete(b.pending, key)
-		return e
+		return slices.Collect(maps.Values(peerSet)), true
 	}
-	return nil
+	return nil, false
 }
 
-// takeDue removes and returns all entries whose deadline has passed.
-func (b *gossipBuffer) takeDue(now time.Time) []*pendingGossip {
-	return b.take(func(e *pendingGossip) bool { return !e.deadline.After(now) })
+// takeAll removes and returns all buffered entries.
+func (b *gossipBuffer) takeAll() []gossipBatch {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+
+	if len(b.pending) == 0 {
+		return nil
+	}
+
+	out := make([]gossipBatch, 0, len(b.pending))
+	for key, peerSet := range b.pending {
+		out = append(out, gossipBatch{
+			addressee: swarm.NewAddress([]byte(key)),
+			peers:     slices.Collect(maps.Values(peerSet)),
+		})
+	}
+	b.pending = make(map[string]map[string]swarm.Address)
+	return out
 }
 
 func (b *gossipBuffer) clearAddressee(addressee swarm.Address) {
 	b.mu.Lock()
 	defer b.mu.Unlock()
-
 	delete(b.pending, addressee.ByteString())
 }
 
 func (b *gossipBuffer) pendingAddressees() int {
 	b.mu.Lock()
 	defer b.mu.Unlock()
-
 	return len(b.pending)
 }
-
-func (b *gossipBuffer) take(match func(*pendingGossip) bool) []*pendingGossip {
-	b.mu.Lock()
-	defer b.mu.Unlock()
-
-	var out []*pendingGossip
-	for key, e := range b.pending {
-		if match(e) {
-			out = append(out, e)
-			delete(b.pending, key)
-		}
-	}
-	return out
-}
-
-func (e *pendingGossip) addresses() []swarm.Address {
-	out := make([]swarm.Address, 0, len(e.peers))
-	for _, p := range e.peers {
-		out = append(out, p)
-	}
-	return out
-}

pkg/hive/gossip_buffer_test.go

@@ -11,55 +11,57 @@ import (
 	"github.com/ethersphere/bee/v2/pkg/swarm"
 )
 
-func TestGossipBufferAddAndDue(t *testing.T) {
+func TestGossipBufferAddAndTakeAll(t *testing.T) {
 	t.Parallel()
 
-	const interval = 100 * time.Millisecond
-
-	b := newGossipBuffer(interval, maxBatchSize)
+	b := newGossipBuffer(time.Second, maxBatchSize)
 	addressee := swarm.RandAddress(t)
 	peer1 := swarm.RandAddress(t)
 	peer2 := swarm.RandAddress(t)
 
-	now := time.Now()
-	if full := b.add(now, addressee, peer1); full != nil {
-		t.Fatal("unexpected immediate flush")
+	if pending := b.takeAll(); len(pending) != 0 {
+		t.Fatalf("want no pending entries, got %d", len(pending))
 	}
 
-	if due := b.takeDue(now); len(due) != 0 {
-		t.Fatalf("want no due entries, got %d", len(due))
+	if _, flush := b.stagePeers(addressee, peer1); flush {
+		t.Fatal("unexpected immediate flush")
 	}
 
-	if full := b.add(now, addressee, peer2); full != nil {
+	if _, flush := b.stagePeers(addressee, peer2); flush {
 		t.Fatal("unexpected immediate flush")
 	}
 
-	afterDeadline := now.Add(interval + defaultGossipCoalesceJitter + time.Millisecond)
-	due := b.takeDue(afterDeadline)
-	if len(due) != 1 {
-		t.Fatalf("want 1 due entry, got %d", len(due))
+	pending := b.takeAll()
+	if len(pending) != 1 {
+		t.Fatalf("want 1 pending entry, got %d", len(pending))
 	}
-	if got := len(due[0].addresses()); got != 2 {
+	if got := len(pending[0].peers); got != 2 {
 		t.Fatalf("want 2 coalesced peers, got %d", got)
 	}
+	if !pending[0].addressee.Equal(addressee) {
+		t.Fatal("unexpected addressee in pending batch")
+	}
+
+	if pending := b.takeAll(); len(pending) != 0 {
+		t.Fatalf("want empty buffer after takeAll, got %d pending", len(pending))
+	}
 }
 
 func TestGossipBufferMaxBatchFlush(t *testing.T) {
 	t.Parallel()
 
 	b := newGossipBuffer(time.Second, 2)
 	addressee := swarm.RandAddress(t)
-	now := time.Now()
 
-	b.add(now, addressee, swarm.RandAddress(t))
-	full := b.add(now, addressee, swarm.RandAddress(t))
-	if full == nil {
+	b.stagePeers(addressee, swarm.RandAddress(t))
+	flushPeers, flush := b.stagePeers(addressee, swarm.RandAddress(t))
+	if !flush {
 		t.Fatal("want immediate flush at maxBatch")
 	}
-	if got := len(full.addresses()); got != 2 {
+	if got := len(flushPeers); got != 2 {
 		t.Fatalf("want 2 peers in full batch, got %d", got)
 	}
-	if due := b.takeDue(now.Add(time.Second)); len(due) != 0 {
-		t.Fatalf("want empty buffer after maxBatch flush, got %d due", len(due))
+	if pending := b.takeAll(); len(pending) != 0 {
+		t.Fatalf("want empty buffer after maxBatch flush, got %d pending", len(pending))
 	}
 }

pkg/hive/hive.go

@@ -163,8 +163,7 @@ func (s *Service) BroadcastPeers(ctx context.Context, addressee swarm.Address, p
 	if len(peers) >= coalesceThreshold {
 		s.metrics.GossipCoalesceImmediatePeers.Add(float64(len(peers)))
 		s.logger.Debug("gossip immediate send", "addressee", addressee, "peer_count", len(peers))
-		_, err := s.broadcastNow(ctx, addressee, peers...)
-		return err
+		return s.broadcastNow(ctx, addressee, false, peers...)
 	}
 
 	select {
@@ -177,52 +176,49 @@ func (s *Service) BroadcastPeers(ctx context.Context, addressee swarm.Address, p
 	s.logger.Debug("gossip buffered", "addressee", addressee, "peer_count", len(peers))
 
 	// Buffer; if it just filled up, flush it synchronously while still in the call
-	if full := s.gossipBuf.add(s.now(), addressee, peers...); full != nil {
-		flushPeers := full.addresses()
+	if flushPeers, flush := s.gossipBuf.stagePeers(addressee, peers...); flush {
 		s.recordCoalesceFlush(coalesceFlushReasonMaxBatch, addressee, flushPeers)
 		s.setCoalesceBufferGauge()
-		sent, err := s.broadcastNow(ctx, addressee, flushPeers...)
-		if dropped := len(flushPeers) - sent; dropped > 0 {
-			s.metrics.GossipCoalesceDropped.Add(float64(dropped))
-		}
-		return err
+		return s.broadcastNow(ctx, addressee, true, flushPeers...)
 	}
 	s.setCoalesceBufferGauge()
 	return nil
 }
 
 // broadcastNow performs the synchronous, rate-limited, batched send.
-// It returns the number of peers successfully sent.
-func (s *Service) broadcastNow(ctx context.Context, addressee swarm.Address, peers ...swarm.Address) (sent int, err error) {
+func (s *Service) broadcastNow(ctx context.Context, addressee swarm.Address, coalesced bool, peers ...swarm.Address) error {
 	maxSize := maxBatchSize
-	total := len(peers)
 
 	for len(peers) > 0 {
 		if maxSize > len(peers) {
 			maxSize = len(peers)
 		}
 
-		// If broadcasting limit is exceeded, return early
 		if !s.outLimiter.Allow(addressee.ByteString(), maxSize) {
-			return total - len(peers), nil
+			if coalesced {
+				s.metrics.GossipCoalesceDropped.Add(float64(len(peers)))
+			}
+			return nil
 		}
 
 		select {
 		case <-ctx.Done():
-			return total - len(peers), ctx.Err()
+			return ctx.Err()
 		case <-s.quit:
-			return total - len(peers), ErrShutdownInProgress
+			return ErrShutdownInProgress
 		default:
 		}
 
 		if err := s.sendPeers(ctx, addressee, peers[:maxSize]); err != nil {
-			return total - len(peers), err
+			if coalesced {
+				s.metrics.GossipCoalesceDropped.Add(float64(len(peers)))
+			}
+			return err
 		}
 
 		peers = peers[maxSize:]
 	}
-
-	return total, nil
+	return nil
 }
 
 func (s *Service) SetAddPeersHandler(h func(addr ...swarm.Address)) {
@@ -360,42 +356,32 @@ func (s *Service) disconnect(peer p2p.Peer) error {
 }
 
 func (s *Service) startGossipCoalescer() {
-	tick := s.gossipBuf.interval / 2
-	if tick <= 0 {
-		tick = s.gossipBuf.interval
-	}
-
 	s.wg.Go(func() {
-		ticker := time.NewTicker(tick)
+		ticker := time.NewTicker(s.gossipBuf.interval)
 		defer ticker.Stop()
 		for {
 			select {
 			case <-ticker.C:
-				s.flushGossipEntries(s.gossipBuf.takeDue(s.now()), coalesceFlushReasonTimer)
+				for _, batch := range s.gossipBuf.takeAll() {
+					s.flushGossipBatch(batch.addressee, batch.peers, coalesceFlushReasonTimer)
+				}
+				s.setCoalesceBufferGauge()
 			case <-s.quit:
 				return
 			}
 		}
 	})
 }
 
-func (s *Service) flushGossipEntries(entries []*pendingGossip, reason string) {
-	s.setCoalesceBufferGauge()
-
-	for _, e := range entries {
-		peers := e.addresses()
-		s.recordCoalesceFlush(reason, e.addressee, peers)
+func (s *Service) flushGossipBatch(addressee swarm.Address, peers []swarm.Address, reason string) {
+	s.recordCoalesceFlush(reason, addressee, peers)
 
-		ctx, cancel := context.WithTimeout(context.Background(), messageTimeout)
-		sent, err := s.broadcastNow(ctx, e.addressee, peers...)
-		if dropped := len(peers) - sent; dropped > 0 {
-			s.metrics.GossipCoalesceDropped.Add(float64(dropped))
-		}
-		if err != nil {
-			s.logger.Debug("coalesced gossip flush failed", "addressee", e.addressee, "reason", reason, "batch_size", len(peers), "error", err)
-		}
-		cancel()
+	ctx, cancel := context.WithTimeout(context.Background(), messageTimeout)
+	err := s.broadcastNow(ctx, addressee, true, peers...)
+	if err != nil {
+		s.logger.Debug("coalesced gossip flush failed", "addressee", addressee, "reason", reason, "batch_size", len(peers), "error", err)
 	}
+	cancel()
 }
 
 func (s *Service) recordCoalesceFlush(reason string, addressee swarm.Address, peers []swarm.Address) {