slot_allocator.go

package bmdb

// SlotAllocator is a small deterministic in-memory allocator for fixed-count
// slot tokens. It is safe for concurrent use by multiple goroutines. The
// implementation uses a buffered channel as the backing store which gives a
// simple, deterministic FIFO-ish behavior and fast non-blocking acquire
// attempts when tokens are available.
type SlotAllocator struct {
	ch chan struct{}
}

// NewSlotAllocator creates an allocator with capacity slots. If cap <= 0,
// it returns nil to indicate no allocation is enforced.
func NewSlotAllocator(cap int64) *SlotAllocator {
	if cap <= 0 {
		return nil
	}
	a := &SlotAllocator{ch: make(chan struct{}, int(cap))}
	// fill channel with tokens
	for i := int64(0); i < cap; i++ {
		a.ch <- struct{}{}
	}
	return a
}

// Acquire attempts to take n tokens from the allocator without blocking.
// If there aren't enough tokens available immediately, it restores any
// partially-taken tokens and returns false.
func (a *SlotAllocator) Acquire(n int64) bool {
	if a == nil || n <= 0 {
		return false
	}
	taken := int64(0)
	for i := int64(0); i < n; i++ {
		select {
		case <-a.ch:
			taken++
		default:
			// not enough tokens: restore what we took and fail
			if taken > 0 {
				for j := int64(0); j < taken; j++ {
					a.ch <- struct{}{}
				}
			}
			return false
		}
	}
	return true
}

// Release returns n tokens back to the allocator. It will not block but if
// more tokens are released than the capacity allows the extra releases will
// be discarded (this should not normally happen if callers properly match
// Acquire/Release semantics).
func (a *SlotAllocator) Release(n int64) {
	if a == nil || n <= 0 {
		return
	}
	for i := int64(0); i < n; i++ {
		select {
		case a.ch <- struct{}{}:
		default:
			// channel full: discard extra
			return
		}
	}
}

// Available returns the number of tokens currently available.
func (a *SlotAllocator) Available() int64 {
	if a == nil {
		return 0
	}
	return int64(len(a.ch))
}