Line data    Source code

       1             : #pragma once
       2             : 
       3             : #include "source/common/upstream/load_balancer_impl.h"
       4             : #include "source/common/upstream/upstream_impl.h"
       5             : 
       6             : namespace Envoy {
       7             : namespace Extensions {
       8             : namespace Clusters {
       9             : namespace Aggregate {
      10             : 
      11             : // AggregateLoadBalancerContext wraps the load balancer context to re-assign priority load
      12             : // according the to host priority selected by the aggregate load balancer.
      13             : class AggregateLoadBalancerContext : public Upstream::LoadBalancerContextBase {
      14             : public:
      15             :   AggregateLoadBalancerContext(Upstream::LoadBalancerContext* context,
      16             :                                Upstream::LoadBalancerBase::HostAvailability host_availability,
      17             :                                uint32_t host_priority)
      18           0 :       : host_availability_(host_availability), host_priority_(host_priority) {
      19           0 :     if (context == nullptr) {
      20           0 :       owned_context_ = std::make_unique<Upstream::LoadBalancerContextBase>();
      21           0 :       context_ = owned_context_.get();
      22           0 :     } else {
      23           0 :       context_ = context;
      24           0 :     }
      25           0 :   }
      26             : 
      27             :   // Upstream::LoadBalancerContext
      28           0 :   absl::optional<uint64_t> computeHashKey() override { return context_->computeHashKey(); }
      29           0 :   const Network::Connection* downstreamConnection() const override {
      30           0 :     return context_->downstreamConnection();
      31           0 :   }
      32           0 :   const Router::MetadataMatchCriteria* metadataMatchCriteria() override {
      33           0 :     return context_->metadataMatchCriteria();
      34           0 :   }
      35           0 :   const Http::RequestHeaderMap* downstreamHeaders() const override {
      36           0 :     return context_->downstreamHeaders();
      37           0 :   }
      38             :   const Upstream::HealthyAndDegradedLoad&
      39             :   determinePriorityLoad(const Upstream::PrioritySet&,
      40             :                         const Upstream::HealthyAndDegradedLoad& original_priority_load,
      41           0 :                         const Upstream::RetryPriority::PriorityMappingFunc&) override {
      42             :     // Re-assign load. Set all traffic to the priority and availability selected in aggregate
      43             :     // cluster.
      44             :     //
      45             :     // Note: context_->determinePriorityLoad() was already called and its result handled in
      46             :     // AggregateClusterLoadBalancer::LoadBalancerImpl::chooseHost().
      47           0 :     const size_t priorities = original_priority_load.healthy_priority_load_.get().size();
      48           0 :     priority_load_.healthy_priority_load_.get().assign(priorities, 0);
      49           0 :     priority_load_.degraded_priority_load_.get().assign(priorities, 0);
      50             : 
      51           0 :     if (host_availability_ == Upstream::LoadBalancerBase::HostAvailability::Healthy) {
      52           0 :       priority_load_.healthy_priority_load_.get()[host_priority_] = 100;
      53           0 :     } else {
      54           0 :       priority_load_.degraded_priority_load_.get()[host_priority_] = 100;
      55           0 :     }
      56           0 :     return priority_load_;
      57           0 :   }
      58           0 :   bool shouldSelectAnotherHost(const Upstream::Host& host) override {
      59           0 :     return context_->shouldSelectAnotherHost(host);
      60           0 :   }
      61           0 :   uint32_t hostSelectionRetryCount() const override { return context_->hostSelectionRetryCount(); }
      62           0 :   Network::Socket::OptionsSharedPtr upstreamSocketOptions() const override {
      63           0 :     return context_->upstreamSocketOptions();
      64           0 :   }
      65           0 :   Network::TransportSocketOptionsConstSharedPtr upstreamTransportSocketOptions() const override {
      66           0 :     return context_->upstreamTransportSocketOptions();
      67           0 :   }
      68             : 
      69             : private:
      70             :   Upstream::HealthyAndDegradedLoad priority_load_;
      71             :   std::unique_ptr<Upstream::LoadBalancerContext> owned_context_;
      72             :   Upstream::LoadBalancerContext* context_{nullptr};
      73             :   const Upstream::LoadBalancerBase::HostAvailability host_availability_;
      74             :   const uint32_t host_priority_;
      75             : };
      76             : 
      77             : } // namespace Aggregate
      78             : } // namespace Clusters
      79             : } // namespace Extensions
      80             : } // namespace Envoy