ShardingSphere自定义分布式主键生成策略、自定义分片规则

主键生成策略源码

开发者手册

KeyGenerateAlgorithm

全限定类名org.apache.shardingsphere.sharding.spi.KeyGenerateAlgorithm

分布式主键生成算法，已知实现

源码入口

package org.apache.shardingsphere.sharding.factory;

/**
 * Key generate algorithm factory.
 */
@NoArgsConstructor(access = AccessLevel.PRIVATE)
public final class KeyGenerateAlgorithmFactory {
    //加载所有的主键生成策略
    static {
        ShardingSphereServiceLoader.register(KeyGenerateAlgorithm.class);
    }
    
    /**
     * 根据配置的主键生成策略，获取一个主键生成算法
     * 例如：spring.shardingsphere.rules.sharding.key-generators.usercourse_keygen.type=SNOWFLAKE
     */
    public static KeyGenerateAlgorithm newInstance(final AlgorithmConfiguration keyGenerateAlgorithmConfig) {
        return ShardingSphereAlgorithmFactory.createAlgorithm(keyGenerateAlgorithmConfig, KeyGenerateAlgorithm.class);
    }
    
    /**
     * 判断是否包含配置的算法
     */
    public static boolean contains(final String keyGenerateAlgorithmType) {
        return TypedSPIRegistry.findRegisteredService(KeyGenerateAlgorithm.class, keyGenerateAlgorithmType).isPresent();
    }
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

先来看主键生成策略是如何加载的：ShardingSphereServiceLoader.register(KeyGenerateAlgorithm.class);

public final class ShardingSphereServiceLoader {
    //线程安全Map，缓存所有主键生成器
    private static final Map<Class<?>, Collection<Object>> SERVICES = new ConcurrentHashMap<>();
    
    // 进入到register()方法中
    public static void register(final Class<?> serviceInterface) {
        if (!SERVICES.containsKey(serviceInterface)) {
            // 调用下方的load()方法
            SERVICES.put(serviceInterface, load(serviceInterface));
        }
    }
    
    //使用java的SPI机制加载接口的所有实现类
    private static <T> Collection<Object> load(final Class<T> serviceInterface) {
        Collection<Object> result = new LinkedList<>();
        for (T each : ServiceLoader.load(serviceInterface)) {
            result.add(each);
        }
        return result;
    }
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

实现

ShardingJDBC是通过SPI机制，加载org.apache.shardingsphere.sharding.spi.KeyGenerateAlgorithm接口的实现类，也就是上方表格中的内容

我们就可以直接在yml配置文件中进行配置分布式主键生成算法

接下来就以SNOWFLAKE雪花算法举例，下方就列举出了几个关键方法

// 实现了KeyGenerateAlgorithm接口
public final class SnowflakeKeyGenerateAlgorithm implements KeyGenerateAlgorithm, InstanceContextAware {
    
    // 在init方法中，会把我们yml配置文件中定义的props的配置项，保存在下面方法的形参中，并赋值给props成员属性
    // 其他地方再用props对象获取我们的配置项
    @Override
    public void init(final Properties props) {
        this.props = props;
        maxVibrationOffset = getMaxVibrationOffset(props);
        maxTolerateTimeDifferenceMilliseconds = getMaxTolerateTimeDifferenceMilliseconds(props);
    }
    
    // 实现KeyGenerateAlgorithm接口中的抽象方法generateKey()
    // 也就是在这个方法中具体生成分布式主键值的
    @Override
    public synchronized Long generateKey() {
        long currentMilliseconds = timeService.getCurrentMillis();
        if (waitTolerateTimeDifferenceIfNeed(currentMilliseconds)) {
            currentMilliseconds = timeService.getCurrentMillis();
        }
        if (lastMilliseconds == currentMilliseconds) {
            if (0L == (sequence = (sequence + 1) & SEQUENCE_MASK)) {
                currentMilliseconds = waitUntilNextTime(currentMilliseconds);
            }
        } else {
            vibrateSequenceOffset();
            sequence = sequenceOffset;
        }
        lastMilliseconds = currentMilliseconds;
        return ((currentMilliseconds - EPOCH) << TIMESTAMP_LEFT_SHIFT_BITS) | (getWorkerId() << WORKER_ID_LEFT_SHIFT_BITS) | sequence;
    }
    
    // getType() 方法中返回的字符串就是我们上方yml配置文件中type配置项填写的值
    @Override
    public String getType() {
        return "SNOWFLAKE";
    }
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38

其他几个实现类也是一样的格式

扩展 自定义分布式主键生成策略

package com.hs.sharding.algorithm;

import com.google.common.base.Preconditions;
import org.apache.shardingsphere.infra.instance.InstanceContext;
import org.apache.shardingsphere.infra.instance.InstanceContextAware;
import org.apache.shardingsphere.sharding.algorithm.keygen.TimeService;
import org.apache.shardingsphere.sharding.spi.KeyGenerateAlgorithm;

import java.util.Calendar;
import java.util.Properties;

/**
 * 改进雪花算法，让他能够 %4 均匀分布。
 * @auth hs
 */
public final class MySnowFlakeAlgorithm implements KeyGenerateAlgorithm, InstanceContextAware {

    public static final long EPOCH;

    private static final String MAX_VIBRATION_OFFSET_KEY = "max-vibration-offset";

    private static final String MAX_TOLERATE_TIME_DIFFERENCE_MILLISECONDS_KEY = "max-tolerate-time-difference-milliseconds";

    private static final long SEQUENCE_BITS = 12L;

    private static final long WORKER_ID_BITS = 10L;

    private static final long SEQUENCE_MASK = (1 << SEQUENCE_BITS) - 1;

    private static final long WORKER_ID_LEFT_SHIFT_BITS = SEQUENCE_BITS;

    private static final long TIMESTAMP_LEFT_SHIFT_BITS = WORKER_ID_LEFT_SHIFT_BITS + WORKER_ID_BITS;

    private static final int DEFAULT_VIBRATION_VALUE = 1;

    private static final int MAX_TOLERATE_TIME_DIFFERENCE_MILLISECONDS = 10;

    private static final long DEFAULT_WORKER_ID = 0;

    private static TimeService timeService = new TimeService();

    public static void setTimeService(TimeService timeService) {
        MySnowFlakeAlgorithm.timeService = timeService;
    }

    private Properties props;

    @Override
    public Properties getProps() {
        return props;
    }

    private int maxVibrationOffset;

    private int maxTolerateTimeDifferenceMilliseconds;

    private volatile int sequenceOffset = -1;

    private volatile long sequence;

    private volatile long lastMilliseconds;

    private volatile InstanceContext instanceContext;

    static {
        Calendar calendar = Calendar.getInstance();
        calendar.set(2016, Calendar.NOVEMBER, 1);
        calendar.set(Calendar.HOUR_OF_DAY, 0);
        calendar.set(Calendar.MINUTE, 0);
        calendar.set(Calendar.SECOND, 0);
        calendar.set(Calendar.MILLISECOND, 0);
        EPOCH = calendar.getTimeInMillis();
    }

    @Override
    public void init(final Properties props) {
        this.props = props;
        maxVibrationOffset = getMaxVibrationOffset(props);
        maxTolerateTimeDifferenceMilliseconds = getMaxTolerateTimeDifferenceMilliseconds(props);
    }

    @Override
    public void setInstanceContext(final InstanceContext instanceContext) {
        this.instanceContext = instanceContext;
        if (null != instanceContext) {
            instanceContext.generateWorkerId(props);
        }
    }

    private int getMaxVibrationOffset(final Properties props) {
        int result = Integer.parseInt(props.getOrDefault(MAX_VIBRATION_OFFSET_KEY, DEFAULT_VIBRATION_VALUE).toString());
        Preconditions.checkArgument(result >= 0 && result <= SEQUENCE_MASK, "Illegal max vibration offset.");
        return result;
    }

    private int getMaxTolerateTimeDifferenceMilliseconds(final Properties props) {
        return Integer.parseInt(props.getOrDefault(MAX_TOLERATE_TIME_DIFFERENCE_MILLISECONDS_KEY, MAX_TOLERATE_TIME_DIFFERENCE_MILLISECONDS).toString());
    }

    @Override
    public synchronized Long generateKey() {
        long currentMilliseconds = timeService.getCurrentMillis();
        if (waitTolerateTimeDifferenceIfNeed(currentMilliseconds)) {
            currentMilliseconds = timeService.getCurrentMillis();
        }
        if (lastMilliseconds == currentMilliseconds) {
//            if (0L == (sequence = (sequence + 1) & SEQUENCE_MASK)) {
                currentMilliseconds = waitUntilNextTime(currentMilliseconds);
//            }
        } else {
            vibrateSequenceOffset();
//            sequence = sequenceOffset;
            sequence = sequence >= SEQUENCE_MASK ? 0:sequence+1;
        }
        lastMilliseconds = currentMilliseconds;
        return ((currentMilliseconds - EPOCH) << TIMESTAMP_LEFT_SHIFT_BITS) | (getWorkerId() << WORKER_ID_LEFT_SHIFT_BITS) | sequence;
    }

    private boolean waitTolerateTimeDifferenceIfNeed(final long currentMilliseconds) {
        if (lastMilliseconds <= currentMilliseconds) {
            return false;
        }
        long timeDifferenceMilliseconds = lastMilliseconds - currentMilliseconds;
        Preconditions.checkState(timeDifferenceMilliseconds < maxTolerateTimeDifferenceMilliseconds,
                "Clock is moving backwards, last time is %d milliseconds, current time is %d milliseconds", lastMilliseconds, currentMilliseconds);
        try {
            Thread.sleep(timeDifferenceMilliseconds);
        } catch (InterruptedException e) {
        }
        return true;
    }

    private long waitUntilNextTime(final long lastTime) {
        long result = timeService.getCurrentMillis();
        while (result <= lastTime) {
            result = timeService.getCurrentMillis();
        }
        return result;
    }

    @SuppressWarnings("NonAtomicOperationOnVolatileField")
    private void vibrateSequenceOffset() {
        sequenceOffset = sequenceOffset >= maxVibrationOffset ? 0 : sequenceOffset + 1;
    }

    private long getWorkerId() {
        return null == instanceContext ? DEFAULT_WORKER_ID : instanceContext.getWorkerId();
    }

    @Override
    public String getType() {
        return "MYSNOWFLAKE";
    }

    @Override
    public boolean isDefault() {
        return true;
    }
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159

使用spi机制加载我们上方定义的类

yml配置文件中使用我们自己定义的类

分片算法

开发者手册

ShardingAlgorithm

全限定类名org.apache.shardingsphere.sharding.spi.ShardingAlgorithm

分片算法，已知实现

实现

这里就拿CLASS_BASED自定义分片策略来举例。我们之前的配置项如下所示。

这里就有一个问题，props的值我怎么知道写什么，我又怎么知道我自定义的类需要实现什么接口？

我们现在进入到CLASS_BASED分片算法的实现类中ClassBasedShardingAlgorithm去看看它的源码

public final class ClassBasedShardingAlgorithm implements StandardShardingAlgorithm<Comparable<?>>, ComplexKeysShardingAlgorithm<Comparable<?>>, HintShardingAlgorithm<Comparable<?>> {
    
    // 定义两个常量，我们会发现这里就是props中我们进行配置的值
    private static final String STRATEGY_KEY = "strategy";
    
    private static final String ALGORITHM_CLASS_NAME_KEY = "algorithmClassName";
    
    @Getter
    private Properties props;
    
    private ClassBasedShardingAlgorithmStrategyType strategy;
    
    private String algorithmClassName;
    
    private StandardShardingAlgorithm standardShardingAlgorithm;
    
    private ComplexKeysShardingAlgorithm complexKeysShardingAlgorithm;
    
    private HintShardingAlgorithm hintShardingAlgorithm;
    
    // init()方法中会获取到props对象，props对象中保存了我们yml配置文件中的配置内容
    // 这里就会取出来，赋值给 strategy  和  algorithmClassName 成员属性
    @Override
    public void init(final Properties props) {
        this.props = props;
        strategy = getStrategy(props);
        algorithmClassName = getAlgorithmClassName(props);
        initAlgorithmInstance(props);
    }
    
    private ClassBasedShardingAlgorithmStrategyType getStrategy(final Properties props) {
        String strategy = props.getProperty(STRATEGY_KEY);
        Preconditions.checkNotNull(strategy, "Properties `%s` can not be null when uses class based sharding strategy.", STRATEGY_KEY);
        return ClassBasedShardingAlgorithmStrategyType.valueOf(strategy.toUpperCase().trim());
    }
    
    private String getAlgorithmClassName(final Properties props) {
        String result = props.getProperty(ALGORITHM_CLASS_NAME_KEY);
        Preconditions.checkNotNull(result, "Properties `%s` can not be null when uses class based sharding strategy.", ALGORITHM_CLASS_NAME_KEY);
        return result;
    }
    
    // 这里就会判断 strategy 属性是哪一个  STANDARD、COMPLEX、HINT
    // 然后在进行具体的实例 StandardShardingAlgorithm、ComplexKeysShardingAlgorithm、HintShardingAlgorithm
    private void initAlgorithmInstance(final Properties props) {
        switch (strategy) {
            case STANDARD:
                standardShardingAlgorithm = ClassBasedShardingAlgorithmFactory.newInstance(algorithmClassName, StandardShardingAlgorithm.class, props);
                break;
            case COMPLEX:
                complexKeysShardingAlgorithm = ClassBasedShardingAlgorithmFactory.newInstance(algorithmClassName, ComplexKeysShardingAlgorithm.class, props);
                break;
            case HINT:
                hintShardingAlgorithm = ClassBasedShardingAlgorithmFactory.newInstance(algorithmClassName, HintShardingAlgorithm.class, props);
                break;
            default:
                break;
        }
    }
    
    
    // doSharding()方法，具体的分片算法逻辑
    @SuppressWarnings("unchecked")
    @Override
    public String doSharding(final Collection<String> availableTargetNames, final PreciseShardingValue<Comparable<?>> shardingValue) {
        return standardShardingAlgorithm.doSharding(availableTargetNames, shardingValue);
    }
    
    @SuppressWarnings("unchecked")
    @Override
    public Collection<String> doSharding(final Collection<String> availableTargetNames, final RangeShardingValue<Comparable<?>> shardingValue) {
        return standardShardingAlgorithm.doSharding(availableTargetNames, shardingValue);
    }
    
    @SuppressWarnings("unchecked")
    @Override
    public Collection<String> doSharding(final Collection<String> availableTargetNames, final ComplexKeysShardingValue<Comparable<?>> shardingValue) {
        return complexKeysShardingAlgorithm.doSharding(availableTargetNames, shardingValue);
    }
    
    @SuppressWarnings("unchecked")
    @Override
    public Collection<String> doSharding(final Collection<String> availableTargetNames, final HintShardingValue<Comparable<?>> shardingValue) {
        return hintShardingAlgorithm.doSharding(availableTargetNames, shardingValue);
    }
    
    // 返回trye为 CLASS_BASED  这里也就是和yml配置文件中的type对应上了
    @Override
    public String getType() {
        return "CLASS_BASED";
    }
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92

其他的分片算法也是类似的实现

扩展 自定义分片算法

自定义一个java类，实现ShardingAlgorithm接口，或者是它的子接口 StandardShardingAlgorithm、ComplexKeysShardingAlgorithm、HintShardingAlgorithm都行，重写其中的doSharding()方法，我们自己指定分片逻辑

重写getType()方法，返回一个字符串，能够让我们在yml配置文件中进行配置

@Override
public String getType() {
    return "MY_COMPLEX_ALGORITHM";
}
1
2
3
4

例如我现在自定义的分片类如下

package com.hs.sharding.algorithm;

import com.google.common.base.Preconditions;
import com.google.common.collect.Range;
import org.apache.shardingsphere.sharding.api.sharding.standard.PreciseShardingValue;
import org.apache.shardingsphere.sharding.api.sharding.standard.RangeShardingValue;
import org.apache.shardingsphere.sharding.api.sharding.standard.StandardShardingAlgorithm;


import java.util.*;

/**
 * 自定义分片策略 , 我们这里实现标准的分片算法接口StandardShardingAlgorithm
 * 我这里是分片逻辑就是按照个数取模，在分发到sys_user1   sys_user2数据表中
 */
public class HsComplexAlgorithm implements StandardShardingAlgorithm<Long> {

    /**
     * 数据库个数
     */
    private final String DB_COUNT = "db-count";
    /**
     * 数据表个数
     */
    private final String TAB_COUNT = "tab-count";
    /**
     * 真实数据表前缀
     */
    private final String PERTAB = "pertab";


    private Integer dbCount;
    private Integer tabCount;
    private String pertab;

    private Properties props;

    @Override
    public void init(Properties props) {
        this.props = props;
        this.dbCount = getDbCount(props);
        this.tabCount = getTabCount(props);
        this.pertab = getPertab(props);
        // 校验条件
        Preconditions.checkState(null != pertab && !pertab.isEmpty(),
                "Inline hsComplex algorithm expression cannot be null or empty.");
    }

    /**
     * 精确查询分片执行接口（对应的sql是where ??=值)
     * @param collection 可用的分片名集合(分库就是库名，分表就是表名)
     * @param preciseShardingValue 分片键
     */
    @Override
    public String doSharding(Collection<String> collection, PreciseShardingValue<Long> preciseShardingValue) {

        Long uid = preciseShardingValue.getValue();
        String resultTableName = pertab + ((uid + 1) % (dbCount * tabCount) / tabCount + 1);
        if (collection.contains(resultTableName)){
            return resultTableName;
        }
        throw new UnsupportedOperationException("route: " + resultTableName + " is not supported, please check your config");
    }

    /**
     * 范围分片规则（对应的是where ??>='XXX' and ??<='XXX')
     * 范围查询分片算法（分片键涉及区间查询时会进入该方法进行分片计算）
     */
    @Override
    public Collection<String> doSharding(Collection<String> collection, RangeShardingValue<Long> rangeShardingValue) {
        List<String> result = new ArrayList<>();
        Range<Long> valueRange = rangeShardingValue.getValueRange();
        Long upperEndpoint = valueRange.upperEndpoint();
        Long aLong = valueRange.lowerEndpoint();

        // TODO 进行相应的分片判断
//        return result;

        return collection;
    }

    private String getPertab(Properties props) {
        return props.getProperty(PERTAB);
    }

    private Integer getDbCount(Properties props) {
        String count = props.getProperty(DB_COUNT);
        return count == null || count.isEmpty() ? 0 : Integer.valueOf(count);
    }

    private Integer getTabCount(Properties props) {
        String count = props.getProperty(TAB_COUNT);
        return count == null || count.isEmpty() ? 0 : Integer.valueOf(count);
    }


    @Override
    public Properties getProps() {
        return props;
    }

    @Override
    public String getType() {
        return "HS";
    }
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107

需要添加一个SPI的配置文件org.apache.shardingsphere.sharding.spi.ShardingAlgorithm，在该文件中指定我们上方创建的java类

yml配置文件中进行相应的更改

踩的坑

我先是自定义的类实现的是ComplexKeysShardingAlgorithm接口，但是我们yml配置类中还是一直按照standard的配置，导致我自定义的类中的doSharding()方法所以就一直没有调用到

之后我修改了complex就能调用了

在配置还是standard时，我通过debug，发现init() 和 getType()方法都能够调用，证明API机制相关的文件没问题。

我就想会不会是单分片键、精确查询、范围查询相关问题导致的？

我修改了实现接口，改为了StandardShardingAlgorithm，然后就进入了其中单分片键的doSharding()方法。最后就一点一点的排查，再到了这上面的配置