View Javadoc
1   /*
2    * Licensed to the Apache Software Foundation (ASF) under one or more
3    * contributor license agreements.  See the NOTICE file distributed with
4    * this work for additional information regarding copyright ownership.
5    * The ASF licenses this file to You under the Apache License, Version 2.0
6    * (the "License"); you may not use this file except in compliance with
7    * the License.  You may obtain a copy of the License at
8    *
9    *     http://www.apache.org/licenses/LICENSE-2.0
10   *
11   * Unless required by applicable law or agreed to in writing, software
12   * distributed under the License is distributed on an "AS IS" BASIS,
13   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14   * See the License for the specific language governing permissions and
15   * limitations under the License.
16   */
17  
18  package org.apache.shardingsphere.traffic.engine;
19  
20  import lombok.RequiredArgsConstructor;
21  import org.apache.shardingsphere.infra.algorithm.loadbalancer.core.LoadBalanceAlgorithm;
22  import org.apache.shardingsphere.infra.instance.InstanceContext;
23  import org.apache.shardingsphere.infra.instance.metadata.InstanceMetaData;
24  import org.apache.shardingsphere.infra.instance.metadata.InstanceType;
25  import org.apache.shardingsphere.infra.session.query.QueryContext;
26  import org.apache.shardingsphere.traffic.rule.TrafficRule;
27  import org.apache.shardingsphere.traffic.rule.TrafficStrategyRule;
28  
29  import java.util.ArrayList;
30  import java.util.Map;
31  import java.util.Optional;
32  
33  /**
34   * Traffic engine.
35   */
36  @RequiredArgsConstructor
37  public final class TrafficEngine {
38      
39      private final TrafficRule trafficRule;
40      
41      private final InstanceContext instanceContext;
42      
43      /**
44       * Dispatch.
45       *
46       * @param queryContext query context
47       * @param inTransaction is in transaction
48       * @return instance id
49       */
50      public Optional<String> dispatch(final QueryContext queryContext, final boolean inTransaction) {
51          Optional<TrafficStrategyRule> strategyRule = trafficRule.findMatchedStrategyRule(queryContext, inTransaction);
52          if (!strategyRule.isPresent() || isInvalidStrategyRule(strategyRule.get())) {
53              return Optional.empty();
54          }
55          Map<String, InstanceMetaData> instances = instanceContext.getAllClusterInstances(InstanceType.PROXY, strategyRule.get().getLabels());
56          if (!instances.isEmpty()) {
57              LoadBalanceAlgorithm loadBalancer = strategyRule.get().getLoadBalancer();
58              String instanceId = 1 == instances.size() ? instances.keySet().iterator().next() : loadBalancer.getTargetName(strategyRule.get().getName(), new ArrayList<>(instances.keySet()));
59              return Optional.of(instanceId);
60          }
61          return Optional.empty();
62      }
63      
64      private boolean isInvalidStrategyRule(final TrafficStrategyRule strategyRule) {
65          return strategyRule.getLabels().isEmpty() || null == strategyRule.getLoadBalancer();
66      }
67  }