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

grpc / grpc-java / #19441

29 Aug 2024 08:13PM UTC coverage: 84.505% (+0.01%) from 84.494%
#19441

push

github

ejona86
Focus MultiChildLB updates around ResolvedAddresses of children

This makes ClusterManagerLB more straight-forward, focusing on just the
things that are relevant to it, and it avoids specialized map key
handling in updateChildrenWithResolvedAddresses().

33387 of 39509 relevant lines covered (84.5%)

0.85 hits per line

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

95.97
/../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.
83
   */
84
  protected Map<Object, ResolvedAddresses> createChildAddressesMap(
85
      ResolvedAddresses resolvedAddresses) {
86
    Map<Object, ResolvedAddresses> childAddresses = new HashMap<>();
1✔
87
    for (EquivalentAddressGroup eag : resolvedAddresses.getAddresses()) {
1✔
88
      ResolvedAddresses addresses = resolvedAddresses.toBuilder()
1✔
89
          .setAddresses(Collections.singletonList(eag))
1✔
90
          .setAttributes(Attributes.newBuilder().set(IS_PETIOLE_POLICY, true).build())
1✔
91
          .setLoadBalancingPolicyConfig(null)
1✔
92
          .build();
1✔
93
      childAddresses.put(new Endpoint(eag), addresses);
1✔
94
    }
1✔
95
    return childAddresses;
1✔
96
  }
97

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

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

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

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

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

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

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

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

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

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

171
    updateChildrenWithResolvedAddresses(newChildAddresses);
1✔
172

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

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

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

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

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

230
  protected final Helper getHelper() {
231
    return helper;
1✔
232
  }
233

234
  @VisibleForTesting
235
  public final Collection<ChildLbState> getChildLbStates() {
236
    return childLbStates.values();
1✔
237
  }
238

239
  @VisibleForTesting
240
  public final ChildLbState getChildLbState(Object key) {
241
    if (key == null) {
1✔
242
      return null;
×
243
    }
244
    if (key instanceof EquivalentAddressGroup) {
1✔
245
      key = new Endpoint((EquivalentAddressGroup) key);
1✔
246
    }
247
    return childLbStates.get(key);
1✔
248
  }
249

250
  @VisibleForTesting
251
  public final ChildLbState getChildLbStateEag(EquivalentAddressGroup eag) {
252
    return getChildLbState(new Endpoint(eag));
1✔
253
  }
254

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

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

283
    private final LoadBalancer lb;
284
    private ConnectivityState currentState;
285
    private SubchannelPicker currentPicker = new FixedResultPicker(PickResult.withNoResult());
1✔
286

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

293
    protected ChildLbStateHelper createChildHelper() {
294
      return new ChildLbStateHelper();
1✔
295
    }
296

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

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

314
    public final Object getKey() {
315
      return key;
1✔
316
    }
317

318
    @VisibleForTesting
319
    public final LoadBalancer getLb() {
320
      return lb;
1✔
321
    }
322

323
    @VisibleForTesting
324
    public final SubchannelPicker getCurrentPicker() {
325
      return currentPicker;
1✔
326
    }
327

328
    public final ConnectivityState getCurrentState() {
329
      return currentState;
1✔
330
    }
331

332
    protected final void setCurrentState(ConnectivityState newState) {
333
      currentState = newState;
1✔
334
    }
1✔
335

336
    protected final void setCurrentPicker(SubchannelPicker newPicker) {
337
      currentPicker = newPicker;
1✔
338
    }
1✔
339

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

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

352
    @VisibleForTesting
353
    public final ResolvedAddresses getResolvedAddresses() {
354
      return resolvedAddresses;
1✔
355
    }
356

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

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

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

386
      @Override
387
      protected Helper delegate() {
388
        return helper;
1✔
389
      }
390
    }
391
  }
392

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

402
    public Endpoint(EquivalentAddressGroup eag) {
1✔
403
      checkNotNull(eag, "eag");
1✔
404

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

418
    @Override
419
    public int hashCode() {
420
      return hashCode;
1✔
421
    }
422

423
    @Override
424
    public boolean equals(Object other) {
425
      if (this == other) {
1✔
426
        return true;
1✔
427
      }
428

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

437
      return o.addrs.containsAll(addrs);
1✔
438
    }
439

440
    @Override
441
    public String toString() {
442
      return addrs.toString();
1✔
443
    }
444
  }
445

446
  protected static class AcceptResolvedAddrRetVal {
447
    public final Status status;
448
    public final List<ChildLbState> removedChildren;
449

450
    public AcceptResolvedAddrRetVal(Status status, List<ChildLbState> removedChildren) {
1✔
451
      this.status = status;
1✔
452
      this.removedChildren = removedChildren;
1✔
453
    }
1✔
454
  }
455
}
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