feat 排队

internal/handler/ws_handler/login/queue_up.go

@@ -0,0 +1,104 @@
+package login
+
+import (
+	"context"
+	"errors"
+	"gateway/internal/handler/ws_handler/client"
+	"sync"
+	"sync/atomic"
+)
+
+var queueUp *QueueUp
+
+// QueueUp 排队队列结构
+type QueueUp struct {
+	queue     chan *QueueUser // 用户队列
+	waiting   sync.Map        // map[usn]*QueueUser
+	minTicket atomic.Int64    // 最小的Ticket
+	maxTicket atomic.Int64    // 最大的Ticket
+	ctx       context.Context
+	cancel    context.CancelFunc
+}
+
+type QueueUser struct {
+	Cli    *client.Client
+	Ticket int64
+}
+
+func NewQueueUp(maxSize int) *QueueUp {
+	ctx, cancel := context.WithCancel(context.Background())
+	queueUp = &QueueUp{
+		queue:  make(chan *QueueUser, maxSize),
+		ctx:    ctx,
+		cancel: cancel,
+	}
+	return queueUp
+}
+
+func GetQueueUp() *QueueUp {
+	return queueUp
+}
+
+// Enqueue 将用户加入排队队列
+func (q *QueueUp) Enqueue(cli *client.Client) error {
+	select {
+	case <-q.ctx.Done():
+		return errors.New("queue stopped")
+	default:
+	}
+
+	ticket := q.maxTicket.Add(1)
+	item := &QueueUser{Cli: cli, Ticket: ticket}
+
+	select {
+	case q.queue <- item:
+		q.waiting.Store(cli.USN, item)
+		return nil
+	default:
+		return errors.New("queue is full")
+	}
+}
+
+// Dequeue 从排队队列中取出下一个有效用户
+func (q *QueueUp) Dequeue() (*client.Client, error) {
+	select {
+	case item, ok := <-q.queue:
+		if ok {
+			q.minTicket.Store(item.Ticket)
+			if _, loaded := q.waiting.LoadAndDelete(item.Cli.USN); loaded {
+				return item.Cli, nil
+			}
+			return q.Dequeue()
+		}
+	case <-q.ctx.Done():
+		return nil, q.ctx.Err()
+	default:
+	}
+	return nil, nil
+}
+
+// GetPosition 返回用户前面还有多少人在排队
+func (q *QueueUp) GetPosition(usn int64) (int64, bool) {
+	val, ok := q.waiting.Load(usn)
+	if !ok {
+		return 0, false
+	}
+	user := val.(*QueueUser)
+	return user.Ticket - q.minTicket.Load() - 1, true
+}
+
+// RemoveUser 安全移除用户（标记为取消）
+func (q *QueueUp) RemoveUser(usn int64) bool {
+	_, loaded := q.waiting.LoadAndDelete(usn)
+	return loaded
+}
+
+// GetQueueSize 获取当前排队人数
+func (q *QueueUp) GetQueueSize() int {
+	return len(q.queue)
+}
+
+// Stop 停止整个队列服务
+func (q *QueueUp) Stop() {
+	q.cancel()
+}