throttling

type BandwidthThrottler interface {
	// Blocks until [nodeID] can read a message of size [msgSize].
	// AddNode([nodeID], ...) must have been called since
	// the last time RemoveNode([nodeID]) was called, if any.
	// It's safe for multiple goroutines to concurrently call Acquire.
	Acquire(msgSize uint64, nodeID ids.ShortID)

	// Add a new node to this throttler.
	// Must be called before Acquire(..., [nodeID]) is called.
	// RemoveNode([nodeID]) must have been called since the last time
	// AddNode([nodeID], ...) was called, if any.
	// Its bandwidth allocation refills at a rate of [refillRate].
	// Its bandwidth allocation can hold up to [maxBurstSize] at a time.
	// [maxBurstSize] must be at least the maximum message size.
	// It's safe for multiple goroutines to concurrently call AddNode.
	AddNode(nodeID ids.ShortID)

	// Remove a node from this throttler.
	// AddNode([nodeID], ...) must have been called since
	// the last time RemoveNode([nodeID]) was called, if any.
	// Must be called when we stop reading messages from [nodeID].
	// It's safe for multiple goroutines to concurrently call RemoveNode.
	RemoveNode(nodeID ids.ShortID)
}

type BandwidthThrottlerConfig struct {
	// Rate at which the inbound bandwidth consumable by a peer replenishes
	RefillRate uint64 `json:"bandwidthRefillRate"`
	// Max amount of consumable bandwidth that can accumulate for a given peer
	MaxBurstSize uint64 `json:"bandwidthMaxBurstRate"`
}

type DialThrottler interface {
	// Block until the event associated with this Acquire can happen.
	// If [ctx] is canceled, gives up and returns an error.
	Acquire(ctx context.Context) error
}

type InboundConnUpgradeThrottler interface {
	// Dispatch starts this InboundConnUpgradeThrottler.
	// Must be called before [ShouldUpgrade].
	// Blocks until [Stop] is called (i.e. should be called in a goroutine.)
	Dispatch()
	// Stop this InboundConnUpgradeThrottler and causes [Dispatch] to return.
	// Should be called when we're done with this InboundConnUpgradeThrottler.
	// This InboundConnUpgradeThrottler must not be used after [Stop] is called.
	Stop()
	// Returns whether we should upgrade an inbound connection from [ipStr].
	// Must only be called after [Dispatch] has been called.
	// If [ip] is a local IP, this method always returns true.
	// Must not be called after [Stop] has been called.
	ShouldUpgrade(ip utils.IPDesc) bool
}

type InboundConnUpgradeThrottlerConfig struct {
	// ShouldUpgrade(ipStr) returns true if it has been at least [UpgradeCooldown]
	// since the last time ShouldUpgrade(ipStr) returned true or if
	// ShouldUpgrade(ipStr) has never been called.
	// If <= 0, inbound connections not rate-limited.
	UpgradeCooldown time.Duration `json:"upgradeCooldown"`
	// Maximum number of inbound connections upgraded within [UpgradeCooldown].
	// (As implemented in inboundConnUpgradeThrottler, may actually upgrade
	// [MaxRecentConnsUpgraded+1] due to a race condition but that's fine.)
	// If <= 0, inbound connections not rate-limited.
	MaxRecentConnsUpgraded int `json:"maxRecentConnsUpgraded"`
}

type InboundMsgThrottler interface {
	BandwidthThrottler

	// Blocks until we can read a message of size [msgSize] from [nodeID].
	// For every call to Acquire([msgSize], [nodeID]), we must (!) call
	// Release([msgSize], [nodeID]) when done processing the message
	// (or when we give up trying to read the message.)
	Acquire(msgSize uint64, nodeID ids.ShortID)

	// Mark that we're done processing a message of size [msgSize]
	// from [nodeID].
	Release(msgSize uint64, nodeID ids.ShortID)
}

type InboundMsgThrottlerConfig struct {
	MsgByteThrottlerConfig
	BandwidthThrottlerConfig
	MaxProcessingMsgsPerNode uint64 `json:"maxProcessingMsgsPerNode"`
}

type MsgByteThrottlerConfig struct {
	VdrAllocSize        uint64 `json:"vdrAllocSize"`
	AtLargeAllocSize    uint64 `json:"atLargeAllocSize"`
	NodeMaxAtLargeBytes uint64 `json:"nodeMaxAtLargeBytes"`
}

type OutboundMsgThrottler interface {
	// Returns true if we can queue a message of size [msgSize] to be sent to node [nodeID].
	// Returns false if the message should be dropped (not sent to [nodeID]).
	// If this method returns true, Release([msgSize], [nodeID]) must be called (!) when
	// the message is sent (or when we give up trying to send the message, if applicable.)
	// If this method returns false, do not make a corresponding call to Release.
	Acquire(msgSize uint64, nodeID ids.ShortID) bool

	// Mark that a message of size [msgSize] has been sent to [nodeID] or we have
	// given up sending the message. Must correspond to a previous call to
	// Acquire([msgSize], [nodeID]) that returned true.
	Release(msgSize uint64, nodeID ids.ShortID)
}