Add support for publishNotReadyAddresses (istio#1231)

* Add support for `publishNotReadyAddresses`

* rename

proto/workload.proto

@@ -137,12 +137,20 @@ message LoadBalancing {
     // 4. Any endpoints
     FAILOVER = 2;
   }
+  enum HealthPolicy {
+    // Only select healthy endpoints
+    ONLY_HEALTHY = 0;
+    // Include all endpoints, even if they are unhealthy.
+    ALLOW_ALL = 1;
+  }
 
   // routing_preference defines what scopes we want to keep traffic within.
   // The `mode` determines how these routing preferences are handled
   repeated Scope routing_preference = 1;
   // mode defines how we should handle the routing preferences.
   Mode mode = 2;
+  // health_policy defines how we should filter endpoints
+  HealthPolicy health_policy = 3;
 }
 
 // Workload represents a workload - an endpoint (or collection behind a hostname).

src/admin.rs

@@ -685,6 +685,7 @@ mod tests {
             load_balancing: Some(XdsLoadBalancing {
                 routing_preference: vec![1, 2],
                 mode: 1,
+                health_policy: 1,
             }), // ..Default::default() // intentionally don't default. we want all fields populated
             ip_families: 0,
         };

src/proxy.rs

@@ -874,6 +874,7 @@ mod tests {
                 },
                 address: addr,
                 port: ports.clone(),
+                status: state::workload::HealthStatus::Healthy,
             },
         );
         Service {

src/proxy/inbound.rs

@@ -515,6 +515,7 @@ mod tests {
         test_helpers,
     };
 
+    use crate::state::workload::HealthStatus;
     use hickory_resolver::config::{ResolverConfig, ResolverOpts};
     use prometheus_client::registry::Registry;
     use test_case::test_case;
@@ -621,6 +622,7 @@ mod tests {
                         },
                         address: ep_addr,
                         port: std::collections::HashMap::new(),
+                        status: HealthStatus::Healthy,
                     },
                 )]
                 .into_iter()

src/state.rs

@@ -349,8 +349,7 @@ impl ProxyState {
         });
 
         match svc.load_balancer {
-            None => endpoints.choose(&mut rand::thread_rng()),
-            Some(ref lb) => {
+            Some(ref lb) if lb.mode != LoadBalancerMode::Standard => {
                 let ranks = endpoints
                     .filter_map(|(ep, wl)| {
                         // Load balancer will define N targets we want to match
@@ -394,6 +393,7 @@ impl ProxyState {
                     .map(|(_, ep, wl)| (ep, wl))
                     .choose(&mut rand::thread_rng())
             }
+            _ => endpoints.choose(&mut rand::thread_rng()),
         }
     }
 }
@@ -908,8 +908,8 @@ impl ProxyStateManager {
 
 #[cfg(test)]
 mod tests {
-    use crate::state::service::LoadBalancer;
-    use crate::state::workload::Locality;
+    use crate::state::service::{LoadBalancer, LoadBalancerHealthPolicy};
+    use crate::state::workload::{HealthStatus, Locality};
     use prometheus_client::registry::Registry;
     use std::{net::Ipv4Addr, net::SocketAddrV4, time::Duration};
 
@@ -1180,6 +1180,7 @@ mod tests {
                         network: "".into(),
                     }),
                     port: tc.endpoint_mapping(),
+                    status: HealthStatus::Healthy,
                 },
             )]),
             ports: tc.service_mapping(),
@@ -1355,6 +1356,7 @@ mod tests {
                         network: "".into(),
                     }),
                     port: HashMap::from([(80u16, 80u16)]),
+                    status: HealthStatus::Healthy,
                 },
             ),
             (
@@ -1370,6 +1372,7 @@ mod tests {
                         network: "".into(),
                     }),
                     port: HashMap::from([(80u16, 80u16)]),
+                    status: HealthStatus::Healthy,
                 },
             ),
             (
@@ -1385,6 +1388,7 @@ mod tests {
                         network: "".into(),
                     }),
                     port: HashMap::from([(80u16, 80u16)]),
+                    status: HealthStatus::Healthy,
                 },
             ),
         ]);
@@ -1397,6 +1401,7 @@ mod tests {
                     LoadBalancerScopes::Region,
                     LoadBalancerScopes::Zone,
                 ],
+                health_policy: LoadBalancerHealthPolicy::OnlyHealthy,
             }),
             ports: HashMap::from([(80u16, 80u16)]),
             ..test_helpers::mock_default_service()
@@ -1410,6 +1415,7 @@ mod tests {
                     LoadBalancerScopes::Region,
                     LoadBalancerScopes::Zone,
                 ],
+                health_policy: LoadBalancerHealthPolicy::OnlyHealthy,
             }),
             ports: HashMap::from([(80u16, 80u16)]),
             ..test_helpers::mock_default_service()

src/state/service.rs

@@ -22,11 +22,11 @@ use tracing::trace;
 
 use xds::istio::workload::Service as XdsService;
 
-use crate::state::workload::is_default;
 use crate::state::workload::{
     byte_to_ip, network_addr, GatewayAddress, NamespacedHostname, NetworkAddress, Workload,
     WorkloadError,
 };
+use crate::state::workload::{is_default, HealthStatus};
 use crate::strng::Strng;
 use crate::xds::istio::workload::load_balancing::Scope as XdsScope;
 use crate::xds::istio::workload::{IpFamilies, PortList};
@@ -59,18 +59,42 @@ pub struct Service {
 
 #[derive(Debug, Eq, PartialEq, Clone, serde::Serialize, serde::Deserialize)]
 pub enum LoadBalancerMode {
+    // Do not consider LoadBalancerScopes when picking endpoints
+    Standard,
+    // Only select endpoints matching all LoadBalancerScopes when picking endpoints; otherwise, fail.
     Strict,
+    // Prefer select endpoints matching all LoadBalancerScopes when picking endpoints but allow mismatches
     Failover,
 }
 
 impl From<xds::istio::workload::load_balancing::Mode> for LoadBalancerMode {
     fn from(value: xds::istio::workload::load_balancing::Mode) -> Self {
         match value {
             xds::istio::workload::load_balancing::Mode::Strict => LoadBalancerMode::Strict,
+            xds::istio::workload::load_balancing::Mode::Failover => LoadBalancerMode::Failover,
             xds::istio::workload::load_balancing::Mode::UnspecifiedMode => {
-                LoadBalancerMode::Failover
+                LoadBalancerMode::Standard
+            }
+        }
+    }
+}
+
+#[derive(Default, Debug, Eq, PartialEq, Clone, serde::Serialize, serde::Deserialize)]
+pub enum LoadBalancerHealthPolicy {
+    #[default]
+    OnlyHealthy,
+    AllowAll,
+}
+
+impl From<xds::istio::workload::load_balancing::HealthPolicy> for LoadBalancerHealthPolicy {
+    fn from(value: xds::istio::workload::load_balancing::HealthPolicy) -> Self {
+        match value {
+            xds::istio::workload::load_balancing::HealthPolicy::OnlyHealthy => {
+                LoadBalancerHealthPolicy::OnlyHealthy
+            }
+            xds::istio::workload::load_balancing::HealthPolicy::AllowAll => {
+                LoadBalancerHealthPolicy::AllowAll
             }
-            xds::istio::workload::load_balancing::Mode::Failover => LoadBalancerMode::Failover,
         }
     }
 }
@@ -105,6 +129,7 @@ impl TryFrom<XdsScope> for LoadBalancerScopes {
 pub struct LoadBalancer {
     pub routing_preferences: Vec<LoadBalancerScopes>,
     pub mode: LoadBalancerMode,
+    pub health_policy: LoadBalancerHealthPolicy,
 }
 
 impl From<xds::istio::workload::IpFamilies> for Option<IpFamily> {
@@ -147,6 +172,15 @@ impl Service {
     pub fn contains_endpoint(&self, wl: &Workload, addr: Option<&NetworkAddress>) -> bool {
         self.endpoints.contains_key(&endpoint_uid(&wl.uid, addr))
     }
+
+    pub fn should_include_endpoint(&self, ep_health: HealthStatus) -> bool {
+        ep_health == HealthStatus::Healthy
+            || self
+                .load_balancer
+                .as_ref()
+                .map(|lb| lb.health_policy == LoadBalancerHealthPolicy::AllowAll)
+                .unwrap_or(false)
+    }
 }
 
 #[derive(Debug, Hash, Eq, PartialEq, Clone, serde::Serialize)]
@@ -181,6 +215,9 @@ pub struct Endpoint {
 
     /// The port mapping.
     pub port: HashMap<u16, u16>,
+
+    /// Health status for the endpoint
+    pub status: HealthStatus,
 }
 
 pub fn endpoint_uid(workload_uid: &str, address: Option<&NetworkAddress>) -> Strng {
@@ -220,6 +257,10 @@ impl TryFrom<&XdsService> for Service {
                     })
                     .collect::<Result<Vec<LoadBalancerScopes>, WorkloadError>>()?,
                 mode: xds::istio::workload::load_balancing::Mode::try_from(lb.mode)?.into(),
+                health_policy: xds::istio::workload::load_balancing::HealthPolicy::try_from(
+                    lb.health_policy,
+                )?
+                .into(),
             })
         } else {
             None
@@ -324,7 +365,12 @@ impl ServiceStore {
     pub fn insert_endpoint(&mut self, ep: Endpoint) {
         let ep_uid = endpoint_uid(&ep.workload_uid, ep.address.as_ref());
         if let Some(svc) = self.get_by_namespaced_host(&ep.service) {
+            // We may or may not accept the endpoint based on it's health
+            if !svc.should_include_endpoint(ep.status) {
+                return;
+            }
             let mut svc = Arc::unwrap_or_clone(svc);
+
             // Clone the service and add the endpoint.
             svc.endpoints.insert(ep_uid, ep);
 
@@ -386,7 +432,9 @@ impl ServiceStore {
         let namespaced_hostname = service.namespaced_hostname();
         if let Some(endpoints) = self.staged_services.remove(&namespaced_hostname) {
             for (wip, ep) in endpoints {
-                service.endpoints.insert(wip.clone(), ep);
+                if service.should_include_endpoint(ep.status) {
+                    service.endpoints.insert(wip.clone(), ep);
+                }
             }
         }