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    private static long twepoch = 1680411660000L;
083    /**
084     * 可容忍的时间偏移量。
085     */
086    private static long offsetPeriod = 5L;
087
088    /**
089     * 工作机器 ID。
090     */
091    private final long workerId;
092    /**
093     * 数据中心 ID。
094     */
095    private final long dataCenterId;
096
097    /**
098     * IP 地址信息，用来生成工作机器 ID 和数据中心 ID。
099     */
100    protected InetAddress address;
101    /**
102     * 同一毫秒内的最新序号，最大值可为（2^12 - 1 = 4095）。
103     */
104    private long sequence;
105    /**
106     * 上次生产 ID 时间戳。
107     */
108    private long lastTimeMillis = -1L;
109
110    /**
111     * 雪花算法 ID 生成器。
112     */
113    public SnowFlakeIDKeyGenerator() {
114        this(null);
115    }
116
117    /**
118     * 根据 IP 地址计算数据中心 ID 和工作机器 ID 生成数据库 ID。
119     *
120     * @param address IP 地址
121     */
122    public SnowFlakeIDKeyGenerator(InetAddress address) {
123        this.address = address;
124        this.dataCenterId = getDataCenterId(MAX_DATA_CENTER_ID);
125        this.workerId = getWorkerId(dataCenterId, MAX_WORKER_ID);
126    }
127
128    /**
129     * 根据数据中心 ID 和工作机器 ID 生成数据库 ID。
130     *
131     * @param workerId     工作机器 ID
132     * @param dataCenterId 数据中心 ID
133     */
134    public SnowFlakeIDKeyGenerator(long workerId, long dataCenterId) {
135        if (workerId > MAX_WORKER_ID || workerId < 0) {
136            throw new IllegalArgumentException(
137                String.format("workerId must be greater than 0 and less than %d.", MAX_WORKER_ID));
138        }
139        if (dataCenterId > MAX_DATA_CENTER_ID || dataCenterId < 0) {
140            throw new IllegalArgumentException(
141                String.format("dataCenterId must be greater than 0 and less than %d.", MAX_DATA_CENTER_ID));
142        }
143        this.workerId = workerId;
144        this.dataCenterId = dataCenterId;
145    }
146
147    /**
148     * 根据 MAC + PID 的 hashCode 获取 16 个低位生成工作机器 ID。
149     */
150    protected long getWorkerId(long dataCenterId, long maxWorkerId) {
151        StringBuilder mpId = new StringBuilder();
152        mpId.append(dataCenterId);
153        String name = ManagementFactory.getRuntimeMXBean().getName();
154        if (StringUtil.isNotBlank(name)) {
155            // GET jvmPid
156            mpId.append(name.split("@")[0]);
157        }
158        // MAC + PID 的 hashCode 获取16个低位
159        return (mpId.toString().hashCode() & 0xffff) % (maxWorkerId + 1);
160    }
161
162    /**
163     * 根据网卡 MAC 地址计算余数作为数据中心 ID。
164     */
165    protected long getDataCenterId(long maxDataCenterId) {
166        long id = 0L;
167        try {
168            if (address == null) {
169                address = InetAddress.getLocalHost();
170            }
171            NetworkInterface network = NetworkInterface.getByInetAddress(address);
172            if (null == network) {
173                id = 1L;
174            } else {
175                byte[] mac = network.getHardwareAddress();
176                if (null != mac) {
177                    id = ((0x000000FF & (long) mac[mac.length - 2]) | (0x0000FF00 & (((long) mac[mac.length - 1]) << 8))) >> 6;
178                    id = id % (maxDataCenterId + 1);
179                }
180            }
181        } catch (Exception e) {
182            throw FlexExceptions.wrap(e, "dataCenterId: %s", e.getMessage());
183        }
184        return id;
185    }
186
187    @Override
188    public Object generate(Object entity, String keyColumn) {
189        return nextId();
190    }
191
192    /**
193     * 获取下一个 ID。
194     */
195    public synchronized long nextId() {
196        long currentTimeMillis = System.currentTimeMillis();
197        // 当前时间小于上一次生成 ID 使用的时间，可能出现服务器时钟回拨问题。
198        if (currentTimeMillis < lastTimeMillis) {
199            long offset = lastTimeMillis - currentTimeMillis;
200            // 在可容忍的时间差值之内等待时间恢复正常
201            if (offset <= offsetPeriod) {
202                try {
203                    wait(offset << 1L);
204                    currentTimeMillis = System.currentTimeMillis();
205                    if (currentTimeMillis < lastTimeMillis) {
206                        throw FlexExceptions.wrap("Clock moved backwards, please check the time. Current timestamp: %d, last used timestamp: %d", currentTimeMillis, lastTimeMillis);
207                    }
208                } catch (InterruptedException e) {
209                    throw FlexExceptions.wrap(e);
210                }
211            } else {
212                throw FlexExceptions.wrap("Clock moved backwards, please check the time. Current timestamp: %d, last used timestamp: %d", currentTimeMillis, lastTimeMillis);
213            }
214        }
215
216        if (currentTimeMillis == lastTimeMillis) {
217            // 相同毫秒内，序列号自增
218            sequence = (sequence + 1) & SEQUENCE_MASK;
219            if (sequence == 0) {
220                // 同一毫秒的序列数已经达到最大
221                currentTimeMillis = tilNextMillis(lastTimeMillis);
222            }
223        } else {
224            // 不同毫秒内，序列号置为 0。
225            sequence = 0L;
226        }
227
228        // 记录最后一次使用的毫秒时间戳
229        lastTimeMillis = currentTimeMillis;
230
231        // 时间戳部分 | 数据中心部分 | 机器标识部分 | 序列号部分
232        return ((currentTimeMillis - twepoch) << TIMESTAMP_SHIFT)
233            | (dataCenterId << DATA_CENTER_ID_SHIFT)
234            | (workerId << WORK_ID_SHIFT)
235            | sequence;
236    }
237
238    /**
239     * 获取指定时间戳的接下来的时间戳。
240     */
241    private long tilNextMillis(long lastTimestamp) {
242        long currentTimeMillis = System.currentTimeMillis();
243        while (currentTimeMillis <= lastTimestamp) {
244            currentTimeMillis = System.currentTimeMillis();
245        }
246        return currentTimeMillis;
247    }
248
249    public static long getTwepoch() {
250        return twepoch;
251    }
252
253    public static void setTwepoch(long twepoch) {
254        SnowFlakeIDKeyGenerator.twepoch = twepoch;
255    }
256
257    public static long getOffsetPeriod() {
258        return offsetPeriod;
259    }
260
261    public static void setOffsetPeriod(long offsetPeriod) {
262        SnowFlakeIDKeyGenerator.offsetPeriod = offsetPeriod;
263    }
264
265    public long getWorkerId() {
266        return workerId;
267    }
268
269    public long getDataCenterId() {
270        return dataCenterId;
271    }
272
273    public InetAddress getAddress() {
274        return address;
275    }
276
277    public void setAddress(InetAddress address) {
278        this.address = address;
279    }
280
281    public long getSequence() {
282        return sequence;
283    }
284
285    public void setSequence(long sequence) {
286        this.sequence = sequence;
287    }
288
289    public long getLastTimeMillis() {
290        return lastTimeMillis;
291    }
292
293    public void setLastTimeMillis(long lastTimeMillis) {
294        this.lastTimeMillis = lastTimeMillis;
295    }
296}