1 /* 2 * Copyright 2019 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.xds.orca; 18 19 import static com.google.common.truth.Truth.assertThat; 20 import static org.mockito.AdditionalAnswers.delegatesTo; 21 import static org.mockito.ArgumentMatchers.any; 22 import static org.mockito.ArgumentMatchers.argThat; 23 import static org.mockito.Mockito.doNothing; 24 import static org.mockito.Mockito.mock; 25 import static org.mockito.Mockito.verify; 26 import static org.mockito.Mockito.verifyNoInteractions; 27 import static org.mockito.Mockito.verifyNoMoreInteractions; 28 import static org.mockito.Mockito.when; 29 30 import com.github.xds.data.orca.v3.OrcaLoadReport; 31 import com.google.common.base.Objects; 32 import io.grpc.ClientStreamTracer; 33 import io.grpc.Metadata; 34 import io.grpc.services.MetricReport; 35 import io.grpc.xds.orca.OrcaPerRequestUtil.OrcaPerRequestReportListener; 36 import io.grpc.xds.orca.OrcaPerRequestUtil.OrcaReportingTracerFactory; 37 import org.junit.Rule; 38 import org.junit.Test; 39 import org.junit.runner.RunWith; 40 import org.junit.runners.JUnit4; 41 import org.mockito.ArgumentCaptor; 42 import org.mockito.ArgumentMatcher; 43 import org.mockito.Mock; 44 import org.mockito.junit.MockitoJUnit; 45 import org.mockito.junit.MockitoRule; 46 47 /** 48 * Unit tests for {@link OrcaPerRequestUtil} class. 49 */ 50 @RunWith(JUnit4.class) 51 public class OrcaPerRequestUtilTest { 52 @Rule public final MockitoRule mocks = MockitoJUnit.rule(); 53 54 private static final ClientStreamTracer.StreamInfo STREAM_INFO = 55 ClientStreamTracer.StreamInfo.newBuilder().build(); 56 57 @Mock 58 private OrcaPerRequestReportListener orcaListener1; 59 @Mock 60 private OrcaPerRequestReportListener orcaListener2; 61 62 /** 63 * Tests a single load balance policy's listener receive per-request ORCA reports upon call 64 * trailer arrives. 65 */ 66 @Test singlePolicyTypicalWorkflow()67 public void singlePolicyTypicalWorkflow() { 68 // Use a mocked noop stream tracer factory as the original stream tracer factory. 69 ClientStreamTracer.Factory fakeDelegateFactory = mock(ClientStreamTracer.Factory.class); 70 ClientStreamTracer fakeTracer = mock(ClientStreamTracer.class); 71 doNothing().when(fakeTracer).inboundTrailers(any(Metadata.class)); 72 when(fakeDelegateFactory.newClientStreamTracer( 73 any(ClientStreamTracer.StreamInfo.class), any(Metadata.class))) 74 .thenReturn(fakeTracer); 75 76 // The OrcaReportingTracerFactory will augment the StreamInfo passed to its 77 // newClientStreamTracer method. The augmented StreamInfo's CallOptions will contain 78 // a OrcaReportBroker, in which has the registered listener. 79 ClientStreamTracer.Factory factory = 80 OrcaPerRequestUtil.getInstance() 81 .newOrcaClientStreamTracerFactory(fakeDelegateFactory, orcaListener1); 82 ClientStreamTracer tracer = factory.newClientStreamTracer(STREAM_INFO, new Metadata()); 83 ArgumentCaptor<ClientStreamTracer.StreamInfo> streamInfoCaptor = 84 ArgumentCaptor.forClass(ClientStreamTracer.StreamInfo.class); 85 verify(fakeDelegateFactory) 86 .newClientStreamTracer(streamInfoCaptor.capture(), any(Metadata.class)); 87 ClientStreamTracer.StreamInfo capturedInfo = streamInfoCaptor.getValue(); 88 assertThat(capturedInfo).isNotEqualTo(STREAM_INFO); 89 90 // When the trailer does not contain ORCA report, listener callback will not be invoked. 91 Metadata trailer = new Metadata(); 92 tracer.inboundTrailers(trailer); 93 verifyNoMoreInteractions(orcaListener1); 94 95 // When the trailer contains an ORCA report, listener callback will be invoked. 96 trailer.put( 97 OrcaReportingTracerFactory.ORCA_ENDPOINT_LOAD_METRICS_KEY, 98 OrcaLoadReport.getDefaultInstance()); 99 tracer.inboundTrailers(trailer); 100 ArgumentCaptor<MetricReport> reportCaptor = ArgumentCaptor.forClass(MetricReport.class); 101 verify(orcaListener1).onLoadReport(reportCaptor.capture()); 102 assertThat(reportEqual(reportCaptor.getValue(), 103 OrcaPerRequestUtil.fromOrcaLoadReport(OrcaLoadReport.getDefaultInstance()))).isTrue(); 104 } 105 106 static final class MetricsReportMatcher implements ArgumentMatcher<MetricReport> { 107 private MetricReport original; 108 MetricsReportMatcher(MetricReport report)109 public MetricsReportMatcher(MetricReport report) { 110 this.original = report; 111 } 112 113 @Override matches(MetricReport argument)114 public boolean matches(MetricReport argument) { 115 return reportEqual(original, argument); 116 } 117 } 118 reportEqual(MetricReport a, MetricReport b)119 static boolean reportEqual(MetricReport a, 120 MetricReport b) { 121 return a.getCpuUtilization() == b.getCpuUtilization() 122 && a.getApplicationUtilization() == b.getApplicationUtilization() 123 && a.getMemoryUtilization() == b.getMemoryUtilization() 124 && a.getQps() == b.getQps() 125 && a.getEps() == b.getEps() 126 && Objects.equal(a.getRequestCostMetrics(), b.getRequestCostMetrics()) 127 && Objects.equal(a.getUtilizationMetrics(), b.getUtilizationMetrics()); 128 } 129 130 /** 131 * Tests parent-child load balance policies' listeners both receive per-request ORCA reports upon 132 * call trailer arrives and ORCA report deserialization happens only once. 133 */ 134 @Test twoLevelPoliciesTypicalWorkflow()135 public void twoLevelPoliciesTypicalWorkflow() { 136 ClientStreamTracer.Factory parentFactory = 137 mock(ClientStreamTracer.Factory.class, 138 delegatesTo( 139 OrcaPerRequestUtil.getInstance().newOrcaClientStreamTracerFactory(orcaListener1))); 140 141 ClientStreamTracer.Factory childFactory = 142 OrcaPerRequestUtil.getInstance() 143 .newOrcaClientStreamTracerFactory(parentFactory, orcaListener2); 144 // Child factory will augment the StreamInfo and pass it to the parent factory. 145 ClientStreamTracer childTracer = 146 childFactory.newClientStreamTracer(STREAM_INFO, new Metadata()); 147 ArgumentCaptor<ClientStreamTracer.StreamInfo> streamInfoCaptor = 148 ArgumentCaptor.forClass(ClientStreamTracer.StreamInfo.class); 149 verify(parentFactory).newClientStreamTracer(streamInfoCaptor.capture(), any(Metadata.class)); 150 ClientStreamTracer.StreamInfo parentStreamInfo = streamInfoCaptor.getValue(); 151 assertThat(parentStreamInfo).isNotEqualTo(STREAM_INFO); 152 153 // When the trailer does not contain ORCA report, no listener callback will be invoked. 154 Metadata trailer = new Metadata(); 155 childTracer.inboundTrailers(trailer); 156 verifyNoMoreInteractions(orcaListener1); 157 verifyNoMoreInteractions(orcaListener2); 158 159 // When the trailer contains an ORCA report, callbacks for both listeners will be invoked. 160 // Both listener will receive the same ORCA report instance, which means deserialization 161 // happens only once. 162 trailer.put( 163 OrcaReportingTracerFactory.ORCA_ENDPOINT_LOAD_METRICS_KEY, 164 OrcaLoadReport.getDefaultInstance()); 165 childTracer.inboundTrailers(trailer); 166 ArgumentCaptor<MetricReport> parentReportCap = ArgumentCaptor.forClass(MetricReport.class); 167 ArgumentCaptor<MetricReport> childReportCap = ArgumentCaptor.forClass(MetricReport.class); 168 verify(orcaListener1).onLoadReport(parentReportCap.capture()); 169 verify(orcaListener2).onLoadReport(childReportCap.capture()); 170 assertThat(reportEqual(parentReportCap.getValue(), 171 OrcaPerRequestUtil.fromOrcaLoadReport(OrcaLoadReport.getDefaultInstance()))).isTrue(); 172 assertThat(childReportCap.getValue()).isSameInstanceAs(parentReportCap.getValue()); 173 } 174 175 /** 176 * Tests the case when parent policy creates its own {@link ClientStreamTracer.Factory}, ORCA 177 * reports are only forwarded to the parent's listener. 178 */ 179 @Test onlyParentPolicyReceivesReportsIfCreatesOwnTracer()180 public void onlyParentPolicyReceivesReportsIfCreatesOwnTracer() { 181 ClientStreamTracer.Factory parentFactory = 182 OrcaPerRequestUtil.getInstance().newOrcaClientStreamTracerFactory(orcaListener1); 183 ClientStreamTracer.Factory childFactory = 184 mock(ClientStreamTracer.Factory.class, 185 delegatesTo(OrcaPerRequestUtil.getInstance() 186 .newOrcaClientStreamTracerFactory(parentFactory, orcaListener2))); 187 ClientStreamTracer parentTracer = 188 parentFactory.newClientStreamTracer(STREAM_INFO, new Metadata()); 189 Metadata trailer = new Metadata(); 190 OrcaLoadReport report = OrcaLoadReport.getDefaultInstance(); 191 trailer.put( 192 OrcaReportingTracerFactory.ORCA_ENDPOINT_LOAD_METRICS_KEY, report); 193 parentTracer.inboundTrailers(trailer); 194 verify(orcaListener1).onLoadReport( 195 argThat(new MetricsReportMatcher(OrcaPerRequestUtil.fromOrcaLoadReport(report)))); 196 verifyNoInteractions(childFactory); 197 verifyNoInteractions(orcaListener2); 198 } 199 } 200