Mybatis Config 雪花算法id,用 redis 管理 workerId

2023-06-06
2 min read

Mybatis Config 雪花算法id,用 redis 自动管理 workerId。

本章代码直接提供了使用 mybatis 配置雪花算法 id,并且利用 redis 自动注册 workerId;
另外提供了一个方法,让 springboot 在启动好之前就初始化好数据库连接池;

@Slf4j
@Configuration
@EnableTransactionManagement
public class MybatisPlusConfig {

    private final String idWorker = "WorkerId";

    @Value("${spring.application.name}")
    private String serviceName;

    @Resource
    private RedisTemplate<String,Long> redisTemplate;

    private final String script =
            "local now = redis.call('TIME')[1]\n" +
                    "local idWordsKey = KEYS[1]\n" +
                    "local sp = ':'\n" +
                    "for i = 0, 1023 do\n" +
                    "    local serviceKey = idWordsKey..sp..i\n" +
                    "    if redis.call('SETNX', serviceKey, now) == 1 then\n" +
                    "        redis.call('Expire', serviceKey, 30)\n" +
                    "        return i;\n" +
                    "    end\n" +
                    "end\n" +
                    "return -1";

    /**
     * 分页插件
     */
    @Bean
    public MybatisPlusInterceptor mybatisPlusInterceptor() {
        MybatisPlusInterceptor interceptor = new MybatisPlusInterceptor();
        interceptor.addInnerInterceptor(new PaginationInnerInterceptor(DbType.MYSQL));
        return interceptor;
    }

    @Bean
    public IdentifierGenerator idGenerator() {
        if (serviceName == null || serviceName.length() == 0){
            log.error("雪花算法初始化失败,【 app.idKey 】配置为空。");
            return null;
        }
        long num = getWorkerIdNum();
        // 获取前  5 位
        long dataCenterId = num >> 5;
        // 获 取 后 5 位
        long workerId = num & (~(-1L << 5L));
        // 自定义初始化雪花算法
        log.info("==== [Init Snowflake suceessfully]: dataCenterId:{},workerId:{}", dataCenterId, workerId);
        return new DefaultIdentifierGenerator(workerId, dataCenterId);
    }

    /**
     * 获取机器标识号
     * param serviceName 服务名称,不再需要,r edis 框架自动添加服务名
     */
    private Long getWorkerIdNum() {
        // 实例化脚本对象
        DefaultRedisScript<Long> lua = new DefaultRedisScript<>();
        lua.setResultType(Long.class);
        lua.setScriptText(script);
        List<String> keys = new ArrayList<>(2);
        keys.add(idWorker);

        // 获取序列号
        Long num = redisTemplate.execute(lua, keys, keys.size());
        String targetKey = String.join(":", keys) + ":" + num;

        // -1 代表机器用完了,重试
        if (num < 0){
            log.erro r( "目前 Id 已用完,请重新启动试试");
            System.exit(0);
        }

        // 自动续期
        this.autoExpire(targetKey);

        return num;
    }

    /**
     * 自动续期
     */
    private void autoExpire(String key) {
        ScheduledExecutorService executorService = new ScheduledThreadPoolExecutor(1,
                new BasicThreadFactory.Builder().namingPattern("id_auto_expire-%d").daemon(true).build());
        executorService.scheduleAtFixedRate(() -> {
            redisTemplate.expire(key, 30, TimeUnit.SECONDS);
            log.debu g( "自动续期 id 成功:{}", key);
        }, 0, 10, TimeUnit.SECONDS);
    }

    public String getServiceName() {
        return serviceName;
    }

    public void setServiceName(String serviceName) {
        this.serviceName = serviceName;
    }

    /**
     * init connection when springboot startup
     */
    @Bean
    public ApplicationRunner connectionInit(DataSource dataSource) {
        return args -> {
            dataSource.getConnection();
            log.info("==== [datasource connection inited]");
        };
    }