diff --git a/pkg/collections/orderedintset.go b/pkg/collections/orderedintset.go index 23abab04d3f58..7442f2e93f560 100644 --- a/pkg/collections/orderedintset.go +++ b/pkg/collections/orderedintset.go @@ -1,12 +1,13 @@ package collections import ( + "sort" "sync" ) // OrderedIntSet is a thread-safe sorted set and a stack. type OrderedIntSet struct { - sync.RWMutex + sync.Mutex set []int } @@ -15,29 +16,22 @@ func NewOrderedIntSet() *OrderedIntSet { return &OrderedIntSet{} } -// Push takes a string and adds it to the set. If the elem aready exists, it has no effect. +// Push takes an int and adds it to the set. If the elem aready exists, it has no effect. func (s *OrderedIntSet) Push(elem int) { - s.RLock() - for _, e := range s.set { - if e == elem { - s.RUnlock() - return - } - } - s.RUnlock() - s.Lock() + if len(s.set) == 0 { + s.set = append(s.set, elem) + s.Unlock() + return + } // Make sure the list is always sorted - for i, e := range s.set { - if elem < e { - s.set = append(s.set[:i], append([]int{elem}, s.set[i:]...)...) - s.Unlock() - return - } + i := sort.SearchInts(s.set, elem) + if i < len(s.set) && s.set[i] == elem { + s.Unlock() + return } - // If we reach here, then elem is the biggest elem of the list. - s.set = append(s.set, elem) + s.set = append(s.set[:i], append([]int{elem}, s.set[i:]...)...) s.Unlock() } @@ -46,28 +40,26 @@ func (s *OrderedIntSet) Pop() int { return s.PopFront() } -// Pop returns the first elemen from the list and removes it. +// Pop returns the first element from the list and removes it. // If the list is empty, it returns 0 func (s *OrderedIntSet) PopFront() int { - s.RLock() - - for i, e := range s.set { - ret := e - s.RUnlock() - s.Lock() - s.set = append(s.set[:i], s.set[i+1:]...) + s.Lock() + if len(s.set) == 0 { s.Unlock() - return ret + return 0 } - s.RUnlock() - - return 0 + ret := s.set[0] + s.set = s.set[1:] + s.Unlock() + return ret } // PullBack retrieve the last element of the list. // The element is not removed. // If the list is empty, an empty element is returned. func (s *OrderedIntSet) PullBack() int { + s.Lock() + defer s.Unlock() if len(s.set) == 0 { return 0 } @@ -76,21 +68,28 @@ func (s *OrderedIntSet) PullBack() int { // Exists checks if the given element present in the list. func (s *OrderedIntSet) Exists(elem int) bool { - for _, e := range s.set { - if e == elem { - return true - } + s.Lock() + if len(s.set) == 0 { + s.Unlock() + return false } - return false + i := sort.SearchInts(s.set, elem) + res := i < len(s.set) && s.set[i] == elem + s.Unlock() + return res } // Remove removes an element from the list. // If the element is not found, it has no effect. func (s *OrderedIntSet) Remove(elem int) { - for i, e := range s.set { - if e == elem { - s.set = append(s.set[:i], s.set[i+1:]...) - return - } + s.Lock() + if len(s.set) == 0 { + s.Unlock() + return } + i := sort.SearchInts(s.set, elem) + if i < len(s.set) && s.set[i] == elem { + s.set = append(s.set[:i], s.set[i+1:]...) + } + s.Unlock() } diff --git a/pkg/collections/orderedintset_test.go b/pkg/collections/orderedintset_test.go new file mode 100644 index 0000000000000..0ac4ca545547b --- /dev/null +++ b/pkg/collections/orderedintset_test.go @@ -0,0 +1,71 @@ +package collections + +import ( + "math/rand" + "testing" +) + +func BenchmarkPush(b *testing.B) { + var testSet []int + for i := 0; i < 1000; i++ { + testSet = append(testSet, rand.Int()) + } + s := NewOrderedIntSet() + b.ResetTimer() + for i := 0; i < b.N; i++ { + for _, elem := range testSet { + s.Push(elem) + } + } +} + +func BenchmarkPop(b *testing.B) { + var testSet []int + for i := 0; i < 1000; i++ { + testSet = append(testSet, rand.Int()) + } + s := NewOrderedIntSet() + for _, elem := range testSet { + s.Push(elem) + } + b.ResetTimer() + for i := 0; i < b.N; i++ { + for j := 0; j < 1000; j++ { + s.Pop() + } + } +} + +func BenchmarkExist(b *testing.B) { + var testSet []int + for i := 0; i < 1000; i++ { + testSet = append(testSet, rand.Intn(2000)) + } + s := NewOrderedIntSet() + for _, elem := range testSet { + s.Push(elem) + } + b.ResetTimer() + for i := 0; i < b.N; i++ { + for j := 0; j < 1000; j++ { + s.Exists(j) + } + } +} + +func BenchmarkRemove(b *testing.B) { + var testSet []int + for i := 0; i < 1000; i++ { + testSet = append(testSet, rand.Intn(2000)) + } + s := NewOrderedIntSet() + for _, elem := range testSet { + s.Push(elem) + } + b.ResetTimer() + for i := 0; i < b.N; i++ { + for j := 0; j < 1000; j++ { + s.Remove(j) + } + } +}