玩家匹配

team_maker.go

package example

import (
	"context"
	"errors"
	"sync/atomic"
	"time"

	"github.com/joyparty/gokit"
)

type teamMember[T any] struct {
	player    T
	ctx       context.Context
	cancelCtx context.CancelFunc
}

// TeamMakerOption 队伍匹配器选项
type TeamMakerOption[T any] struct {
	// 队伍人数
	Number int

	// 等待超时时间
	WaitTimeout time.Duration

	// 成功回调
	OnSuccess func(players []T)

	// 超时回调，可选
	OnTimeout func(player T)
}

// TeamMaker 队伍匹配
//
// T is player type
type TeamMaker[T any] struct {
	opt     TeamMakerOption[T]
	memberC chan *teamMember[T]
	done    chan struct{}
	waiting int32
}

// NewTeamMaker 构造队伍匹配器
func NewTeamMaker[T any](opt TeamMakerOption[T]) (*TeamMaker[T], error) {
	if opt.Number < 2 {
		return nil, errors.New("number must be greater than 1")
	} else if opt.WaitTimeout <= 0 {
		return nil, errors.New("wait timeout must be greater than 0")
	} else if opt.OnSuccess == nil {
		return nil, errors.New("OnSuccess must not be nil")
	}

	maker := &TeamMaker[T]{
		opt:     opt,
		memberC: make(chan *teamMember[T]),
		done:    make(chan struct{}),
	}

	go maker.run()
	return maker, nil
}

// Submit 提交匹配
//
// 调用返回的cancel即可放弃匹配
func (tm *TeamMaker[T]) Submit(player T) (cancel context.CancelFunc) {
	ctx, cancel := context.WithTimeout(context.Background(), tm.opt.WaitTimeout)

	member := &teamMember[T]{
		player:    player,
		ctx:       ctx,
		cancelCtx: cancel,
	}

	go func() {
		select {
		case tm.memberC <- member:
		default:
			select {
			case tm.memberC <- member:
			case <-tm.done:
				member.cancelCtx()
			case <-ctx.Done():
				tm.onTimeout(ctx.Err(), player)
			}
		}
	}()

	return cancel
}

// Waiting 正在等待的人数
func (tm *TeamMaker[T]) Waiting() int32 {
	return tm.waiting
}

func (tm *TeamMaker[T]) run() {
	team := gokit.NewListOf[*teamMember[T]]()

	defer func() {
		for el := team.Front(); el != nil; el = el.Next() {
			team.Remove(el)
			el.Value().cancelCtx()
		}
	}()

	// 删除超时或放弃的玩家
	removeInvalid := func() {
		for el := team.Front(); el != nil; el = el.Next() {
			if err := el.Value().ctx.Err(); err != nil {
				team.Remove(el)

				tm.onTimeout(err, el.Value().player)
			}
		}
	}

	handle := func(member *teamMember[T]) {
		removeInvalid()

		// 如果不是N缺1，就一起等待
		if team.Len() < tm.opt.Number-1 {
			team.PushBack(member)
			return
		}

		// 队伍匹配成功
		players := make([]T, 0, tm.opt.Number)
		for el := team.Front(); el != nil; el = el.Next() {
			team.Remove(el)
			el.Value().cancelCtx()

			players = append(players, el.Value().player)
		}

		member.cancelCtx()
		players = append(players, member.player)

		go tm.opt.OnSuccess(players)
	}

	for {
		if team.Len() == 0 {
			select {
			case <-tm.done:
				return
			case member := <-tm.memberC:
				handle(member)
			}
		} else {
			first := team.Front().Value()

			select {
			case <-tm.done:
				return
			case member := <-tm.memberC:
				handle(member)
			case <-first.ctx.Done():
				removeInvalid()
			}
		}

		atomic.StoreInt32(&tm.waiting, int32(team.Len()))
	}
}

// Close 关闭
func (tm *TeamMaker[T]) Close() {
	close(tm.done)
	close(tm.memberC)
}

func (tm *TeamMaker[T]) onTimeout(err error, player T) {
	if onTimeout := tm.opt.OnTimeout; onTimeout != nil && err == context.DeadlineExceeded {
		go onTimeout(player)
	}
}