Samuel McBroom
97 lines
2.2 KiB
Go
97 lines
2.2 KiB
Go
package server
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/sirupsen/logrus"
|
|
)
|
|
|
|
type IDownScaler interface {
|
|
Reset()
|
|
Begin(serverAddress string)
|
|
Cancel(serverAddress string)
|
|
}
|
|
|
|
var DownScaler IDownScaler
|
|
|
|
func NewDownScaler(ctx context.Context, enabled bool, delay time.Duration) IDownScaler {
|
|
ds := &downScalerImpl{
|
|
enabled: enabled,
|
|
delay: delay,
|
|
parentContext: ctx,
|
|
contextCancellations: make(map[string]context.CancelFunc),
|
|
}
|
|
|
|
return ds
|
|
}
|
|
|
|
type downScalerImpl struct {
|
|
sync.RWMutex
|
|
enabled bool
|
|
delay time.Duration
|
|
parentContext context.Context
|
|
contextCancellations map[string]context.CancelFunc
|
|
}
|
|
|
|
func (ds *downScalerImpl) Reset() {
|
|
// Cancel all existing scale down routines
|
|
for _, scaleDownCancel := range ds.contextCancellations {
|
|
scaleDownCancel()
|
|
}
|
|
ds.contextCancellations = make(map[string]context.CancelFunc)
|
|
}
|
|
|
|
func (ds *downScalerImpl) Begin(serverAddress string) {
|
|
ds.Lock()
|
|
defer ds.Unlock()
|
|
|
|
if !ds.enabled {
|
|
return
|
|
}
|
|
|
|
// If an existing scale down routine exists, cancel it
|
|
if scaleDownCancel, ok := ds.contextCancellations[serverAddress]; ok {
|
|
scaleDownCancel()
|
|
}
|
|
|
|
logrus.WithField("serverAddress", serverAddress).Debug("Beginning scale down")
|
|
scaleDownContext, scaleDownContextCancellation := context.WithCancel(ds.parentContext)
|
|
ds.contextCancellations[serverAddress] = scaleDownContextCancellation
|
|
go ds.scaleDown(scaleDownContext, serverAddress)
|
|
}
|
|
|
|
func (ds *downScalerImpl) Cancel(serverAddress string) {
|
|
ds.Lock()
|
|
defer ds.Unlock()
|
|
|
|
if !ds.enabled {
|
|
return
|
|
}
|
|
|
|
if scaleDownContextCancellation, ok := ds.contextCancellations[serverAddress]; ok {
|
|
logrus.WithField("serverAddress", serverAddress).Debug("Canceling scale down")
|
|
scaleDownContextCancellation()
|
|
delete(ds.contextCancellations, serverAddress)
|
|
}
|
|
}
|
|
|
|
func (ds *downScalerImpl) scaleDown(ctx context.Context, serverAddress string) {
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case <-time.After(ds.delay):
|
|
_, _, _, sleeper := Routes.FindBackendForServerAddress(ctx, serverAddress)
|
|
if sleeper == nil {
|
|
return
|
|
}
|
|
if err := sleeper(ctx); err != nil {
|
|
logrus.WithField("serverAddress", serverAddress).WithError(err).Error("failed to scale down backend")
|
|
}
|
|
return
|
|
}
|
|
}
|
|
}
|