From 0a1efe73f3181538a16860f85602bcdcd43002e3 Mon Sep 17 00:00:00 2001
From: Eric Haberkorn <erichaberkorn@gmail.com>
Date: Tue, 13 Jun 2023 11:03:30 -0400
Subject: [PATCH] Refactor disco chain prioritize by locality structs (#17696)

This includes prioritize by localities on disco chain targets rather than
resolvers, allowing different targets within the same partition to have
different policies.
---
 agent/consul/discoverychain/compile.go      | 22 ++++++++++-------
 agent/consul/discoverychain/compile_test.go |  1 +
 agent/structs/discovery_chain.go            | 26 +++++++++++----------
 agent/structs/structs.deepcopy.go           |  8 +++----
 4 files changed, 33 insertions(+), 24 deletions(-)

diff --git a/agent/consul/discoverychain/compile.go b/agent/consul/discoverychain/compile.go
index 634ab4a6ff76..20227db3ef19 100644
--- a/agent/consul/discoverychain/compile.go
+++ b/agent/consul/discoverychain/compile.go
@@ -1009,19 +1009,25 @@ RESOLVE_AGAIN:
 		Type: structs.DiscoveryGraphNodeTypeResolver,
 		Name: target.ID,
 		Resolver: &structs.DiscoveryResolver{
-			Default:              resolver.IsDefault(),
-			Target:               target.ID,
-			ConnectTimeout:       connectTimeout,
-			RequestTimeout:       resolver.RequestTimeout,
-			PrioritizeByLocality: resolver.PrioritizeByLocality.ToDiscovery(),
+			Default:        resolver.IsDefault(),
+			Target:         target.ID,
+			ConnectTimeout: connectTimeout,
+			RequestTimeout: resolver.RequestTimeout,
 		},
 		LoadBalancer: resolver.LoadBalancer,
 	}
 
-	// Merge default values from the proxy defaults
 	proxyDefault := c.entries.GetProxyDefaults(targetID.PartitionOrDefault())
-	if proxyDefault != nil && node.Resolver.PrioritizeByLocality == nil {
-		node.Resolver.PrioritizeByLocality = proxyDefault.PrioritizeByLocality.ToDiscovery()
+
+	// Only set PrioritizeByLocality for targets in the same partition.
+	if target.Partition == c.evaluateInPartition && target.Peer == "" {
+		if resolver.PrioritizeByLocality != nil {
+			target.PrioritizeByLocality = resolver.PrioritizeByLocality.ToDiscovery()
+		}
+
+		if target.PrioritizeByLocality == nil && proxyDefault != nil {
+			target.PrioritizeByLocality = proxyDefault.PrioritizeByLocality.ToDiscovery()
+		}
 	}
 
 	target.Subset = resolver.Subsets[target.ServiceSubset]
diff --git a/agent/consul/discoverychain/compile_test.go b/agent/consul/discoverychain/compile_test.go
index 692c0c4fdaca..ca39aa236fc9 100644
--- a/agent/consul/discoverychain/compile_test.go
+++ b/agent/consul/discoverychain/compile_test.go
@@ -3301,6 +3301,7 @@ func newTarget(opts structs.DiscoveryTargetOpts, modFn func(t *structs.Discovery
 	t.SNI = connect.TargetSNI(t, "trustdomain.consul")
 	t.Name = t.SNI
 	t.ConnectTimeout = 5 * time.Second // default
+	t.PrioritizeByLocality = opts.PrioritizeByLocality
 	if modFn != nil {
 		modFn(t)
 	}
diff --git a/agent/structs/discovery_chain.go b/agent/structs/discovery_chain.go
index edbc7384a383..029fc3b0a8be 100644
--- a/agent/structs/discovery_chain.go
+++ b/agent/structs/discovery_chain.go
@@ -121,12 +121,11 @@ func (s *DiscoveryGraphNode) MapKey() string {
 
 // compiled form of ServiceResolverConfigEntry
 type DiscoveryResolver struct {
-	Default              bool                           `json:",omitempty"`
-	ConnectTimeout       time.Duration                  `json:",omitempty"`
-	RequestTimeout       time.Duration                  `json:",omitempty"`
-	Target               string                         `json:",omitempty"`
-	Failover             *DiscoveryFailover             `json:",omitempty"`
-	PrioritizeByLocality *DiscoveryPrioritizeByLocality `json:",omitempty"`
+	Default        bool               `json:",omitempty"`
+	ConnectTimeout time.Duration      `json:",omitempty"`
+	RequestTimeout time.Duration      `json:",omitempty"`
+	Target         string             `json:",omitempty"`
+	Failover       *DiscoveryFailover `json:",omitempty"`
 }
 
 func (r *DiscoveryResolver) MarshalJSON() ([]byte, error) {
@@ -238,6 +237,8 @@ type DiscoveryTarget struct {
 	// balancer objects.  This has a structure similar to SNI, but will not be
 	// affected by SNI customizations.
 	Name string `json:",omitempty"`
+
+	PrioritizeByLocality *DiscoveryPrioritizeByLocality `json:",omitempty"`
 }
 
 func (t *DiscoveryTarget) MarshalJSON() ([]byte, error) {
@@ -277,12 +278,13 @@ func (t *DiscoveryTarget) UnmarshalJSON(data []byte) error {
 }
 
 type DiscoveryTargetOpts struct {
-	Service       string
-	ServiceSubset string
-	Namespace     string
-	Partition     string
-	Datacenter    string
-	Peer          string
+	Service              string
+	ServiceSubset        string
+	Namespace            string
+	Partition            string
+	Datacenter           string
+	Peer                 string
+	PrioritizeByLocality *DiscoveryPrioritizeByLocality
 }
 
 func MergeDiscoveryTargetOpts(opts ...DiscoveryTargetOpts) DiscoveryTargetOpts {
diff --git a/agent/structs/structs.deepcopy.go b/agent/structs/structs.deepcopy.go
index 59a5c7224952..cdd007c26cdb 100644
--- a/agent/structs/structs.deepcopy.go
+++ b/agent/structs/structs.deepcopy.go
@@ -129,6 +129,10 @@ func (o *CompiledDiscoveryChain) DeepCopy() *CompiledDiscoveryChain {
 					cp_Targets_v2.Locality = new(Locality)
 					*cp_Targets_v2.Locality = *v2.Locality
 				}
+				if v2.PrioritizeByLocality != nil {
+					cp_Targets_v2.PrioritizeByLocality = new(DiscoveryPrioritizeByLocality)
+					*cp_Targets_v2.PrioritizeByLocality = *v2.PrioritizeByLocality
+				}
 			}
 			cp.Targets[k2] = cp_Targets_v2
 		}
@@ -240,10 +244,6 @@ func (o *DiscoveryResolver) DeepCopy() *DiscoveryResolver {
 	if o.Failover != nil {
 		cp.Failover = o.Failover.DeepCopy()
 	}
-	if o.PrioritizeByLocality != nil {
-		cp.PrioritizeByLocality = new(DiscoveryPrioritizeByLocality)
-		*cp.PrioritizeByLocality = *o.PrioritizeByLocality
-	}
 	return &cp
 }