github.com/Jeffail/benthos/v3@v3.65.0/lib/input/reader/async_cut_off.go (about)

     1  package reader
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  
     7  	"github.com/Jeffail/benthos/v3/lib/types"
     8  )
     9  
    10  //------------------------------------------------------------------------------
    11  
    12  type asyncCutOffMsg struct {
    13  	msg   types.Message
    14  	ackFn AsyncAckFn
    15  }
    16  
    17  // AsyncCutOff is a wrapper for reader.Async implementations that exits from
    18  // WaitForClose immediately. This is only useful when the underlying readable
    19  // resource cannot be closed reliably and can block forever.
    20  type AsyncCutOff struct {
    21  	msgChan chan asyncCutOffMsg
    22  	errChan chan error
    23  	ctx     context.Context
    24  	close   func()
    25  
    26  	r Async
    27  }
    28  
    29  // NewAsyncCutOff returns a new AsyncCutOff wrapper around a reader.Async.
    30  func NewAsyncCutOff(r Async) *AsyncCutOff {
    31  	ctx, done := context.WithCancel(context.Background())
    32  	return &AsyncCutOff{
    33  		msgChan: make(chan asyncCutOffMsg),
    34  		errChan: make(chan error),
    35  		ctx:     ctx,
    36  		close:   done,
    37  		r:       r,
    38  	}
    39  }
    40  
    41  //------------------------------------------------------------------------------
    42  
    43  // ConnectWithContext attempts to establish a connection to the source, if
    44  // unsuccessful returns an error. If the attempt is successful (or not
    45  // necessary) returns nil.
    46  func (c *AsyncCutOff) ConnectWithContext(ctx context.Context) error {
    47  	return c.r.ConnectWithContext(ctx)
    48  }
    49  
    50  // ReadWithContext attempts to read a new message from the source.
    51  func (c *AsyncCutOff) ReadWithContext(ctx context.Context) (types.Message, AsyncAckFn, error) {
    52  	go func() {
    53  		msg, ackFn, err := c.r.ReadWithContext(ctx)
    54  		if err == nil {
    55  			select {
    56  			case c.msgChan <- asyncCutOffMsg{
    57  				msg:   msg,
    58  				ackFn: ackFn,
    59  			}:
    60  			case <-ctx.Done():
    61  			}
    62  		} else {
    63  			select {
    64  			case c.errChan <- err:
    65  			case <-ctx.Done():
    66  			}
    67  		}
    68  	}()
    69  	select {
    70  	case m := <-c.msgChan:
    71  		return m.msg, m.ackFn, nil
    72  	case e := <-c.errChan:
    73  		return nil, nil, e
    74  	case <-ctx.Done():
    75  		c.r.CloseAsync()
    76  	case <-c.ctx.Done():
    77  	}
    78  	return nil, nil, types.ErrTypeClosed
    79  }
    80  
    81  // CloseAsync triggers the asynchronous closing of the reader.
    82  func (c *AsyncCutOff) CloseAsync() {
    83  	c.r.CloseAsync()
    84  	c.close()
    85  }
    86  
    87  // WaitForClose blocks until either the reader is finished closing or a timeout
    88  // occurs.
    89  func (c *AsyncCutOff) WaitForClose(tout time.Duration) error {
    90  	return nil // We don't block here.
    91  }
    92  
    93  //------------------------------------------------------------------------------