github.com/nsqio/nsq@v1.3.0/nsqd/in_flight_pqueue.go (about)

     1  package nsqd
     2  
     3  type inFlightPqueue []*Message
     4  
     5  func newInFlightPqueue(capacity int) inFlightPqueue {
     6  	return make(inFlightPqueue, 0, capacity)
     7  }
     8  
     9  func (pq inFlightPqueue) Swap(i, j int) {
    10  	pq[i], pq[j] = pq[j], pq[i]
    11  	pq[i].index = i
    12  	pq[j].index = j
    13  }
    14  
    15  func (pq *inFlightPqueue) Push(x *Message) {
    16  	n := len(*pq)
    17  	c := cap(*pq)
    18  	if n+1 > c {
    19  		npq := make(inFlightPqueue, n, c*2)
    20  		copy(npq, *pq)
    21  		*pq = npq
    22  	}
    23  	*pq = (*pq)[0 : n+1]
    24  	x.index = n
    25  	(*pq)[n] = x
    26  	pq.up(n)
    27  }
    28  
    29  func (pq *inFlightPqueue) Pop() *Message {
    30  	n := len(*pq)
    31  	c := cap(*pq)
    32  	pq.Swap(0, n-1)
    33  	pq.down(0, n-1)
    34  	if n < (c/2) && c > 25 {
    35  		npq := make(inFlightPqueue, n, c/2)
    36  		copy(npq, *pq)
    37  		*pq = npq
    38  	}
    39  	x := (*pq)[n-1]
    40  	x.index = -1
    41  	*pq = (*pq)[0 : n-1]
    42  	return x
    43  }
    44  
    45  func (pq *inFlightPqueue) Remove(i int) *Message {
    46  	n := len(*pq)
    47  	if n-1 != i {
    48  		pq.Swap(i, n-1)
    49  		pq.down(i, n-1)
    50  		pq.up(i)
    51  	}
    52  	x := (*pq)[n-1]
    53  	x.index = -1
    54  	*pq = (*pq)[0 : n-1]
    55  	return x
    56  }
    57  
    58  func (pq *inFlightPqueue) PeekAndShift(max int64) (*Message, int64) {
    59  	if len(*pq) == 0 {
    60  		return nil, 0
    61  	}
    62  
    63  	x := (*pq)[0]
    64  	if x.pri > max {
    65  		return nil, x.pri - max
    66  	}
    67  	pq.Pop()
    68  
    69  	return x, 0
    70  }
    71  
    72  func (pq *inFlightPqueue) up(j int) {
    73  	for {
    74  		i := (j - 1) / 2 // parent
    75  		if i == j || (*pq)[j].pri >= (*pq)[i].pri {
    76  			break
    77  		}
    78  		pq.Swap(i, j)
    79  		j = i
    80  	}
    81  }
    82  
    83  func (pq *inFlightPqueue) down(i, n int) {
    84  	for {
    85  		j1 := 2*i + 1
    86  		if j1 >= n || j1 < 0 { // j1 < 0 after int overflow
    87  			break
    88  		}
    89  		j := j1 // left child
    90  		if j2 := j1 + 1; j2 < n && (*pq)[j1].pri >= (*pq)[j2].pri {
    91  			j = j2 // = 2*i + 2  // right child
    92  		}
    93  		if (*pq)[j].pri >= (*pq)[i].pri {
    94  			break
    95  		}
    96  		pq.Swap(i, j)
    97  		i = j
    98  	}
    99  }