github.com/songzhibin97/go-baseutils@v0.0.2-0.20240302024150-487d8ce9c082/sys/syncx/poolqueue.go (about)

     1  //go:build !race
     2  // +build !race
     3  
     4  package syncx
     5  
     6  import (
     7  	"sync/atomic"
     8  	"unsafe"
     9  )
    10  
    11  // poolDequeue is a lock-free fixed-size single-producer,
    12  // multi-consumer queue. The single producer can both push and pop
    13  // from the head, and consumers can pop from the tail.
    14  //
    15  // It has the added feature that it nils out unused slots to avoid
    16  // unnecessary retention of objects. This is important for sync.Pool,
    17  // but not typically a property considered in the literature.
    18  type poolDequeue struct {
    19  	// headTail packs together a 32-bit head index and a 32-bit
    20  	// tail index. Both are indexes into vals modulo len(vals)-1.
    21  	//
    22  	// tail = index of oldest data in queue
    23  	// head = index of next slot to fill
    24  	//
    25  	// Slots in the range [tail, head) are owned by consumers.
    26  	// A consumer continues to own a slot outside this range until
    27  	// it nils the slot, at which point ownership passes to the
    28  	// producer.
    29  	//
    30  	// The head index is stored in the most-significant bits so
    31  	// that we can atomically add to it and the overflow is
    32  	// harmless.
    33  	headTail uint64
    34  
    35  	// vals is a ring buffer of interface{} values stored in this
    36  	// dequeue. The size of this must be a power of 2.
    37  	//
    38  	// vals[i].typ is nil if the slot is empty and non-nil
    39  	// otherwise. A slot is still in use until *both* the tail
    40  	// index has moved beyond it and typ has been set to nil. This
    41  	// is set to nil atomically by the consumer and read
    42  	// atomically by the producer.
    43  	vals []eface
    44  }
    45  
    46  type eface struct {
    47  	typ, val unsafe.Pointer
    48  }
    49  
    50  const dequeueBits = 32
    51  
    52  // dequeueLimit is the maximum size of a poolDequeue.
    53  //
    54  // This must be at most (1<<dequeueBits)/2 because detecting fullness
    55  // depends on wrapping around the ring buffer without wrapping around
    56  // the index. We divide by 4 so this fits in an int on 32-bit.
    57  const dequeueLimit = (1 << dequeueBits) / 4
    58  
    59  // dequeueNil is used in poolDeqeue to represent interface{}(nil).
    60  // Since we use nil to represent empty slots, we need a sentinel value
    61  // to represent nil.
    62  type dequeueNil *struct{}
    63  
    64  func (d *poolDequeue) unpack(ptrs uint64) (head, tail uint32) {
    65  	const mask = 1<<dequeueBits - 1
    66  	head = uint32((ptrs >> dequeueBits) & mask)
    67  	tail = uint32(ptrs & mask)
    68  	return
    69  }
    70  
    71  func (d *poolDequeue) pack(head, tail uint32) uint64 {
    72  	const mask = 1<<dequeueBits - 1
    73  	return (uint64(head) << dequeueBits) |
    74  		uint64(tail&mask)
    75  }
    76  
    77  // pushHead adds val at the head of the queue. It returns false if the
    78  // queue is full. It must only be called by a single producer.
    79  func (d *poolDequeue) pushHead(val *block) bool {
    80  	ptrs := atomic.LoadUint64(&d.headTail)
    81  	head, tail := d.unpack(ptrs)
    82  	if (tail+uint32(len(d.vals)))&(1<<dequeueBits-1) == head {
    83  		// Queue is full.
    84  		return false
    85  	}
    86  	slot := &d.vals[head&uint32(len(d.vals)-1)]
    87  
    88  	// Check if the head slot has been released by popTail.
    89  	typ := atomic.LoadPointer(&slot.typ)
    90  	if typ != nil {
    91  		// Another goroutine is still cleaning up the tail, so
    92  		// the queue is actually still full.
    93  		return false
    94  	}
    95  
    96  	// The head slot is free, so we own it.
    97  	*(**block)(unsafe.Pointer(slot)) = val
    98  
    99  	// Increment head. This passes ownership of slot to popTail
   100  	// and acts as a store barrier for writing the slot.
   101  	atomic.AddUint64(&d.headTail, 1<<dequeueBits)
   102  	return true
   103  }
   104  
   105  // popHead removes and returns the element at the head of the queue.
   106  // It returns false if the queue is empty. It must only be called by a
   107  // single producer.
   108  func (d *poolDequeue) popHead() (*block, bool) {
   109  	var slot *eface
   110  	for {
   111  		ptrs := atomic.LoadUint64(&d.headTail)
   112  		head, tail := d.unpack(ptrs)
   113  		if tail == head {
   114  			// Queue is empty.
   115  			return nil, false
   116  		}
   117  
   118  		// Confirm tail and decrement head. We do this before
   119  		// reading the value to take back ownership of this
   120  		// slot.
   121  		head--
   122  		ptrs2 := d.pack(head, tail)
   123  		if atomic.CompareAndSwapUint64(&d.headTail, ptrs, ptrs2) {
   124  			// We successfully took back slot.
   125  			slot = &d.vals[head&uint32(len(d.vals)-1)]
   126  			break
   127  		}
   128  	}
   129  
   130  	val := *(**block)(unsafe.Pointer(slot))
   131  	// Zero the slot. Unlike popTail, this isn't racing with
   132  	// pushHead, so we don't need to be careful here.
   133  	*slot = eface{}
   134  	return val, true
   135  }
   136  
   137  // popTail removes and returns the element at the tail of the queue.
   138  // It returns false if the queue is empty. It may be called by any
   139  // number of consumers.
   140  func (d *poolDequeue) popTail() (*block, bool) {
   141  	var slot *eface
   142  	for {
   143  		ptrs := atomic.LoadUint64(&d.headTail)
   144  		head, tail := d.unpack(ptrs)
   145  		if tail == head {
   146  			// Queue is empty.
   147  			return nil, false
   148  		}
   149  
   150  		// Confirm head and tail (for our speculative check
   151  		// above) and increment tail. If this succeeds, then
   152  		// we own the slot at tail.
   153  		ptrs2 := d.pack(head, tail+1)
   154  		if atomic.CompareAndSwapUint64(&d.headTail, ptrs, ptrs2) {
   155  			// Success.
   156  			slot = &d.vals[tail&uint32(len(d.vals)-1)]
   157  			break
   158  		}
   159  	}
   160  
   161  	// We now own slot.
   162  	val := *(**block)(unsafe.Pointer(slot))
   163  
   164  	// Tell pushHead that we're done with this slot. Zeroing the
   165  	// slot is also important so we don't leave behind references
   166  	// that could keep this object live longer than necessary.
   167  	//
   168  	// We write to val first and then publish that we're done with
   169  	// this slot by atomically writing to typ.
   170  	slot.val = nil
   171  	atomic.StorePointer(&slot.typ, nil)
   172  	// At this point pushHead owns the slot.
   173  	return val, true
   174  }
   175  
   176  // poolChain is a dynamically-sized version of poolDequeue.
   177  //
   178  // This is implemented as a doubly-linked list queue of poolDequeues
   179  // where each dequeue is double the size of the previous one. Once a
   180  // dequeue fills up, this allocates a new one and only ever pushes to
   181  // the latest dequeue. Pops happen from the other end of the list and
   182  // once a dequeue is exhausted, it gets removed from the list.
   183  type poolChain struct {
   184  	size int32
   185  
   186  	// head is the poolDequeue to push to. This is only accessed
   187  	// by the producer, so doesn't need to be synchronized.
   188  	head *poolChainElt
   189  
   190  	// tail is the poolDequeue to popTail from. This is accessed
   191  	// by consumers, so reads and writes must be atomic.
   192  	tail *poolChainElt
   193  }
   194  
   195  type poolChainElt struct {
   196  	poolDequeue
   197  
   198  	// next and prev link to the adjacent poolChainElts in this
   199  	// poolChain.
   200  	//
   201  	// next is written atomically by the producer and read
   202  	// atomically by the consumer. It only transitions from nil to
   203  	// non-nil.
   204  	//
   205  	// prev is written atomically by the consumer and read
   206  	// atomically by the producer. It only transitions from
   207  	// non-nil to nil.
   208  	next, prev *poolChainElt
   209  }
   210  
   211  func storePoolChainElt(pp **poolChainElt, v *poolChainElt) {
   212  	atomic.StorePointer((*unsafe.Pointer)(unsafe.Pointer(pp)), unsafe.Pointer(v))
   213  }
   214  
   215  func loadPoolChainElt(pp **poolChainElt) *poolChainElt {
   216  	return (*poolChainElt)(atomic.LoadPointer((*unsafe.Pointer)(unsafe.Pointer(pp))))
   217  }
   218  
   219  func (c *poolChain) pushHead(val *block) {
   220  	atomic.AddInt32(&c.size, 1)
   221  	d := c.head
   222  	if d == nil {
   223  		// Initialize the chain.
   224  		const initSize = 8 // Must be a power of 2
   225  		d = new(poolChainElt)
   226  		d.vals = make([]eface, initSize)
   227  		c.head = d
   228  		storePoolChainElt(&c.tail, d)
   229  	}
   230  
   231  	if d.pushHead(val) {
   232  		return
   233  	}
   234  
   235  	// The current dequeue is full. Allocate a new one of twice
   236  	// the size.
   237  	newSize := len(d.vals) * 2
   238  	if newSize >= dequeueLimit {
   239  		// Can't make it any bigger.
   240  		newSize = dequeueLimit
   241  	}
   242  
   243  	d2 := &poolChainElt{prev: d}
   244  	d2.vals = make([]eface, newSize)
   245  	c.head = d2
   246  	storePoolChainElt(&d.next, d2)
   247  	d2.pushHead(val)
   248  }
   249  
   250  func (c *poolChain) popHead() (*block, bool) {
   251  	d := c.head
   252  	for d != nil {
   253  		if val, ok := d.popHead(); ok {
   254  			atomic.AddInt32(&c.size, -1)
   255  			return val, ok
   256  		}
   257  		// There may still be unconsumed elements in the
   258  		// previous dequeue, so try backing up.
   259  		d = loadPoolChainElt(&d.prev)
   260  	}
   261  	return nil, false
   262  }
   263  
   264  func (c *poolChain) popTail() (*block, bool) {
   265  	d := loadPoolChainElt(&c.tail)
   266  	if d == nil {
   267  		return nil, false
   268  	}
   269  
   270  	for {
   271  		// It's important that we load the next pointer
   272  		// *before* popping the tail. In general, d may be
   273  		// transiently empty, but if next is non-nil before
   274  		// the pop and the pop fails, then d is permanently
   275  		// empty, which is the only condition under which it's
   276  		// safe to drop d from the chain.
   277  		d2 := loadPoolChainElt(&d.next)
   278  
   279  		if val, ok := d.popTail(); ok {
   280  			atomic.AddInt32(&c.size, -1)
   281  			return val, ok
   282  		}
   283  
   284  		if d2 == nil {
   285  			// This is the only dequeue. It's empty right
   286  			// now, but could be pushed to in the future.
   287  			return nil, false
   288  		}
   289  
   290  		// The tail of the chain has been drained, so move on
   291  		// to the next dequeue. Try to drop it from the chain
   292  		// so the next pop doesn't have to look at the empty
   293  		// dequeue again.
   294  		if atomic.CompareAndSwapPointer((*unsafe.Pointer)(unsafe.Pointer(&c.tail)), unsafe.Pointer(d), unsafe.Pointer(d2)) {
   295  			// We won the race. Clear the prev pointer so
   296  			// the garbage collector can collect the empty
   297  			// dequeue and so popHead doesn't back up
   298  			// further than necessary.
   299  			storePoolChainElt(&d2.prev, nil)
   300  		}
   301  		d = d2
   302  	}
   303  }