fix random pod spread

Signed-off-by: Areeb Ahmed <areebahmed0709@gmail.com>

build/build.go

@@ -1030,7 +1030,11 @@ func detectSharedMounts(ctx context.Context, reqs map[string][]*reqForNode) (_ m
 	m := map[string]map[fsKey]*fsTracker{}
 	for _, reqs := range reqs {
 		for _, req := range reqs {
-			nodeName := req.ResolvedNode.Node().Name
+			nodeName := req.Node().Name
+			// skip shared-session optimisation: targets may connect to different replicas.
+			if req.Node().Driver != nil && req.Node().Driver.RequiresUncachedClient() {
+				continue
+			}
 			if _, ok := m[nodeName]; !ok {
 				m[nodeName] = map[fsKey]*fsTracker{}
 			}

build/resolver/driver.go

@@ -57,6 +57,14 @@ func (dp ResolvedNode) Platforms() []ocispecs.Platform {
 }
 
 func (dp ResolvedNode) Client(ctx context.Context) (*client.Client, error) {
+	node := dp.resolver.nodes[dp.driverIndex]
+	// loadbalance=random requires a fresh connection per call so each target lands on a different pod.
+	if node.Driver != nil && node.Driver.RequiresUncachedClient() {
+		if _, err := dp.resolver.boot(ctx, []int{dp.driverIndex}, nil); err != nil {
+			return nil, err
+		}
+		return node.Driver.UncachedClient(ctx)
+	}
 	clients, err := dp.resolver.boot(ctx, []int{dp.driverIndex}, nil)
 	if err != nil {
 		return nil, err

driver/driver.go

@@ -57,6 +57,10 @@ type Info struct {
 	DynamicNodes []store.Node
 }
 
+type UncachedClientDriver interface {
+	RequiresUncachedClient() bool
+}
+
 type Driver interface {
 	Factory() Factory
 	Bootstrap(context.Context, progress.Logger) error

driver/kubernetes/driver.go

@@ -45,6 +45,7 @@ type Driver struct {
 	// if you add fields, remember to update docs:
 	// https://github.com/docker/docs/blob/main/content/build/drivers/kubernetes.md
 	minReplicas       int
+	loadbalance       string
 	deployment        *appsv1.Deployment
 	statefulSet       *appsv1.StatefulSet
 	configMaps        []*corev1.ConfigMap
@@ -61,6 +62,10 @@ func (d *Driver) IsMobyDriver() bool {
 	return false
 }
 
+func (d *Driver) RequiresUncachedClient() bool {
+	return d.loadbalance == LoadbalanceRandom
+}
+
 func (d *Driver) Config() driver.InitConfig {
 	return d.InitConfig
 }

driver/kubernetes/factory.go

@@ -162,6 +162,7 @@ func (f *factory) New(ctx context.Context, cfg driver.InitConfig) (driver.Driver
 			StatefulSet: d.statefulSet,
 		}
 	}
+	d.loadbalance = loadbalance
 	return d, nil
 }
 

driver/kubernetes/factory_test.go

@@ -268,3 +268,41 @@ func TestFactory_processDriverOpts(t *testing.T) {
 		},
 	)
 }
+
+func TestRequiresUncachedClient(t *testing.T) {
+	f := factory{
+		cc: &mockClientConfig{
+			clientConfig: &rest.Config{},
+		},
+	}
+	baseCfg := driver.InitConfig{
+		Name: driver.BuilderName("test"),
+	}
+
+	t.Run("RandomLoadbalance", func(t *testing.T) {
+		cfg := baseCfg
+		cfg.DriverOpts = map[string]string{"loadbalance": "random"}
+		d, err := f.New(t.Context(), cfg)
+		require.NoError(t, err)
+		require.True(t, d.(*Driver).RequiresUncachedClient(),
+			"expected RequiresUncachedClient=true for loadbalance=random")
+	})
+
+	t.Run("StickyLoadbalance", func(t *testing.T) {
+		cfg := baseCfg
+		cfg.DriverOpts = map[string]string{"loadbalance": "sticky"}
+		d, err := f.New(t.Context(), cfg)
+		require.NoError(t, err)
+		require.False(t, d.(*Driver).RequiresUncachedClient(),
+			"expected RequiresUncachedClient=false for loadbalance=sticky")
+	})
+
+	t.Run("DefaultLoadbalance", func(t *testing.T) {
+		cfg := baseCfg
+		cfg.DriverOpts = map[string]string{}
+		d, err := f.New(t.Context(), cfg)
+		require.NoError(t, err)
+		require.False(t, d.(*Driver).RequiresUncachedClient(),
+			"expected RequiresUncachedClient=false for default (sticky) loadbalance")
+	})
+}

driver/kubernetes/podchooser/podchooser.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"math/rand"
 	"sort"
-	"time"
 
 	"github.com/docker/buildx/driver/kubernetes/kubeclient"
 	"github.com/pkg/errors"
@@ -20,7 +19,6 @@ type PodChooser interface {
 }
 
 type RandomPodChooser struct {
-	RandSource  rand.Source
 	PodClient   kubeclient.PodClient
 	Deployment  *appsv1.Deployment
 	StatefulSet *appsv1.StatefulSet
@@ -34,12 +32,7 @@ func (pc *RandomPodChooser) ChoosePod(ctx context.Context) (*corev1.Pod, error)
 	if len(pods) == 0 {
 		return nil, errors.New("no running buildkit pods found")
 	}
-	randSource := pc.RandSource
+	n := rand.Intn(len(pods)) // #nosec G404 -- no strong seeding required
-	if randSource == nil {
-		randSource = rand.NewSource(time.Now().Unix())
-	}
-	rnd := rand.New(randSource) // #nosec G404 -- 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
 }

driver/manager.go

@@ -128,6 +128,17 @@ func (d *DriverHandle) Client(ctx context.Context, opt ...client.ClientOpt) (*cl
 	return d.client, d.err
 }
 
+func (d *DriverHandle) UncachedClient(ctx context.Context) (*client.Client, error) {
+	return d.Driver.Client(ctx, d.getClientOptions()...)
+}
+
+func (d *DriverHandle) RequiresUncachedClient() bool {
+	if p, ok := d.Driver.(UncachedClientDriver); ok {
+		return p.RequiresUncachedClient()
+	}
+	return false
+}
+
 func (d *DriverHandle) getClientOptions() []client.ClientOpt {
 	return []client.ClientOpt{
 		client.WithTracerDelegate(delegated.DefaultExporter),