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  }