github.com/unionj-cloud/go-doudou/v2@v2.3.5/toolkit/memberlist/queue.go (about) 1 package memberlist 2 3 import ( 4 "math" 5 "sync" 6 7 "github.com/google/btree" 8 ) 9 10 //go:generate mockgen -destination ./mock/mock_queue.go -package mock -source=./queue.go 11 12 // TransmitLimitedQueue is used to queue messages to broadcast to 13 // the cluster (via gossip) but limits the number of transmits per 14 // message. It also prioritizes messages with lower transmit counts 15 // (hence newer messages). 16 type TransmitLimitedQueue struct { 17 // NumNodes returns the number of nodes in the cluster. This is 18 // used to determine the retransmit count, which is calculated 19 // based on the log of this. 20 NumNodes func() int 21 22 // RetransmitMult is the multiplier used to determine the maximum 23 // number of retransmissions attempted. 24 RetransmitMult int 25 RetransmitMultGetter func() int 26 27 mu sync.Mutex 28 tq *btree.BTree // stores *limitedBroadcast as btree.Item 29 tm map[string]*limitedBroadcast 30 idGen int64 31 } 32 33 type limitedBroadcast struct { 34 transmits int // btree-key[0]: Number of transmissions attempted. 35 msgLen int64 // btree-key[1]: copied from len(b.Message()) 36 id int64 // btree-key[2]: unique incrementing id stamped at submission time 37 b Broadcast 38 39 name string // set if Broadcast is a NamedBroadcast 40 } 41 42 // Less tests whether the current item is less than the given argument. 43 // 44 // This must provide a strict weak ordering. 45 // If !a.Less(b) && !b.Less(a), we treat this to mean a == b (i.e. we can only 46 // hold one of either a or b in the tree). 47 // 48 // default ordering is 49 // - [transmits=0, ..., transmits=inf] 50 // - [transmits=0:len=999, ..., transmits=0:len=2, ...] 51 // - [transmits=0:len=999,id=999, ..., transmits=0:len=999:id=1, ...] 52 func (b *limitedBroadcast) Less(than btree.Item) bool { 53 o := than.(*limitedBroadcast) 54 if b.transmits < o.transmits { 55 return true 56 } else if b.transmits > o.transmits { 57 return false 58 } 59 if b.msgLen > o.msgLen { 60 return true 61 } else if b.msgLen < o.msgLen { 62 return false 63 } 64 return b.id > o.id 65 } 66 67 // for testing; emits in transmit order if reverse=false 68 func (q *TransmitLimitedQueue) orderedView(reverse bool) []*limitedBroadcast { 69 q.mu.Lock() 70 defer q.mu.Unlock() 71 72 out := make([]*limitedBroadcast, 0, q.lenLocked()) 73 q.walkReadOnlyLocked(reverse, func(cur *limitedBroadcast) bool { 74 out = append(out, cur) 75 return true 76 }) 77 78 return out 79 } 80 81 // walkReadOnlyLocked calls f for each item in the queue traversing it in 82 // natural order (by Less) when reverse=false and the opposite when true. You 83 // must hold the mutex. 84 // 85 // This method panics if you attempt to mutate the item during traversal. The 86 // underlying btree should also not be mutated during traversal. 87 func (q *TransmitLimitedQueue) walkReadOnlyLocked(reverse bool, f func(*limitedBroadcast) bool) { 88 if q.lenLocked() == 0 { 89 return 90 } 91 92 iter := func(item btree.Item) bool { 93 cur := item.(*limitedBroadcast) 94 95 prevTransmits := cur.transmits 96 prevMsgLen := cur.msgLen 97 prevID := cur.id 98 99 keepGoing := f(cur) 100 101 if prevTransmits != cur.transmits || prevMsgLen != cur.msgLen || prevID != cur.id { 102 panic("edited queue while walking read only") 103 } 104 105 return keepGoing 106 } 107 108 if reverse { 109 q.tq.Descend(iter) // end with transmit 0 110 } else { 111 q.tq.Ascend(iter) // start with transmit 0 112 } 113 } 114 115 // Broadcast is something that can be broadcasted via gossip to 116 // the memberlist cluster. 117 type Broadcast interface { 118 // Invalidates checks if enqueuing the current broadcast 119 // invalidates a previous broadcast 120 Invalidates(b Broadcast) bool 121 122 // Returns a byte form of the message 123 Message() []byte 124 125 // Finished is invoked when the message will no longer 126 // be broadcast, either due to invalidation or to the 127 // transmit limit being reached 128 Finished() 129 } 130 131 // NamedBroadcast is an optional extension of the Broadcast interface that 132 // gives each message a unique string name, and that is used to optimize 133 // 134 // You shoud ensure that Invalidates() checks the same uniqueness as the 135 // example below: 136 // 137 // func (b *foo) Invalidates(other Broadcast) bool { 138 // nb, ok := other.(NamedBroadcast) 139 // if !ok { 140 // return false 141 // } 142 // return b.Name() == nb.Name() 143 // } 144 // 145 // Invalidates() isn't currently used for NamedBroadcasts, but that may change 146 // in the future. 147 type NamedBroadcast interface { 148 Broadcast 149 // The unique identity of this broadcast message. 150 Name() string 151 } 152 153 // UniqueBroadcast is an optional interface that indicates that each message is 154 // intrinsically unique and there is no need to scan the broadcast queue for 155 // duplicates. 156 // 157 // You should ensure that Invalidates() always returns false if implementing 158 // this interface. Invalidates() isn't currently used for UniqueBroadcasts, but 159 // that may change in the future. 160 type UniqueBroadcast interface { 161 Broadcast 162 // UniqueBroadcast is just a marker method for this interface. 163 UniqueBroadcast() 164 } 165 166 // QueueBroadcast is used to enqueue a broadcast 167 func (q *TransmitLimitedQueue) QueueBroadcast(b Broadcast) { 168 q.queueBroadcast(b, 0) 169 } 170 171 // lazyInit initializes internal data structures the first time they are 172 // needed. You must already hold the mutex. 173 func (q *TransmitLimitedQueue) lazyInit() { 174 if q.tq == nil { 175 q.tq = btree.New(32) 176 } 177 if q.tm == nil { 178 q.tm = make(map[string]*limitedBroadcast) 179 } 180 } 181 182 // queueBroadcast is like QueueBroadcast but you can use a nonzero value for 183 // the initial transmit tier assigned to the message. This is meant to be used 184 // for unit testing. 185 func (q *TransmitLimitedQueue) queueBroadcast(b Broadcast, initialTransmits int) { 186 q.mu.Lock() 187 defer q.mu.Unlock() 188 189 q.lazyInit() 190 191 if q.idGen == math.MaxInt64 { 192 // it's super duper unlikely to wrap around within the retransmit limit 193 q.idGen = 1 194 } else { 195 q.idGen++ 196 } 197 id := q.idGen 198 199 lb := &limitedBroadcast{ 200 transmits: initialTransmits, 201 msgLen: int64(len(b.Message())), 202 id: id, 203 b: b, 204 } 205 unique := false 206 if nb, ok := b.(NamedBroadcast); ok { 207 lb.name = nb.Name() 208 } else if _, ok := b.(UniqueBroadcast); ok { 209 unique = true 210 } 211 212 // Check if this message invalidates another. 213 if lb.name != "" { 214 if old, ok := q.tm[lb.name]; ok { 215 old.b.Finished() 216 q.deleteItem(old) 217 } 218 } else if !unique { 219 // Slow path, hopefully nothing hot hits this. 220 var remove []*limitedBroadcast 221 q.tq.Ascend(func(item btree.Item) bool { 222 cur := item.(*limitedBroadcast) 223 224 // Special Broadcasts can only invalidate each other. 225 switch cur.b.(type) { 226 case NamedBroadcast: 227 // noop 228 case UniqueBroadcast: 229 // noop 230 default: 231 if b.Invalidates(cur.b) { 232 cur.b.Finished() 233 remove = append(remove, cur) 234 } 235 } 236 return true 237 }) 238 for _, cur := range remove { 239 q.deleteItem(cur) 240 } 241 } 242 243 // Append to the relevant queue. 244 q.addItem(lb) 245 } 246 247 // deleteItem removes the given item from the overall datastructure. You 248 // must already hold the mutex. 249 func (q *TransmitLimitedQueue) deleteItem(cur *limitedBroadcast) { 250 _ = q.tq.Delete(cur) 251 if cur.name != "" { 252 delete(q.tm, cur.name) 253 } 254 255 if q.tq.Len() == 0 { 256 // At idle there's no reason to let the id generator keep going 257 // indefinitely. 258 q.idGen = 0 259 } 260 } 261 262 // addItem adds the given item into the overall datastructure. You must already 263 // hold the mutex. 264 func (q *TransmitLimitedQueue) addItem(cur *limitedBroadcast) { 265 _ = q.tq.ReplaceOrInsert(cur) 266 if cur.name != "" { 267 q.tm[cur.name] = cur 268 } 269 } 270 271 // getTransmitRange returns a pair of min/max values for transmit values 272 // represented by the current queue contents. Both values represent actual 273 // transmit values on the interval [0, len). You must already hold the mutex. 274 func (q *TransmitLimitedQueue) getTransmitRange() (minTransmit, maxTransmit int) { 275 if q.lenLocked() == 0 { 276 return 0, 0 277 } 278 minItem, maxItem := q.tq.Min(), q.tq.Max() 279 if minItem == nil || maxItem == nil { 280 return 0, 0 281 } 282 283 min := minItem.(*limitedBroadcast).transmits 284 max := maxItem.(*limitedBroadcast).transmits 285 286 return min, max 287 } 288 289 // GetBroadcasts is used to get a number of broadcasts, up to a byte limit 290 // and applying a per-message overhead as provided. 291 func (q *TransmitLimitedQueue) GetBroadcasts(overhead, limit int) [][]byte { 292 q.mu.Lock() 293 defer q.mu.Unlock() 294 295 // Fast path the default case 296 if q.lenLocked() == 0 { 297 return nil 298 } 299 300 mult := q.RetransmitMult 301 if q.RetransmitMultGetter != nil { 302 mult = q.RetransmitMultGetter() 303 } 304 transmitLimit := retransmitLimit(mult, q.NumNodes()) 305 306 var ( 307 bytesUsed int 308 toSend [][]byte 309 reinsert []*limitedBroadcast 310 ) 311 312 // Visit fresher items first, but only look at stuff that will fit. 313 // We'll go tier by tier, grabbing the largest items first. 314 minTr, maxTr := q.getTransmitRange() 315 for transmits := minTr; transmits <= maxTr; /*do not advance automatically*/ { 316 free := int64(limit - bytesUsed - overhead) 317 if free <= 0 { 318 break // bail out early 319 } 320 321 // Search for the least element on a given tier (by transmit count) as 322 // defined in the limitedBroadcast.Less function that will fit into our 323 // remaining space. 324 greaterOrEqual := &limitedBroadcast{ 325 transmits: transmits, 326 msgLen: free, 327 id: math.MaxInt64, 328 } 329 lessThan := &limitedBroadcast{ 330 transmits: transmits + 1, 331 msgLen: math.MaxInt64, 332 id: math.MaxInt64, 333 } 334 var keep *limitedBroadcast 335 q.tq.AscendRange(greaterOrEqual, lessThan, func(item btree.Item) bool { 336 cur := item.(*limitedBroadcast) 337 // Check if this is within our limits 338 if int64(len(cur.b.Message())) > free { 339 // If this happens it's a bug in the datastructure or 340 // surrounding use doing something like having len(Message()) 341 // change over time. There's enough going on here that it's 342 // probably sane to just skip it and move on for now. 343 return true 344 } 345 keep = cur 346 return false 347 }) 348 if keep == nil { 349 // No more items of an appropriate size in the tier. 350 transmits++ 351 continue 352 } 353 354 msg := keep.b.Message() 355 356 // Add to slice to send 357 bytesUsed += overhead + len(msg) 358 toSend = append(toSend, msg) 359 360 // Check if we should stop transmission 361 q.deleteItem(keep) 362 if keep.transmits+1 >= transmitLimit { 363 keep.b.Finished() 364 } else { 365 // We need to bump this item down to another transmit tier, but 366 // because it would be in the same direction that we're walking the 367 // tiers, we will have to delay the reinsertion until we are 368 // finished our search. Otherwise we'll possibly re-add the message 369 // when we ascend to the next tier. 370 keep.transmits++ 371 reinsert = append(reinsert, keep) 372 } 373 } 374 375 for _, cur := range reinsert { 376 q.addItem(cur) 377 } 378 379 return toSend 380 } 381 382 // NumQueued returns the number of queued messages 383 func (q *TransmitLimitedQueue) NumQueued() int { 384 q.mu.Lock() 385 defer q.mu.Unlock() 386 return q.lenLocked() 387 } 388 389 // lenLocked returns the length of the overall queue datastructure. You must 390 // hold the mutex. 391 func (q *TransmitLimitedQueue) lenLocked() int { 392 if q.tq == nil { 393 return 0 394 } 395 return q.tq.Len() 396 } 397 398 // Reset clears all the queued messages. Should only be used for tests. 399 func (q *TransmitLimitedQueue) Reset() { 400 q.mu.Lock() 401 defer q.mu.Unlock() 402 403 q.walkReadOnlyLocked(false, func(cur *limitedBroadcast) bool { 404 cur.b.Finished() 405 return true 406 }) 407 408 q.tq = nil 409 q.tm = nil 410 q.idGen = 0 411 } 412 413 // Prune will retain the maxRetain latest messages, and the rest 414 // will be discarded. This can be used to prevent unbounded queue sizes 415 func (q *TransmitLimitedQueue) Prune(maxRetain int) { 416 q.mu.Lock() 417 defer q.mu.Unlock() 418 419 // Do nothing if queue size is less than the limit 420 for q.tq.Len() > maxRetain { 421 item := q.tq.Max() 422 if item == nil { 423 break 424 } 425 cur := item.(*limitedBroadcast) 426 cur.b.Finished() 427 q.deleteItem(cur) 428 } 429 }