001/*
002 *  Copyright (c) 2022-2025, 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;
021import org.apache.ibatis.logging.LogFactory;
022
023import java.lang.management.ManagementFactory;
024import java.net.InetAddress;
025import java.net.NetworkInterface;
026
027/**
028 * <p>雪花算法 ID 生成器。
029 *
030 * <ul>
031 *     <li>最高 1 位固定值 0，因为生成的 ID 是正整数；
032 *     <li>接下来 41 位存储毫秒级时间戳，2 ^ 41 / ( 1000 * 60 * 60 * 24 * 365) = 69，大概可以使用 69 年；
033 *     <li>再接下 10 位存储机器码，包括 5 位 dataCenterId 和 5 位 workerId，最多可以部署 2 ^ 10 = 1024 台机器；
034 *     <li>最后 12 位存储序列号，同一毫秒时间戳时，通过这个递增的序列号来区分，即对于同一台机器而言，同一毫秒时间戳下，可以生成 2 ^ 12 = 4096 个不重复 ID。
035 * </ul>
036 *
037 * <p>优化自开源项目：<a href="https://gitee.com/yu120/sequence">Sequence</a>
038 *
039 * @author 王帅
040 * @since 2023-05-12
041 */
042public class SnowFlakeIDKeyGenerator implements IKeyGenerator {
043
044    /**
045     * 工作机器 ID 占用的位数（5bit）。
046     */
047    private static final long WORKER_ID_BITS = 5L;
048    /**
049     * 数据中心 ID 占用的位数（5bit）。
050     */
051    private static final long DATA_CENTER_ID_BITS = 5L;
052    /**
053     * 序号占用的位数（12bit）。
054     */
055    private static final long SEQUENCE_BITS = 12L;
056    /**
057     * 工作机器 ID 占用 5bit 时的最大值 31。
058     */
059    private static final long MAX_WORKER_ID = ~(-1L << WORKER_ID_BITS);
060    /**
061     * 数据中心 ID 占用 5bit 时的最大值 31。
062     */
063    private static final long MAX_DATA_CENTER_ID = ~(-1L << DATA_CENTER_ID_BITS);
064    /**
065     * 序号掩码，用于与自增后的序列号进行位“与”操作，如果值为 0，则代表自增后的序列号超过了 4095。
066     */
067    private static final long SEQUENCE_MASK = ~(-1L << SEQUENCE_BITS);
068    /**
069     * 工作机器 ID 位需要左移的位数（12bit）。
070     */
071    private static final long WORK_ID_SHIFT = SEQUENCE_BITS;
072    /**
073     * 数据中心 ID 位需要左移的位数（12bit + 5bit）。
074     */
075    private static final long DATA_CENTER_ID_SHIFT = SEQUENCE_BITS + WORKER_ID_BITS;
076    /**
077     * 时间戳需要左移的位数（12bit + 5bit + 5bit）。
078     */
079    private static final long TIMESTAMP_SHIFT = SEQUENCE_BITS + WORKER_ID_BITS + DATA_CENTER_ID_BITS;
080    /**
081     * 时间起始标记点，一旦确定不能变动（2023-04-02 13:01:00）。
082     */
083    private static long twepoch = 1680411660000L;
084    /**
085     * 可容忍的时间偏移量。
086     */
087    private static long offsetPeriod = 5L;
088
089    /**
090     * 工作机器 ID。
091     */
092    private final long workerId;
093    /**
094     * 数据中心 ID。
095     */
096    private final long dataCenterId;
097
098    /**
099     * IP 地址信息，用来生成工作机器 ID 和数据中心 ID。
100     */
101    protected InetAddress address;
102    /**
103     * 同一毫秒内的最新序号，最大值可为（2^12 - 1 = 4095）。
104     */
105    private long sequence;
106    /**
107     * 上次生产 ID 时间戳。
108     */
109    private long lastTimeMillis = -1L;
110
111    /**
112     * 雪花算法 ID 生成器。
113     */
114    public SnowFlakeIDKeyGenerator() {
115        this(null);
116    }
117
118    /**
119     * 根据 IP 地址计算数据中心 ID 和工作机器 ID 生成数据库 ID。
120     *
121     * @param address IP 地址
122     */
123    public SnowFlakeIDKeyGenerator(InetAddress address) {
124        this.address = address;
125        this.dataCenterId = getDataCenterId(MAX_DATA_CENTER_ID);
126        this.workerId = getWorkerId(dataCenterId, MAX_WORKER_ID);
127    }
128
129    /**
130     * 根据数据中心 ID 和工作机器 ID 生成数据库 ID。
131     *
132     * @param workerId     工作机器 ID
133     * @param dataCenterId 数据中心 ID
134     */
135    public SnowFlakeIDKeyGenerator(long workerId, long dataCenterId) {
136        if (workerId > MAX_WORKER_ID || workerId < 0) {
137            throw new IllegalArgumentException(
138                String.format("workerId must be greater than 0 and less than %d.", MAX_WORKER_ID));
139        }
140        if (dataCenterId > MAX_DATA_CENTER_ID || dataCenterId < 0) {
141            throw new IllegalArgumentException(
142                String.format("dataCenterId must be greater than 0 and less than %d.", MAX_DATA_CENTER_ID));
143        }
144        this.workerId = workerId;
145        this.dataCenterId = dataCenterId;
146    }
147
148    /**
149     * 根据 MAC + PID 的 hashCode 获取 16 个低位生成工作机器 ID。
150     */
151    protected long getWorkerId(long dataCenterId, long maxWorkerId) {
152        StringBuilder mpId = new StringBuilder();
153        mpId.append(dataCenterId);
154        String name = ManagementFactory.getRuntimeMXBean().getName();
155        if (StringUtil.isNotBlank(name)) {
156            // GET jvmPid
157            mpId.append(name.split("@")[0]);
158        }
159        // MAC + PID 的 hashCode 获取16个低位
160        return (mpId.toString().hashCode() & 0xffff) % (maxWorkerId + 1);
161    }
162
163    /**
164     * 根据网卡 MAC 地址计算余数作为数据中心 ID。
165     */
166    protected long getDataCenterId(long maxDataCenterId) {
167        long id = 1L;
168        try {
169            if (address == null) {
170                address = InetAddress.getLocalHost();
171            }
172            NetworkInterface network = NetworkInterface.getByInetAddress(address);
173            if (null != network) {
174                byte[] mac = network.getHardwareAddress();
175                if (null != mac) {
176                    id = ((0x000000FF & (long) mac[mac.length - 2]) | (0x0000FF00 & (((long) mac[mac.length - 1]) << 8))) >> 6;
177                    id = id % (maxDataCenterId + 1);
178                }
179            }
180        } catch (Exception e) {
181            LogFactory.getLog(SnowFlakeIDKeyGenerator.class).error(e.toString(), e);
182        }
183        return id;
184    }
185
186    @Override
187    public Object generate(Object entity, String keyColumn) {
188        return nextId();
189    }
190
191    /**
192     * 获取下一个 ID。
193     */
194    public synchronized long nextId() {
195        long currentTimeMillis = System.currentTimeMillis();
196        // 当前时间小于上一次生成 ID 使用的时间，可能出现服务器时钟回拨问题。
197        if (currentTimeMillis < lastTimeMillis) {
198            long offset = lastTimeMillis - currentTimeMillis;
199            // 在可容忍的时间差值之内等待时间恢复正常
200            if (offset <= offsetPeriod) {
201                try {
202                    wait(offset << 1L);
203                    currentTimeMillis = System.currentTimeMillis();
204                    if (currentTimeMillis < lastTimeMillis) {
205                        throw FlexExceptions.wrap("Clock moved backwards, please check the time. Current timestamp: %d, last used timestamp: %d", currentTimeMillis, lastTimeMillis);
206                    }
207                } catch (InterruptedException e) {
208                    throw FlexExceptions.wrap(e);
209                }
210            } else {
211                throw FlexExceptions.wrap("Clock moved backwards, please check the time. Current timestamp: %d, last used timestamp: %d", currentTimeMillis, lastTimeMillis);
212            }
213        }
214
215        if (currentTimeMillis == lastTimeMillis) {
216            // 相同毫秒内，序列号自增
217            sequence = (sequence + 1) & SEQUENCE_MASK;
218            if (sequence == 0) {
219                // 同一毫秒的序列数已经达到最大
220                currentTimeMillis = tilNextMillis(lastTimeMillis);
221            }
222        } else {
223            // 不同毫秒内，序列号置为 0。
224            sequence = 0L;
225        }
226
227        // 记录最后一次使用的毫秒时间戳
228        lastTimeMillis = currentTimeMillis;
229
230        // 时间戳部分 | 数据中心部分 | 机器标识部分 | 序列号部分
231        return ((currentTimeMillis - twepoch) << TIMESTAMP_SHIFT)
232            | (dataCenterId << DATA_CENTER_ID_SHIFT)
233            | (workerId << WORK_ID_SHIFT)
234            | sequence;
235    }
236
237    /**
238     * 获取指定时间戳的接下来的时间戳。
239     */
240    private long tilNextMillis(long lastTimestamp) {
241        long currentTimeMillis = System.currentTimeMillis();
242        while (currentTimeMillis <= lastTimestamp) {
243            currentTimeMillis = System.currentTimeMillis();
244        }
245        return currentTimeMillis;
246    }
247
248    public static long getTwepoch() {
249        return twepoch;
250    }
251
252    public static void setTwepoch(long twepoch) {
253        SnowFlakeIDKeyGenerator.twepoch = twepoch;
254    }
255
256    public static long getOffsetPeriod() {
257        return offsetPeriod;
258    }
259
260    public static void setOffsetPeriod(long offsetPeriod) {
261        SnowFlakeIDKeyGenerator.offsetPeriod = offsetPeriod;
262    }
263
264    public long getWorkerId() {
265        return workerId;
266    }
267
268    public long getDataCenterId() {
269        return dataCenterId;
270    }
271
272    public InetAddress getAddress() {
273        return address;
274    }
275
276    public void setAddress(InetAddress address) {
277        this.address = address;
278    }
279
280    public long getSequence() {
281        return sequence;
282    }
283
284    public void setSequence(long sequence) {
285        this.sequence = sequence;
286    }
287
288    public long getLastTimeMillis() {
289        return lastTimeMillis;
290    }
291
292    public void setLastTimeMillis(long lastTimeMillis) {
293        this.lastTimeMillis = lastTimeMillis;
294    }
295}