package entity

import (
	"context"
	"sync"
	"time"
)

//并发限制

type reqLimit struct {
	doPool   chan struct{}
	waitPool chan struct{}
}

var (
	ReqLimitInit *reqLimit
	ReqLimitLock = &sync.Mutex{}
)

func NewLimit(num int) {
	doPool := make(chan struct{}, num)
	for i := 0; i < num; i++ {
		doPool <- struct{}{}
	}
	waitPool := make(chan struct{}, num)
	for i := 0; i < num; i++ {
		waitPool <- struct{}{}
	}
	ReqLimitInit = &reqLimit{
		doPool:   doPool,
		waitPool: waitPool,
	}
}

// -2 等待池已满
// -1 超时
// 1：可以执行查询
func (r *reqLimit) Limit(ctx context.Context) int {
	ctx, cancel := context.WithTimeout(ctx, 30*time.Second) //30秒
	defer cancel()
	select {
	case <-r.waitPool:
		defer func() {
			r.waitPool <- struct{}{}
		}()
		select {
		case <-r.doPool:
			return 1
		case <-ctx.Done(): //超时
			return -1
		}
	default:
		return -2
	}
}

func (r *reqLimit) Release() {
	r.doPool <- struct{}{}
}