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

grpc / grpc-java / #19419

12 Aug 2024 11:39PM CUT coverage: 84.476% (-0.01%) from 84.489%
#19419

push

github

ejona86
util: MultiChildLb children should always start with a NoResult picker

That's the obvious default, and all current usages use (something
equivalent to) that default.

33389 of 39525 relevant lines covered (84.48%)

0.84 hits per line

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

96.05
/../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.checkArgument;
20
import static com.google.common.base.Preconditions.checkNotNull;
21
import static io.grpc.ConnectivityState.CONNECTING;
22
import static io.grpc.ConnectivityState.IDLE;
23
import static io.grpc.ConnectivityState.READY;
24
import static io.grpc.ConnectivityState.SHUTDOWN;
25
import static io.grpc.ConnectivityState.TRANSIENT_FAILURE;
26

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

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

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

65
  protected final LoadBalancerProvider pickFirstLbProvider = new PickFirstLoadBalancerProvider();
1✔
66

67
  protected ConnectivityState currentConnectivityState;
68

69

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

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

82
  /**
83
   * Override to utilize parsing of the policy configuration or alternative helper/lb generation.
84
   */
85
  protected Map<Object, ChildLbState> createChildLbMap(ResolvedAddresses resolvedAddresses) {
86
    Map<Object, ChildLbState> childLbMap = new HashMap<>();
1✔
87
    List<EquivalentAddressGroup> addresses = resolvedAddresses.getAddresses();
1✔
88
    for (EquivalentAddressGroup eag : addresses) {
1✔
89
      Endpoint endpoint = new Endpoint(eag); // keys need to be just addresses
1✔
90
      ChildLbState existingChildLbState = childLbStates.get(endpoint);
1✔
91
      if (existingChildLbState != null) {
1✔
92
        childLbMap.put(endpoint, existingChildLbState);
1✔
93
      } else {
94
        childLbMap.put(endpoint, createChildLbState(endpoint, null, resolvedAddresses));
1✔
95
      }
96
    }
1✔
97
    return childLbMap;
1✔
98
  }
99

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

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

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

122
      // Update the picker and our connectivity state
123
      updateOverallBalancingState();
1✔
124

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

133
  /**
134
   * Override this if your keys are not of type Endpoint.
135
   * @param key Key to identify the ChildLbState
136
   * @param resolvedAddresses list of addresses which include attributes
137
   * @param childConfig a load balancing policy config. This field is optional.
138
   * @return a fully loaded ResolvedAddresses object for the specified key
139
   */
140
  protected ResolvedAddresses getChildAddresses(Object key, ResolvedAddresses resolvedAddresses,
141
      Object childConfig) {
142
    Endpoint endpointKey;
143
    if (key instanceof EquivalentAddressGroup) {
1✔
144
      endpointKey = new Endpoint((EquivalentAddressGroup) key);
×
145
    } else {
146
      checkArgument(key instanceof Endpoint, "key is wrong type");
1✔
147
      endpointKey = (Endpoint) key;
1✔
148
    }
149

150
    // Retrieve the non-stripped version
151
    EquivalentAddressGroup eagToUse = null;
1✔
152
    for (EquivalentAddressGroup currEag : resolvedAddresses.getAddresses()) {
1✔
153
      if (endpointKey.equals(new Endpoint(currEag))) {
1✔
154
        eagToUse = currEag;
1✔
155
        break;
1✔
156
      }
157
    }
1✔
158

159
    checkNotNull(eagToUse, key + " no longer present in load balancer children");
1✔
160

161
    return resolvedAddresses.toBuilder()
1✔
162
        .setAddresses(Collections.singletonList(eagToUse))
1✔
163
        .setAttributes(Attributes.newBuilder().set(IS_PETIOLE_POLICY, true).build())
1✔
164
        .setLoadBalancingPolicyConfig(childConfig)
1✔
165
        .build();
1✔
166
  }
167

168
  /**
169
   * Handle the name resolution error.
170
   *
171
   * <p/>Override if you need special handling.
172
   */
173
  @Override
174
  public void handleNameResolutionError(Status error) {
175
    if (currentConnectivityState != READY)  {
1✔
176
      helper.updateBalancingState(TRANSIENT_FAILURE, getErrorPicker(error));
1✔
177
    }
178
  }
1✔
179

180
  /**
181
   * Handle the name resolution error only for the specified child.
182
   *
183
   * <p/>Override if you need special handling.
184
   */
185
  protected void handleNameResolutionError(ChildLbState child, Status error) {
186
    child.lb.handleNameResolutionError(error);
1✔
187
  }
1✔
188

189
  /**
190
   * Creates a new picker representing an error status.
191
   *
192
   * <p/>Override to produce a custom picker when there are errors.
193
   */
194
  protected SubchannelPicker getErrorPicker(Status error)  {
195
    return new FixedResultPicker(PickResult.withError(error));
1✔
196
  }
197

198
  @Override
199
  public void shutdown() {
200
    logger.log(Level.FINE, "Shutdown");
1✔
201
    for (ChildLbState state : childLbStates.values()) {
1✔
202
      state.shutdown();
1✔
203
    }
1✔
204
    childLbStates.clear();
1✔
205
  }
1✔
206

207
  /**
208
   *   This does the work to update the child map and calculate which children have been removed.
209
   *   You must call {@link #updateOverallBalancingState} to update the picker
210
   *   and call {@link #shutdownRemoved(List)} to shutdown the endpoints that have been removed.
211
    */
212
  protected final AcceptResolvedAddrRetVal acceptResolvedAddressesInternal(
213
      ResolvedAddresses resolvedAddresses) {
214
    logger.log(Level.FINE, "Received resolution result: {0}", resolvedAddresses);
1✔
215

216
    // Subclass handles any special manipulation to create appropriate types of keyed ChildLbStates
217
    Map<Object, ChildLbState> newChildren = createChildLbMap(resolvedAddresses);
1✔
218

219
    // Handle error case
220
    if (newChildren.isEmpty()) {
1✔
221
      Status unavailableStatus = Status.UNAVAILABLE.withDescription(
1✔
222
          "NameResolver returned no usable address. " + resolvedAddresses);
223
      handleNameResolutionError(unavailableStatus);
1✔
224
      return new AcceptResolvedAddrRetVal(unavailableStatus, null);
1✔
225
    }
226

227
    addMissingChildren(newChildren);
1✔
228

229
    updateChildrenWithResolvedAddresses(resolvedAddresses, newChildren);
1✔
230

231
    return new AcceptResolvedAddrRetVal(Status.OK, getRemovedChildren(newChildren.keySet()));
1✔
232
  }
233

234
  protected final void addMissingChildren(Map<Object, ChildLbState> newChildren) {
235
    // Do adds and identify reused children
236
    for (Map.Entry<Object, ChildLbState> entry : newChildren.entrySet()) {
1✔
237
      final Object key = entry.getKey();
1✔
238
      if (!childLbStates.containsKey(key)) {
1✔
239
        childLbStates.put(key, entry.getValue());
1✔
240
      }
241
    }
1✔
242
  }
1✔
243

244
  protected final void updateChildrenWithResolvedAddresses(ResolvedAddresses resolvedAddresses,
245
                                                     Map<Object, ChildLbState> newChildren) {
246
    for (Map.Entry<Object, ChildLbState> entry : newChildren.entrySet()) {
1✔
247
      Object childConfig = entry.getValue().getConfig();
1✔
248
      ChildLbState childLbState = childLbStates.get(entry.getKey());
1✔
249
      ResolvedAddresses childAddresses =
1✔
250
          getChildAddresses(entry.getKey(), resolvedAddresses, childConfig);
1✔
251
      childLbState.setResolvedAddresses(childAddresses); // update child
1✔
252
      childLbState.lb.handleResolvedAddresses(childAddresses); // update child LB
1✔
253
    }
1✔
254
  }
1✔
255

256
  /**
257
   * Identifies which children have been removed (are not part of the newChildKeys).
258
   */
259
  protected final List<ChildLbState> getRemovedChildren(Set<Object> newChildKeys) {
260
    List<ChildLbState> removedChildren = new ArrayList<>();
1✔
261
    // Do removals
262
    for (Object key : ImmutableList.copyOf(childLbStates.keySet())) {
1✔
263
      if (!newChildKeys.contains(key)) {
1✔
264
        ChildLbState childLbState = childLbStates.remove(key);
1✔
265
        removedChildren.add(childLbState);
1✔
266
      }
267
    }
1✔
268
    return removedChildren;
1✔
269
  }
270

271
  protected final void shutdownRemoved(List<ChildLbState> removedChildren) {
272
    // Do shutdowns after updating picker to reduce the chance of failing an RPC by picking a
273
    // subchannel that has been shutdown.
274
    for (ChildLbState childLbState : removedChildren) {
1✔
275
      childLbState.shutdown();
1✔
276
    }
1✔
277
  }
1✔
278

279
  @Nullable
280
  protected static ConnectivityState aggregateState(
281
      @Nullable ConnectivityState overallState, ConnectivityState childState) {
282
    if (overallState == null) {
1✔
283
      return childState;
1✔
284
    }
285
    if (overallState == READY || childState == READY) {
1✔
286
      return READY;
1✔
287
    }
288
    if (overallState == CONNECTING || childState == CONNECTING) {
1✔
289
      return CONNECTING;
1✔
290
    }
291
    if (overallState == IDLE || childState == IDLE) {
1✔
292
      return IDLE;
×
293
    }
294
    return overallState;
1✔
295
  }
296

297
  protected final Helper getHelper() {
298
    return helper;
1✔
299
  }
300

301
  @VisibleForTesting
302
  public final ImmutableMap<Object, ChildLbState> getImmutableChildMap() {
303
    return ImmutableMap.copyOf(childLbStates);
1✔
304
  }
305

306
  @VisibleForTesting
307
  public final Collection<ChildLbState> getChildLbStates() {
308
    return childLbStates.values();
1✔
309
  }
310

311
  @VisibleForTesting
312
  public final ChildLbState getChildLbState(Object key) {
313
    if (key == null) {
1✔
314
      return null;
×
315
    }
316
    if (key instanceof EquivalentAddressGroup) {
1✔
317
      key = new Endpoint((EquivalentAddressGroup) key);
1✔
318
    }
319
    return childLbStates.get(key);
1✔
320
  }
321

322
  @VisibleForTesting
323
  public final ChildLbState getChildLbStateEag(EquivalentAddressGroup eag) {
324
    return getChildLbState(new Endpoint(eag));
1✔
325
  }
326

327
  /**
328
   * Filters out non-ready child load balancers (subchannels).
329
   */
330
  protected final List<ChildLbState> getReadyChildren() {
331
    List<ChildLbState> activeChildren = new ArrayList<>();
1✔
332
    for (ChildLbState child : getChildLbStates()) {
1✔
333
      if (child.getCurrentState() == READY) {
1✔
334
        activeChildren.add(child);
1✔
335
      }
336
    }
1✔
337
    return activeChildren;
1✔
338
  }
339

340
  /**
341
   * This represents the state of load balancer children.  Each endpoint (represented by an
342
   * EquivalentAddressGroup or EDS string) will have a separate ChildLbState which in turn will
343
   * have a single child LoadBalancer created from the provided factory.
344
   *
345
   * <p>A ChildLbStateHelper is the glue between ChildLbState and the helpers associated with the
346
   * petiole policy above and the PickFirstLoadBalancer's helper below.
347
   *
348
   * <p>If you wish to store additional state information related to each subchannel, then extend
349
   * this class.
350
   */
351
  public class ChildLbState {
352
    private final Object key;
353
    private ResolvedAddresses resolvedAddresses;
354
    private final Object config;
355

356
    private final LoadBalancer lb;
357
    private ConnectivityState currentState;
358
    private SubchannelPicker currentPicker = new FixedResultPicker(PickResult.withNoResult());
1✔
359

360
    public ChildLbState(Object key, LoadBalancer.Factory policyFactory, Object childConfig) {
1✔
361
      this.key = key;
1✔
362
      this.config = childConfig;
1✔
363
      this.lb = policyFactory.newLoadBalancer(createChildHelper());
1✔
364
      this.currentState = CONNECTING;
1✔
365
    }
1✔
366

367
    protected ChildLbStateHelper createChildHelper() {
368
      return new ChildLbStateHelper();
1✔
369
    }
370

371
    /**
372
     * Override for unique behavior such as delayed shutdowns of subchannels.
373
     */
374
    protected void shutdown() {
375
      lb.shutdown();
1✔
376
      this.currentState = SHUTDOWN;
1✔
377
      logger.log(Level.FINE, "Child balancer {0} deleted", key);
1✔
378
    }
1✔
379

380
    @Override
381
    public String toString() {
382
      return "Address = " + key
×
383
          + ", state = " + currentState
384
          + ", picker type: " + currentPicker.getClass()
×
385
          + ", lb: " + lb;
386
    }
387

388
    public final Object getKey() {
389
      return key;
1✔
390
    }
391

392
    @VisibleForTesting
393
    public final LoadBalancer getLb() {
394
      return lb;
1✔
395
    }
396

397
    @VisibleForTesting
398
    public final SubchannelPicker getCurrentPicker() {
399
      return currentPicker;
1✔
400
    }
401

402
    public final ConnectivityState getCurrentState() {
403
      return currentState;
1✔
404
    }
405

406
    protected final void setCurrentState(ConnectivityState newState) {
407
      currentState = newState;
1✔
408
    }
1✔
409

410
    protected final void setCurrentPicker(SubchannelPicker newPicker) {
411
      currentPicker = newPicker;
1✔
412
    }
1✔
413

414
    public final EquivalentAddressGroup getEag() {
415
      if (resolvedAddresses == null || resolvedAddresses.getAddresses().isEmpty()) {
1✔
416
        return null;
×
417
      }
418
      return resolvedAddresses.getAddresses().get(0);
1✔
419
    }
420

421
    protected final void setResolvedAddresses(ResolvedAddresses newAddresses) {
422
      checkNotNull(newAddresses, "Missing address list for child");
1✔
423
      resolvedAddresses = newAddresses;
1✔
424
    }
1✔
425

426
    private Object getConfig() {
427
      return config;
1✔
428
    }
429

430
    @VisibleForTesting
431
    public final ResolvedAddresses getResolvedAddresses() {
432
      return resolvedAddresses;
1✔
433
    }
434

435
    /**
436
     * ChildLbStateHelper is the glue between ChildLbState and the helpers associated with the
437
     * petiole policy above and the PickFirstLoadBalancer's helper below.
438
     *
439
     * <p>The ChildLbState updates happen during updateBalancingState.  Otherwise, it is doing
440
     * simple forwarding.
441
     */
442
    protected class ChildLbStateHelper extends ForwardingLoadBalancerHelper {
1✔
443

444
      /**
445
       * Update current state and picker for this child and then use
446
       * {@link #updateOverallBalancingState()} for the parent LB.
447
       */
448
      @Override
449
      public void updateBalancingState(final ConnectivityState newState,
450
          final SubchannelPicker newPicker) {
451
        if (!childLbStates.containsKey(key)) {
1✔
452
          return;
×
453
        }
454

455
        currentState = newState;
1✔
456
        currentPicker = newPicker;
1✔
457
        // If we are already in the process of resolving addresses, the overall balancing state
458
        // will be updated at the end of it, and we don't need to trigger that update here.
459
        if (!resolvingAddresses) {
1✔
460
          updateOverallBalancingState();
1✔
461
        }
462
      }
1✔
463

464
      @Override
465
      protected Helper delegate() {
466
        return helper;
1✔
467
      }
468
    }
469
  }
470

471
  /**
472
   * Endpoint is an optimization to quickly lookup and compare EquivalentAddressGroup address sets.
473
   * It ignores the attributes. Is used as a key for ChildLbState for most load balancers
474
   * (ClusterManagerLB uses a String).
475
   */
476
  protected static class Endpoint {
477
    final Collection<SocketAddress> addrs;
478
    final int hashCode;
479

480
    public Endpoint(EquivalentAddressGroup eag) {
1✔
481
      checkNotNull(eag, "eag");
1✔
482

483
      if (eag.getAddresses().size() < 10) {
1✔
484
        addrs = eag.getAddresses();
1✔
485
      } else {
486
        // This is expected to be very unlikely in practice
487
        addrs = new HashSet<>(eag.getAddresses());
1✔
488
      }
489
      int sum = 0;
1✔
490
      for (SocketAddress address : eag.getAddresses()) {
1✔
491
        sum += address.hashCode();
1✔
492
      }
1✔
493
      hashCode = sum;
1✔
494
    }
1✔
495

496
    @Override
497
    public int hashCode() {
498
      return hashCode;
1✔
499
    }
500

501
    @Override
502
    public boolean equals(Object other) {
503
      if (this == other) {
1✔
504
        return true;
1✔
505
      }
506

507
      if (!(other instanceof Endpoint)) {
1✔
508
        return false;
1✔
509
      }
510
      Endpoint o = (Endpoint) other;
1✔
511
      if (o.hashCode != hashCode || o.addrs.size() != addrs.size()) {
1✔
512
        return false;
1✔
513
      }
514

515
      return o.addrs.containsAll(addrs);
1✔
516
    }
517

518
    @Override
519
    public String toString() {
520
      return addrs.toString();
1✔
521
    }
522
  }
523

524
  protected static class AcceptResolvedAddrRetVal {
525
    public final Status status;
526
    public final List<ChildLbState> removedChildren;
527

528
    public AcceptResolvedAddrRetVal(Status status, List<ChildLbState> removedChildren) {
1✔
529
      this.status = status;
1✔
530
      this.removedChildren = removedChildren;
1✔
531
    }
1✔
532
  }
533
}
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

© 2025 Coveralls, Inc