github.com/Jeffail/benthos/v3@v3.65.0/internal/bundle/tracing/input.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 tracedInput struct {
    12  	e       *events
    13  	ctr     *uint64
    14  	wrapped types.Input
    15  	tChan   chan types.Transaction
    16  	shutSig *shutdown.Signaller
    17  }
    18  
    19  func traceInput(e *events, counter *uint64, i types.Input) types.Input {
    20  	t := &tracedInput{
    21  		e:       e,
    22  		ctr:     counter,
    23  		wrapped: i,
    24  		tChan:   make(chan types.Transaction),
    25  		shutSig: shutdown.NewSignaller(),
    26  	}
    27  	go t.loop()
    28  	return t
    29  }
    30  
    31  func (t *tracedInput) loop() {
    32  	defer close(t.tChan)
    33  	readChan := t.wrapped.TransactionChan()
    34  	for {
    35  		tran, open := <-readChan
    36  		if !open {
    37  			return
    38  		}
    39  		_ = tran.Payload.Iter(func(i int, part types.Part) error {
    40  			_ = atomic.AddUint64(t.ctr, 1)
    41  			t.e.Add(EventProduce, string(part.Get()))
    42  			return nil
    43  		})
    44  		select {
    45  		case t.tChan <- tran:
    46  		case <-t.shutSig.CloseNowChan():
    47  			// Stop flushing if we fully timed out
    48  			return
    49  		}
    50  	}
    51  }
    52  
    53  func (t *tracedInput) TransactionChan() <-chan types.Transaction {
    54  	return t.tChan
    55  }
    56  
    57  func (t *tracedInput) Connected() bool {
    58  	return t.wrapped.Connected()
    59  }
    60  
    61  func (t *tracedInput) CloseAsync() {
    62  	t.wrapped.CloseAsync()
    63  }
    64  
    65  func (t *tracedInput) WaitForClose(timeout time.Duration) error {
    66  	err := t.wrapped.WaitForClose(timeout)
    67  	t.shutSig.CloseNow()
    68  	return err
    69  }