001/*
002 *  Copyright (c) 2022-2023, Mybatis-Flex (fuhai999@gmail.com).
003 *  <p>
004 *  Licensed under the Apache License, Version 2.0 (the "License");
005 *  you may not use this file except in compliance with the License.
006 *  You may obtain a copy of the License at
007 *  <p>
008 *  http://www.apache.org/licenses/LICENSE-2.0
009 *  <p>
010 *  Unless required by applicable law or agreed to in writing, software
011 *  distributed under the License is distributed on an "AS IS" BASIS,
012 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 *  See the License for the specific language governing permissions and
014 *  limitations under the License.
015 */
016package com.mybatisflex.core.keygen.impl;
017
018import com.mybatisflex.core.exception.FlexExceptions;
019import com.mybatisflex.core.keygen.IKeyGenerator;
020import com.mybatisflex.core.util.StringUtil;
021
022import java.lang.management.ManagementFactory;
023import java.net.InetAddress;
024import java.net.NetworkInterface;
025
026/**
027 * <p>雪花算法 ID 生成器。
028 *
029 * <ul>
030 *     <li>最高 1 位固定值 0，因为生成的 ID 是正整数；
031 *     <li>接下来 41 位存储毫秒级时间戳，2 ^ 41 / ( 1000 * 60 * 60 * 24 * 365) = 69，大概可以使用 69 年；
032 *     <li>再接下 10 位存储机器码，包括 5 位 dataCenterId 和 5 位 workerId，最多可以部署 2 ^ 10 = 1024 台机器；
033 *     <li>最后 12 位存储序列号，同一毫秒时间戳时，通过这个递增的序列号来区分，即对于同一台机器而言，同一毫秒时间戳下，可以生成 2 ^ 12 = 4096 个不重复 ID。
034 * </ul>
035 *
036 * <p>优化自开源项目：<a href="https://gitee.com/yu120/sequence">Sequence</a>
037 *
038 * @author 王帅
039 * @since 2023-05-12
040 */
041public class SnowFlakeIDKeyGenerator implements IKeyGenerator {
042
043    /**
044     * 工作机器 ID 占用的位数（5bit）。
045     */
046    private static final long WORKER_ID_BITS = 5L;
047    /**
048     * 数据中心 ID 占用的位数（5bit）。
049     */
050    private static final long DATA_CENTER_ID_BITS = 5L;
051    /**
052     * 序号占用的位数（12bit）。
053     */
054    private static final long SEQUENCE_BITS = 12L;
055    /**
056     * 工作机器 ID 占用 5bit 时的最大值 31。
057     */
058    private static final long MAX_WORKER_ID = ~(-1L << WORKER_ID_BITS);
059    /**
060     * 数据中心 ID 占用 5bit 时的最大值 31。
061     */
062    private static final long MAX_DATA_CENTER_ID = ~(-1L << DATA_CENTER_ID_BITS);
063    /**
064     * 序号掩码，用于与自增后的序列号进行位“与”操作，如果值为 0，则代表自增后的序列号超过了 4095。
065     */
066    private static final long SEQUENCE_MASK = ~(-1L << SEQUENCE_BITS);
067    /**
068     * 工作机器 ID 位需要左移的位数（12bit）。
069     */
070    private static final long WORK_ID_SHIFT = SEQUENCE_BITS;
071    /**
072     * 数据中心 ID 位需要左移的位数（12bit + 5bit）。
073     */
074    private static final long DATA_CENTER_ID_SHIFT = SEQUENCE_BITS + WORKER_ID_BITS;
075    /**
076     * 时间戳需要左移的位数（12bit + 5bit + 5bit）。
077     */
078    private static final long TIMESTAMP_SHIFT = SEQUENCE_BITS + WORKER_ID_BITS + DATA_CENTER_ID_BITS;
079    /**
080     * 时间起始标记点，一旦确定不能变动（2023-04-02 13:01:00）。
081     */
082    protected static long twepoch = 1680411660000L;
083    /**
084     * 可容忍的时间偏移量。
085     */
086    protected static long offsetPeriod = 5L;
087    /**
088     * 工作机器 ID。
089     */
090    private final long workerId;
091    /**
092     * 数据中心 ID。
093     */
094    private final long dataCenterId;
095    /**
096     * IP 地址信息，用来生成工作机器 ID 和数据中心 ID。
097     */
098    protected InetAddress address;
099    /**
100     * 同一毫秒内的最新序号，最大值可为（2^12 - 1 = 4095）。
101     */
102    private long sequence;
103    /**
104     * 上次生产 ID 时间戳。
105     */
106    private long lastTimeMillis = -1L;
107
108    /**
109     * 雪花算法 ID 生成器。
110     */
111    public SnowFlakeIDKeyGenerator() {
112        this(null);
113    }
114
115    /**
116     * 根据 IP 地址计算数据中心 ID 和工作机器 ID 生成数据库 ID。
117     *
118     * @param address IP 地址
119     */
120    public SnowFlakeIDKeyGenerator(InetAddress address) {
121        this.address = address;
122        this.dataCenterId = getDataCenterId(MAX_DATA_CENTER_ID);
123        this.workerId = getWorkerId(dataCenterId, MAX_WORKER_ID);
124    }
125
126    /**
127     * 根据数据中心 ID 和工作机器 ID 生成数据库 ID。
128     *
129     * @param workerId     工作机器 ID
130     * @param dataCenterId 数据中心 ID
131     */
132    public SnowFlakeIDKeyGenerator(long workerId, long dataCenterId) {
133        if (workerId > MAX_WORKER_ID || workerId < 0) {
134            throw new IllegalArgumentException(
135                String.format("workerId must be greater than 0 and less than %d.", MAX_WORKER_ID));
136        }
137        if (dataCenterId > MAX_DATA_CENTER_ID || dataCenterId < 0) {
138            throw new IllegalArgumentException(
139                String.format("dataCenterId must be greater than 0 and less than %d.", MAX_DATA_CENTER_ID));
140        }
141        this.workerId = workerId;
142        this.dataCenterId = dataCenterId;
143    }
144
145    /**
146     * 根据 MAC + PID 的 hashCode 获取 16 个低位生成工作机器 ID。
147     */
148    protected long getWorkerId(long dataCenterId, long maxWorkerId) {
149        StringBuilder mpId = new StringBuilder();
150        mpId.append(dataCenterId);
151        String name = ManagementFactory.getRuntimeMXBean().getName();
152        if (StringUtil.isNotBlank(name)) {
153            // GET jvmPid
154            mpId.append(name.split("@")[0]);
155        }
156        // MAC + PID 的 hashCode 获取16个低位
157        return (mpId.toString().hashCode() & 0xffff) % (maxWorkerId + 1);
158    }
159
160    /**
161     * 根据网卡 MAC 地址计算余数作为数据中心 ID。
162     */
163    protected long getDataCenterId(long maxDataCenterId) {
164        long id = 0L;
165        try {
166            if (address == null) {
167                address = InetAddress.getLocalHost();
168            }
169            NetworkInterface network = NetworkInterface.getByInetAddress(address);
170            if (null == network) {
171                id = 1L;
172            } else {
173                byte[] mac = network.getHardwareAddress();
174                if (null != mac) {
175                    id = ((0x000000FF & (long) mac[mac.length - 2]) | (0x0000FF00 & (((long) mac[mac.length - 1]) << 8))) >> 6;
176                    id = id % (maxDataCenterId + 1);
177                }
178            }
179        } catch (Exception e) {
180            throw FlexExceptions.wrap(e, "dataCenterId: %s", e.getMessage());
181        }
182        return id;
183    }
184
185    @Override
186    public Object generate(Object entity, String keyColumn) {
187        return nextId();
188    }
189
190    /**
191     * 获取下一个 ID。
192     */
193    public synchronized long nextId() {
194        long currentTimeMillis = System.currentTimeMillis();
195        // 当前时间小于上一次生成 ID 使用的时间，可能出现服务器时钟回拨问题。
196        if (currentTimeMillis < lastTimeMillis) {
197            long offset = lastTimeMillis - currentTimeMillis;
198            // 在可容忍的时间差值之内等待时间恢复正常
199            if (offset <= offsetPeriod) {
200                try {
201                    wait(offset << 1L);
202                    currentTimeMillis = System.currentTimeMillis();
203                    if (currentTimeMillis < lastTimeMillis) {
204                        throw FlexExceptions.wrap("Clock moved backwards, please check the time. Current timestamp: %d, last used timestamp: %d", currentTimeMillis, lastTimeMillis);
205                    }
206                } catch (InterruptedException e) {
207                    throw FlexExceptions.wrap(e);
208                }
209            } else {
210                throw FlexExceptions.wrap("Clock moved backwards, please check the time. Current timestamp: %d, last used timestamp: %d", currentTimeMillis, lastTimeMillis);
211            }
212        }
213
214        if (currentTimeMillis == lastTimeMillis) {
215            // 相同毫秒内，序列号自增
216            sequence = (sequence + 1) & SEQUENCE_MASK;
217            if (sequence == 0) {
218                // 同一毫秒的序列数已经达到最大
219                currentTimeMillis = tilNextMillis(lastTimeMillis);
220            }
221        } else {
222            // 不同毫秒内，序列号置为 0。
223            sequence = 0L;
224        }
225
226        // 记录最后一次使用的毫秒时间戳
227        lastTimeMillis = currentTimeMillis;
228
229        // 时间戳部分 | 数据中心部分 | 机器标识部分 | 序列号部分
230        return ((currentTimeMillis - twepoch) << TIMESTAMP_SHIFT)
231            | (dataCenterId << DATA_CENTER_ID_SHIFT)
232            | (workerId << WORK_ID_SHIFT)
233            | sequence;
234    }
235
236    /**
237     * 获取指定时间戳的接下来的时间戳。
238     */
239    private long tilNextMillis(long lastTimestamp) {
240        long currentTimeMillis = System.currentTimeMillis();
241        while (currentTimeMillis <= lastTimestamp) {
242            currentTimeMillis = System.currentTimeMillis();
243        }
244        return currentTimeMillis;
245    }
246
247}