github.com/andy2046/gopie@v0.7.0/pkg/spsc/spsc.go (about) 1 // Package spsc implements a Single-Producer / Single-Consumer queue. 2 package spsc 3 4 import ( 5 "golang.org/x/sys/cpu" 6 "runtime" 7 "sync/atomic" 8 "unsafe" 9 ) 10 11 const ( 12 // CacheLinePadSize is the size of OS CacheLine. 13 CacheLinePadSize = unsafe.Sizeof(cpu.CacheLinePad{}) 14 // DefaultMaxBatch is the default max batch size. 15 DefaultMaxBatch uint32 = (1 << 8) - 1 16 ) 17 18 // SPSC is the SPSC queue structure. 19 type SPSC struct { 20 _ [CacheLinePadSize]byte 21 writeCache int64 22 _ [CacheLinePadSize - 8%CacheLinePadSize]byte 23 writeIndex int64 24 _ [CacheLinePadSize - 8%CacheLinePadSize]byte 25 readIndex int64 26 _ [CacheLinePadSize - 8%CacheLinePadSize]byte 27 readCache int64 28 _ [CacheLinePadSize - 8%CacheLinePadSize]byte 29 done int32 30 _ [CacheLinePadSize - 4%CacheLinePadSize]byte 31 data []unsafe.Pointer 32 mask int64 33 maxbatch int64 34 } 35 36 // New create a new SPSC with bounded `size`. 37 func New(size uint32, batchSize ...uint32) *SPSC { 38 sp := SPSC{} 39 sp.data = make([]unsafe.Pointer, nextPowerOf2(size)) 40 sp.mask = int64(len(sp.data) - 1) 41 bSize := DefaultMaxBatch 42 if len(batchSize) > 0 && batchSize[0] != 0 { 43 bSize = batchSize[0] 44 } 45 sp.maxbatch = int64(nextPowerOf2(bSize) - 1) 46 return &sp 47 } 48 49 // Close the SPSC, it shall NOT be called before `Offer()` or `Put()`. 50 func (sp *SPSC) Close() { 51 atomic.AddInt32(&sp.done, 1) 52 } 53 54 // Poll the value at the head of the queue to given variable, 55 // non-blocking, return false if the queue is empty / closed. 56 func (sp *SPSC) Poll(i interface{}) bool { 57 readCache := sp.readCache 58 if writeIndex := atomic.LoadInt64(&sp.writeIndex); readCache >= writeIndex { 59 if readCache > sp.readIndex { 60 atomic.StoreInt64(&sp.readIndex, readCache) 61 } 62 if atomic.LoadInt32(&sp.done) > 0 { 63 // queue is closed 64 return false 65 } 66 if writeIndex = atomic.LoadInt64(&sp.writeCache); readCache >= writeIndex { 67 // queue is empty 68 return false 69 } 70 } 71 72 inject(i, sp.data[readCache&sp.mask]) 73 readCache++ 74 sp.readCache = readCache 75 if readCache-sp.readIndex > sp.maxbatch { 76 atomic.StoreInt64(&sp.readIndex, readCache) 77 } 78 return true 79 } 80 81 // Get the value at the head of the queue to given variable, 82 // blocking, return false if the queue is closed. 83 func (sp *SPSC) Get(i interface{}) bool { 84 readCache := sp.readCache 85 if writeIndex := atomic.LoadInt64(&sp.writeIndex); readCache >= writeIndex { 86 if readCache > sp.readIndex { 87 atomic.StoreInt64(&sp.readIndex, readCache) 88 } 89 for readCache >= writeIndex { 90 if atomic.LoadInt32(&sp.done) > 0 { 91 // queue is closed 92 return false 93 } 94 sp.spin() 95 writeIndex = atomic.LoadInt64(&sp.writeCache) 96 } 97 } 98 99 inject(i, sp.data[readCache&sp.mask]) 100 readCache++ 101 sp.readCache = readCache 102 if readCache-sp.readIndex > sp.maxbatch { 103 atomic.StoreInt64(&sp.readIndex, readCache) 104 } 105 return true 106 } 107 108 // Offer given variable at the tail of the queue, 109 // non-blocking, return false if the queue is full. 110 func (sp *SPSC) Offer(i interface{}) bool { 111 writeCache := sp.writeCache 112 if masked, readIndex := writeCache-sp.mask, 113 atomic.LoadInt64(&sp.readIndex); masked >= readIndex { 114 if writeCache > sp.writeIndex { 115 atomic.StoreInt64(&sp.writeIndex, writeCache) 116 } 117 if readIndex = atomic.LoadInt64(&sp.readIndex); masked >= readIndex { 118 // queue is full 119 return false 120 } 121 } 122 123 sp.data[writeCache&sp.mask] = extractptr(i) 124 writeCache = atomic.AddInt64(&sp.writeCache, 1) 125 if writeCache-sp.writeIndex > sp.maxbatch { 126 atomic.StoreInt64(&sp.writeIndex, writeCache) 127 } 128 return true 129 } 130 131 // Put given variable at the tail of the queue, 132 // blocking. 133 func (sp *SPSC) Put(i interface{}) { 134 writeCache := sp.writeCache 135 if masked, readIndex := writeCache-sp.mask, 136 atomic.LoadInt64(&sp.readIndex); masked >= readIndex { 137 if writeCache > sp.writeIndex { 138 atomic.StoreInt64(&sp.writeIndex, writeCache) 139 } 140 for masked >= readIndex { 141 sp.spin() 142 readIndex = atomic.LoadInt64(&sp.readIndex) 143 } 144 } 145 146 sp.data[writeCache&sp.mask] = extractptr(i) 147 writeCache = atomic.AddInt64(&sp.writeCache, 1) 148 if writeCache-sp.writeIndex > sp.maxbatch { 149 atomic.StoreInt64(&sp.writeIndex, writeCache) 150 } 151 } 152 153 func (sp *SPSC) spin() { 154 runtime.Gosched() 155 }