storj.io/uplink@v1.13.0/private/storage/streams/pieceupload/manager_test.go (about)

     1  // Copyright (C) 2023 Storj Labs, Inc.
     2  // See LICENSE for copying information.
     3  
     4  package pieceupload
     5  
     6  import (
     7  	"context"
     8  	"errors"
     9  	"testing"
    10  
    11  	"github.com/stretchr/testify/assert"
    12  	"github.com/stretchr/testify/require"
    13  
    14  	"storj.io/common/pb"
    15  )
    16  
    17  func TestManager(t *testing.T) {
    18  	t.Run("results returned in piece order", func(t *testing.T) {
    19  		manager := newManager(2)
    20  
    21  		done0 := requireNextPiece(t, manager, piecenum{0}, revision{0})
    22  		requireNextPieceAndFinish(t, manager, piecenum{1}, revision{0}, true)
    23  		done0(true)
    24  
    25  		requireDone(t, manager)
    26  
    27  		assertResults(t, manager, revision{0},
    28  			makeResult(piecenum{0}, revision{0}),
    29  			makeResult(piecenum{1}, revision{0}),
    30  		)
    31  	})
    32  
    33  	t.Run("piece retried on failure", func(t *testing.T) {
    34  		manager := newManager(3)
    35  
    36  		// 0(0) fails
    37  		// 1(0) succeeds
    38  		// 2(0) fails
    39  		requireNextPieceAndFinish(t, manager, piecenum{0}, revision{0}, false)
    40  		requireNextPieceAndFinish(t, manager, piecenum{1}, revision{0}, true)
    41  		requireNextPieceAndFinish(t, manager, piecenum{2}, revision{0}, false)
    42  
    43  		// Now that the all uploads have been attempted, we expect the next
    44  		// call to exchange the upload limits and return a piece to retry.
    45  
    46  		// Retries happen in first-failed order, so:
    47  		// 2(1) is retried and succeeds
    48  		// 0(1) is retried and fails
    49  		requireNextPieceAndFinish(t, manager, piecenum{0}, revision{1}, true)
    50  		requireNextPieceAndFinish(t, manager, piecenum{2}, revision{1}, false)
    51  
    52  		// 0(2) is retried again (after another limit exchange) and succeeds
    53  		requireNextPieceAndFinish(t, manager, piecenum{2}, revision{2}, true)
    54  
    55  		requireDone(t, manager)
    56  
    57  		assertResults(t, manager, revision{2},
    58  			makeResult(piecenum{0}, revision{1}),
    59  			makeResult(piecenum{1}, revision{0}),
    60  			makeResult(piecenum{2}, revision{2}),
    61  		)
    62  	})
    63  
    64  	t.Run("piece retry fails if exchange fails", func(t *testing.T) {
    65  		manager := newManagerWithExchanger(1, failExchange{})
    66  		requireNextPieceAndFinish(t, manager, piecenum{0}, revision{0}, false)
    67  
    68  		_, _, _, err := manager.NextPiece(context.Background())
    69  		require.EqualError(t, err, "piece limit exchange failed: oh no")
    70  	})
    71  }
    72  
    73  func makeResult(num piecenum, rev revision) *pb.SegmentPieceUploadResult {
    74  	return &pb.SegmentPieceUploadResult{PieceNum: int32(num.value), Hash: hash(num), NodeId: nodeID(num, rev)}
    75  }
    76  
    77  func assertResults(t *testing.T, manager *Manager, expectedRev revision, expectedResults ...*pb.SegmentPieceUploadResult) {
    78  	t.Helper()
    79  	actualID, actualResults := manager.Results()
    80  	assert.Equal(t, makeSegmentID(expectedRev), actualID)
    81  	assert.Equal(t, expectedResults, actualResults)
    82  }
    83  
    84  func requireNextPiece(t *testing.T, manager *Manager, expectedNum piecenum, expectedRev revision) func(bool) {
    85  	reader, limit, done, err := manager.NextPiece(context.Background())
    86  	require.NoError(t, err, "expected next piece")
    87  	require.Equal(t, expectedNum, pieceReaderNum(reader), "unexpected next piece number")
    88  	require.Equal(t, makeLimit(expectedNum, expectedRev), limit, "unexpected next piece number limit")
    89  	return func(uploaded bool) {
    90  		if uploaded {
    91  			done(hash(expectedNum), uploaded)
    92  		} else {
    93  			done(nil, false)
    94  		}
    95  	}
    96  }
    97  
    98  func requireDone(t *testing.T, manager *Manager) {
    99  	reader, _, _, err := manager.NextPiece(context.Background())
   100  	if err == nil {
   101  		require.FailNowf(t, "expected done", "next piece %d", pieceReaderNum(reader))
   102  		return
   103  	}
   104  	require.True(t, errors.Is(err, ErrDone), "expected done but got %q", err)
   105  }
   106  
   107  func requireNextPieceAndFinish(t *testing.T, manager *Manager, expectedNum piecenum, expectedRev revision, uploaded bool) {
   108  	requireNextPiece(t, manager, expectedNum, expectedRev)(uploaded)
   109  }