github.com/pharosnet/flyline@v1.0.2/queue.go (about)

     1  package flyline
     2  
     3  import (
     4  	"runtime"
     5  	"sync"
     6  	"sync/atomic"
     7  	"unsafe"
     8  )
     9  
    10  type node struct {
    11  	value interface{}
    12  	prev  *node
    13  	next  *node
    14  }
    15  
    16  type queue struct {
    17  	wg      sync.WaitGroup
    18  	in      *node
    19  	inAvail *node
    20  	out     *node
    21  	outTail *node
    22  	outFree *node
    23  }
    24  
    25  func (q *queue) newNode() *node {
    26  	if q.inAvail != nil {
    27  		n := q.inAvail
    28  		q.inAvail = q.inAvail.next
    29  		return n
    30  	}
    31  	times := 10
    32  	for {
    33  		q.inAvail = (*node)(atomic.LoadPointer(
    34  			(*unsafe.Pointer)(unsafe.Pointer(&q.outFree)),
    35  		))
    36  		if q.inAvail == nil {
    37  			return &node{}
    38  		}
    39  		if atomic.CompareAndSwapPointer(
    40  			(*unsafe.Pointer)(unsafe.Pointer(&q.outFree)),
    41  			unsafe.Pointer(q.inAvail), nil) {
    42  			return q.newNode()
    43  		}
    44  		times--
    45  		if times <= 0 {
    46  			runtime.Gosched()
    47  			times = 10
    48  		}
    49  	}
    50  }
    51  
    52  func (q *queue) peek() *node {
    53  	return (*node)(atomic.LoadPointer((*unsafe.Pointer)(unsafe.Pointer(&q.in))))
    54  }
    55  
    56  func (q *queue) swapNode(old *node, new *node) bool {
    57  	return atomic.CompareAndSwapPointer(
    58  		(*unsafe.Pointer)(unsafe.Pointer(&q.in)),
    59  		unsafe.Pointer(old),
    60  		unsafe.Pointer(new),
    61  	)
    62  }
    63  
    64  func (q *queue) free(recv *node) {
    65  	times := 10
    66  	for {
    67  		freed := (*node)(atomic.LoadPointer(
    68  			(*unsafe.Pointer)(unsafe.Pointer(&q.outFree)),
    69  		))
    70  		q.outTail.next = freed
    71  		if atomic.CompareAndSwapPointer(
    72  			(*unsafe.Pointer)(unsafe.Pointer(&q.outFree)),
    73  			unsafe.Pointer(freed), unsafe.Pointer(recv)) {
    74  			return
    75  		}
    76  		times--
    77  		if times <= 0 {
    78  			runtime.Gosched()
    79  			times = 10
    80  		}
    81  	}
    82  }
    83  
    84  var emptyNode = &node{}
    85  
    86  func (q *queue) add(value interface{}) {
    87  	n := q.newNode()
    88  	n.value = value
    89  	hasRemains := false
    90  	times := 10
    91  	for {
    92  		n.next = q.peek()
    93  		if n.next == emptyNode {
    94  			if q.swapNode(n.next, n.next.next) {
    95  				hasRemains = true
    96  			}
    97  		} else {
    98  			if q.swapNode(n.next, n) {
    99  				break
   100  			}
   101  		}
   102  		times--
   103  		if times <= 0 {
   104  			runtime.Gosched()
   105  			times = 10
   106  		}
   107  	}
   108  	if hasRemains {
   109  		q.wg.Done()
   110  	}
   111  }
   112  
   113  func (q *queue) poll() interface{} {
   114  	if q.out != nil {
   115  		v := q.out.value
   116  		if q.out.prev == nil {
   117  			q.free(q.out)
   118  			q.out = nil
   119  		} else {
   120  			q.out = q.out.prev
   121  			q.out.next.prev = nil
   122  		}
   123  		return v
   124  	}
   125  	var n *node
   126  	times := 10
   127  	for {
   128  		n = q.peek()
   129  		if n == nil {
   130  			q.wg.Add(1)
   131  			if q.swapNode(n, emptyNode) {
   132  				q.wg.Wait()
   133  			} else {
   134  				q.wg.Done()
   135  			}
   136  		} else if q.swapNode(n, nil) {
   137  			break
   138  		}
   139  		times--
   140  		if times <= 0 {
   141  			runtime.Gosched()
   142  			times = 10
   143  		}
   144  	}
   145  	for n.next != nil {
   146  		n.next.prev = n
   147  		n = n.next
   148  	}
   149  	q.out = n
   150  	q.outTail = n
   151  	return q.poll()
   152  }