github.com/pdmccormick/importable-docker-buildx@v0.0.0-20240426161518-e47091289030/driver/kubernetes/podchooser/podchooser.go (about)

     1  package podchooser
     2  
     3  import (
     4  	"context"
     5  	"math/rand"
     6  	"sort"
     7  	"time"
     8  
     9  	"github.com/pkg/errors"
    10  	"github.com/serialx/hashring"
    11  	"github.com/sirupsen/logrus"
    12  	appsv1 "k8s.io/api/apps/v1"
    13  	corev1 "k8s.io/api/core/v1"
    14  	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
    15  	clientcorev1 "k8s.io/client-go/kubernetes/typed/core/v1"
    16  )
    17  
    18  type PodChooser interface {
    19  	ChoosePod(ctx context.Context) (*corev1.Pod, error)
    20  }
    21  
    22  type RandomPodChooser struct {
    23  	RandSource rand.Source
    24  	PodClient  clientcorev1.PodInterface
    25  	Deployment *appsv1.Deployment
    26  }
    27  
    28  func (pc *RandomPodChooser) ChoosePod(ctx context.Context) (*corev1.Pod, error) {
    29  	pods, err := ListRunningPods(ctx, pc.PodClient, pc.Deployment)
    30  	if err != nil {
    31  		return nil, err
    32  	}
    33  	if len(pods) == 0 {
    34  		return nil, errors.New("no running buildkit pods found")
    35  	}
    36  	randSource := pc.RandSource
    37  	if randSource == nil {
    38  		randSource = rand.NewSource(time.Now().Unix())
    39  	}
    40  	rnd := rand.New(randSource) //nolint:gosec // no strong seeding required
    41  	n := rnd.Int() % len(pods)
    42  	logrus.Debugf("RandomPodChooser.ChoosePod(): len(pods)=%d, n=%d", len(pods), n)
    43  	return pods[n], nil
    44  }
    45  
    46  type StickyPodChooser struct {
    47  	Key        string
    48  	PodClient  clientcorev1.PodInterface
    49  	Deployment *appsv1.Deployment
    50  }
    51  
    52  func (pc *StickyPodChooser) ChoosePod(ctx context.Context) (*corev1.Pod, error) {
    53  	pods, err := ListRunningPods(ctx, pc.PodClient, pc.Deployment)
    54  	if err != nil {
    55  		return nil, err
    56  	}
    57  	var podNames []string
    58  	podMap := make(map[string]*corev1.Pod, len(pods))
    59  	for _, pod := range pods {
    60  		podNames = append(podNames, pod.Name)
    61  		podMap[pod.Name] = pod
    62  	}
    63  	ring := hashring.New(podNames)
    64  	chosen, ok := ring.GetNode(pc.Key)
    65  	if !ok {
    66  		// NOTREACHED
    67  		logrus.Errorf("no pod found for key %q", pc.Key)
    68  		rpc := &RandomPodChooser{
    69  			PodClient:  pc.PodClient,
    70  			Deployment: pc.Deployment,
    71  		}
    72  		return rpc.ChoosePod(ctx)
    73  	}
    74  	return podMap[chosen], nil
    75  }
    76  
    77  func ListRunningPods(ctx context.Context, client clientcorev1.PodInterface, depl *appsv1.Deployment) ([]*corev1.Pod, error) {
    78  	selector, err := metav1.LabelSelectorAsSelector(depl.Spec.Selector)
    79  	if err != nil {
    80  		return nil, err
    81  	}
    82  	listOpts := metav1.ListOptions{
    83  		LabelSelector: selector.String(),
    84  	}
    85  	podList, err := client.List(ctx, listOpts)
    86  	if err != nil {
    87  		return nil, err
    88  	}
    89  	var runningPods []*corev1.Pod
    90  	for i := range podList.Items {
    91  		pod := &podList.Items[i]
    92  		if pod.Status.Phase == corev1.PodRunning {
    93  			logrus.Debugf("pod runnning: %q", pod.Name)
    94  			runningPods = append(runningPods, pod)
    95  		}
    96  	}
    97  	sort.Slice(runningPods, func(i, j int) bool {
    98  		return runningPods[i].Name < runningPods[j].Name
    99  	})
   100  	return runningPods, nil
   101  }