Merge pull request #17617 from enj/enj/f/ranked_set

Automatic merge from submit-queue.

Implement a ranked set based on key and rank

This change adds a ranked set that uses a key (string) to determine uniqueness and a rank (int64) to determine ordering.  For example, a structure with a name and time could be ranked based on the time and deduped based on the name.

Signed-off-by: Monis Khan <[email protected]>

---

/assign @simo5

In regards to using this in the token timeout PR:

1. Implement Key using token.Name and Rank using token.timeout().Unix()
2. Use RankedSet.Insert to add tokens
3. Use RankedSet.LessThan(flushInterval.Unix(), true) to get all tokens that are about to expire (and remove them from the set)
4. Process the tokens as needed

Author: openshift-merge-robot

pkg/util/rankedset/rankedset.go

@@ -0,0 +1,162 @@
+package rankedset
+
+import "github.com/google/btree"
+
+// Item represents a single object in a RankedSet.
+type Item interface {
+	// Key returns the unique identifier for this item.
+	Key() string
+	// Rank is used to sort items.
+	// Items with the same rank are sorted lexicographically based on Key.
+	Rank() int64
+}
+
+// RankedSet stores Items based on Key (uniqueness) and Rank (sorting).
+type RankedSet struct {
+	rank *btree.BTree
+	set  map[string]*treeItem
+}
+
+// StringItem implements Item using a string.
+// It has two main uses:
+// 1. If all items in a RankedSet are StringItems, the set becomes a store of unique strings sorted lexicographically.
+// 2. It serves as a Key item that can be passed into methods that ignore Rank such as RankedSet.Delete.
+type StringItem string
+
+func (s StringItem) Key() string {
+	return string(s)
+}
+
+func (s StringItem) Rank() int64 {
+	return 0
+}
+
+func New() *RankedSet {
+	return &RankedSet{
+		rank: btree.New(32),
+		set:  make(map[string]*treeItem),
+	}
+}
+
+// Insert adds the item into the set.
+// If an item with the same Key existed in the set, it is deleted and returned.
+func (s *RankedSet) Insert(item Item) Item {
+	old := s.Delete(item)
+
+	key := item.Key()
+	value := &treeItem{item: item}
+
+	s.rank.ReplaceOrInsert(value) // should always return nil because we call Delete first
+	s.set[key] = value
+
+	return old
+}
+
+// Delete removes the item from the set based on Key (Rank is ignored).
+// The removed item is returned if it existed in the set.
+func (s *RankedSet) Delete(item Item) Item {
+	key := item.Key()
+	value, ok := s.set[key]
+	if !ok {
+		return nil
+	}
+
+	s.rank.Delete(value) // should always return the same data as value (non-nil)
+	delete(s.set, key)
+
+	return value.item
+}
+
+func (s *RankedSet) Min() Item {
+	if min := s.rank.Min(); min != nil {
+		return min.(*treeItem).item
+	}
+	return nil
+}
+
+func (s *RankedSet) Max() Item {
+	if max := s.rank.Max(); max != nil {
+		return max.(*treeItem).item
+	}
+	return nil
+}
+
+func (s *RankedSet) Len() int {
+	return len(s.set)
+}
+
+func (s *RankedSet) Get(item Item) Item {
+	if value, ok := s.set[item.Key()]; ok {
+		return value.item
+	}
+	return nil
+}
+
+func (s *RankedSet) Has(item Item) bool {
+	_, ok := s.set[item.Key()]
+	return ok
+}
+
+// List returns all items in the set in ranked order.
+// If delete is set to true, the returned items are removed from the set.
+func (s *RankedSet) List(delete bool) []Item {
+	return s.ascend(
+		func(item Item) bool {
+			return true
+		},
+		delete,
+	)
+}
+
+// LessThan returns all items less than the given rank in ranked order.
+// If delete is set to true, the returned items are removed from the set.
+func (s *RankedSet) LessThan(rank int64, delete bool) []Item {
+	return s.ascend(
+		func(item Item) bool {
+			return item.Rank() < rank
+		},
+		delete,
+	)
+}
+
+// setItemIterator allows callers of ascend to iterate in-order over the set.
+// When this function returns false, iteration will stop.
+type setItemIterator func(item Item) bool
+
+func (s *RankedSet) ascend(iterator setItemIterator, delete bool) []Item {
+	var items []Item
+	s.rank.Ascend(func(i btree.Item) bool {
+		item := i.(*treeItem).item
+		if !iterator(item) {
+			return false
+		}
+		items = append(items, item)
+		return true
+	})
+	// delete after Ascend since it is probably not safe to remove while iterating
+	if delete {
+		for _, item := range items {
+			s.Delete(item)
+		}
+	}
+	return items
+}
+
+var _ btree.Item = &treeItem{}
+
+type treeItem struct {
+	item Item
+}
+
+func (i *treeItem) Less(than btree.Item) bool {
+	other := than.(*treeItem).item
+
+	selfRank := i.item.Rank()
+	otherRank := other.Rank()
+
+	if selfRank == otherRank {
+		return i.item.Key() < other.Key()
+	}
+
+	return selfRank < otherRank
+}