Plug more gaps in forced shutdown cleanup

internal/bundle/tracing/input.go

@@ -37,8 +37,14 @@ func (t *tracedInput) loop() {
 	defer close(t.tChan)
 	readChan := t.wrapped.TransactionChan()
 	for {
-		tran, open := <-readChan
-		if !open {
+		var tran message.Transaction
+		var open bool
+		select {
+		case tran, open = <-readChan:
+			if !open {
+				return
+			}
+		case <-t.shutSig.CloseNowChan():
 			return
 		}
 		if t.e.IsEnabled() {

internal/bundle/tracing/output.go

@@ -72,6 +72,7 @@ func (t *tracedOutput) Connected() bool {
 
 func (t *tracedOutput) TriggerCloseNow() {
 	t.wrapped.TriggerCloseNow()
+	t.shutSig.CloseNow()
 }
 
 func (t *tracedOutput) WaitForClose(ctx context.Context) error {

internal/impl/pure/input_broker_fan_in.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/benthosdev/benthos/v4/internal/component/input"
 	"github.com/benthosdev/benthos/v4/internal/message"
+	"github.com/benthosdev/benthos/v4/internal/shutdown"
 )
 
 type fanInInputBroker struct {
@@ -17,7 +18,7 @@ type fanInInputBroker struct {
 	remainingMap    map[int]struct{}
 	remainingMapMut sync.Mutex
 
-	closedChan chan struct{}
+	shutSig *shutdown.Signaller
 }
 
 func newFanInInputBroker(inputs []input.Streamed) (*fanInInputBroker, error) {
@@ -31,8 +32,8 @@ func newFanInInputBroker(inputs []input.Streamed) (*fanInInputBroker, error) {
 		inputClosedChan: make(chan int),
 		remainingMap:    make(map[int]struct{}),
 
-		closables:  []input.Streamed{},
-		closedChan: make(chan struct{}),
+		closables: []input.Streamed{},
+		shutSig:   shutdown.NewSignaller(),
 	}
 
 	for n, input := range inputs {
@@ -48,11 +49,21 @@ func newFanInInputBroker(inputs []input.Streamed) (*fanInInputBroker, error) {
 				i.inputClosedChan <- index
 			}()
 			for {
-				in, open := <-inputs[index].TransactionChan()
-				if !open {
+				var in message.Transaction
+				var open bool
+				select {
+				case in, open = <-inputs[index].TransactionChan():
+					if !open {
+						return
+					}
+				case <-i.shutSig.CloseNowChan():
+					return
+				}
+				select {
+				case i.transactions <- in:
+				case <-i.shutSig.CloseNowChan():
 					return
 				}
-				i.transactions <- in
 			}
 		}(n)
 	}
@@ -85,7 +96,7 @@ func (i *fanInInputBroker) loop() {
 	defer func() {
 		close(i.inputClosedChan)
 		close(i.transactions)
-		close(i.closedChan)
+		i.shutSig.ShutdownComplete()
 	}()
 
 	for {
@@ -112,11 +123,12 @@ func (i *fanInInputBroker) TriggerCloseNow() {
 	for _, closable := range i.closables {
 		closable.TriggerCloseNow()
 	}
+	i.shutSig.CloseNow()
 }
 
 func (i *fanInInputBroker) WaitForClose(ctx context.Context) error {
 	select {
-	case <-i.closedChan:
+	case <-i.shutSig.HasClosedChan():
 	case <-ctx.Done():
 		return ctx.Err()
 	}

internal/pipeline/pool.go

@@ -149,6 +149,7 @@ func (p *Pool) TriggerCloseNow() {
 	for _, w := range p.workers {
 		w.TriggerCloseNow()
 	}
+	p.shutSig.CloseNow()
 }
 
 // WaitForClose blocks until the component has closed down or the context is

internal/pipeline/processor.go

@@ -93,8 +93,7 @@ func (p *Processor) dispatchMessages(ctx context.Context, msgs []message.Batch,
 
 	pending := msgs
 	for len(pending) > 0 {
-		wg := sync.WaitGroup{}
-		wg.Add(len(pending))
+		doneChan := make(chan struct{}, len(pending))
 
 		var newPending []message.Batch
 		var newPendingMut sync.Mutex
@@ -107,18 +106,31 @@ func (p *Processor) dispatchMessages(ctx context.Context, msgs []message.Batch,
 					newPending = append(newPending, b)
 					newPendingMut.Unlock()
 				}
-				wg.Done()
+				select {
+				case doneChan <- struct{}{}:
+				default:
+				}
 				return nil
 			})
 
 			select {
 			case p.messagesOut <- transac:
 			case <-ctx.Done():
-				wg.Done()
+				select {
+				case doneChan <- struct{}{}:
+				default:
+				}
+				return
+			}
+		}
+
+		for i := 0; i < len(pending); i++ {
+			select {
+			case <-doneChan:
+			case <-ctx.Done():
 				return
 			}
 		}
-		wg.Wait()
 
 		if pending = newPending; len(pending) > 0 && !throt.Retry() {
 			return