storj.io/uplink@v1.13.0/private/storage/streams/pieceupload/manager.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  	"io"
     9  	"sort"
    10  	"sync"
    11  
    12  	"github.com/zeebo/errs"
    13  
    14  	"storj.io/common/pb"
    15  	"storj.io/common/storj"
    16  )
    17  
    18  var (
    19  	// ErrDone is returned from the Manager NextPiece when all of the piece
    20  	// uploads have completed.
    21  	ErrDone = errs.New("all pieces have been uploaded")
    22  )
    23  
    24  // PieceReader provides a reader for a piece with the given number.
    25  type PieceReader interface {
    26  	PieceReader(num int) io.Reader
    27  }
    28  
    29  // LimitsExchanger exchanges piece upload limits.
    30  type LimitsExchanger interface {
    31  	ExchangeLimits(ctx context.Context, segmentID storj.SegmentID, pieceNumbers []int) (storj.SegmentID, []*pb.AddressedOrderLimit, error)
    32  }
    33  
    34  // Manager tracks piece uploads for a segment. It provides callers with piece
    35  // data and limits and tracks which uploads have been successful (or not). It
    36  // also manages obtaining new piece upload limits for failed uploads to
    37  // add resiliency to segment uploads. The manager also keeps track of the
    38  // upload results, which the caller can use to commit the segment, including
    39  // the segment ID, which is updated as limits are exchanged.
    40  type Manager struct {
    41  	exchanger   LimitsExchanger
    42  	pieceReader PieceReader
    43  
    44  	mu         sync.Mutex
    45  	retries    int
    46  	segmentID  storj.SegmentID
    47  	limits     []*pb.AddressedOrderLimit
    48  	next       chan int
    49  	exchange   chan struct{}
    50  	done       chan struct{}
    51  	xchgFailed chan struct{}
    52  	xchgError  error
    53  	failed     []int
    54  	results    []*pb.SegmentPieceUploadResult
    55  }
    56  
    57  // NewManager returns a new piece upload manager.
    58  func NewManager(exchanger LimitsExchanger, pieceReader PieceReader, segmentID storj.SegmentID, limits []*pb.AddressedOrderLimit) *Manager {
    59  	next := make(chan int, len(limits))
    60  	for num := 0; num < len(limits); num++ {
    61  		next <- num
    62  	}
    63  	return &Manager{
    64  		exchanger:   exchanger,
    65  		pieceReader: pieceReader,
    66  		segmentID:   segmentID,
    67  		limits:      limits,
    68  		next:        next,
    69  		exchange:    make(chan struct{}, 1),
    70  		done:        make(chan struct{}),
    71  		xchgFailed:  make(chan struct{}),
    72  	}
    73  }
    74  
    75  // NextPiece returns a reader and limit for the next piece to upload. It also
    76  // returns a callback that the caller uses to indicate success (along with the
    77  // results) or not. NextPiece may return data with a new limit for a piece that
    78  // was previously attempted but failed. It will return ErrDone when enough
    79  // pieces have finished successfully to satisfy the optimal threshold. If
    80  // NextPiece is unable to exchange limits for failed pieces, it will return
    81  // an error.
    82  func (mgr *Manager) NextPiece(ctx context.Context) (_ io.Reader, _ *pb.AddressedOrderLimit, _ func(hash *pb.PieceHash, uploaded bool), err error) {
    83  	var num int
    84  	for acquired := false; !acquired; {
    85  		// If NextPiece is called with a cancelled context, we want to ensure
    86  		// that we return before hitting the select and possibly picking up
    87  		// another piece to upload.
    88  		if err := ctx.Err(); err != nil {
    89  			return nil, nil, nil, err
    90  		}
    91  
    92  		select {
    93  		case num = <-mgr.next:
    94  			acquired = true
    95  		case <-mgr.exchange:
    96  			if err := mgr.exchangeLimits(ctx); err != nil {
    97  				return nil, nil, nil, err
    98  			}
    99  		case <-ctx.Done():
   100  			return nil, nil, nil, ctx.Err()
   101  		case <-mgr.done:
   102  			return nil, nil, nil, ErrDone
   103  		case <-mgr.xchgFailed:
   104  			return nil, nil, nil, mgr.xchgError
   105  		}
   106  	}
   107  
   108  	mgr.mu.Lock()
   109  	defer mgr.mu.Unlock()
   110  
   111  	limit := mgr.limits[num]
   112  	piece := mgr.pieceReader.PieceReader(num)
   113  
   114  	invoked := false
   115  	done := func(hash *pb.PieceHash, uploaded bool) {
   116  		mgr.mu.Lock()
   117  		defer mgr.mu.Unlock()
   118  
   119  		// Protect against the callback being invoked twice.
   120  		if invoked {
   121  			return
   122  		}
   123  		invoked = true
   124  
   125  		if uploaded {
   126  			mgr.results = append(mgr.results, &pb.SegmentPieceUploadResult{
   127  				PieceNum: int32(num),
   128  				NodeId:   limit.Limit.StorageNodeId,
   129  				Hash:     hash,
   130  			})
   131  		} else {
   132  			mgr.failed = append(mgr.failed, num)
   133  		}
   134  
   135  		if len(mgr.results)+len(mgr.failed) < len(mgr.limits) {
   136  			return
   137  		}
   138  
   139  		// All of the uploads are accounted for. If there are failed pieces
   140  		// then signal that an exchange should take place so that the
   141  		// uploads can hopefully continue.
   142  		if len(mgr.failed) > 0 {
   143  			select {
   144  			case mgr.exchange <- struct{}{}:
   145  			default:
   146  			}
   147  			return
   148  		}
   149  
   150  		// Otherwise, all piece uploads have finished and we can signal the
   151  		// other callers that we are done.
   152  		close(mgr.done)
   153  	}
   154  
   155  	return piece, limit, done, nil
   156  }
   157  
   158  // Results returns the results of each piece successfully updated as well as
   159  // the segment ID, which may differ from that passed into NewManager if piece
   160  // limits needed to be exchanged for failed piece uploads.
   161  func (mgr *Manager) Results() (storj.SegmentID, []*pb.SegmentPieceUploadResult) {
   162  	mgr.mu.Lock()
   163  	segmentID := mgr.segmentID
   164  	results := append([]*pb.SegmentPieceUploadResult(nil), mgr.results...)
   165  	mgr.mu.Unlock()
   166  
   167  	// The satellite expects the results to be sorted by piece number...
   168  	sort.Slice(results, func(i, j int) bool {
   169  		return results[i].PieceNum < results[j].PieceNum
   170  	})
   171  
   172  	return segmentID, results
   173  }
   174  
   175  func (mgr *Manager) exchangeLimits(ctx context.Context) (err error) {
   176  	mgr.mu.Lock()
   177  	defer mgr.mu.Unlock()
   178  
   179  	// any error in exchangeLimits is permanently fatal because the
   180  	// api call should have retries in it already.
   181  	defer func() {
   182  		if err != nil && mgr.xchgError == nil {
   183  			mgr.xchgError = err
   184  			close(mgr.xchgFailed)
   185  		}
   186  	}()
   187  
   188  	if len(mgr.failed) == 0 {
   189  		// purely defensive: shouldn't happen.
   190  		return errs.New("failed piece list is empty")
   191  	}
   192  
   193  	if mgr.retries > 10 {
   194  		return errs.New("too many retries: are any nodes reachable?")
   195  	}
   196  	mgr.retries++
   197  
   198  	segmentID, limits, err := mgr.exchanger.ExchangeLimits(ctx, mgr.segmentID, mgr.failed)
   199  	if err != nil {
   200  		return errs.New("piece limit exchange failed: %w", err)
   201  	}
   202  	mgr.segmentID = segmentID
   203  	mgr.limits = limits
   204  	for _, num := range mgr.failed {
   205  		mgr.next <- num
   206  	}
   207  	mgr.failed = mgr.failed[:0]
   208  	return nil
   209  }