s_cherry_e/sr/internal/extend/maps/ordered_map.go

package maps

import "sync/atomic"

type (
	OrderedMap[K comparable, V any] struct {
		batchLoop int           // 批量处理循环次数
		writeChan chan any      // 写入缓冲区
		stopChan  chan struct{} // 停止信号
		closed    int32         // 原子关闭标记
	}

	queueCmdSet[K comparable, V any] struct {
		key   K
		value V
	}

	queueCmdPop[K comparable, V any] struct {
		limit  int
		result chan []V
	}

	queueCmdLen struct {
		result chan int
	}
	// internal data structure
	// orderQueue 封装了有序队列的底层实现，包含 Map 和 Slice 的管理
	orderQueue[K comparable, V any] struct {
		pairs map[K]V // 所有数据
		keys  []K     // 所有key（保证先进先出)
		head  int     // 记录当前有效数据的起始索引
	}
)

func newOrderQueue[K comparable, V any]() *orderQueue[K, V] {
	return &orderQueue[K, V]{
		pairs: make(map[K]V),
		keys:  make([]K, 0, 1024),
		head:  0,
	}
}

func (q *orderQueue[K, V]) set(key K, value V) {
	if _, found := q.pairs[key]; found {
		q.pairs[key] = value
	} else {
		q.keys = append(q.keys, key)
		q.pairs[key] = value
	}
}

func (q *orderQueue[K, V]) pop(limit int) []V {
	total := len(q.keys) - q.head
	if total <= 0 {
		return nil
	}

	batchSize := limit
	if batchSize > total {
		batchSize = total
	}

	if batchSize <= 0 {
		return nil
	}

	result := make([]V, 0, batchSize)
	for i := 0; i < batchSize; i++ {
		k := q.keys[q.head]
		v := q.pairs[k]
		result = append(result, v)

		delete(q.pairs, k)
		var zero K
		q.keys[q.head] = zero // free reference
		q.head++
	}

	// incremental compaction
	if q.head > 1024 && q.head > len(q.keys)/2 {
		newKeys := make([]K, len(q.keys)-q.head)
		copy(newKeys, q.keys[q.head:])
		q.keys = newKeys
		q.head = 0
	}

	return result
}

func (q *orderQueue[K, V]) len() int {
	return len(q.pairs)
}

func NewOrderMap[K comparable, V any](chanSize, batchLoop int) *OrderedMap[K, V] {
	if chanSize <= 0 {
		chanSize = 10000
	}

	if batchLoop <= 0 {
		batchLoop = 10
	}

	om := &OrderedMap[K, V]{
		batchLoop: batchLoop,
		writeChan: make(chan any, chanSize),
		stopChan:  make(chan struct{}),
	}

	go om.loop()
	return om
}

func (om *OrderedMap[K, V]) loop() {
	queue := newOrderQueue[K, V]()

	// processCmd 处理单个命令
	processCmd := func(cmd any) {
		switch c := cmd.(type) {
		case *queueCmdSet[K, V]:
			queue.set(c.key, c.value)

		case *queueCmdPop[K, V]:
			c.result <- queue.pop(c.limit)

		case *queueCmdLen:
			c.result <- queue.len()
		}
	}

	for {
		// 1. 阻塞等待：如果没有任务，在这里让出 CPU
		select {
		case <-om.stopChan:
			return
		case cmd := <-om.writeChan:
			// 收到第一个任务，立即处理
			processCmd(cmd)
		}

		// 2. 批量消费：唤醒后，贪婪地消费所有积压任务
		// 设置一个最大批量限制，防止单个协程长期霸占 CPU
	BatchLoop:
		for i := 0; i < om.batchLoop; i++ {
			select {
			case cmd := <-om.writeChan:
				processCmd(cmd)
			case <-om.stopChan:
				return
			default:
				// Channel 空了，回到外层阻塞
				break BatchLoop
			}
		}
	}
}

func (om *OrderedMap[K, V]) isClosed() bool {
	return atomic.LoadInt32(&om.closed) == 1
}

// Set 是异步的 (fire and forget)
func (om *OrderedMap[K, V]) Set(key K, value V) {
	if om.isClosed() {
		return
	}
	// 直接发送，如果满则阻塞（背压）
	om.writeChan <- &queueCmdSet[K, V]{key: key, value: value}
}

func (om *OrderedMap[K, V]) Len() int {
	if om.isClosed() {
		return 0
	}
	resultChan := make(chan int, 1)
	om.writeChan <- &queueCmdLen{result: resultChan}
	return <-resultChan
}

// PopBatch 高效批量弹出
func (om *OrderedMap[K, V]) PopBatch(limit int) []V {
	if limit <= 0 || om.isClosed() {
		return nil
	}
	resultChan := make(chan []V, 1)
	om.writeChan <- &queueCmdPop[K, V]{limit: limit, result: resultChan}
	return <-resultChan
}

func (om *OrderedMap[K, V]) Pop() (V, bool) {
	batch := om.PopBatch(1)
	if len(batch) > 0 {
		return batch[0], true
	}
	var zero V
	return zero, false
}

func (om *OrderedMap[K, V]) Close() {
	if !atomic.CompareAndSwapInt32(&om.closed, 0, 1) {
		return
	}
	close(om.stopChan)
}

func (om *OrderedMap[K, V]) GetStopChan() <-chan struct{} {
	return om.stopChan
}