Repository: grpc/grpc-go
File: xds/internal/balancer/clusterimpl/picker.go (Go)

Code before:
}

func (d *dropper) drop() (ret bool) {
	return d.w.Next().(bool)
}

// loadReporter wraps the methods from the loadStore that are used here.
type loadReporter interface {
	CallStarted(locality string)
	CallFinished(locality string, err error)
	CallServerLoad(locality, name string, val float64)
	CallDropped(locality string)
}

// Picker implements RPC drop, circuit breaking drop and load reporting.
type picker struct {
	drops           []*dropper
	s               balancer.State
	loadStore       loadReporter
	counter         *xdsclient.ClusterRequestsCounter
	countMax        uint32
	telemetryLabels map[string]string
}

func (b *clusterImplBalancer) newPicker(config *dropConfigs) *picker {
	b.childStateMu.Lock()
	defer b.childStateMu.Unlock()
	b.configMu.Lock()
	defer b.configMu.Unlock()

	return &picker{
		drops:           config.drops,
		s:               b.childState,
		loadStore:       b.loadWrapper,
		counter:         config.requestCounter,
		countMax:        config.requestCountMax,
		telemetryLabels: b.telemetryLabels,
	}
}

func telemetryLabels(ctx context.Context) map[string]string {
	if ctx == nil {
		return nil
	}
	labels := stats.GetLabels(ctx)
	if labels == nil {
		return nil
	}
	return labels.TelemetryLabels
}


Review comment (suggestion):
I would get rid of this method here, and would instead add a `newPickerLocked` in `clusterimpl.go` that is always called with the mutex held. In the process I would also get rid of the `dropConfigs` struct and would change `handleDropAndRequestCount()` to instead return a boolean.

```golang
// handleDropAndRequestCount compares drop and request counter in newConfig with
// the one currently used by picker. It returns a boolean indicating if a new
// picker needs to be generated.
func (b *clusterImplBalancer) handleDropAndRequestCount(newConfig *LBConfig) bool {
	var updatePicker bool
	if !equalDropCategories(b.dropCategories, newConfig.DropCategories) {
		b.dropCategories = newConfig.DropCategories
		b.drops = make([]*dropper, 0, len(newConfig.DropCategories))
		for _, c := range newConfig.DropCategories {
			b.drops = append(b.drops, newDropper(c))
		}
		updatePicker = true
	}

	if b.requestCounterCluster != newConfig.Cluster || b.requestCounterService != newConfig.EDSServiceName {
		b.requestCounterCluster = newConfig.Cluster
		b.requestCounterService = newConfig.EDSServiceName
		b.requestCounter = xdsclient.GetClusterRequestsCounter(newConfig.Cluster, newConfig.EDSServiceName)
		updatePicker = true
	}
	var newRequestCountMax uint32 = 1024
	if newConfig.MaxConcurrentRequests != nil {
		newRequestCountMax = *newConfig.MaxConcurrentRequests
	}
	if b.requestCountMax != newRequestCountMax {
		b.requestCountMax = newRequestCountMax
		updatePicker = true
	}

	return updatePicker
}

func (b *clusterImplBalancer) newPickerLocked() *picker {
	return &picker{
		drops:           b.drops,
		s:               b.childState,
		loadStore:       b.loadWrapper,
		counter:         b.requestCounter,
		countMax:        b.requestCountMax,
		telemetryLabels: b.telemetryLabels,
	}
}
```

Code after:
}

func (d *dropper) drop() (ret bool) {
	return d.w.Next().(bool)
}

// loadReporter wraps the methods from the loadStore that are used here.
type loadReporter interface {
	CallStarted(locality string)
	CallFinished(locality string, err error)
	CallServerLoad(locality, name string, val float64)
	CallDropped(locality string)
}

// Picker implements RPC drop, circuit breaking drop and load reporting.
type picker struct {
	drops           []*dropper
	s               balancer.State
	loadStore       loadReporter
	counter         *xdsclient.ClusterRequestsCounter
	countMax        uint32
	telemetryLabels map[string]string
}

func telemetryLabels(ctx context.Context) map[string]string {
	if ctx == nil {
		return nil
	}
	labels := stats.GetLabels(ctx)
	if labels == nil {
		return nil
	}
	return labels.TelemetryLabels
}

func (d *picker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
	// Unconditionally set labels if present, even dropped or queued RPC's can
	// use these labels.
	if labels := telemetryLabels(info.Ctx); labels != nil {
		for key, value := range d.telemetryLabels {
			labels[key] = value
		}
	}

	// Don't drop unless the inner picker is READY. Similar to
	// https://github.com/grpc/grpc-go/issues/2622.
	if d.s.ConnectivityState == connectivity.Ready {
		// Check if this RPC should be dropped by category.
		for _, dp := range d.drops {
			if dp.drop() {
				if d.loadStore != nil {