fast and safe sort slice like queue

sortslice_list.go

package sortslice

import (
	"container/heap"
	"sort"
	"sync"
)

type Item interface {
	Compare(other Item) int
	Key() interface{}
}

type ItemsHeap []Item

func (p ItemsHeap) Len() int {
	return len(p)
}
func (p ItemsHeap) Less(i, j int) bool {
	// We want Pop to give us the highest, not lowest, priority so we use greater than here.
	return p[i].Compare(p[j]) > 0
}
func (p ItemsHeap) Swap(i, j int) {
	p[i], p[j] = p[j], p[i]
}

func (p *ItemsHeap) Push(x interface{}) {
	*p = append(*p, x.(Item))
}
func (p ItemsHeap) Get(number int) []Item {
	if number > p.Len() {
		return p[:]
	}
	return p[:number]
}

func (h *ItemsHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

type List struct {
	items    ItemsHeap
	itemKeys map[interface{}]struct{}
	lock     sync.RWMutex
}

func (l *List) Add(item Item) bool {
	l.lock.Lock()
	defer l.lock.Unlock()

	k := item.Key()
	if _, ok := l.itemKeys[k]; ok {
		return false
	}
	heap.Push(&l.items, item)
	l.itemKeys[k] = struct{}{}
	return true
}
func (l *List) Get(index int) Item {
	l.lock.RLock()
	defer l.lock.RUnlock()

	if l.items.Len()-1 < index {
		return nil
	}
	return l.items[index]
}

func (l *List) List() []Item {
	if l.items.Len() == 0 {
		return nil
	}
	l.lock.Lock()
	defer l.lock.Unlock()
	sort.Sort(l.items)

	cpy := make([]Item, l.items.Len())
	copy(cpy, l.items)
	return cpy
}
func (l *List) Range(f func(item Item) bool) {
	l.lock.Lock()
	defer l.lock.Unlock()
	for i := 0; i < l.items.Len(); i++ {
		if !f(l.items[i]) {
			break
		}
	}
}

func (l *List) Pop() Item {
	if l.items.Len() == 0 {
		return nil
	}
	l.lock.Lock()
	defer l.lock.Unlock()
	if l.items.Len() == 0 {
		return nil
	}

	return l.pop()
}

func (l *List) pop() Item {
	top := heap.Pop(&l.items)
	delete(l.itemKeys, top.(Item).Key())
	return top.(Item)
}
func (l *List) Exist(key interface{}) bool {
	if l.items.Len() == 0 {
		return false
	}
	l.lock.RLock()
	_, ok := l.itemKeys[key]
	l.lock.RUnlock()
	return ok
}
func (l *List) Remove(key interface{}) bool {
	if l.items.Len() == 0 {
		return false
	}
	l.lock.Lock()
	defer l.lock.Unlock()
	if l.items.Len() == 0 {
		return false
	}

	if _, ok := l.itemKeys[key]; !ok {
		return false
	}

	//遍历删除
	for i := 0; i < l.items.Len(); i++ {
		if l.items[i].Key() == key {
			l.remove(i)
			return true
		}
	}
	return false
}

//移除指定位置的项，这里是根据情况采用高效算法。
func (l *List) remove(index int) {
	//find,如果刚好是在第一个，则直接取出
	if index == 0 {
		l.pop()
		return
	}
	//如果不是末尾单元，则将最后一个单元和当前位置交换
	//且不需要重新排序
	lastIndex := l.items.Len() - 1
	if index != lastIndex {
		l.items.Swap(index, lastIndex)
	}
	//此时需要删除的项，在末尾
	item := l.items[lastIndex]
	delete(l.itemKeys, item.Key())
	l.items[lastIndex] = nil      //释放
	l.items = l.items[:lastIndex] //截断掉最后一项
}

func (l *List) Len() int {
	l.lock.RLock()
	defer l.lock.RUnlock()
	return l.items.Len()
}

func (l *List) Empty() bool {
	return l.Len() == 0
}

func (l *List) Release() {
	l.lock.Lock()
	defer l.lock.Unlock()
	for i := 0; i < l.items.Len(); i++ {
		l.items[i] = nil
	}
	for k := range l.itemKeys {
		delete(l.itemKeys, k)
	}
	l.items = l.items[:0]
}

func NewList(hint int) *List {
	list := List{
		items:    make(ItemsHeap, 0, hint),
		itemKeys: make(map[interface{}]struct{}, hint),
	}
	heap.Init(&list.items)
	return &list
}

sortslice_list_test.go

package sortslice

import (
	"fmt"
	"math/rand"
	"sync"
	"testing"

	"github.com/stretchr/testify/require"
)

func TestList_Add(t *testing.T) {
	list := NewList(10)

	items := []int{5, 4, 2}

	for _, v := range items {
		ok := list.Add(mockItem(v))
		require.True(t, ok)
	}
	require.Equal(t, len(items), list.Len())

	//不允许重复
	t.Run("ignoreRepeat", func(t *testing.T) {
		for _, v := range items {
			ok := list.Add(mockItem(v))
			require.False(t, ok)
		}
	})
}

func TestList_Pop(t *testing.T) {
	list := NewList(10)

	items := []int{5, 4, 2, 6, 10}
	for _, v := range items {
		list.Add(mockItem(v))
	}
	require.Equal(t, mockItem(10), list.Pop())
	require.Equal(t, mockItem(6), list.Pop())
	require.Equal(t, mockItem(5), list.Pop())
	require.Equal(t, mockItem(4), list.Pop())
	require.Equal(t, mockItem(2), list.Pop())

	require.True(t, list.Empty())
	require.Nil(t, list.Pop(), "should be nil when is empty")
}

func TestList_Remove(t *testing.T) {
	list := NewList(10)
	items := []int{5, 4, 2, 6, 10}
	for _, v := range items {
		list.Add(mockItem(v))
	}

	removeOrders := []struct {
		remove int
		top    int
	}{
		{5, 10},
		{10, 6},
		{2, 6},
		{6, 4},
	}

	for _, c := range removeOrders {
		ok := list.Remove(fmt.Sprintf("key-%d", c.remove))
		require.True(t, ok)
		require.Equal(t, mockItem(c.top), list.Get(0))
	}
}

func TestList_Penter(t *testing.T) {
	list := NewList(10)

	randV := func() int {
		return rand.Intn(10000)
	}
	randItem := func() mockItem {
		return mockItem(randV())
	}
	//随机写、读、删
	var (
		wg   sync.WaitGroup
		quit = make(chan struct{})
		add  int
		get  int
		pop  int
	)

	wg.Add(1000)

	//写
	go func() {
		for {
			select {
			case <-quit:
				return
			default:
				if list.Add(randItem()) {
					add++
				}
			}
		}
	}()
	//读
	go func() {
		for {
			select {
			case <-quit:
				return
			default:
				if list.Get(randV()) != nil {
					get++
				}
			}
		}
	}()
	//删
	go func() {
		var removed int
		for {
			select {
			case <-quit:
				return
			default:
				ok := list.Remove(randItem().Key())
				if ok {
					removed++
					wg.Done()
				}
			}
		}
	}()
	//取
	go func() {
		for {
			select {
			case <-quit:
				return
			default:
				if list.Pop() != nil {
					pop++
				}
			}
		}

	}()

	wg.Wait()
	close(quit)

	t.Logf("add:%d,get:%d,pop:%d", add, get, pop)
}

func BenchmarkList_Pop(b *testing.B) {
	list := NewList(1000)
	for i := 0; i < b.N; i++ {
		list.Add(mockItem(rand.Intn(b.N)))
	}

	var wait sync.WaitGroup
	wait.Add(1)
	go func() {
		wait.Done()
		for i := 0; i < b.N; i++ {
			list.Add(mockItem(rand.Intn(b.N)))
		}
	}()
	wait.Wait()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		list.Pop()
	}
}

sortslice_mock_test.go

package sortslice

import "fmt"

type mockItem int

func (m mockItem) Key() interface{} {
	return fmt.Sprintf("key-%d", m)
}

func (mi mockItem) Compare(other Item) int {
	omi := other.(mockItem)
	if mi > omi {
		return 1
	} else if mi == omi {
		return 0
	}
	return -1
}

sortslice_pool.go

package sortslice

import "sync"

var defPool = sync.Pool{
	New: func() interface{} {
		return NewList(0)
	},
}

func Get() *List { return defPool.Get().(*List) }
func Put(list *List) {
	list.Release()
	defPool.Put(list)
}