proxy-pool/internal/domain/proxy/capacity.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
	reservationEnabled atomic.Bool
	reserving          atomic.Int64
	configuredObserver *activityObserver
	observer           atomic.Pointer[activityObserver]
}

type activityObserver struct{ notify func(bool) }

func NewCapacity(max int64) *Capacity {
	return NewCapacityWithActivityObserver(max, nil)
}

// NewCapacityWithActivityObserver reports successful zero-to-nonzero and
// nonzero-to-zero transitions. The observer must tolerate concurrent calls.
func NewCapacityWithActivityObserver(max int64, observer func(nonzero bool)) *Capacity {
	capacity := &Capacity{}
	if max < 0 || max > int64(counterMask) {
		max = 0
	}
	capacity.max.Store(uint32(max))
	capacity.reservationEnabled.Store(true)
	if observer != nil {
		capacity.configuredObserver = &activityObserver{notify: observer}
		capacity.observer.Store(capacity.configuredObserver)
	}
	return capacity
}

// SetActivityObservationEnabled lets snapshot ownership disable callbacks for
// current Proxies and enable them only while a runtime is retired and draining.
func (c *Capacity) SetActivityObservationEnabled(enabled bool) {
	if c == nil || c.configuredObserver == nil {
		return
	}
	if enabled {
		c.observer.Store(c.configuredObserver)
		return
	}
	c.observer.Store(nil)
}

// SetReservationEnabled gates new reservations without changing existing
// Active or Reserved counters. It is used while a snapshot retires a Proxy.
func (c *Capacity) SetReservationEnabled(enabled bool) {
	if c == nil {
		return
	}
	c.reservationEnabled.Store(enabled)
}

// Reclaimable reports whether a retired capacity can be forgotten without
// losing an in-flight reservation or draining runtime counter.
func (c *Capacity) Reclaimable() bool {
	if c == nil || c.reservationEnabled.Load() || c.reserving.Load() != 0 {
		return false
	}
	active, reserved, _ := c.Counters()
	return active == 0 && reserved == 0
}

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) Counters() (active, reserved, maximum int64) {
	if c == nil {
		return 0, 0, 0
	}
	packed := c.counters.Load()
	activeCounter, reservedCounter := unpack(packed)
	return int64(activeCounter), int64(reservedCounter), int64(c.max.Load())
}

func (c *Capacity) Reserve() (*Reservation, bool) {
	if c == nil || !c.beginReservation() {
		return nil, false
	}
	defer c.reserving.Add(-1)
	for {
		if !c.reservationEnabled.Load() {
			return nil, false
		}
		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) {
			if !c.reservationEnabled.Load() {
				c.cancel()
				return nil, false
			}
			if active+reserved == 0 {
				if observer := c.observer.Load(); observer != nil {
					observer.notify(true)
				}
			}
			return &Reservation{capacity: c}, true
		}
	}
}

func (c *Capacity) beginReservation() bool {
	if !c.reservationEnabled.Load() {
		return false
	}
	c.reserving.Add(1)
	if c.reservationEnabled.Load() {
		return true
	}
	c.reserving.Add(-1)
	return false
}

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 {
			return
		}
		next := pack(active, reserved-1)
		if c.counters.CompareAndSwap(current, next) {
			if active+reserved == 1 {
				if observer := c.observer.Load(); observer != nil {
					observer.notify(false)
				}
			}
			return
		}
	}
}

func (c *Capacity) release() {
	for {
		current := c.counters.Load()
		active, reserved := unpack(current)
		if active == 0 {
			return
		}
		next := pack(active-1, reserved)
		if c.counters.CompareAndSwap(current, next) {
			if active+reserved == 1 {
				if observer := c.observer.Load(); observer != nil {
					observer.notify(false)
				}
			}
			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
}