github.com/treeverse/lakefs@v1.24.1-0.20240520134607-95648127bfb0/pkg/graveler/committed/batch.go (about)

     1  package committed
     2  
     3  import (
     4  	"errors"
     5  	"sync"
     6  )
     7  
     8  type ResultCloser interface {
     9  	Close() (*WriteResult, error)
    10  }
    11  
    12  type BatchCloser struct {
    13  	// mu protects results and error
    14  	mu      sync.Mutex
    15  	results []WriteResult
    16  	err     error
    17  
    18  	wg sync.WaitGroup
    19  	ch chan ResultCloser
    20  }
    21  
    22  // NewBatchCloser returns a new BatchCloser
    23  func NewBatchCloser(numClosers int) *BatchCloser {
    24  	ret := &BatchCloser{
    25  		// Block when all closer goroutines are busy.
    26  		ch: make(chan ResultCloser),
    27  	}
    28  
    29  	ret.wg.Add(numClosers)
    30  	for i := 0; i < numClosers; i++ {
    31  		go ret.handleClose()
    32  	}
    33  
    34  	return ret
    35  }
    36  
    37  var ErrMultipleWaitCalls = errors.New("wait has already been called")
    38  
    39  // CloseWriterAsync adds RangeWriter instance for the BatchCloser to handle.
    40  // Any writes executed to the writer after this call are not guaranteed to succeed.
    41  // If Wait() has already been called, returns an error.
    42  func (bc *BatchCloser) CloseWriterAsync(w ResultCloser) error {
    43  	bc.mu.Lock()
    44  
    45  	if bc.err != nil {
    46  		// Don't accept new writers if previous error occurred.
    47  		// In particular, if Wait has started then this is errMultipleWaitCalls.
    48  		bc.mu.Unlock()
    49  		return bc.err
    50  	}
    51  
    52  	bc.mu.Unlock()
    53  	bc.ch <- w
    54  
    55  	return nil
    56  }
    57  
    58  func (bc *BatchCloser) handleClose() {
    59  	for w := range bc.ch {
    60  		bc.closeWriter(w)
    61  	}
    62  	bc.wg.Done()
    63  }
    64  
    65  func (bc *BatchCloser) closeWriter(w ResultCloser) {
    66  	res, err := w.Close()
    67  
    68  	// long operation is over, we can lock to have synchronized access to err and results
    69  	bc.mu.Lock()
    70  	defer bc.mu.Unlock()
    71  
    72  	if err != nil {
    73  		if bc.nilErrOrMultipleCalls() {
    74  			// keeping first error is enough
    75  			bc.err = err
    76  		}
    77  		return
    78  	}
    79  
    80  	bc.results = append(bc.results, *res)
    81  }
    82  
    83  // Wait returns when all Writers finished.  Returns a nil results slice and an error if *any*
    84  // RangeWriter failed to close and upload.
    85  func (bc *BatchCloser) Wait() ([]WriteResult, error) {
    86  	bc.mu.Lock()
    87  	if bc.err != nil {
    88  		defer bc.mu.Unlock()
    89  		return nil, bc.err
    90  	}
    91  	bc.err = ErrMultipleWaitCalls
    92  	bc.mu.Unlock()
    93  
    94  	close(bc.ch)
    95  
    96  	bc.wg.Wait()
    97  
    98  	// all writers finished
    99  	bc.mu.Lock()
   100  	defer bc.mu.Unlock()
   101  	if !bc.nilErrOrMultipleCalls() {
   102  		return nil, bc.err
   103  	}
   104  	return bc.results, nil
   105  }
   106  
   107  func (bc *BatchCloser) nilErrOrMultipleCalls() bool {
   108  	return bc.err == nil || errors.Is(bc.err, ErrMultipleWaitCalls)
   109  }