github.com/isyscore/isc-gobase@v1.5.3-0.20231218061332-cbc7451899e9/isc/queue.go (about)

     1  package isc
     2  
     3  import (
     4  	"sync"
     5  	"sync/atomic"
     6  	t0 "time"
     7  )
     8  
     9  type element struct {
    10  	next  *element
    11  	value interface{}
    12  }
    13  
    14  func newElement(value interface{}) *element {
    15  	return &element{value: value}
    16  }
    17  
    18  type Queue struct {
    19  	input   *element
    20  	output  *element
    21  	ch      chan bool
    22  	waitNum int32
    23  	num     int32
    24  	lock    *sync.RWMutex
    25  	rLock   *sync.Mutex
    26  	timer   *timer
    27  }
    28  
    29  type timer struct {
    30  	t     *t0.Timer
    31  	isEnd chan bool
    32  }
    33  
    34  func newTimer() *timer {
    35  	return &timer{t: t0.NewTimer(t0.Second * 10), isEnd: make(chan bool)}
    36  }
    37  
    38  func (timer *timer) wait() bool {
    39  	select {
    40  	case <-timer.t.C:
    41  		return false
    42  	case fa := <-timer.isEnd:
    43  		return fa
    44  	}
    45  	// return false
    46  }
    47  
    48  func (timer *timer) end() {
    49  	fa := timer.t.Stop()
    50  	if fa {
    51  		timer.isEnd <- true
    52  	}
    53  }
    54  
    55  func (timer *timer) reset(duration t0.Duration) {
    56  	timer.t.Reset(duration)
    57  }
    58  
    59  var poolTimer = &sync.Pool{
    60  	New: func() interface{} {
    61  		return newTimer()
    62  	},
    63  }
    64  
    65  func getTimer(duration t0.Duration) *timer {
    66  	ti := poolTimer.Get().(*timer)
    67  	ti.reset(duration)
    68  	return ti
    69  }
    70  
    71  func freeTimer(timer *timer) {
    72  	poolTimer.Put(timer)
    73  }
    74  
    75  func NewQueue() *Queue {
    76  	return &Queue{ch: make(chan bool), waitNum: 0, num: 0, lock: new(sync.RWMutex), rLock: new(sync.Mutex)}
    77  }
    78  
    79  func (queue *Queue) Offer(value interface{}) (num int32) {
    80  	ele := newElement(value)
    81  	queue.lock.Lock()
    82  	if queue.num == 0 {
    83  		queue.input = ele
    84  		queue.output = ele
    85  	} else {
    86  		queue.input.next = ele
    87  		queue.input = ele
    88  	}
    89  	num = atomic.AddInt32(&queue.num, 1)
    90  	if queue.waitNum > 0 {
    91  		queue.waitNum--
    92  		queue.lock.Unlock()
    93  		queue.ch <- true
    94  	} else {
    95  		queue.lock.Unlock()
    96  	}
    97  	return
    98  }
    99  
   100  func (queue *Queue) Num() int32 {
   101  	return queue.num
   102  }
   103  
   104  func (queue *Queue) Peek() (value interface{}, num int32) {
   105  	queue.lock.RLock()
   106  	num = queue.num
   107  	if queue.num > 0 {
   108  		value = queue.output.value
   109  		queue.lock.RUnlock()
   110  		return
   111  	} else {
   112  		queue.lock.RUnlock()
   113  		return nil, 0
   114  	}
   115  }
   116  
   117  func (queue *Queue) Poll() (value interface{}, num int32) {
   118  	for {
   119  		queue.lock.Lock()
   120  		if queue.num > 0 {
   121  			if queue.num == 1 {
   122  				value, num = queue.readOne()
   123  				queue.lock.Unlock()
   124  				return
   125  			} else {
   126  				queue.lock.Unlock()
   127  				queue.rLock.Lock()
   128  				val, n, last := queue.readGtOne()
   129  				if last {
   130  					queue.rLock.Unlock()
   131  				} else {
   132  					queue.rLock.Unlock()
   133  					return val, n
   134  				}
   135  			}
   136  		} else {
   137  			queue.waitNum++
   138  			queue.lock.Unlock()
   139  			<-queue.ch
   140  		}
   141  	}
   142  }
   143  
   144  func (queue *Queue) readOne() (value interface{}, num int32) {
   145  	var ele = queue.output
   146  	value = ele.value
   147  	num = atomic.AddInt32(&queue.num, -1)
   148  	return value, num
   149  }
   150  
   151  func (queue *Queue) readGtOne() (value interface{}, num int32, isLast bool) {
   152  	var ele = queue.output
   153  	if ele.next == nil {
   154  		return nil, 0, true
   155  	}
   156  	value = ele.value
   157  	queue.output = ele.next
   158  	ele.next = nil
   159  	num = atomic.AddInt32(&queue.num, -1)
   160  	return value, num, false
   161  }
   162  
   163  func (queue *Queue) Take(duration t0.Duration) (value interface{}, num int32) {
   164  	for {
   165  		queue.lock.Lock()
   166  		if queue.num > 0 {
   167  			if queue.num == 1 {
   168  				value, num = queue.readOne()
   169  				queue.lock.Unlock()
   170  				return
   171  			} else {
   172  				queue.lock.Unlock()
   173  				queue.rLock.Lock()
   174  				val, n, last := queue.readGtOne()
   175  				if last {
   176  					queue.rLock.Unlock()
   177  				} else {
   178  					queue.rLock.Unlock()
   179  					return val, n
   180  				}
   181  			}
   182  		} else {
   183  			queue.waitNum++
   184  			queue.lock.Unlock()
   185  			tm := getTimer(duration)
   186  			go func() {
   187  				fa := tm.wait()
   188  				if !fa {
   189  					queue.lock.Lock()
   190  					if queue.waitNum > 0 {
   191  						queue.waitNum--
   192  						queue.lock.Unlock()
   193  						queue.ch <- false
   194  					} else {
   195  						queue.lock.Unlock()
   196  					}
   197  				}
   198  			}()
   199  			flag := <-queue.ch
   200  			tm.end()
   201  			freeTimer(tm)
   202  			if !flag {
   203  				return nil, 0
   204  			}
   205  		}
   206  	}
   207  }