github.com/egonelbre/exp@v0.0.0-20240430123955-ed1d3aa93911/combiner/tbb.go (about)

     1  package combiner
     2  
     3  import (
     4  	"sync/atomic"
     5  	"unsafe"
     6  )
     7  
     8  // based on https://software.intel.com/en-us/blogs/2013/02/22/combineraggregator-synchronization-primitive
     9  type TBB struct {
    10  	head    unsafe.Pointer // *tbbNode
    11  	_       [7]uint64
    12  	batcher Batcher
    13  	busy    int64
    14  }
    15  
    16  type tbbNode struct {
    17  	next     unsafe.Pointer // *tbbNode
    18  	argument interface{}
    19  }
    20  
    21  func NewTBB(batcher Batcher) *TBB {
    22  	return &TBB{
    23  		batcher: batcher,
    24  		head:    nil,
    25  	}
    26  }
    27  
    28  func (c *TBB) DoAsync(op interface{}) { c.do(op, true) }
    29  func (c *TBB) Do(op interface{})      { c.do(op, false) }
    30  
    31  func (c *TBB) do(arg interface{}, async bool) {
    32  	node := &tbbNode{argument: arg}
    33  
    34  	var cmp unsafe.Pointer
    35  	for {
    36  		cmp = atomic.LoadPointer(&c.head)
    37  		node.next = cmp
    38  		if atomic.CompareAndSwapPointer(&c.head, cmp, unsafe.Pointer(node)) {
    39  			break
    40  		}
    41  	}
    42  
    43  	if cmp != nil {
    44  		if async {
    45  			return
    46  		}
    47  
    48  		// 2. If we are not the combiner, wait for arg.next to become nil
    49  		// (which means the operation is finished).
    50  		for try := 0; atomic.LoadPointer(&node.next) != nil; spin(&try) {
    51  		}
    52  	} else {
    53  		// 3. We are the combiner.
    54  
    55  		// wait for previous combiner to finish
    56  		for try := 0; atomic.LoadInt64(&c.busy) == 1; spin(&try) {
    57  		}
    58  		atomic.StoreInt64(&c.busy, 1)
    59  
    60  		// First, execute own operation.
    61  		c.batcher.Start()
    62  		defer c.batcher.Finish()
    63  
    64  		// Grab the batch of operations only once
    65  		for {
    66  			cmp = atomic.LoadPointer(&c.head)
    67  			if atomic.CompareAndSwapPointer(&c.head, cmp, nil) {
    68  				break
    69  			}
    70  		}
    71  
    72  		node = (*tbbNode)(cmp)
    73  		// Execute the list of operations.
    74  		for node != nil {
    75  			next := (*tbbNode)(node.next)
    76  			c.batcher.Include(node.argument)
    77  			// Mark completion.
    78  			atomic.StorePointer(&node.next, nil)
    79  			node = next
    80  		}
    81  
    82  		// allow next combiner to proceed
    83  		atomic.StoreInt64(&c.busy, 0)
    84  	}
    85  }