oras.land/oras-go/v2@v2.5.1-0.20240520045656-aef90e4d04c4/internal/syncutil/merge.go (about)

     1  /*
     2  Copyright The ORAS Authors.
     3  Licensed under the Apache License, Version 2.0 (the "License");
     4  you may not use this file except in compliance with the License.
     5  You may obtain a copy of the License at
     6  
     7  http://www.apache.org/licenses/LICENSE-2.0
     8  
     9  Unless required by applicable law or agreed to in writing, software
    10  distributed under the License is distributed on an "AS IS" BASIS,
    11  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  See the License for the specific language governing permissions and
    13  limitations under the License.
    14  */
    15  
    16  package syncutil
    17  
    18  import "sync"
    19  
    20  // mergeStatus represents the merge status of an item.
    21  type mergeStatus struct {
    22  	// main indicates if items are being merged by the current go-routine.
    23  	main bool
    24  	// err represents the error of the merge operation.
    25  	err error
    26  }
    27  
    28  // Merge represents merge operations on items.
    29  // The state transfer is shown as below:
    30  //
    31  //	           +----------+
    32  //	           |  Start   +--------+-------------+
    33  //	           +----+-----+        |             |
    34  //	                |              |             |
    35  //	                v              v             v
    36  //	           +----+-----+   +----+----+   +----+----+
    37  //	   +-------+ Prepare  +<--+ Pending +-->+ Waiting |
    38  //	   |       +----+-----+   +---------+   +----+----+
    39  //	   |            |                            |
    40  //	   |            v                            |
    41  //	   |       + ---+---- +                      |
    42  //	On Error   | Resolve  |                      |
    43  //	   |       + ---+---- +                      |
    44  //	   |            |                            |
    45  //	   |            v                            |
    46  //	   |       +----+-----+                      |
    47  //	   +------>+ Complete +<---------------------+
    48  //	           +----+-----+
    49  //	                |
    50  //	                v
    51  //	           +----+-----+
    52  //	           |   End    |
    53  //	           +----------+
    54  type Merge[T any] struct {
    55  	lock          sync.Mutex
    56  	committed     bool
    57  	items         []T
    58  	status        chan mergeStatus
    59  	pending       []T
    60  	pendingStatus chan mergeStatus
    61  }
    62  
    63  // Do merges concurrent operations of items into a single call of prepare and
    64  // resolve.
    65  // If Do is called multiple times concurrently, only one of the calls will be
    66  // selected to invoke prepare and resolve.
    67  func (m *Merge[T]) Do(item T, prepare func() error, resolve func(items []T) error) error {
    68  	status := <-m.assign(item)
    69  	if status.main {
    70  		err := prepare()
    71  		items := m.commit()
    72  		if err == nil {
    73  			err = resolve(items)
    74  		}
    75  		m.complete(err)
    76  		return err
    77  	}
    78  	return status.err
    79  }
    80  
    81  // assign adds a new item into the item list.
    82  func (m *Merge[T]) assign(item T) <-chan mergeStatus {
    83  	m.lock.Lock()
    84  	defer m.lock.Unlock()
    85  
    86  	if m.committed {
    87  		if m.pendingStatus == nil {
    88  			m.pendingStatus = make(chan mergeStatus, 1)
    89  		}
    90  		m.pending = append(m.pending, item)
    91  		return m.pendingStatus
    92  	}
    93  
    94  	if m.status == nil {
    95  		m.status = make(chan mergeStatus, 1)
    96  		m.status <- mergeStatus{main: true}
    97  	}
    98  	m.items = append(m.items, item)
    99  	return m.status
   100  }
   101  
   102  // commit closes the assignment window, and the assigned items will be ready
   103  // for resolve.
   104  func (m *Merge[T]) commit() []T {
   105  	m.lock.Lock()
   106  	defer m.lock.Unlock()
   107  
   108  	m.committed = true
   109  	return m.items
   110  }
   111  
   112  // complete completes the previous merge, and moves the pending items to the
   113  // stage for the next merge.
   114  func (m *Merge[T]) complete(err error) {
   115  	// notify results
   116  	if err == nil {
   117  		close(m.status)
   118  	} else {
   119  		remaining := len(m.items) - 1
   120  		status := m.status
   121  		for remaining > 0 {
   122  			status <- mergeStatus{err: err}
   123  			remaining--
   124  		}
   125  	}
   126  
   127  	// move pending items to the stage
   128  	m.lock.Lock()
   129  	defer m.lock.Unlock()
   130  
   131  	m.committed = false
   132  	m.items = m.pending
   133  	m.status = m.pendingStatus
   134  	m.pending = nil
   135  	m.pendingStatus = nil
   136  
   137  	if m.status != nil {
   138  		m.status <- mergeStatus{main: true}
   139  	}
   140  }