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 }