github.com/Jeffail/benthos/v3@v3.65.0/public/service/input.go (about)

     1  package service
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"time"
     7  
     8  	"github.com/Jeffail/benthos/v3/internal/shutdown"
     9  	"github.com/Jeffail/benthos/v3/lib/input/reader"
    10  	"github.com/Jeffail/benthos/v3/lib/message"
    11  	"github.com/Jeffail/benthos/v3/lib/response"
    12  	"github.com/Jeffail/benthos/v3/lib/types"
    13  )
    14  
    15  // AckFunc is a common function returned by inputs that must be called once for
    16  // each message consumed. This function ensures that the source of the message
    17  // receives either an acknowledgement (err is nil) or an error that can either
    18  // be propagated upstream as a nack, or trigger a reattempt at delivering the
    19  // same message.
    20  //
    21  // If your input implementation doesn't have a specific mechanism for dealing
    22  // with a nack then you can wrap your input implementation with AutoRetryNacks
    23  // to get automatic retries.
    24  type AckFunc func(ctx context.Context, err error) error
    25  
    26  // Input is an interface implemented by Benthos inputs. Calls to Read should
    27  // block until either a message has been received, the connection is lost, or
    28  // the provided context is cancelled.
    29  type Input interface {
    30  	// Establish a connection to the upstream service. Connect will always be
    31  	// called first when a reader is instantiated, and will be continuously
    32  	// called with back off until a nil error is returned.
    33  	//
    34  	// The provided context remains open only for the duration of the connecting
    35  	// phase, and should not be used to establish the lifetime of the connection
    36  	// itself.
    37  	//
    38  	// Once Connect returns a nil error the Read method will be called until
    39  	// either ErrNotConnected is returned, or the reader is closed.
    40  	Connect(context.Context) error
    41  
    42  	// Read a single message from a source, along with a function to be called
    43  	// once the message can be either acked (successfully sent or intentionally
    44  	// filtered) or nacked (failed to be processed or dispatched to the output).
    45  	//
    46  	// The AckFunc will be called for every message at least once, but there are
    47  	// no guarantees as to when this will occur. If your input implementation
    48  	// doesn't have a specific mechanism for dealing with a nack then you can
    49  	// wrap your input implementation with AutoRetryNacks to get automatic
    50  	// retries.
    51  	//
    52  	// If this method returns ErrNotConnected then Read will not be called again
    53  	// until Connect has returned a nil error. If ErrEndOfInput is returned then
    54  	// Read will no longer be called and the pipeline will gracefully terminate.
    55  	Read(context.Context) (*Message, AckFunc, error)
    56  
    57  	Closer
    58  }
    59  
    60  //------------------------------------------------------------------------------
    61  
    62  // BatchInput is an interface implemented by Benthos inputs that produce
    63  // messages in batches, where there is a desire to process and send the batch as
    64  // a logical group rather than as individual messages.
    65  //
    66  // Calls to ReadBatch should block until either a message batch is ready to
    67  // process, the connection is lost, or the provided context is cancelled.
    68  type BatchInput interface {
    69  	// Establish a connection to the upstream service. Connect will always be
    70  	// called first when a reader is instantiated, and will be continuously
    71  	// called with back off until a nil error is returned.
    72  	//
    73  	// The provided context remains open only for the duration of the connecting
    74  	// phase, and should not be used to establish the lifetime of the connection
    75  	// itself.
    76  	//
    77  	// Once Connect returns a nil error the Read method will be called until
    78  	// either ErrNotConnected is returned, or the reader is closed.
    79  	Connect(context.Context) error
    80  
    81  	// Read a message batch from a source, along with a function to be called
    82  	// once the entire batch can be either acked (successfully sent or
    83  	// intentionally filtered) or nacked (failed to be processed or dispatched
    84  	// to the output).
    85  	//
    86  	// The AckFunc will be called for every message batch at least once, but
    87  	// there are no guarantees as to when this will occur. If your input
    88  	// implementation doesn't have a specific mechanism for dealing with a nack
    89  	// then you can wrap your input implementation with AutoRetryNacksBatched to
    90  	// get automatic retries.
    91  	//
    92  	// If this method returns ErrNotConnected then ReadBatch will not be called
    93  	// again until Connect has returned a nil error. If ErrEndOfInput is
    94  	// returned then Read will no longer be called and the pipeline will
    95  	// gracefully terminate.
    96  	ReadBatch(context.Context) (MessageBatch, AckFunc, error)
    97  
    98  	Closer
    99  }
   100  
   101  //------------------------------------------------------------------------------
   102  
   103  // Implements input.AsyncReader
   104  type airGapReader struct {
   105  	r Input
   106  
   107  	sig *shutdown.Signaller
   108  }
   109  
   110  func newAirGapReader(r Input) reader.Async {
   111  	return &airGapReader{r, shutdown.NewSignaller()}
   112  }
   113  
   114  func (a *airGapReader) ConnectWithContext(ctx context.Context) error {
   115  	err := a.r.Connect(ctx)
   116  	if err != nil && errors.Is(err, ErrEndOfInput) {
   117  		err = types.ErrTypeClosed
   118  	}
   119  	return err
   120  }
   121  
   122  func (a *airGapReader) ReadWithContext(ctx context.Context) (types.Message, reader.AsyncAckFn, error) {
   123  	msg, ackFn, err := a.r.Read(ctx)
   124  	if err != nil {
   125  		if errors.Is(err, ErrNotConnected) {
   126  			err = types.ErrNotConnected
   127  		} else if errors.Is(err, ErrEndOfInput) {
   128  			err = types.ErrTypeClosed
   129  		}
   130  		return nil, nil, err
   131  	}
   132  	tMsg := message.New(nil)
   133  	tMsg.Append(msg.part)
   134  	return tMsg, func(c context.Context, r types.Response) error {
   135  		return ackFn(c, r.Error())
   136  	}, nil
   137  }
   138  
   139  func (a *airGapReader) CloseAsync() {
   140  	go func() {
   141  		// TODO: Determine whether to continue trying or log/exit.
   142  		_ = a.r.Close(context.Background())
   143  		a.sig.ShutdownComplete()
   144  	}()
   145  }
   146  
   147  func (a *airGapReader) WaitForClose(tout time.Duration) error {
   148  	select {
   149  	case <-a.sig.HasClosedChan():
   150  	case <-time.After(tout):
   151  		return types.ErrTimeout
   152  	}
   153  	return nil
   154  }
   155  
   156  //------------------------------------------------------------------------------
   157  
   158  // Implements input.AsyncReader
   159  type airGapBatchReader struct {
   160  	r BatchInput
   161  
   162  	sig *shutdown.Signaller
   163  }
   164  
   165  func newAirGapBatchReader(r BatchInput) reader.Async {
   166  	return &airGapBatchReader{r, shutdown.NewSignaller()}
   167  }
   168  
   169  func (a *airGapBatchReader) ConnectWithContext(ctx context.Context) error {
   170  	err := a.r.Connect(ctx)
   171  	if err != nil && errors.Is(err, ErrEndOfInput) {
   172  		err = types.ErrTypeClosed
   173  	}
   174  	return err
   175  }
   176  
   177  func (a *airGapBatchReader) ReadWithContext(ctx context.Context) (types.Message, reader.AsyncAckFn, error) {
   178  	batch, ackFn, err := a.r.ReadBatch(ctx)
   179  	if err != nil {
   180  		if errors.Is(err, ErrNotConnected) {
   181  			err = types.ErrNotConnected
   182  		} else if errors.Is(err, ErrEndOfInput) {
   183  			err = types.ErrTypeClosed
   184  		}
   185  		return nil, nil, err
   186  	}
   187  	tMsg := message.New(nil)
   188  	for _, msg := range batch {
   189  		tMsg.Append(msg.part)
   190  	}
   191  	return tMsg, func(c context.Context, r types.Response) error {
   192  		return ackFn(c, r.Error())
   193  	}, nil
   194  }
   195  
   196  func (a *airGapBatchReader) CloseAsync() {
   197  	go func() {
   198  		if err := a.r.Close(context.Background()); err == nil {
   199  			a.sig.ShutdownComplete()
   200  		}
   201  	}()
   202  }
   203  
   204  func (a *airGapBatchReader) WaitForClose(tout time.Duration) error {
   205  	select {
   206  	case <-a.sig.HasClosedChan():
   207  	case <-time.After(tout):
   208  		return types.ErrTimeout
   209  	}
   210  	return nil
   211  }
   212  
   213  //------------------------------------------------------------------------------
   214  
   215  // OwnedInput provides direct ownership of an input extracted from a plugin
   216  // config. Connectivity of the input is handled internally, and so the consumer
   217  // of this type should only be concerned with reading messages and eventually
   218  // calling Close to terminate the input.
   219  type OwnedInput struct {
   220  	i types.Input
   221  }
   222  
   223  // ReadBatch attemps to read a message batch from the input, along with a
   224  // function to be called once the entire batch can be either acked (successfully
   225  // sent or intentionally filtered) or nacked (failed to be processed or
   226  // dispatched to the output).
   227  //
   228  // If this method returns ErrEndOfInput then that indicates that the input has
   229  // finished and will no longer yield new messages.
   230  func (o *OwnedInput) ReadBatch(ctx context.Context) (MessageBatch, AckFunc, error) {
   231  	var tran types.Transaction
   232  	var open bool
   233  	select {
   234  	case tran, open = <-o.i.TransactionChan():
   235  	case <-ctx.Done():
   236  		return nil, nil, ctx.Err()
   237  	}
   238  	if !open {
   239  		return nil, nil, ErrEndOfInput
   240  	}
   241  
   242  	var b MessageBatch
   243  	_ = tran.Payload.Iter(func(i int, part types.Part) error {
   244  		b = append(b, newMessageFromPart(part))
   245  		return nil
   246  	})
   247  
   248  	return b, func(actx context.Context, err error) error {
   249  		var res types.Response
   250  		if err != nil {
   251  			res = response.NewError(err)
   252  		} else {
   253  			res = response.NewAck()
   254  		}
   255  		select {
   256  		case tran.ResponseChan <- res:
   257  		case <-actx.Done():
   258  			return actx.Err()
   259  		}
   260  		return nil
   261  	}, nil
   262  }
   263  
   264  // Close the input.
   265  func (o *OwnedInput) Close(ctx context.Context) error {
   266  	o.i.CloseAsync()
   267  	for {
   268  		// Gross but will do for now until we replace these with context params.
   269  		if err := o.i.WaitForClose(time.Millisecond * 100); err == nil {
   270  			return nil
   271  		}
   272  		select {
   273  		case <-ctx.Done():
   274  			return ctx.Err()
   275  		default:
   276  		}
   277  	}
   278  
   279  }