Spring Boot項目下JPA自定義雪花算法ID生成器詳解

作者: 白石(https://github.com/wjw465150)

本文分享下Spring boot項目下使用JPA操作數據庫時關于雪花ID生成器的相關實現代碼。

SnowFlake 算法（雪花算法），是 Twitter 開源的分布式 id 生成算法。其核心思想就是：使用一個 64 bit 的 long 型的數字作為全局唯一 id。在分布式系統中的應用十分廣泛，且ID 引入了時間戳，基本上保持自增。

在JPA中一個數據表必須要有主鍵，主鍵類型一般是推薦使用Long類型，那么在分布式微服務下需要保證ID的唯一性，此時往往需要自定義主鍵生成策略。

首先實現一個實體類的基類，在基類中定義ID的生成策略，子類繼承其實現，這樣就不用每個實體類都去寫一遍了。

實體類的基類AbstractBaseEntity

package org.wjw.jpa.snowflake;

import java.io.Serializable;

import javax.persistence.Column;
import javax.persistence.GeneratedValue;
import javax.persistence.Id;
import javax.persistence.MappedSuperclass;

import org.hibernate.annotations.GenericGenerator;

@MappedSuperclass
public abstract class AbstractBaseEntity implements Serializable {
  @Id
  @GenericGenerator(name = "snowFlakeIdGenerator", strategy = "org.wjw.jpa.snowflake.SnowFlakeIdGenerator")
  @GeneratedValue(generator = "snowFlakeIdGenerator")
  @Column(name = "id", length=18)
  private String id;

    /**
     * 獲取主鍵id
     * 
     * @return id 因為如果是Long類型,前端js能處理的長度低于Java，防止精度丟失;java的Long類型是18位， 而 js的Long類型(雖然沒有明確定義的Long類型)是16位， 所以會造成丟失精度，
     */
    public String getId() {
      return id;
    }

    /**
     * 設置主鍵id
     * 
     * @param id 主鍵id
     */
    public void setId(String id) {
      this.id = id;
    }
}

strategy 表示生成策略實現類的全路徑名。

使用雪花算法要注意的是，保證機器的時鐘是一直增加的，也就是說不可以將時鐘往前調，不然就不能保證ID的自增，并且有可能發生ID沖突（產生了重復的ID）。因此，上面的代碼中，在檢查到時鐘異常時會拋出異常。

雪花算法類SnowFlake:

package org.wjw.jpa.snowflake;

import javax.annotation.PostConstruct;

import org.hibernate.HibernateException;
import org.hibernate.internal.CoreLogging;
import org.hibernate.internal.CoreMessageLogger;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.stereotype.Component;

/**
 * 雪花算法
 */
@Component
public class SnowFlake implements ApplicationContextAware {
  private static final CoreMessageLogger LOG = CoreLogging.messageLogger(SnowFlake.class);

  private static ApplicationContext _applicationContext;

  @Override
  public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
    SnowFlake._applicationContext = applicationContext;
  }
  
  public static SnowFlake getBean() {
    return _applicationContext.getBean(SnowFlake.class);
  }
  
  /**
   * 起始的時間戳
   */
  private final long twepoch = 1557825652094L;

  /**
   * 每一部分占用的位數
   */
  private final long workerIdBits = 5L;
  private final long datacenterIdBits = 5L;
  private final long sequenceBits = 12L;

  /**
   * 每一部分的最大值
   */
  private final long maxWorkerId = -1L ^ (-1L << workerIdBits);
  private final long maxDatacenterId = -1L ^ (-1L << datacenterIdBits);
  private final long maxSequence = -1L ^ (-1L << sequenceBits);

  /**
   * 每一部分向左的位移
   */
  private final long workerIdShift = sequenceBits;
  private final long datacenterIdShift = sequenceBits + workerIdBits;
  private final long timestampShift = sequenceBits + workerIdBits + datacenterIdBits;

  @Value("${snowflake.datacenter-id:1}")
  private long datacenterId; // 數據中心ID

  @Value("${snowflake.worker-id:0}")
  private long workerId; // 機器ID

  private long sequence = 0L; // 序列號
  private long lastTimestamp = -1L; // 上一次時間戳
  
  @PostConstruct
  public void init() {
    String msg;
    if (workerId > maxWorkerId || workerId < 0) {
      msg = String.format("worker Id can't be greater than %d or less than 0", maxWorkerId);
      LOG.unsuccessful(msg);
    }
    if (datacenterId > maxDatacenterId || datacenterId < 0) {
      msg = String.format("datacenter Id can't be greater than %d or less than 0", maxDatacenterId);
      LOG.unsuccessful(msg);
    }
  }

  public String nextIdString() {
    long id = nextId();
    
    return String.valueOf(id);
  }
  
  public synchronized long nextId() {
    long timestamp = timeGen();
    if (timestamp < lastTimestamp) {
      throw new HibernateException(String.format("Clock moved backwards.  Refusing to generate id for %d milliseconds", lastTimestamp - timestamp));
    }
    if (timestamp == lastTimestamp) {
      sequence = (sequence + 1) & maxSequence;
      if (sequence == 0L) {
        timestamp = tilNextMillis();
      }
    } else {
      sequence = 0L;
    }
    lastTimestamp = timestamp;

    return (timestamp - twepoch) << timestampShift // 時間戳部分
        | datacenterId << datacenterIdShift // 數據中心部分
        | workerId << workerIdShift // 機器標識部分
        | sequence; // 序列號部分
  }

  private long tilNextMillis() {
    long timestamp = timeGen();
    while (timestamp <= lastTimestamp) {
      timestamp = timeGen();
    }
    return timestamp;
  }

  private long timeGen() {
    return System.currentTimeMillis();
  }

}

代碼中使用***@Value(“$snowflake.datacenter?id:1")???和???@Value("${snowflake.worker-id:0}”)***注解從環境配置中讀取當前的數據中心id機器id。

集成JPA配置SnowFlakeIdGenerator

package org.wjw.jpa.snowflake;

import java.io.Serializable;
import java.util.Properties;

import org.hibernate.HibernateException;
import org.hibernate.MappingException;
import org.hibernate.engine.spi.SharedSessionContractImplementor;
import org.hibernate.id.Configurable;
import org.hibernate.id.IdentifierGenerator;
import org.hibernate.service.ServiceRegistry;
import org.hibernate.type.Type;

/**
 * JPA 雪花算法ID生成器
 */
public class SnowFlakeIdGenerator implements IdentifierGenerator, Configurable {

  @Override
  public Serializable generate(SharedSessionContractImplementor session, Object o) throws HibernateException {
    return SnowFlake.getBean().nextIdString();
  }

  @Override
  public void configure(Type type, Properties params, ServiceRegistry serviceRegistry) throws MappingException {
  }
}

測試

創建一個 @Entity

好的，接下來就是正常實體類繼承基類就可以了，如下User實體類：

package org.wjw.jpa.entity;

import javax.persistence.Entity;

import org.wjw.jpa.snowflake.AbstractBaseEntity;

@Entity
public class User extends AbstractBaseEntity {
  private String name;

  private String email;

  public String getName() {
    return name;
  }

  public void setName(String name) {
    this.name = name;
  }

  public String getEmail() {
    return email;
  }

  public void setEmail(String email) {
    this.email = email;
  }

  @Override
  public String toString() {
    return "User [id=" + id + ", name=" + name + ", email=" + email + "]";
  }
  
}

創建一個 Repository

倉庫類UserJpaRepository:

package org.wjw.jpa.repository;

import org.springframework.data.jpa.repository.JpaRepository;
import org.wjw.jpa.entity.User;

public interface UserJpaRepository extends JpaRepository<User,Long> {
}

進行單元測試

單元測試代碼:

package org.wjw.jpa;

import org.junit.jupiter.api.Test;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.context.SpringBootTest;
import org.wjw.jpa.entity.User;
import org.wjw.jpa.repository.UserJpaRepository;

@SpringBootTest
public class MyTests {
  @Autowired
  private UserJpaRepository userRepository;
  
  @Test
  public void testAddNewUser() {
    User n = new User();
    n.setName("Querydsl");
    n.setEmail("Querydsl@gamil.com");
    userRepository.save(n);
  }
}