/*
 *
 * Copyright 2020 gRPC authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

// Package clusterimpl implements the xds_cluster_impl balancing policy. It
// handles the cluster features (e.g. circuit_breaking, RPC dropping).
//
// Note that it doesn't handle name resolution, which is done by policy
// xds_cluster_resolver.
package clusterimpl

import (
	"context"
	"encoding/json"
	"fmt"
	"net"
	"slices"
	"sync"
	"sync/atomic"
	"time"

	"google.golang.org/grpc/balancer"
	"google.golang.org/grpc/balancer/weightedroundrobin"
	"google.golang.org/grpc/connectivity"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/credentials/tls/certprovider"
	"google.golang.org/grpc/internal/balancer/gracefulswitch"
	"google.golang.org/grpc/internal/credentials/xds"
	"google.golang.org/grpc/internal/grpclog"
	"google.golang.org/grpc/internal/grpcsync"
	"google.golang.org/grpc/internal/pretty"
	xdsinternal "google.golang.org/grpc/internal/xds"

	"google.golang.org/grpc/internal/xds/balancer/clusterimpl/internal"
	"google.golang.org/grpc/internal/xds/balancer/loadstore"
	"google.golang.org/grpc/internal/xds/bootstrap"
	"google.golang.org/grpc/internal/xds/clients"
	"google.golang.org/grpc/internal/xds/clients/lrsclient"
	"google.golang.org/grpc/internal/xds/xdsclient"
	"google.golang.org/grpc/internal/xds/xdsclient/xdsresource"
	"google.golang.org/grpc/resolver"
	"google.golang.org/grpc/serviceconfig"
)

const (
	// Name is the name of the cluster_impl balancer.
	Name                   = "xds_cluster_impl_experimental"
	defaultRequestCountMax = 1024
	loadStoreStopTimeout   = 1 * time.Second
)

var (
	// Below function is no-op in actual code, but can be overridden in
	// tests to give tests visibility into exactly when certain events happen.
	clientConnUpdateHook = func() {}
	pickerUpdateHook     = func() {}
	buildProvider        = buildProviderFunc
)

func init() {
	balancer.Register(bb{})
}

type bb struct{}

func (bb) Build(cc balancer.ClientConn, bOpts balancer.BuildOptions) balancer.Balancer {
	b := &clusterImplBalancer{
		ClientConn:      cc,
		loadWrapper:     loadstore.NewWrapper(),
		requestCountMax: defaultRequestCountMax,
	}
	b.xdsHIPtr.Store(xds.NewHandshakeInfo(nil, nil, nil, "", false, false, false))
	b.logger = prefixLogger(b)
	b.child = gracefulswitch.NewBalancer(b, bOpts)
	b.logger.Infof("Created")

	var creds credentials.TransportCredentials
	switch {
	case bOpts.DialCreds != nil:
		creds = bOpts.DialCreds
	case bOpts.CredsBundle != nil:
		creds = bOpts.CredsBundle.TransportCredentials()
	}
	if xc, ok := creds.(interface{ UsesXDS() bool }); ok && xc.UsesXDS() {
		b.xdsCredsInUse = true
	}
	b.logger.Infof("xDS credentials in use: %v", b.xdsCredsInUse)
	return b
}

func (bb) Name() string {
	return Name
}

func (bb) ParseConfig(c json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
	return parseConfig(c)
}

type clusterImplBalancer struct {
	balancer.ClientConn

	// The following fields are set at creation time, and are read-only after
	// that, and therefore need not be protected by a mutex.
	logger *grpclog.PrefixLogger
	// TODO: #8366 -  Refactor usage of loadWrapper to easily plugin a test
	// load reporter from tests.
	loadWrapper *loadstore.Wrapper

	// The following fields are only accessed from balancer API methods, which
	// are guaranteed to be called serially by gRPC.
	xdsClient        xdsclient.XDSClient     // Sent down in ResolverState attributes.
	cancelLoadReport func(context.Context)   // To stop reporting load through the above xDS client.
	edsServiceName   string                  // EDS service name to report load for.
	lrsServer        *bootstrap.ServerConfig // Load reporting server configuration.
	dropCategories   []DropConfig            // The categories for drops.
	child            *gracefulswitch.Balancer
	xdsHIPtr         atomic.Pointer[grpcsync.RefCounted[*xds.HandshakeInfo]] // Accessed atomically as it is shared between the balancer and the transport.
	xdsCredsInUse    bool

	// The active security configuration received from the control plane.
	securityConfig *xdsresource.SecurityConfig

	// The following fields are protected by mu, since they are accessed in
	// balancer API methods and in methods called from the child policy.
	mu                       sync.Mutex
	clusterName              string                                      // The cluster name for credentials handshaking.
	inhibitPickerUpdates     bool                                        // Inhibits state updates from child policy when processing an update from the parent.
	pendingPickerUpdates     bool                                        // True if a picker update from the child policy was inhibited when processing an update from the parent.
	childState               balancer.State                              // Most recent state update from the child policy.
	drops                    []*dropper                                  // Drops implementation.
	requestCounterCluster    string                                      // The cluster name for the request counter, from LB config.
	requestCounterService    string                                      // The service name for the request counter, from LB config.
	requestCountMax          uint32                                      // Max concurrent requests, from LB config.
	requestCounter           *xdsclient.ClusterRequestsCounter           // Tracks total inflight requests for a given service.
	telemetryLabels          map[string]string                           // Telemetry labels to set on picks, from LB config.
	lrsReportEndpointMetrics *xdsresource.LRSReportEndpointMetricsConfig // LRS metrics to propagate.
}

// dropRequestsPerMillion scales a drop overload's numerator and denominator to
// a number of requests to drop per million. numerator can be as large as a
// million (a 100% drop), so the multiplication is done in uint64 to avoid
// overflowing a uint32, and the result is capped at a million because a drop
// ratio above 100% is treated as 100%.
func dropRequestsPerMillion(numerator, denominator uint32) uint32 {
	rpm := uint64(numerator) * million / uint64(denominator)
	if rpm > million {
		rpm = million
	}
	return uint32(rpm)
}

// handleClusterConfigLocked updates the internal state of the balancer with the
// new cluster configuration. It returns true if a new picker needs to be
// generated as a result of these changes. It must be called with b.mu held.
func (b *clusterImplBalancer) handleClusterConfigLocked(clusterConfig xdsresource.ClusterConfig) bool {
	clusterUpdate := clusterConfig.Cluster
	var updatePicker bool

	b.telemetryLabels = clusterUpdate.TelemetryLabels

	if !b.lrsReportEndpointMetrics.Equal(clusterUpdate.LRSReportEndpointMetrics) {
		b.lrsReportEndpointMetrics = clusterUpdate.LRSReportEndpointMetrics
		updatePicker = true
	}

	var newDrops []DropConfig
	if clusterUpdate.ClusterType == xdsresource.ClusterTypeEDS {
		edsUpdate := clusterConfig.EndpointConfig.EDSUpdate
		newDrops = make([]DropConfig, 0, len(edsUpdate.Drops))
		for _, d := range edsUpdate.Drops {
			newDrops = append(newDrops, DropConfig{
				Category:           d.Category,
				RequestsPerMillion: dropRequestsPerMillion(d.Numerator, d.Denominator),
			})
		}
		if !slices.Equal(b.dropCategories, newDrops) {
			b.dropCategories = newDrops
			b.drops = make([]*dropper, 0, len(newDrops))
			for _, c := range newDrops {
				b.drops = append(b.drops, newDropper(c))
			}
			updatePicker = true
		}
	}

	if b.requestCounterCluster != clusterUpdate.ClusterName || b.requestCounterService != clusterUpdate.EDSServiceName {
		b.requestCounterCluster = clusterUpdate.ClusterName
		b.requestCounterService = clusterUpdate.EDSServiceName
		b.requestCounter = xdsclient.GetClusterRequestsCounter(clusterUpdate.ClusterName, clusterUpdate.EDSServiceName)
		updatePicker = true
	}
	var newRequestCountMax uint32 = 1024
	if clusterUpdate.MaxRequests != nil {
		newRequestCountMax = *clusterUpdate.MaxRequests
	}
	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,
		clusterName:     b.clusterName,
		metrics:         b.lrsReportEndpointMetrics,
	}
}

// updateLoadStore checks the config for load store, and decides whether it
// needs to restart the load reporting stream.
func (b *clusterImplBalancer) updateLoadStore(clusterUpdate *xdsresource.ClusterUpdate) error {
	var updateLoadClusterAndService bool

	// ClusterName is different, restart. ClusterName is from ClusterName and
	// EDSServiceName.
	clusterName := b.getClusterName()
	if clusterName != clusterUpdate.ClusterName {
		updateLoadClusterAndService = true
		b.setClusterName(clusterUpdate.ClusterName)
		clusterName = clusterUpdate.ClusterName
	}
	if b.edsServiceName != clusterUpdate.EDSServiceName {
		updateLoadClusterAndService = true
		b.edsServiceName = clusterUpdate.EDSServiceName
	}
	if updateLoadClusterAndService {
		// This updates the clusterName and serviceName that will be reported
		// for the loads. The update here is too early, the perfect timing is
		// when the picker is updated with the new connection. But from this
		// balancer's point of view, it's impossible to tell.
		//
		// On the other hand, this will almost never happen. Each LRS policy
		// shouldn't get updated config. The parent should do a graceful switch
		// when the clusterName or serviceName is changed.
		b.loadWrapper.UpdateClusterAndService(clusterName, b.edsServiceName)
	}

	var (
		stopOldLoadReport  bool
		startNewLoadReport bool
	)

	// Check if it's necessary to restart load report.
	if b.lrsServer == nil {
		if clusterUpdate.LRSServerConfig != nil {
			// Old is nil, new is not nil, start new LRS.
			b.lrsServer = clusterUpdate.LRSServerConfig
			startNewLoadReport = true
		}
		// Old is nil, new is nil, do nothing.
	} else if clusterUpdate.LRSServerConfig == nil {
		// Old is not nil, new is nil, stop old, don't start new.
		b.lrsServer = nil
		stopOldLoadReport = true
	} else {
		// Old is not nil, new is not nil, compare string values, if
		// different, stop old and start new.
		if !b.lrsServer.Equal(clusterUpdate.LRSServerConfig) {
			b.lrsServer = clusterUpdate.LRSServerConfig
			stopOldLoadReport = true
			startNewLoadReport = true
		}
	}

	if stopOldLoadReport {
		if b.cancelLoadReport != nil {
			stopCtx, stopCancel := context.WithTimeout(context.Background(), loadStoreStopTimeout)
			defer stopCancel()
			b.cancelLoadReport(stopCtx)
			b.cancelLoadReport = nil
			if !startNewLoadReport {
				// If a new LRS stream will be started later, no need to update
				// it to nil here.
				b.loadWrapper.UpdateLoadStore(nil)
			}
		}
	}
	if startNewLoadReport {
		var loadStore *lrsclient.LoadStore
		if b.xdsClient != nil {
			loadStore, b.cancelLoadReport = b.xdsClient.ReportLoad(b.lrsServer)
		}
		b.loadWrapper.UpdateLoadStore(loadStore)
	}

	return nil
}

func buildProviderFunc(configs map[string]*certprovider.BuildableConfig, instanceName, certName string, wantIdentity, wantRoot bool) (certprovider.Provider, error) {
	cfg := configs[instanceName]
	provider, err := cfg.Build(certprovider.BuildOptions{
		CertName:     certName,
		WantIdentity: wantIdentity,
		WantRoot:     wantRoot,
	})
	if err != nil {
		// This error is not expected since the bootstrap process parses the
		// config and makes sure that it is acceptable to the plugin. Still, it
		// is possible that the plugin parses the config successfully, but its
		// Build() method errors out.
		return nil, fmt.Errorf("xds: failed to get security plugin instance (%+v): %v", cfg, err)
	}
	return provider, nil
}

func (b *clusterImplBalancer) buildProviders(config *xdsresource.SecurityConfig) (certprovider.Provider, certprovider.Provider, error) {
	cpc := b.xdsClient.BootstrapConfig().CertProviderConfigs()
	var rootProvider certprovider.Provider
	if config.UseSystemRootCerts {
		rootProvider = systemRootCertsProvider{}
	} else {
		rp, err := buildProvider(cpc, config.RootInstanceName, config.RootCertName, false, true)
		if err != nil {
			return nil, nil, err
		}
		rootProvider = rp
	}

	var identityProvider certprovider.Provider
	if name, cert := config.IdentityInstanceName, config.IdentityCertName; name != "" {
		ip, err := buildProvider(cpc, name, cert, true, false)
		if err != nil {
			if rootProvider != nil {
				rootProvider.Close()
			}
			return nil, nil, err
		}
		identityProvider = ip
	}
	return rootProvider, identityProvider, nil
}

// handleSecurityConfig processes the security configuration received from the
// management server, creates appropriate certificate provider plugins, and
// updates the HandshakeInfo which is added as an address attribute in
// NewSubConn() calls.
func (b *clusterImplBalancer) handleSecurityConfig(config *xdsresource.SecurityConfig) error {
	// If xdsCredentials are not in use, i.e, the user did not want to get
	// security configuration from an xDS server, we should not be acting on the
	// received security config here. Doing so poses a security threat.
	if !b.xdsCredsInUse || config.Equal(b.securityConfig) {
		return nil
	}

	// Security config being nil is a valid case where the management server has
	// not sent any security configuration. The xdsCredentials implementation
	// handles this by delegating to its fallback credentials.
	if config == nil {
		b.securityConfig = nil
		oldHI := b.xdsHIPtr.Swap(xds.NewHandshakeInfo(nil, nil, nil, "", false, false, false))
		if oldHI != nil {
			oldHI.Decrement()
		}
		return nil
	}

	rootProvider, identityProvider, err := b.buildProviders(config)
	if err != nil {
		return err
	}

	newHI := xds.NewHandshakeInfo(rootProvider, identityProvider, config.SubjectAltNameMatchers, config.SNI, false, config.AutoSNISANValidation, config.UseAutoHostSNI)
	b.securityConfig = config
	oldHI := b.xdsHIPtr.Swap(newHI)
	if oldHI != nil {
		oldHI.Decrement()
	}
	return nil
}

func (b *clusterImplBalancer) UpdateClientConnState(s balancer.ClientConnState) error {
	defer clientConnUpdateHook()

	b.mu.Lock()
	b.inhibitPickerUpdates = true
	b.mu.Unlock()
	if b.logger.V(2) {
		b.logger.Infof("Received configuration: %s", pretty.ToJSON(s.BalancerConfig))
	}
	newConfig, ok := s.BalancerConfig.(*LBConfig)
	if !ok {
		return fmt.Errorf("unexpected balancer config with type: %T", s.BalancerConfig)
	}

	// Need to check for potential errors at the beginning of this function, so
	// that on errors, we reject the whole config, instead of applying part of
	// it.
	bb := balancer.Get(newConfig.ChildPolicy.Name)
	if bb == nil {
		return fmt.Errorf("child policy %q not registered", newConfig.ChildPolicy.Name)
	}

	if b.xdsClient == nil {
		c := xdsclient.FromResolverState(s.ResolverState)
		if c == nil {
			return balancer.ErrBadResolverState
		}
		b.xdsClient = c
	}

	xdsConfig := xdsresource.XDSConfigFromResolverState(s.ResolverState)
	if xdsConfig == nil {
		b.logger.Warningf("Received balancer config with no xDS config")
		return balancer.ErrBadResolverState
	}
	clusterCfg := xdsConfig.Clusters[newConfig.Cluster]
	clusterUpdate := clusterCfg.Config.Cluster
	if err := b.handleSecurityConfig(clusterUpdate.SecurityCfg); err != nil {
		// If the security config is invalid, for example, if the provider
		// instance is not found in the bootstrap config, we need to put the
		// channel in transient failure.
		return fmt.Errorf("received Cluster resource that contains invalid security config: %v", err)

	}
	// Update load reporting config. This needs to be done before updating the
	// child policy because we need the loadStore from the updated client to be
	// passed to the ccWrapper, so that the next picker from the child policy
	// will pick up the new loadStore.
	if err := b.updateLoadStore(clusterUpdate); err != nil {
		return err
	}

	// Build config for the gracefulswitch balancer. It is safe to ignore JSON
	// marshaling errors here, since the config was already validated as part of
	// ParseConfig().
	cfg := []map[string]any{{newConfig.ChildPolicy.Name: newConfig.ChildPolicy.Config}}
	cfgJSON, _ := json.Marshal(cfg)
	parsedCfg, err := gracefulswitch.ParseConfig(cfgJSON)
	if err != nil {
		return err
	}

	// Addresses and sub-balancer config are sent to sub-balancer.
	err = b.child.UpdateClientConnState(balancer.ClientConnState{
		ResolverState:  weightedroundrobin.SetBackendService(s.ResolverState, b.clusterName),
		BalancerConfig: parsedCfg,
	})

	b.mu.Lock()
	updatePicker := b.handleClusterConfigLocked(clusterCfg.Config)
	// We want to send a picker update to the parent if one of the two
	// conditions are met:
	// - drop/request count config or LRS metrics config has changed *and* there
	//   is already a picker from the child, or
	// - there is a pending picker update from the child (and this covers the
	//   case where the drop/request config has not changed, but the child sent
	//   a picker update while we were still processing config from our parent).
	if (updatePicker && b.childState.Picker != nil) || b.pendingPickerUpdates {
		b.pendingPickerUpdates = false
		b.ClientConn.UpdateState(balancer.State{
			ConnectivityState: b.childState.ConnectivityState,
			Picker:            b.newPickerLocked(),
		})
	}
	b.inhibitPickerUpdates = false
	b.mu.Unlock()
	pickerUpdateHook()
	return err
}

func (b *clusterImplBalancer) ResolverError(err error) {
	b.child.ResolverError(err)
}

func (b *clusterImplBalancer) updateSubConnState(_ balancer.SubConn, s balancer.SubConnState, cb func(balancer.SubConnState)) {
	// Trigger re-resolution when a SubConn turns transient failure. This is
	// necessary for the LogicalDNS in cluster_resolver policy to re-resolve.
	//
	// Note that this happens not only for the addresses from DNS, but also for
	// EDS (cluster_impl doesn't know if it's DNS or EDS, only the parent
	// knows). The parent priority policy is configured to ignore re-resolution
	// signal from the EDS children.
	if s.ConnectivityState == connectivity.TransientFailure {
		b.ClientConn.ResolveNow(resolver.ResolveNowOptions{})
	}

	if cb != nil {
		cb(s)
	}
}

func (b *clusterImplBalancer) UpdateSubConnState(sc balancer.SubConn, s balancer.SubConnState) {
	b.logger.Errorf("UpdateSubConnState(%v, %+v) called unexpectedly", sc, s)
}

func (b *clusterImplBalancer) Close() {
	b.child.Close()
	b.childState = balancer.State{}

	if b.cancelLoadReport != nil {
		stopCtx, stopCancel := context.WithTimeout(context.Background(), loadStoreStopTimeout)
		defer stopCancel()
		b.cancelLoadReport(stopCtx)
		b.cancelLoadReport = nil
	}
	oldHI := b.xdsHIPtr.Swap(nil)
	if oldHI != nil {
		oldHI.Decrement()
	}
	b.securityConfig = nil
	b.logger.Infof("Shutdown")
}

func (b *clusterImplBalancer) ExitIdle() {
	b.child.ExitIdle()
}

// Override methods to accept updates from the child LB.

func (b *clusterImplBalancer) UpdateState(state balancer.State) {
	b.mu.Lock()
	defer b.mu.Unlock()

	// Inhibit sending a picker update to our parent as part of handling new
	// state from the child, if we are currently handling an update from our
	// parent. Update the childState field regardless.
	b.childState = state
	if b.inhibitPickerUpdates {
		b.pendingPickerUpdates = true
		if b.logger.V(2) {
			b.logger.Infof("Received a picker update from the child when processing an update from the parent")
		}
		return
	}

	b.ClientConn.UpdateState(balancer.State{
		ConnectivityState: state.ConnectivityState,
		Picker:            b.newPickerLocked(),
	})
	pickerUpdateHook()
}

func (b *clusterImplBalancer) setClusterName(n string) {
	b.mu.Lock()
	defer b.mu.Unlock()
	b.clusterName = n
}

func (b *clusterImplBalancer) getClusterName() string {
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.clusterName
}

// scWrapper is a wrapper of SubConn with locality ID. The locality ID can be
// retrieved from the addresses when creating SubConn.
//
// All SubConns passed to the child policies are wrapped in this, so that the
// picker can get the localityID from the picked SubConn, and do load reporting.
//
// After wrapping, all SubConns to and from the parent ClientConn (e.g. for
// SubConn state update, update/remove SubConn) must be the original SubConns.
// All SubConns to and from the child policy (NewSubConn, forwarding SubConn
// state update) must be the wrapper. The balancer keeps a map from the original
// SubConn to the wrapper for this purpose.
type scWrapper struct {
	balancer.SubConn

	localityID clients.Locality

	hostname string
}

func (b *clusterImplBalancer) NewSubConn(addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
	clusterName := b.getClusterName()
	newAddrs := make([]resolver.Address, len(addrs))
	for i, addr := range addrs {
		newAddrs[i] = xdsinternal.SetXDSHandshakeClusterName(addr, clusterName)
		newAddrs[i] = xds.SetHandshakeInfo(newAddrs[i], &b.xdsHIPtr)

		host := xdsresource.Hostname(addr)
		// If the hostname contains a port, strip it. Per [RFC 6066, Section
		// 3](https://www.rfc-editor.org/rfc/rfc6066.html#section-3), the SNI
		// may only contain a qualified DNS hostname, which excludes port
		// numbers.
		h, _, err := net.SplitHostPort(host)
		if err == nil {
			host = h
		}
		// Store hostname in the address attributes, so that it can be used in
		// the client handshake.
		newAddrs[i] = xds.SetAddressHostname(newAddrs[i], host)
	}
	var sc balancer.SubConn
	scw := &scWrapper{}
	if len(addrs) > 0 {
		scw.hostname = xdsresource.Hostname(addrs[0])
		scw.localityID = xdsinternal.GetLocalityID(addrs[0])
	}
	oldListener := opts.StateListener
	opts.StateListener = func(state balancer.SubConnState) {
		b.updateSubConnState(sc, state, oldListener)
	}
	sc, err := b.ClientConn.NewSubConn(newAddrs, opts)
	if err != nil {
		return nil, err
	}
	scw.SubConn = sc
	return scw, nil
}

func (b *clusterImplBalancer) RemoveSubConn(sc balancer.SubConn) {
	b.logger.Errorf("RemoveSubConn(%v) called unexpectedly", sc)
}

func (b *clusterImplBalancer) UpdateAddresses(sc balancer.SubConn, _ []resolver.Address) {
	b.logger.Errorf("UpdateAddresses(%v) called unexpectedly", sc)
}

// systemRootCertsProvider implements a certprovider.Provider that returns the
// system default root certificates for validation.
type systemRootCertsProvider struct{}

func (systemRootCertsProvider) Close() {}

func (systemRootCertsProvider) KeyMaterial(context.Context) (*certprovider.KeyMaterial, error) {
	rootCAs, err := internal.X509SystemCertPoolFunc()
	if err != nil {
		return nil, err
	}
	return &certprovider.KeyMaterial{Roots: rootCAs}, nil
}