tinymq/mapx.go

package tinymq

/* 优化连接方式

- 高效匹配字符串算法
- 随机匹配算法
- 排序算法
- 支持万级连接数
*/

import (
	"math/rand"
	"sync"
)

type Mapx struct {
	sync.RWMutex
	id        int // 生成 id 号
	lineMap   map[int]*Line
	randLines []*Line
}

func NewMapx() *Mapx {
	return &Mapx{
		lineMap:   make(map[int]*Line),
		randLines: make([]*Line, 0),
	}
}

// 存储连接
func (m *Mapx) Store(line *Line) {
	m.Lock()
	defer m.Unlock()
	m.id++
	m.lineMap[m.id] = line
	m.randLines = append(m.randLines, line)
	if len(m.randLines) < 3 {
		return
	}
	// 交换一下位置，随机数组
	i := rand.Intn(len(m.randLines) - 1)
	t := m.randLines[i]
	m.randLines[i] = line
	m.randLines[len(m.randLines)-1] = t
}

// 删除连接
func (m *Mapx) Delete(line *Line) {
	m.Lock()
	defer m.Unlock()
	for k, v := range m.lineMap {
		if v == line {
			delete(m.lineMap, k)
			break
		}
	}
	for i, v := range m.randLines {
		if v == line {
			l := len(m.randLines) - 1
			m.randLines[i] = m.randLines[l]
			m.randLines = m.randLines[0:l]
			break
		}
	}
}

// ID 方式删除连接
func (m *Mapx) DeleteById(id int) {
	m.Lock()
	defer m.Unlock()
	line, ok := m.lineMap[id]
	if !ok {
		return
	}
	delete(m.lineMap, id)
	for i, v := range m.randLines {
		if v == line {
			l := len(m.randLines) - 1
			m.randLines[i] = m.randLines[l]
			m.randLines = m.randLines[0:l]
			break
		}
	}
}

// 删除不可用连接
func (m *Mapx) DeleteInvalidLines(expired int64) {
	m.Lock()
	defer m.Unlock()
	for id, line := range m.lineMap {
		if line.state != Connected && line.updated.UnixMilli() < expired {
			delete(m.lineMap, id)
		}
	}
	for i := len(m.randLines) - 1; i >= 0; i-- {
		line := m.randLines[i]
		if line.state != Connected && line.updated.UnixMilli() < expired {
			l := len(m.randLines) - 1
			m.randLines[i] = m.randLines[l]
			m.randLines = m.randLines[0:l]
		}
	}
}

// 循环遍历
func (m *Mapx) Range(fn func(id int, line *Line) bool) {
	m.RLock()
	defer m.RUnlock()
	for k, v := range m.lineMap {
		if !fn(k, v) {
			break
		}
	}
}

// 打乱随机数组
func (m *Mapx) Rand() {
	m.Lock()
	defer m.Unlock()
	l := len(m.randLines) / 2
	for i := range l {
		k := rand.Intn(l) + l
		v := m.randLines[i]
		m.randLines[i] = m.randLines[k]
		m.randLines[k] = v
	}
}

// 随机遍历
func (m *Mapx) RandRange(fn func(i int, v *Line) bool, rand bool) {
	if rand {
		m.Rand()
	}
	m.RLock()
	defer m.RUnlock()
	for i, v := range m.randLines {
		if !fn(i, v) {
			break
		}
	}
}

// 是否存在
func (m *Mapx) Exist(line *Line) bool {
	for _, v := range m.lineMap {
		if v == line {
			return true
		}
	}
	return false
}