fix: address review comments

Author: sbackend123

pkg/hive/export_test.go

@@ -15,6 +15,7 @@ var (
 	MaxBatchSize      = maxBatchSize
 	LimitBurst        = limitBurst
 	CoalesceThreshold = coalesceThreshold
+	MessageTimeout    = messageTimeout
 )
 
 func (s *Service) SetTimeFunc(f func() time.Time) {

pkg/hive/hive.go

@@ -87,6 +87,8 @@ type Service struct {
 	inLimiter         *ratelimit.Limiter
 	outLimiter        *ratelimit.Limiter
 	quit              chan struct{}
+	bgCtx             context.Context
+	bgCancel          context.CancelFunc
 	wg                sync.WaitGroup
 	peersChan         chan pb.Peers
 	sem               *semaphore.Weighted
@@ -103,6 +105,7 @@ type Service struct {
 }
 
 func New(streamer p2p.Streamer, addressbook addressbook.GetPutSeener, networkID uint64, overlay swarm.Address, logger log.Logger, o Options) *Service {
+	bgCtx, bgCancel := context.WithCancel(context.Background())
 	svc := &Service{
 		streamer:           streamer,
 		logger:             logger.WithName(loggerName).Register(),
@@ -112,6 +115,8 @@ func New(streamer p2p.Streamer, addressbook addressbook.GetPutSeener, networkID
 		inLimiter:          ratelimit.New(limitRate, limitBurst),
 		outLimiter:         ratelimit.New(limitRate, limitBurst),
 		quit:               make(chan struct{}),
+		bgCtx:              bgCtx,
+		bgCancel:           bgCancel,
 		peersChan:          make(chan pb.Peers),
 		sem:                semaphore.NewWeighted(int64(swarm.MaxBins)),
 		bootnode:           o.BootnodeMode,
@@ -200,6 +205,7 @@ func (s *Service) broadcastNow(ctx context.Context, addressee swarm.Address, coa
 			if coalesced {
 				s.metrics.GossipCoalesceDropped.Add(float64(len(peers)))
 			}
+			s.logger.Debug("gossip dropped by outbound rate limiter", "addressee", addressee, "dropped", len(peers), "coalesced", coalesced)
 			return nil
 		}
 
@@ -229,6 +235,7 @@ func (s *Service) SetAddPeersHandler(h func(addr ...swarm.Address)) {
 
 func (s *Service) Close() error {
 	close(s.quit)
+	s.bgCancel()
 
 	stopped := make(chan struct{})
 	go func() {
@@ -365,9 +372,9 @@ func (s *Service) startGossipCoalescer() {
 			select {
 			case <-ticker.C:
 				for _, batch := range s.gossipBuf.takeAll() {
-					go func(batch gossipBatch) {
+					s.wg.Go(func() {
 						s.flushGossipBatch(batch.addressee, batch.peers, coalesceFlushReasonTimer)
-					}(batch)
+					})
 				}
 
 			case <-s.quit:
@@ -380,7 +387,7 @@ func (s *Service) startGossipCoalescer() {
 func (s *Service) flushGossipBatch(addressee swarm.Address, peers []swarm.Address, reason string) {
 	s.recordCoalesceFlush(reason, addressee, peers)
 
-	ctx, cancel := context.WithTimeout(context.Background(), messageTimeout)
+	ctx, cancel := context.WithTimeout(s.bgCtx, messageTimeout)
 	defer cancel()
 
 	err := s.broadcastNow(ctx, addressee, true, peers...)
@@ -406,21 +413,15 @@ func (s *Service) setCoalesceBufferGauge() {
 }
 
 func (s *Service) startCheckPeersHandler() {
-	ctx, cancel := context.WithCancel(context.Background())
-	s.wg.Go(func() {
-		<-s.quit
-		cancel()
-	})
-
 	s.wg.Go(func() {
 		for {
 			select {
-			case <-ctx.Done():
+			case <-s.bgCtx.Done():
 				return
 			case newPeers := <-s.peersChan:
 				s.wg.Go(func() {
 					safe.Run(s.logger, "hive-check-and-add-peers", func() {
-						s.checkAndAddPeers(ctx, newPeers)
+						s.checkAndAddPeers(s.bgCtx, newPeers)
 					})
 				})
 			}

pkg/hive/hive_test.go

@@ -1763,3 +1763,87 @@ func TestHiveGossipBuffering(t *testing.T) {
 		})
 	})
 }
+
+// blockUntilCancelStreamer blocks NewStream until ctx is cancelled, then
+// reports that the blocked call has returned.
+type blockUntilCancelStreamer struct {
+	entered  chan struct{}
+	finished chan struct{}
+}
+
+func (s *blockUntilCancelStreamer) NewStream(ctx context.Context, _ swarm.Address, _ p2p.Headers, _, _, _ string) (p2p.Stream, error) {
+	select {
+	case <-s.entered:
+	default:
+		close(s.entered)
+	}
+	<-ctx.Done()
+	select {
+	case <-s.finished:
+	default:
+		close(s.finished)
+	}
+	return nil, ctx.Err()
+}
+
+// TestCoalescedFlushCancelsOnClose verifies that an in-flight async flush does
+// not keep running until messageTimeout after Service.Close. The flush must
+// observe bgCtx cancellation and exit promptly.
+func TestCoalescedFlushCancelsOnClose(t *testing.T) {
+	t.Parallel()
+
+	synctest.Test(t, func(t *testing.T) {
+		const peerCount = 3
+
+		addressbook := ab.New(mock.NewStateStore())
+		networkID := uint64(1)
+		overlays := addTestOverlays(t, addressbook, networkID, peerCount, 5000)
+
+		entered := make(chan struct{})
+		finished := make(chan struct{})
+		streamer := &blockUntilCancelStreamer{entered: entered, finished: finished}
+
+		clientAddress := swarm.RandAddress(t)
+		addressee := swarm.RandAddress(t)
+		client := hive.New(streamer, addressbook, networkID, clientAddress, log.Noop, hive.Options{
+			AllowPrivateCIDRs:      true,
+			GossipCoalesceInterval: hiveGossipBufferingInterval,
+		})
+
+		ctx := context.Background()
+		for _, overlay := range overlays {
+			if err := client.BroadcastPeers(ctx, addressee, overlay); err != nil {
+				t.Fatal(err)
+			}
+		}
+
+		time.Sleep(hiveGossipBufferingInterval + 200*time.Millisecond)
+		synctest.Wait()
+
+		select {
+		case <-entered:
+		default:
+			t.Fatal("expected coalesced flush to block in NewStream before Close")
+		}
+
+		closeAt := time.Now()
+		if err := client.Close(); err != nil {
+			t.Fatalf("Close during coalesced flush: %v", err)
+		}
+
+		// synctest.Wait advances the fake clock only as far as needed for all
+		// goroutines to become durably blocked. If the flush were still waiting
+		// on messageTimeout alone, Wait would advance by a full minute.
+		synctest.Wait()
+
+		select {
+		case <-finished:
+		default:
+			t.Fatal("flush goroutine still blocked in NewStream after Close")
+		}
+
+		if elapsed := time.Since(closeAt); elapsed >= hive.MessageTimeout {
+			t.Fatalf("flush lasted %v after Close; want exit via bgCtx cancel before messageTimeout (%v)", elapsed, hive.MessageTimeout)
+		}
+	})
+}