io.github.dengchen2020.id.snowflake.SnowflakeSmartLifecycle Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of dc-spring-boot-starter-id-generator Show documentation
id生成器
There is a newer version: 0.0.35
Show newest version
package io.github.dengchen2020.id.snowflake;

import io.github.dengchen2020.id.IdHelper;
import io.github.dengchen2020.id.event.IdGeneratorInitCompleteEvent;
import io.github.dengchen2020.id.exception.IdGeneratorException;
import jakarta.annotation.Resource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.context.ApplicationEventPublisher;
import org.springframework.context.SmartLifecycle;
import org.springframework.data.redis.core.StringRedisTemplate;

import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * 雪花算法自动配置
 * 使用方式：{@link IdHelper#nextId()}
 *
 * @author dengchen
 * @since 2022/10/16
 */
public class SnowflakeSmartLifecycle implements SmartLifecycle {

    private static final Logger log = LoggerFactory.getLogger(SnowflakeSmartLifecycle.class);

    AtomicBoolean running = new AtomicBoolean(false);

    private final StringRedisTemplate stringRedisTemplate;

    private final SnowflakeIdGeneratorOptions options;

    @Resource
    private ApplicationEventPublisher applicationEventPublisher;

    public SnowflakeSmartLifecycle(StringRedisTemplate stringRedisTemplate, SnowflakeIdGeneratorOptions options) {
        this.stringRedisTemplate = stringRedisTemplate;
        this.options = options;
    }

    /**
     * 雪花算法机器id集合key
     */
    private static final String SNOWFLAKE_WORKERID_LIST_KEY = "dc:snowflake:workerid:list";

    /**
     * 从redis中获取机器id，并设置机器id
     */
    private void setWorkerIdFromRedis(Short workerId, int maxWorkerIdNumber) throws Exception {
        if (workerId == null || workerId < 0 || workerId > maxWorkerIdNumber) {
            workerId = (short) ThreadLocalRandom.current().nextInt(0, maxWorkerIdNumber);
        }
        for (; workerId <= maxWorkerIdNumber; workerId++) {
            if (!Boolean.TRUE.equals(stringRedisTemplate.opsForSet().isMember(SNOWFLAKE_WORKERID_LIST_KEY, workerId.toString()))) {
                stringRedisTemplate.opsForSet().add(SNOWFLAKE_WORKERID_LIST_KEY, workerId.toString());
                options.setWorkerId(workerId);
                return;
            }
        }
        for (int i = 0; i <= maxWorkerIdNumber; i++) {
            if (!Boolean.TRUE.equals(stringRedisTemplate.opsForSet().isMember(SNOWFLAKE_WORKERID_LIST_KEY, String.valueOf(i)))) {
                stringRedisTemplate.opsForSet().add(SNOWFLAKE_WORKERID_LIST_KEY, String.valueOf(i));
                options.setWorkerId((short) i);
                return;
            }
        }
        throw new Exception("机器id已用完！");
    }

    @Override
    public void start() {
        try {
            //智能设置雪花算法机器id
            setWorkerIdFromRedis(options.getWorkerId(), (1 << options.getWorkerIdBitLength()) - 1);
        } catch (Exception e) {
            running.compareAndSet(true, false);
            throw new IdGeneratorException("雪花算法初始化失败！" + e.getMessage());
        }
        options.setBaseTime(1657209600000L);  //该值不允许更改
        IdHelper.setIdGenerator(new SnowflakeIdGenerator(options));
        if (log.isInfoEnabled()) log.info("雪花算法生成器初始化完成，机器id：{}", options.getWorkerId());
        running.compareAndSet(false, true);
        applicationEventPublisher.publishEvent(new IdGeneratorInitCompleteEvent());
    }

    /**
     * 从redis中移除机器id
     */
    @Override
    public void stop() {
        stringRedisTemplate.opsForSet().remove(SNOWFLAKE_WORKERID_LIST_KEY, String.valueOf(options.getWorkerId()));
        if (log.isInfoEnabled()) log.info("雪花算法移除机器id：{}", options.getWorkerId());
        running.compareAndSet(true, false);
    }

    @Override
    public boolean isRunning() {
        return running.get();
    }
}