github.com/songzhibin97/gkit@v1.2.13/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 }