• Home
  • Features
  • Pricing
  • Docs
  • Announcements
  • Sign In

grpc / grpc-java / #19482

30 Sep 2024 03:17PM UTC coverage: 84.578% (+0.02%) from 84.554%
#19482

push

github

ejona86
util: Simplify MultiChildLB.getChildLbState()

Tests were converted to use getChildLbStateEag() if the argument was an
EAG, so the instanceof was no longer necessary.

33651 of 39787 relevant lines covered (84.58%)

0.85 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

96.58
/../util/src/main/java/io/grpc/util/MultiChildLoadBalancer.java
1
/*
2
 * Copyright 2023 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.util;
18

19
import static com.google.common.base.Preconditions.checkNotNull;
20
import static io.grpc.ConnectivityState.CONNECTING;
21
import static io.grpc.ConnectivityState.IDLE;
22
import static io.grpc.ConnectivityState.READY;
23
import static io.grpc.ConnectivityState.SHUTDOWN;
24
import static io.grpc.ConnectivityState.TRANSIENT_FAILURE;
25

26
import com.google.common.annotations.VisibleForTesting;
27
import com.google.common.collect.ImmutableList;
28
import io.grpc.Attributes;
29
import io.grpc.ConnectivityState;
30
import io.grpc.EquivalentAddressGroup;
31
import io.grpc.Internal;
32
import io.grpc.LoadBalancer;
33
import io.grpc.LoadBalancerProvider;
34
import io.grpc.Status;
35
import io.grpc.internal.PickFirstLoadBalancerProvider;
36
import java.net.SocketAddress;
37
import java.util.ArrayList;
38
import java.util.Collection;
39
import java.util.Collections;
40
import java.util.HashMap;
41
import java.util.HashSet;
42
import java.util.LinkedHashMap;
43
import java.util.List;
44
import java.util.Map;
45
import java.util.Set;
46
import java.util.logging.Level;
47
import java.util.logging.Logger;
48
import javax.annotation.Nullable;
49

50
/**
51
 * A base load balancing policy for those policies which has multiple children such as
52
 * ClusterManager or the petiole policies.  For internal use only.
53
 */
54
@Internal
55
public abstract class MultiChildLoadBalancer extends LoadBalancer {
56

57
  private static final Logger logger = Logger.getLogger(MultiChildLoadBalancer.class.getName());
1✔
58
  private final Map<Object, ChildLbState> childLbStates = new LinkedHashMap<>();
1✔
59
  private final Helper helper;
60
  // Set to true if currently in the process of handling resolved addresses.
61
  protected boolean resolvingAddresses;
62

63
  protected final LoadBalancerProvider pickFirstLbProvider = new PickFirstLoadBalancerProvider();
1✔
64

65
  protected ConnectivityState currentConnectivityState;
66

67

68
  protected MultiChildLoadBalancer(Helper helper) {
1✔
69
    this.helper = checkNotNull(helper, "helper");
1✔
70
    logger.log(Level.FINE, "Created");
1✔
71
  }
1✔
72

73
  /**
74
   * Using the state of all children will calculate the current connectivity state,
75
   * update fields, generate a picker and then call
76
   * {@link Helper#updateBalancingState(ConnectivityState, SubchannelPicker)}.
77
   */
78
  protected abstract void updateOverallBalancingState();
79

80
  /**
81
   * Override to utilize parsing of the policy configuration or alternative helper/lb generation.
82
   * Override this if keys are not Endpoints or if child policies have configuration. Null map
83
   * values preserve the child without delivering the child an update.
84
   */
85
  protected Map<Object, ResolvedAddresses> createChildAddressesMap(
86
      ResolvedAddresses resolvedAddresses) {
87
    Map<Object, ResolvedAddresses> childAddresses = new HashMap<>();
1✔
88
    for (EquivalentAddressGroup eag : resolvedAddresses.getAddresses()) {
1✔
89
      ResolvedAddresses addresses = resolvedAddresses.toBuilder()
1✔
90
          .setAddresses(Collections.singletonList(eag))
1✔
91
          .setAttributes(Attributes.newBuilder().set(IS_PETIOLE_POLICY, true).build())
1✔
92
          .setLoadBalancingPolicyConfig(null)
1✔
93
          .build();
1✔
94
      childAddresses.put(new Endpoint(eag), addresses);
1✔
95
    }
1✔
96
    return childAddresses;
1✔
97
  }
98

99
  /**
100
   * Override to create an instance of a subclass.
101
   */
102
  protected ChildLbState createChildLbState(Object key) {
103
    return new ChildLbState(key, pickFirstLbProvider);
1✔
104
  }
105

106
  /**
107
   *   Override to completely replace the default logic or to do additional activities.
108
   */
109
  @Override
110
  public Status acceptResolvedAddresses(ResolvedAddresses resolvedAddresses) {
111
    try {
112
      resolvingAddresses = true;
1✔
113

114
      // process resolvedAddresses to update children
115
      AcceptResolvedAddrRetVal acceptRetVal = acceptResolvedAddressesInternal(resolvedAddresses);
1✔
116
      if (!acceptRetVal.status.isOk()) {
1✔
117
        return acceptRetVal.status;
1✔
118
      }
119

120
      // Update the picker and our connectivity state
121
      updateOverallBalancingState();
1✔
122

123
      // shutdown removed children
124
      shutdownRemoved(acceptRetVal.removedChildren);
1✔
125
      return acceptRetVal.status;
1✔
126
    } finally {
127
      resolvingAddresses = false;
1✔
128
    }
129
  }
130

131
  /**
132
   * Handle the name resolution error.
133
   *
134
   * <p/>Override if you need special handling.
135
   */
136
  @Override
137
  public void handleNameResolutionError(Status error) {
138
    if (currentConnectivityState != READY)  {
1✔
139
      helper.updateBalancingState(
1✔
140
          TRANSIENT_FAILURE, new FixedResultPicker(PickResult.withError(error)));
1✔
141
    }
142
  }
1✔
143

144
  @Override
145
  public void shutdown() {
146
    logger.log(Level.FINE, "Shutdown");
1✔
147
    for (ChildLbState state : childLbStates.values()) {
1✔
148
      state.shutdown();
1✔
149
    }
1✔
150
    childLbStates.clear();
1✔
151
  }
1✔
152

153
  /**
154
   *   This does the work to update the child map and calculate which children have been removed.
155
   *   You must call {@link #updateOverallBalancingState} to update the picker
156
   *   and call {@link #shutdownRemoved(List)} to shutdown the endpoints that have been removed.
157
    */
158
  protected final AcceptResolvedAddrRetVal acceptResolvedAddressesInternal(
159
      ResolvedAddresses resolvedAddresses) {
160
    logger.log(Level.FINE, "Received resolution result: {0}", resolvedAddresses);
1✔
161

162
    Map<Object, ResolvedAddresses> newChildAddresses = createChildAddressesMap(resolvedAddresses);
1✔
163

164
    // Handle error case
165
    if (newChildAddresses.isEmpty()) {
1✔
166
      Status unavailableStatus = Status.UNAVAILABLE.withDescription(
1✔
167
          "NameResolver returned no usable address. " + resolvedAddresses);
168
      handleNameResolutionError(unavailableStatus);
1✔
169
      return new AcceptResolvedAddrRetVal(unavailableStatus, null);
1✔
170
    }
171

172
    updateChildrenWithResolvedAddresses(newChildAddresses);
1✔
173

174
    return new AcceptResolvedAddrRetVal(Status.OK, getRemovedChildren(newChildAddresses.keySet()));
1✔
175
  }
176

177
  private void updateChildrenWithResolvedAddresses(
178
      Map<Object, ResolvedAddresses> newChildAddresses) {
179
    for (Map.Entry<Object, ResolvedAddresses> entry : newChildAddresses.entrySet()) {
1✔
180
      ChildLbState childLbState = childLbStates.get(entry.getKey());
1✔
181
      if (childLbState == null) {
1✔
182
        childLbState = createChildLbState(entry.getKey());
1✔
183
        childLbStates.put(entry.getKey(), childLbState);
1✔
184
      }
185
      if (entry.getValue() != null) {
1✔
186
        childLbState.setResolvedAddresses(entry.getValue()); // update child
1✔
187
        childLbState.lb.handleResolvedAddresses(entry.getValue()); // update child LB
1✔
188
      }
189
    }
1✔
190
  }
1✔
191

192
  /**
193
   * Identifies which children have been removed (are not part of the newChildKeys).
194
   */
195
  private List<ChildLbState> getRemovedChildren(Set<Object> newChildKeys) {
196
    List<ChildLbState> removedChildren = new ArrayList<>();
1✔
197
    // Do removals
198
    for (Object key : ImmutableList.copyOf(childLbStates.keySet())) {
1✔
199
      if (!newChildKeys.contains(key)) {
1✔
200
        ChildLbState childLbState = childLbStates.remove(key);
1✔
201
        removedChildren.add(childLbState);
1✔
202
      }
203
    }
1✔
204
    return removedChildren;
1✔
205
  }
206

207
  protected final void shutdownRemoved(List<ChildLbState> removedChildren) {
208
    // Do shutdowns after updating picker to reduce the chance of failing an RPC by picking a
209
    // subchannel that has been shutdown.
210
    for (ChildLbState childLbState : removedChildren) {
1✔
211
      childLbState.shutdown();
1✔
212
    }
1✔
213
  }
1✔
214

215
  @Nullable
216
  protected static ConnectivityState aggregateState(
217
      @Nullable ConnectivityState overallState, ConnectivityState childState) {
218
    if (overallState == null) {
1✔
219
      return childState;
1✔
220
    }
221
    if (overallState == READY || childState == READY) {
1✔
222
      return READY;
1✔
223
    }
224
    if (overallState == CONNECTING || childState == CONNECTING) {
1✔
225
      return CONNECTING;
1✔
226
    }
227
    if (overallState == IDLE || childState == IDLE) {
1✔
228
      return IDLE;
×
229
    }
230
    return overallState;
1✔
231
  }
232

233
  protected final Helper getHelper() {
234
    return helper;
1✔
235
  }
236

237
  @VisibleForTesting
238
  public final Collection<ChildLbState> getChildLbStates() {
239
    return childLbStates.values();
1✔
240
  }
241

242
  @VisibleForTesting
243
  public final ChildLbState getChildLbState(Object key) {
244
    return childLbStates.get(key);
1✔
245
  }
246

247
  @VisibleForTesting
248
  public final ChildLbState getChildLbStateEag(EquivalentAddressGroup eag) {
249
    return getChildLbState(new Endpoint(eag));
1✔
250
  }
251

252
  /**
253
   * Filters out non-ready child load balancers (subchannels).
254
   */
255
  protected final List<ChildLbState> getReadyChildren() {
256
    List<ChildLbState> activeChildren = new ArrayList<>();
1✔
257
    for (ChildLbState child : getChildLbStates()) {
1✔
258
      if (child.getCurrentState() == READY) {
1✔
259
        activeChildren.add(child);
1✔
260
      }
261
    }
1✔
262
    return activeChildren;
1✔
263
  }
264

265
  /**
266
   * This represents the state of load balancer children.  Each endpoint (represented by an
267
   * EquivalentAddressGroup or EDS string) will have a separate ChildLbState which in turn will
268
   * have a single child LoadBalancer created from the provided factory.
269
   *
270
   * <p>A ChildLbStateHelper is the glue between ChildLbState and the helpers associated with the
271
   * petiole policy above and the PickFirstLoadBalancer's helper below.
272
   *
273
   * <p>If you wish to store additional state information related to each subchannel, then extend
274
   * this class.
275
   */
276
  public class ChildLbState {
277
    private final Object key;
278
    private ResolvedAddresses resolvedAddresses;
279

280
    private final LoadBalancer lb;
281
    private ConnectivityState currentState;
282
    private SubchannelPicker currentPicker = new FixedResultPicker(PickResult.withNoResult());
1✔
283

284
    public ChildLbState(Object key, LoadBalancer.Factory policyFactory) {
1✔
285
      this.key = key;
1✔
286
      this.lb = policyFactory.newLoadBalancer(createChildHelper());
1✔
287
      this.currentState = CONNECTING;
1✔
288
    }
1✔
289

290
    protected ChildLbStateHelper createChildHelper() {
291
      return new ChildLbStateHelper();
1✔
292
    }
293

294
    /**
295
     * Override for unique behavior such as delayed shutdowns of subchannels.
296
     */
297
    protected void shutdown() {
298
      lb.shutdown();
1✔
299
      this.currentState = SHUTDOWN;
1✔
300
      logger.log(Level.FINE, "Child balancer {0} deleted", key);
1✔
301
    }
1✔
302

303
    @Override
304
    public String toString() {
305
      return "Address = " + key
×
306
          + ", state = " + currentState
307
          + ", picker type: " + currentPicker.getClass()
×
308
          + ", lb: " + lb;
309
    }
310

311
    public final Object getKey() {
312
      return key;
1✔
313
    }
314

315
    @VisibleForTesting
316
    public final LoadBalancer getLb() {
317
      return lb;
1✔
318
    }
319

320
    @VisibleForTesting
321
    public final SubchannelPicker getCurrentPicker() {
322
      return currentPicker;
1✔
323
    }
324

325
    public final ConnectivityState getCurrentState() {
326
      return currentState;
1✔
327
    }
328

329
    protected final void setCurrentState(ConnectivityState newState) {
330
      currentState = newState;
1✔
331
    }
1✔
332

333
    protected final void setCurrentPicker(SubchannelPicker newPicker) {
334
      currentPicker = newPicker;
1✔
335
    }
1✔
336

337
    public final EquivalentAddressGroup getEag() {
338
      if (resolvedAddresses == null || resolvedAddresses.getAddresses().isEmpty()) {
1✔
339
        return null;
×
340
      }
341
      return resolvedAddresses.getAddresses().get(0);
1✔
342
    }
343

344
    protected final void setResolvedAddresses(ResolvedAddresses newAddresses) {
345
      checkNotNull(newAddresses, "Missing address list for child");
1✔
346
      resolvedAddresses = newAddresses;
1✔
347
    }
1✔
348

349
    @VisibleForTesting
350
    public final ResolvedAddresses getResolvedAddresses() {
351
      return resolvedAddresses;
1✔
352
    }
353

354
    /**
355
     * ChildLbStateHelper is the glue between ChildLbState and the helpers associated with the
356
     * petiole policy above and the PickFirstLoadBalancer's helper below.
357
     *
358
     * <p>The ChildLbState updates happen during updateBalancingState.  Otherwise, it is doing
359
     * simple forwarding.
360
     */
361
    protected class ChildLbStateHelper extends ForwardingLoadBalancerHelper {
1✔
362

363
      /**
364
       * Update current state and picker for this child and then use
365
       * {@link #updateOverallBalancingState()} for the parent LB.
366
       */
367
      @Override
368
      public void updateBalancingState(final ConnectivityState newState,
369
          final SubchannelPicker newPicker) {
370
        if (currentState == SHUTDOWN) {
1✔
371
          return;
×
372
        }
373

374
        currentState = newState;
1✔
375
        currentPicker = newPicker;
1✔
376
        // If we are already in the process of resolving addresses, the overall balancing state
377
        // will be updated at the end of it, and we don't need to trigger that update here.
378
        if (!resolvingAddresses) {
1✔
379
          updateOverallBalancingState();
1✔
380
        }
381
      }
1✔
382

383
      @Override
384
      protected Helper delegate() {
385
        return helper;
1✔
386
      }
387
    }
388
  }
389

390
  /**
391
   * Endpoint is an optimization to quickly lookup and compare EquivalentAddressGroup address sets.
392
   * It ignores the attributes. Is used as a key for ChildLbState for most load balancers
393
   * (ClusterManagerLB uses a String).
394
   */
395
  protected static class Endpoint {
396
    final Collection<SocketAddress> addrs;
397
    final int hashCode;
398

399
    public Endpoint(EquivalentAddressGroup eag) {
1✔
400
      checkNotNull(eag, "eag");
1✔
401

402
      if (eag.getAddresses().size() < 10) {
1✔
403
        addrs = eag.getAddresses();
1✔
404
      } else {
405
        // This is expected to be very unlikely in practice
406
        addrs = new HashSet<>(eag.getAddresses());
1✔
407
      }
408
      int sum = 0;
1✔
409
      for (SocketAddress address : eag.getAddresses()) {
1✔
410
        sum += address.hashCode();
1✔
411
      }
1✔
412
      hashCode = sum;
1✔
413
    }
1✔
414

415
    @Override
416
    public int hashCode() {
417
      return hashCode;
1✔
418
    }
419

420
    @Override
421
    public boolean equals(Object other) {
422
      if (this == other) {
1✔
423
        return true;
1✔
424
      }
425

426
      if (!(other instanceof Endpoint)) {
1✔
427
        return false;
1✔
428
      }
429
      Endpoint o = (Endpoint) other;
1✔
430
      if (o.hashCode != hashCode || o.addrs.size() != addrs.size()) {
1✔
431
        return false;
1✔
432
      }
433

434
      return o.addrs.containsAll(addrs);
1✔
435
    }
436

437
    @Override
438
    public String toString() {
439
      return addrs.toString();
1✔
440
    }
441
  }
442

443
  protected static class AcceptResolvedAddrRetVal {
444
    public final Status status;
445
    public final List<ChildLbState> removedChildren;
446

447
    public AcceptResolvedAddrRetVal(Status status, List<ChildLbState> removedChildren) {
1✔
448
      this.status = status;
1✔
449
      this.removedChildren = removedChildren;
1✔
450
    }
1✔
451
  }
452
}
STATUS · Troubleshooting · Open an Issue · Sales · Support · CAREERS · ENTERPRISE · START FREE · SCHEDULE DEMO
ANNOUNCEMENTS · TWITTER · TOS & SLA · Supported CI Services · What's a CI service? · Automated Testing

© 2026 Coveralls, Inc