-
Notifications
You must be signed in to change notification settings - Fork 455
/
podchooser.go
101 lines (93 loc) · 2.57 KB
/
podchooser.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
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
package podchooser
import (
"context"
"math/rand"
"sort"
"time"
"github.com/pkg/errors"
"github.com/serialx/hashring"
"github.com/sirupsen/logrus"
appsv1 "k8s.io/api/apps/v1"
corev1 "k8s.io/api/core/v1"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
clientcorev1 "k8s.io/client-go/kubernetes/typed/core/v1"
)
type PodChooser interface {
ChoosePod(ctx context.Context) (*corev1.Pod, error)
}
type RandomPodChooser struct {
RandSource rand.Source
PodClient clientcorev1.PodInterface
Deployment *appsv1.Deployment
}
func (pc *RandomPodChooser) ChoosePod(ctx context.Context) (*corev1.Pod, error) {
pods, err := ListRunningPods(ctx, pc.PodClient, pc.Deployment)
if err != nil {
return nil, err
}
if len(pods) == 0 {
return nil, errors.New("no running buildkit pods found")
}
randSource := pc.RandSource
if randSource == nil {
randSource = rand.NewSource(time.Now().Unix())
}
rnd := rand.New(randSource) //nolint:gosec // no strong seeding required
n := rnd.Int() % len(pods)
logrus.Debugf("RandomPodChooser.ChoosePod(): len(pods)=%d, n=%d", len(pods), n)
return pods[n], nil
}
type StickyPodChooser struct {
Key string
PodClient clientcorev1.PodInterface
Deployment *appsv1.Deployment
}
func (pc *StickyPodChooser) ChoosePod(ctx context.Context) (*corev1.Pod, error) {
pods, err := ListRunningPods(ctx, pc.PodClient, pc.Deployment)
if err != nil {
return nil, err
}
var podNames []string
podMap := make(map[string]*corev1.Pod, len(pods))
for _, pod := range pods {
podNames = append(podNames, pod.Name)
podMap[pod.Name] = pod
}
ring := hashring.New(podNames)
chosen, ok := ring.GetNode(pc.Key)
if !ok {
// NOTREACHED
logrus.Errorf("no pod found for key %q", pc.Key)
rpc := &RandomPodChooser{
PodClient: pc.PodClient,
Deployment: pc.Deployment,
}
return rpc.ChoosePod(ctx)
}
return podMap[chosen], nil
}
func ListRunningPods(ctx context.Context, client clientcorev1.PodInterface, depl *appsv1.Deployment) ([]*corev1.Pod, error) {
selector, err := metav1.LabelSelectorAsSelector(depl.Spec.Selector)
if err != nil {
return nil, err
}
listOpts := metav1.ListOptions{
LabelSelector: selector.String(),
}
podList, err := client.List(ctx, listOpts)
if err != nil {
return nil, err
}
var runningPods []*corev1.Pod
for i := range podList.Items {
pod := &podList.Items[i]
if pod.Status.Phase == corev1.PodRunning {
logrus.Debugf("pod runnning: %q", pod.Name)
runningPods = append(runningPods, pod)
}
}
sort.Slice(runningPods, func(i, j int) bool {
return runningPods[i].Name < runningPods[j].Name
})
return runningPods, nil
}