雪花id改多workerID依赖redis

package com.meisoo.clearingplat.config;

import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Configuration;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;

import javax.annotation.PostConstruct;
import java.util.Collections;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

/**
 * IdUtil.getSnowflake(SnowflakeConfig.staticWorkerId).nextId()
 * @author Hj
 * @date 2025/11/20
 */
@Slf4j
@Configuration
public class SnowflakeConfig {

    public static Integer staticWorkerId;

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

    @Autowired
    private RedisTemplate<String, Object> redisTemplate;

    private String WORKER_ID_KEY = ":snowflake:workerids";
    private String WORKER_ID_LOCK_KEY = ":snowflake:workerid:lock";
    private final long expireTTL = 20 * 60 * 1000; // 20分钟过期

    @PostConstruct
    public void init() {
        this.staticWorkerId = generateDistributedWorkerId();
    }

    public int generateDistributedWorkerId() {
        Integer workerId = null;
        String appNameMap = appName + WORKER_ID_KEY;
        String lockKey = appName + WORKER_ID_LOCK_KEY;

        // 生成唯一的锁标识，防止误删其他线程的锁
        String lockValue = UUID.randomUUID().toString();

        try {
            // 尝试获取分布式锁
            boolean locked = tryLock(lockKey, lockValue);
            if (!locked) {
                // 如果获取锁失败，重试几次
                for (int i = 0; i < 1000; i++) {
                    Thread.sleep(100); // 等待100ms后重试
                    locked = tryLock(lockKey, lockValue);
                    if (locked) break;
                }
            }

            if (!locked) {
                throw new RuntimeException("无法获取分布式锁，workerId分配失败");
            }

            // 获取锁成功，分配workerId
            workerId = allocateWorkerId(appNameMap);

        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new RuntimeException("workerId分配过程被中断", e);
        } finally {
            // 释放锁
            unlock(lockKey, lockValue);
        }

        if (workerId == null || workerId > 1023) {
            throw new RuntimeException("workerId分配失败，没有可用的workerId");
        }

        // 启动心跳更新任务
        startHeartbeatTask(appNameMap, workerId);

        return workerId;
    }

    /**
     * 尝试获取分布式锁
     */
    private boolean tryLock(String lockKey, String lockValue) {
        Boolean success = redisTemplate.opsForValue().setIfAbsent(lockKey, lockValue);
        return Boolean.TRUE.equals(success);
    }

    /**
     * 释放分布式锁
     */
    private void unlock(String lockKey, String lockValue) {
        try {
            // 使用Lua脚本保证原子性，只有锁的值匹配时才删除
            String luaScript =
                    "if redis.call('get', KEYS[1]) == ARGV[1] then " +
                            "    return redis.call('del', KEYS[1]) " +
                            "else " +
                            "    return 0 " +
                            "end";

            RedisScript<Long> script = new DefaultRedisScript<>(luaScript, Long.class);
            Long result = redisTemplate.execute(script, Collections.singletonList(lockKey), lockValue);

            if (result == null || result == 0) {
                log.warn("释放锁失败，锁可能已被其他线程释放或已过期");
            }
        } catch (Exception e) {
            log.error("释放锁时发生异常", e);
        }
    }

    /**
     * 分配workerId（在锁的保护下执行）
     */
    private Integer allocateWorkerId(String appNameMap) {
        for (int workerId = 0; workerId <= 1023; workerId++) {
            if (!redisTemplate.opsForHash().hasKey(appNameMap, String.valueOf(workerId))) {
                // 找到可用的workerId，立即占用
                redisTemplate.opsForHash().put(appNameMap, String.valueOf(workerId), System.currentTimeMillis());
                log.info("成功分配workerId: {} for application: {}", workerId, appName);
                return workerId;
            }
        }
        return null;
    }

    /**
     * 启动心跳更新任务
     */
    private void startHeartbeatTask(String appNameMap, Integer workerId) {
        ScheduledExecutorService scheduledExecutorService = Executors.newSingleThreadScheduledExecutor();

        scheduledExecutorService.scheduleAtFixedRate(() -> {
            try {
                log.debug("定期更新workerId: {} 的心跳", workerId);
                // 更新自己的心跳
                redisTemplate.opsForHash().put(appNameMap, String.valueOf(workerId), System.currentTimeMillis());

                // 清理过期节点（这里也需要加锁，避免并发清理）
                cleanupExpiredWorkers(appNameMap);

            } catch (Exception e) {
                log.error("心跳更新任务执行异常", e);
            }
        }, 2, 10, TimeUnit.MINUTES); // 每10分钟执行一次，比过期时间短
    }

    /**
     * 清理过期的worker节点（需要加锁）
     */
    private void cleanupExpiredWorkers(String appNameMap) {
        String lockKey = appName + WORKER_ID_LOCK_KEY + ":cleanup";
        String lockValue = UUID.randomUUID().toString();

        try {
            // 尝试获取清理锁，避免多个实例同时清理
            boolean locked = tryLock(lockKey, lockValue); // 10秒锁超时
            if (!locked) {
                return; // 如果没有获取到锁，跳过本次清理
            }

            long currentTime = System.currentTimeMillis();
            Set<Object> workerIds = redisTemplate.opsForHash().keys(appNameMap);

            for (Object workerIdObj : workerIds) {
                try {
                    String workerIdStr = String.valueOf(workerIdObj);
                    Long lastHeartbeat = (Long) redisTemplate.opsForHash().get(appNameMap, workerIdStr);

                    if (lastHeartbeat != null && currentTime - lastHeartbeat > expireTTL) {
                        log.info("清理过期的workerId: {}, 最后心跳时间: {}", workerIdStr, lastHeartbeat);
                        redisTemplate.opsForHash().delete(appNameMap, workerIdStr);
                    }
                } catch (Exception e) {
                    log.warn("清理workerId: {} 时发生异常", workerIdObj, e);
                }
            }

        } finally {
            unlock(lockKey, lockValue);
        }
    }
}