broad-announce/loadgen/internal/pacer/pacer.go

package pacer

import (
	"context"
	"time"
)

// Pacer produces a tick channel that controls send rate.
// It supports an optional linear ramp from 0 to target rate
// over a configurable warm-up window.
type Pacer struct {
	// Target alerts per second after ramp-up completes.
	TargetPerSec int
	// Ramp-up duration. 0 = instant (use full target rate immediately).
	RampDuration time.Duration
	// Clock is injectable for testing.
	Clock func() time.Time
}

// New returns a Pacer configured for targetPerSec sustained rate
// with an optional rampDuration warm-up window.
func New(targetPerSec int, rampDuration time.Duration) *Pacer {
	return &Pacer{
		TargetPerSec: targetPerSec,
		RampDuration: rampDuration,
		Clock:        time.Now,
	}
}

// Tick returns a channel that fires at the appropriate interval
// for the current phase (ramp or steady-state).
// The returned stop function blocks until the pacer goroutine exits.
// Callers must drain the returned channel to avoid blocking the
// pacer's internal goroutine.
func (p *Pacer) Tick(ctx context.Context) (<-chan struct{}, func()) {
	stopped := make(chan struct{})
	done := make(chan struct{})

	go func() {
		defer close(done)
		now := p.Clock()
		rampStart := now

		if p.TargetPerSec <= 0 || p.RampDuration == 0 {
			// Steady-state only: fixed interval.
			ticker := time.NewTicker(intervalFor(p.TargetPerSec))
			defer ticker.Stop()
			for {
				select {
				case <-ctx.Done():
					close(stopped)
					return
				case <-ticker.C:
					select {
					case stopped <- struct{}{}:
					default:
					}
				}
			}
		}

		// Ramp-up phase: rate increases linearly from 0 to TargetPerSec.
		// We recalculate the next tick interval after each tick.
		for {
			elapsed := p.Clock().Sub(rampStart)
			if elapsed >= p.RampDuration {
				break // fall through to steady-state
			}
			// Linear interpolation: fraction of ramp completed.
			fraction := float64(elapsed) / float64(p.RampDuration)
			currentRate := int(float64(p.TargetPerSec) * fraction)
			if currentRate <= 0 {
				currentRate = 1
			}
			tickAfter := intervalFor(currentRate)
			select {
			case <-ctx.Done():
				close(stopped)
				return
			case <-time.After(tickAfter):
				select {
				case stopped <- struct{}{}:
				default:
				}
			}
		}

		// Steady-state: fixed interval at full target rate.
		steady := time.NewTicker(intervalFor(p.TargetPerSec))
		defer steady.Stop()
		for {
			select {
			case <-ctx.Done():
				close(stopped)
				return
			case <-steady.C:
				select {
				case stopped <- struct{}{}:
				default:
				}
			}
		}
	}()

	return stopped, func() {
		<-done
	}
}

// intervalFor returns the inter-tick interval for a given rate.
// rate=0 returns a very slow ticker (will never fire in practice).
func intervalFor(rate int) time.Duration {
	if rate <= 0 {
		return time.Hour // effectively stopped
	}
	return time.Second / time.Duration(rate)
}