142 lines
3.8 KiB
Go
142 lines
3.8 KiB
Go
package prizepool
|
|
|
|
import (
|
|
"context"
|
|
"crypto/rand"
|
|
"errors"
|
|
"math/big"
|
|
)
|
|
|
|
var (
|
|
// ErrNoAvailableItem means every configured item was filtered out by switch,
|
|
// weight, total limit, or user limit.
|
|
ErrNoAvailableItem = errors.New("prizepool: no available item")
|
|
|
|
// ErrUserIssueCounterRequired means at least one item has a user limit, but no
|
|
// counter was provided to check how many times the user has already received it.
|
|
ErrUserIssueCounterRequired = errors.New("prizepool: user issue counter required")
|
|
)
|
|
|
|
// Item is the activity-neutral shape used by the prize pool picker.
|
|
// TotalLimit/UserLimit use 0 to mean unlimited.
|
|
type Item struct {
|
|
ID int64
|
|
Enabled bool
|
|
Weight int
|
|
TotalLimit int64
|
|
UserLimit int64
|
|
IssuedCount int64
|
|
}
|
|
|
|
// PickResult returns the selected item and its original index in the input slice.
|
|
type PickResult struct {
|
|
Item Item
|
|
Index int
|
|
}
|
|
|
|
// UserIssueCounter returns how many cap-occupying records one user already has
|
|
// for a prize item.
|
|
type UserIssueCounter interface {
|
|
CountUserIssues(ctx context.Context, itemID int64, userID int64) (int64, error)
|
|
}
|
|
|
|
// UserIssueCounterFunc adapts a function to UserIssueCounter.
|
|
type UserIssueCounterFunc func(ctx context.Context, itemID int64, userID int64) (int64, error)
|
|
|
|
func (f UserIssueCounterFunc) CountUserIssues(ctx context.Context, itemID int64, userID int64) (int64, error) {
|
|
return f(ctx, itemID, userID)
|
|
}
|
|
|
|
// RandomIntn returns a number in [0, max).
|
|
type RandomIntn func(max int) (int, error)
|
|
|
|
// Picker owns the randomness used by the prize pool. A zero-value Picker uses
|
|
// crypto/rand so production callers do not need to seed a global RNG.
|
|
type Picker struct {
|
|
RandomIntn RandomIntn
|
|
}
|
|
|
|
// PickAvailable filters unavailable items and selects one by Weight.
|
|
func (p Picker) PickAvailable(ctx context.Context, items []Item, userID int64, counter UserIssueCounter) (PickResult, error) {
|
|
available := make([]PickResult, 0, len(items))
|
|
for index, item := range items {
|
|
if !item.Enabled || item.Weight <= 0 {
|
|
continue
|
|
}
|
|
if item.TotalLimit > 0 && item.IssuedCount >= item.TotalLimit {
|
|
continue
|
|
}
|
|
if item.UserLimit > 0 {
|
|
if counter == nil {
|
|
return PickResult{}, ErrUserIssueCounterRequired
|
|
}
|
|
used, err := counter.CountUserIssues(ctx, item.ID, userID)
|
|
if err != nil {
|
|
return PickResult{}, err
|
|
}
|
|
if used >= item.UserLimit {
|
|
continue
|
|
}
|
|
}
|
|
available = append(available, PickResult{Item: item, Index: index})
|
|
}
|
|
if len(available) == 0 {
|
|
return PickResult{}, ErrNoAvailableItem
|
|
}
|
|
return p.PickWeighted(available)
|
|
}
|
|
|
|
// PickWeighted selects one already-filtered item by Weight.
|
|
func (p Picker) PickWeighted(items []PickResult) (PickResult, error) {
|
|
total := 0
|
|
for _, item := range items {
|
|
if item.Item.Enabled && item.Item.Weight > 0 {
|
|
total += item.Item.Weight
|
|
}
|
|
}
|
|
if total <= 0 {
|
|
return PickResult{}, ErrNoAvailableItem
|
|
}
|
|
|
|
n, err := p.randomIntn(total)
|
|
if err != nil {
|
|
return PickResult{}, err
|
|
}
|
|
target := n + 1
|
|
accumulated := 0
|
|
for _, item := range items {
|
|
if !item.Item.Enabled || item.Item.Weight <= 0 {
|
|
continue
|
|
}
|
|
accumulated += item.Item.Weight
|
|
if target <= accumulated {
|
|
return item, nil
|
|
}
|
|
}
|
|
return PickResult{}, ErrNoAvailableItem
|
|
}
|
|
|
|
func (p Picker) randomIntn(max int) (int, error) {
|
|
if max <= 0 {
|
|
return 0, ErrNoAvailableItem
|
|
}
|
|
if p.RandomIntn != nil {
|
|
return p.RandomIntn(max)
|
|
}
|
|
n, err := rand.Int(rand.Reader, big.NewInt(int64(max)))
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
return int(n.Int64()), nil
|
|
}
|
|
|
|
// PickAvailable uses a zero-value Picker.
|
|
func PickAvailable(ctx context.Context, items []Item, userID int64, counter UserIssueCounter) (PickResult, error) {
|
|
return Picker{}.PickAvailable(ctx, items, userID, counter)
|
|
}
|
|
|
|
// PickWeighted uses a zero-value Picker.
|
|
func PickWeighted(items []PickResult) (PickResult, error) {
|
|
return Picker{}.PickWeighted(items)
|
|
}
|