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

     1  package service
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"sync"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/Jeffail/benthos/v3/lib/types"
    11  	"github.com/stretchr/testify/assert"
    12  	"github.com/stretchr/testify/require"
    13  )
    14  
    15  type mockInput struct {
    16  	msgsToSnd []*Message
    17  	ackRcvd   []error
    18  
    19  	connChan  chan error
    20  	readChan  chan error
    21  	ackChan   chan error
    22  	closeChan chan error
    23  }
    24  
    25  func newMockInput() *mockInput {
    26  	return &mockInput{
    27  		connChan:  make(chan error),
    28  		readChan:  make(chan error),
    29  		ackChan:   make(chan error),
    30  		closeChan: make(chan error),
    31  	}
    32  }
    33  
    34  func (i *mockInput) Connect(ctx context.Context) error {
    35  	cerr, open := <-i.connChan
    36  	if !open {
    37  		return types.ErrNotConnected
    38  	}
    39  	return cerr
    40  }
    41  
    42  func (i *mockInput) Read(ctx context.Context) (*Message, AckFunc, error) {
    43  	select {
    44  	case <-ctx.Done():
    45  		return nil, nil, types.ErrTimeout
    46  	case err, open := <-i.readChan:
    47  		if !open {
    48  			return nil, nil, types.ErrNotConnected
    49  		}
    50  		if err != nil {
    51  			return nil, nil, err
    52  		}
    53  	}
    54  	i.ackRcvd = append(i.ackRcvd, errors.New("ack not received"))
    55  	index := len(i.ackRcvd) - 1
    56  
    57  	nextMsg := NewMessage(nil)
    58  	if len(i.msgsToSnd) > 0 {
    59  		nextMsg = i.msgsToSnd[0]
    60  		i.msgsToSnd = i.msgsToSnd[1:]
    61  	}
    62  
    63  	return nextMsg.Copy(), func(ctx context.Context, res error) error {
    64  		i.ackRcvd[index] = res
    65  		return <-i.ackChan
    66  	}, nil
    67  }
    68  
    69  func (i *mockInput) Close(ctx context.Context) error {
    70  	return <-i.closeChan
    71  }
    72  
    73  func TestAutoRetryClose(t *testing.T) {
    74  	t.Parallel()
    75  
    76  	ctx, cancel := context.WithTimeout(context.Background(), time.Second*2)
    77  	defer cancel()
    78  
    79  	readerImpl := newMockInput()
    80  	pres := AutoRetryNacks(readerImpl)
    81  
    82  	expErr := errors.New("foo error")
    83  
    84  	wg := sync.WaitGroup{}
    85  	wg.Add(1)
    86  
    87  	go func() {
    88  		defer wg.Done()
    89  
    90  		err := pres.Connect(ctx)
    91  		require.NoError(t, err)
    92  
    93  		assert.Equal(t, expErr, pres.Close(ctx))
    94  	}()
    95  
    96  	select {
    97  	case readerImpl.connChan <- nil:
    98  	case <-time.After(time.Second):
    99  		t.Error("Timed out")
   100  	}
   101  
   102  	select {
   103  	case readerImpl.closeChan <- expErr:
   104  	case <-time.After(time.Second):
   105  		t.Error("Timed out")
   106  	}
   107  
   108  	wg.Wait()
   109  }
   110  
   111  func TestAutoRetryHappy(t *testing.T) {
   112  	t.Parallel()
   113  
   114  	ctx, cancel := context.WithTimeout(context.Background(), time.Second*2)
   115  	defer cancel()
   116  
   117  	readerImpl := newMockInput()
   118  	readerImpl.msgsToSnd = append(readerImpl.msgsToSnd, NewMessage([]byte("foo")))
   119  
   120  	pres := AutoRetryNacks(readerImpl)
   121  
   122  	go func() {
   123  		select {
   124  		case readerImpl.connChan <- nil:
   125  		case <-time.After(time.Second):
   126  			t.Error("Timed out")
   127  		}
   128  		select {
   129  		case readerImpl.readChan <- nil:
   130  		case <-time.After(time.Second):
   131  			t.Error("Timed out")
   132  		}
   133  	}()
   134  
   135  	require.NoError(t, pres.Connect(ctx))
   136  
   137  	msg, _, err := pres.Read(ctx)
   138  	require.NoError(t, err)
   139  
   140  	act, err := msg.AsBytes()
   141  	require.NoError(t, err)
   142  
   143  	assert.Equal(t, "foo", string(act))
   144  }
   145  
   146  func TestAutoRetryErrorProp(t *testing.T) {
   147  	t.Parallel()
   148  
   149  	ctx, cancel := context.WithTimeout(context.Background(), time.Second*2)
   150  	defer cancel()
   151  
   152  	readerImpl := newMockInput()
   153  	pres := AutoRetryNacks(readerImpl)
   154  
   155  	expErr := errors.New("foo")
   156  
   157  	go func() {
   158  		select {
   159  		case readerImpl.connChan <- expErr:
   160  		case <-time.After(time.Second):
   161  			t.Error("Timed out")
   162  		}
   163  		select {
   164  		case readerImpl.readChan <- expErr:
   165  		case <-time.After(time.Second):
   166  			t.Error("Timed out")
   167  		}
   168  		select {
   169  		case readerImpl.readChan <- nil:
   170  		case <-time.After(time.Second):
   171  			t.Error("Timed out")
   172  		}
   173  		select {
   174  		case readerImpl.ackChan <- expErr:
   175  		case <-time.After(time.Second):
   176  			t.Error("Timed out")
   177  		}
   178  	}()
   179  
   180  	assert.Equal(t, expErr, pres.Connect(ctx))
   181  
   182  	_, _, err := pres.Read(ctx)
   183  	assert.Equal(t, expErr, err)
   184  
   185  	_, aFn, err := pres.Read(ctx)
   186  	require.NoError(t, err)
   187  
   188  	assert.Equal(t, expErr, aFn(ctx, nil))
   189  }
   190  
   191  func TestAutoRetryErrorBackoff(t *testing.T) {
   192  	t.Parallel()
   193  
   194  	readerImpl := newMockInput()
   195  	pres := AutoRetryNacks(readerImpl)
   196  
   197  	go func() {
   198  		select {
   199  		case readerImpl.connChan <- nil:
   200  		case <-time.After(time.Second):
   201  			t.Error("Timed out")
   202  		}
   203  		select {
   204  		case readerImpl.readChan <- nil:
   205  		case <-time.After(time.Second):
   206  			t.Error("Timed out")
   207  		}
   208  		select {
   209  		case readerImpl.closeChan <- nil:
   210  		case <-time.After(time.Second):
   211  			t.Error("Timed out")
   212  		}
   213  	}()
   214  
   215  	ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*500)
   216  	defer cancel()
   217  
   218  	require.NoError(t, pres.Connect(ctx))
   219  
   220  	i := 0
   221  	for {
   222  		_, aFn, actErr := pres.Read(ctx)
   223  		if actErr != nil {
   224  			assert.EqualError(t, actErr, "context deadline exceeded")
   225  			break
   226  		}
   227  		require.NoError(t, aFn(ctx, errors.New("no thanks")))
   228  		i++
   229  		if i == 10 {
   230  			t.Error("Expected backoff to prevent this")
   231  			break
   232  		}
   233  	}
   234  
   235  	require.NoError(t, pres.Close(context.Background()))
   236  }
   237  
   238  func TestAutoRetryBuffer(t *testing.T) {
   239  	t.Parallel()
   240  
   241  	ctx, cancel := context.WithTimeout(context.Background(), time.Second*2)
   242  	defer cancel()
   243  
   244  	readerImpl := newMockInput()
   245  	pres := AutoRetryNacks(readerImpl)
   246  
   247  	sendMsg := func(content string) {
   248  		readerImpl.msgsToSnd = []*Message{
   249  			NewMessage([]byte(content)),
   250  		}
   251  		select {
   252  		case readerImpl.readChan <- nil:
   253  		case <-time.After(time.Second):
   254  			t.Error("Timed out")
   255  		}
   256  	}
   257  	sendAck := func() {
   258  		select {
   259  		case readerImpl.ackChan <- nil:
   260  		case <-time.After(time.Second):
   261  			t.Error("Timed out")
   262  		}
   263  	}
   264  
   265  	// Send message normally.
   266  	exp := "msg 1"
   267  	exp2 := "msg 2"
   268  	exp3 := "msg 3"
   269  
   270  	go sendMsg(exp)
   271  	msg, aFn, err := pres.Read(ctx)
   272  	require.NoError(t, err)
   273  
   274  	b, err := msg.AsBytes()
   275  	require.NoError(t, err)
   276  	assert.Equal(t, exp, string(b))
   277  
   278  	// Prime second message.
   279  	go sendMsg(exp2)
   280  
   281  	// Fail previous message, expecting it to be resent.
   282  	_ = aFn(ctx, errors.New("failed"))
   283  	msg, aFn, err = pres.Read(ctx)
   284  	require.NoError(t, err)
   285  
   286  	b, err = msg.AsBytes()
   287  	require.NoError(t, err)
   288  	assert.Equal(t, exp, string(b))
   289  
   290  	// Read the primed message.
   291  	var aFn2 AckFunc
   292  	msg, aFn2, err = pres.Read(ctx)
   293  	require.NoError(t, err)
   294  
   295  	b, err = msg.AsBytes()
   296  	require.NoError(t, err)
   297  	assert.Equal(t, exp2, string(b))
   298  
   299  	// Fail both messages, expecting them to be resent.
   300  	_ = aFn(ctx, errors.New("failed again"))
   301  	_ = aFn2(ctx, errors.New("failed again"))
   302  
   303  	// Read both messages.
   304  	msg, aFn, err = pres.Read(ctx)
   305  	require.NoError(t, err)
   306  
   307  	b, err = msg.AsBytes()
   308  	require.NoError(t, err)
   309  	assert.Equal(t, exp, string(b))
   310  
   311  	msg, aFn2, err = pres.Read(ctx)
   312  	require.NoError(t, err)
   313  
   314  	b, err = msg.AsBytes()
   315  	require.NoError(t, err)
   316  	assert.Equal(t, exp2, string(b))
   317  
   318  	// Prime a new message and also an acknowledgement.
   319  	go sendMsg(exp3)
   320  	go sendAck()
   321  	go sendAck()
   322  
   323  	// Ack all messages.
   324  	_ = aFn(ctx, nil)
   325  	_ = aFn2(ctx, nil)
   326  
   327  	msg, _, err = pres.Read(ctx)
   328  	require.NoError(t, err)
   329  
   330  	b, err = msg.AsBytes()
   331  	require.NoError(t, err)
   332  	assert.Equal(t, exp3, string(b))
   333  }