This repository has been archived by the owner on Oct 9, 2023. It is now read-only.
/
controller.go
64 lines (52 loc) · 1.84 KB
/
controller.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
package backoff
import (
"context"
"fmt"
"time"
stdAtomic "github.com/lyft/flytestdlib/atomic"
"github.com/lyft/flytestdlib/logger"
"github.com/lyft/flyteplugins/go/tasks/pluginmachinery/k8s"
v1 "k8s.io/api/core/v1"
"k8s.io/apimachinery/pkg/util/clock"
)
// Controller is a name-spaced collection of back-off handlers
type Controller struct {
// Controller.Clock allows the use of fake clock when testing
Clock clock.Clock
backOffHandlerMap HandlerMap
}
func (m *Controller) GetOrCreateHandler(ctx context.Context, key string, backOffBaseSecond int, maxBackOffDuration time.Duration) *ComputeResourceAwareBackOffHandler {
h, loaded := m.backOffHandlerMap.LoadOrStore(key, &ComputeResourceAwareBackOffHandler{
SimpleBackOffBlocker: &SimpleBackOffBlocker{
Clock: m.Clock,
BackOffBaseSecond: backOffBaseSecond,
BackOffExponent: stdAtomic.NewUint32(0),
NextEligibleTime: NewAtomicTime(m.Clock.Now()),
MaxBackOffDuration: maxBackOffDuration,
}, ComputeResourceCeilings: &ComputeResourceCeilings{
computeResourceCeilings: v1.ResourceList{},
},
})
if loaded {
logger.Infof(ctx, "The back-off handler for [%v] has been loaded.\n", key)
} else {
logger.Infof(ctx, "The back-off handler for [%v] has been created.\n", key)
}
if ret, casted := h.(*ComputeResourceAwareBackOffHandler); casted {
return ret
}
return nil
}
func (m *Controller) GetBackOffHandler(key string) (*ComputeResourceAwareBackOffHandler, bool) {
return m.backOffHandlerMap.Get(key)
}
func ComposeResourceKey(o k8s.Resource) string {
return fmt.Sprintf("%v,%v", o.GroupVersionKind().String(), o.GetNamespace())
}
func NewController(ctx context.Context) *Controller {
logger.Infof(ctx, "Initializing the back-off controller.\n")
return &Controller{
Clock: clock.RealClock{},
backOffHandlerMap: HandlerMap{},
}
}