1 /* 2 * Copyright 2022 The gRPC Authors 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 package io.grpc.examples.orca; 18 19 import io.grpc.ConnectivityState; 20 import io.grpc.LoadBalancer; 21 import io.grpc.LoadBalancerProvider; 22 import io.grpc.LoadBalancerRegistry; 23 import io.grpc.services.MetricReport; 24 import io.grpc.util.ForwardingLoadBalancer; 25 import io.grpc.util.ForwardingLoadBalancerHelper; 26 import io.grpc.xds.orca.OrcaOobUtil; 27 import io.grpc.xds.orca.OrcaPerRequestUtil; 28 import java.util.concurrent.TimeUnit; 29 30 /** 31 * Implements a test LB policy that receives ORCA load reports. 32 * The load balancer mostly delegates to {@link io.grpc.internal.PickFirstLoadBalancerProvider}, 33 * in addition, it installs {@link OrcaOobUtil.OrcaOobReportListener} and 34 * {@link OrcaPerRequestUtil.OrcaPerRequestReportListener} to be notified with backend metrics. 35 */ 36 final class CustomBackendMetricsLoadBalancerProvider extends LoadBalancerProvider { 37 38 static final String EXAMPLE_LOAD_BALANCER = "example_backend_metrics_load_balancer"; 39 40 @Override newLoadBalancer(LoadBalancer.Helper helper)41 public LoadBalancer newLoadBalancer(LoadBalancer.Helper helper) { 42 return new CustomBackendMetricsLoadBalancer(helper); 43 } 44 45 @Override isAvailable()46 public boolean isAvailable() { 47 return true; 48 } 49 50 @Override getPriority()51 public int getPriority() { 52 return 5; 53 } 54 55 @Override getPolicyName()56 public String getPolicyName() { 57 return EXAMPLE_LOAD_BALANCER; 58 } 59 60 private final class CustomBackendMetricsLoadBalancer extends ForwardingLoadBalancer { 61 private LoadBalancer delegate; 62 CustomBackendMetricsLoadBalancer(LoadBalancer.Helper helper)63 public CustomBackendMetricsLoadBalancer(LoadBalancer.Helper helper) { 64 this.delegate = LoadBalancerRegistry.getDefaultRegistry() 65 .getProvider("pick_first") 66 .newLoadBalancer(new CustomBackendMetricsLoadBalancerHelper(helper)); 67 } 68 69 @Override delegate()70 public LoadBalancer delegate() { 71 return delegate; 72 } 73 74 private final class CustomBackendMetricsLoadBalancerHelper 75 extends ForwardingLoadBalancerHelper { 76 private final LoadBalancer.Helper orcaHelper; 77 CustomBackendMetricsLoadBalancerHelper(LoadBalancer.Helper helper)78 public CustomBackendMetricsLoadBalancerHelper(LoadBalancer.Helper helper) { 79 this.orcaHelper = OrcaOobUtil.newOrcaReportingHelper(helper); 80 } 81 82 @Override createSubchannel(LoadBalancer.CreateSubchannelArgs args)83 public LoadBalancer.Subchannel createSubchannel(LoadBalancer.CreateSubchannelArgs args) { 84 LoadBalancer.Subchannel subchannel = super.createSubchannel(args); 85 // Installs ORCA OOB metrics reporting listener and configures to receive report every 1s. 86 // The interval can not be smaller than server minimum report interval configuration, 87 // otherwise it is treated as server minimum report interval. 88 OrcaOobUtil.setListener(subchannel, new OrcaOobUtil.OrcaOobReportListener() { 89 @Override 90 public void onLoadReport(MetricReport orcaLoadReport) { 91 System.out.println("Example load balancer received OOB metrics report:\n" 92 + orcaLoadReport); 93 } 94 }, 95 OrcaOobUtil.OrcaReportingConfig.newBuilder() 96 .setReportInterval(1, TimeUnit.SECONDS) 97 .build() 98 ); 99 return subchannel; 100 } 101 102 @Override updateBalancingState(ConnectivityState newState, LoadBalancer.SubchannelPicker newPicker)103 public void updateBalancingState(ConnectivityState newState, LoadBalancer.SubchannelPicker newPicker) { 104 delegate().updateBalancingState(newState, new MayReportLoadPicker(newPicker)); 105 } 106 107 @Override delegate()108 public LoadBalancer.Helper delegate() { 109 return orcaHelper; 110 } 111 } 112 113 private final class MayReportLoadPicker extends LoadBalancer.SubchannelPicker { 114 private LoadBalancer.SubchannelPicker delegate; 115 MayReportLoadPicker(LoadBalancer.SubchannelPicker delegate)116 public MayReportLoadPicker(LoadBalancer.SubchannelPicker delegate) { 117 this.delegate = delegate; 118 } 119 120 @Override pickSubchannel(LoadBalancer.PickSubchannelArgs args)121 public LoadBalancer.PickResult pickSubchannel(LoadBalancer.PickSubchannelArgs args) { 122 LoadBalancer.PickResult result = delegate.pickSubchannel(args); 123 if (result.getSubchannel() == null) { 124 return result; 125 } 126 // Installs ORCA per-query metrics reporting listener. 127 final OrcaPerRequestUtil.OrcaPerRequestReportListener orcaListener = 128 new OrcaPerRequestUtil.OrcaPerRequestReportListener() { 129 @Override 130 public void onLoadReport(MetricReport orcaLoadReport) { 131 System.out.println("Example load balancer received per-rpc metrics report:\n" 132 + orcaLoadReport); 133 } 134 }; 135 if (result.getStreamTracerFactory() == null) { 136 return LoadBalancer.PickResult.withSubchannel( 137 result.getSubchannel(), 138 OrcaPerRequestUtil.getInstance().newOrcaClientStreamTracerFactory(orcaListener)); 139 } else { 140 return LoadBalancer.PickResult.withSubchannel( 141 result.getSubchannel(), 142 OrcaPerRequestUtil.getInstance().newOrcaClientStreamTracerFactory( 143 result.getStreamTracerFactory(), orcaListener)); 144 } 145 } 146 } 147 } 148 } 149