135 lines
2.6 KiB
Go
135 lines
2.6 KiB
Go
package proxy
|
|
|
|
import (
|
|
"errors"
|
|
"sync/atomic"
|
|
)
|
|
|
|
const counterMask = uint64(1<<32 - 1)
|
|
|
|
var (
|
|
ErrReservationCommitted = errors.New("reservation is already committed")
|
|
ErrReservationFinished = errors.New("reservation is already finished")
|
|
)
|
|
|
|
type Capacity struct {
|
|
max atomic.Uint32
|
|
counters atomic.Uint64
|
|
}
|
|
|
|
func NewCapacity(max int64) *Capacity {
|
|
capacity := &Capacity{}
|
|
if max < 0 || max > int64(counterMask) {
|
|
max = 0
|
|
}
|
|
capacity.max.Store(uint32(max))
|
|
return capacity
|
|
}
|
|
|
|
func (c *Capacity) SetMax(max int64) bool {
|
|
if max < 0 || max > int64(counterMask) {
|
|
return false
|
|
}
|
|
c.max.Store(uint32(max))
|
|
return true
|
|
}
|
|
|
|
func (c *Capacity) Max() int64 { return int64(c.max.Load()) }
|
|
|
|
func (c *Capacity) Reserve() (*Reservation, bool) {
|
|
for {
|
|
current := c.counters.Load()
|
|
active, reserved := unpack(current)
|
|
if active+reserved >= c.max.Load() {
|
|
return nil, false
|
|
}
|
|
next := pack(active, reserved+1)
|
|
if c.counters.CompareAndSwap(current, next) {
|
|
return &Reservation{capacity: c}, true
|
|
}
|
|
}
|
|
}
|
|
|
|
func (c *Capacity) Active() int64 {
|
|
active, _ := unpack(c.counters.Load())
|
|
return int64(active)
|
|
}
|
|
|
|
func (c *Capacity) Reserved() int64 {
|
|
_, reserved := unpack(c.counters.Load())
|
|
return int64(reserved)
|
|
}
|
|
|
|
func (c *Capacity) commit() {
|
|
for {
|
|
current := c.counters.Load()
|
|
active, reserved := unpack(current)
|
|
if reserved == 0 {
|
|
return
|
|
}
|
|
if c.counters.CompareAndSwap(current, pack(active+1, reserved-1)) {
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
func (c *Capacity) cancel() {
|
|
for {
|
|
current := c.counters.Load()
|
|
active, reserved := unpack(current)
|
|
if reserved == 0 || c.counters.CompareAndSwap(current, pack(active, reserved-1)) {
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
func (c *Capacity) release() {
|
|
for {
|
|
current := c.counters.Load()
|
|
active, reserved := unpack(current)
|
|
if active == 0 || c.counters.CompareAndSwap(current, pack(active-1, reserved)) {
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
func pack(active, reserved uint32) uint64 {
|
|
return uint64(reserved)<<32 | uint64(active)
|
|
}
|
|
|
|
func unpack(value uint64) (active, reserved uint32) {
|
|
return uint32(value & counterMask), uint32(value >> 32)
|
|
}
|
|
|
|
type Reservation struct {
|
|
capacity *Capacity
|
|
state atomic.Uint32
|
|
}
|
|
|
|
func (r *Reservation) Commit() error {
|
|
if !r.state.CompareAndSwap(0, 1) {
|
|
if r.state.Load() == 1 {
|
|
return ErrReservationCommitted
|
|
}
|
|
return ErrReservationFinished
|
|
}
|
|
r.capacity.commit()
|
|
return nil
|
|
}
|
|
|
|
func (r *Reservation) Cancel() error {
|
|
if !r.state.CompareAndSwap(0, 2) {
|
|
return ErrReservationFinished
|
|
}
|
|
r.capacity.cancel()
|
|
return nil
|
|
}
|
|
|
|
func (r *Reservation) Release() error {
|
|
if !r.state.CompareAndSwap(1, 2) {
|
|
return ErrReservationFinished
|
|
}
|
|
r.capacity.release()
|
|
return nil
|
|
}
|