github.com/Jeffail/benthos/v3@v3.65.0/internal/bundle/tracing/output.go (about)

     1  package tracing
     2  
     3  import (
     4  	"sync/atomic"
     5  	"time"
     6  
     7  	"github.com/Jeffail/benthos/v3/internal/shutdown"
     8  	"github.com/Jeffail/benthos/v3/lib/types"
     9  )
    10  
    11  type tracedOutput struct {
    12  	e       *events
    13  	ctr     *uint64
    14  	wrapped types.Output
    15  	tChan   chan types.Transaction
    16  	shutSig *shutdown.Signaller
    17  }
    18  
    19  func traceOutput(e *events, ctr *uint64, i types.Output) types.Output {
    20  	t := &tracedOutput{
    21  		e:       e,
    22  		ctr:     ctr,
    23  		wrapped: i,
    24  		tChan:   make(chan types.Transaction),
    25  		shutSig: shutdown.NewSignaller(),
    26  	}
    27  	return t
    28  }
    29  
    30  func (t *tracedOutput) loop(inChan <-chan types.Transaction) {
    31  	defer close(t.tChan)
    32  	for {
    33  		tran, open := <-inChan
    34  		if !open {
    35  			return
    36  		}
    37  		_ = tran.Payload.Iter(func(i int, part types.Part) error {
    38  			_ = atomic.AddUint64(t.ctr, 1)
    39  			t.e.Add(EventConsume, string(part.Get()))
    40  			return nil
    41  		})
    42  		select {
    43  		case t.tChan <- tran:
    44  		case <-t.shutSig.CloseNowChan():
    45  			// Stop flushing if we fully timed out
    46  			return
    47  		}
    48  	}
    49  }
    50  
    51  func (t *tracedOutput) Consume(inChan <-chan types.Transaction) error {
    52  	go t.loop(inChan)
    53  	return t.wrapped.Consume(t.tChan)
    54  }
    55  
    56  func (t *tracedOutput) Connected() bool {
    57  	return t.wrapped.Connected()
    58  }
    59  
    60  func (t *tracedOutput) CloseAsync() {
    61  	t.wrapped.CloseAsync()
    62  }
    63  
    64  func (t *tracedOutput) WaitForClose(timeout time.Duration) error {
    65  	err := t.wrapped.WaitForClose(timeout)
    66  	t.shutSig.CloseNow()
    67  	return err
    68  }