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  }