基于美团Leaf-Segment的双buffer方案实现序列号生成器

业务背景

       有时项目中对于流水号有一些特殊的需求。比如，和业务A有关数据，我们在落库时想要给每条数据添加一个流水号字段，用于作为全局唯一标识。流水号格式规则如下，如：BTA(业务A代号)+年月日(20221208)+序列号。并且对序列号的长度有要求，如序列号要求为5位，即从00001到99999，当序列号达到99999后，再次获取则继续从00001开始累加循环。流水号的形式如TX2022120800001。在此之前需要对业务A有关数据每日的数据量进行评估，以上述为例，若一天的单据量超过99999，再次循环可能会造成流水号重复，以致流水号不唯一，所以序列号最大值可以设的稍大一位。

初期方案

最开始实现的方案是每次需要序列号，则通过累加DB中对应业务的序列号的值并获取来实现。这种方案的缺点十分明显，效率低，也会造成某段时间DB压力巨大，当DB宕机时将无法获取到序列号，造成整个系统瘫痪。

基于美团Leaf的改进方案

接触到美团Leaf生成分布式ID的方案后决定将其用于生成序列号的实现方案。但原生的美团Leaf并不太适用于上面的需求。原因如下

1、美团Leaf是独立部署的、作为内部公共的基础技术设施的一个分布式ID的proxy-server。对于一些没有使用病独立部署leaf服务的公司，若某个项目使用该方案，除项目服务外。还需要部署leaf服务节点，代价有点大，得不偿失。

2、美团Leaf的segment方案生成的ID是趋势递增的不支持若序列号要求为5位，当序列号达到99999后，再次获取则继续从00001开始这样的一个需求。

针对以上两点，站在美团leaf的肩膀上，将其改造为项目中的工具类，并且满足我们对于序列号的长度的要求。

代码如下：

建表语句

drop table if exists psm_serial_no_record;
create table psm_serial_no_record
(
   id                   bigint(20) not null auto_increment comment '主键',
   biz_source           varchar(64) not null default '' comment '业务类型',
   max_id               int(11) not null default 1 comment '最大值',
   is_delete            tinyint(1) not null default 0 comment '删除标记:0未删除 1已删除',
   create_time          datetime not null default CURRENT_TIMESTAMP comment '创建时间',
   update_time          datetime not null default CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP comment '修改时间',
   primary key (id),
   unique key `idx_unique_biz_source` (`biz_source`) using btree
)ENGINE=InnoDB AUTO_INCREMENT=6 DEFAULT CHARSET=utf8 COMMENT='序列号表';

PO

import lombok.Data;

import java.io.Serializable;
import java.util.Date;

/**
 *@Description 序列号记录
 **/
@Data
public class PsmSerialNoRecord implements Serializable {

    private static final long serialVersionUID = -42765629647798182L;

    /**
     * ID
     */
    private Long id;

    /**
     * 业务类型
     */
    private String bizSource;

    /**
     * 最大值
     */
    private long maxId;

    /**
     * 删除标记
     */
    private Integer isDelete;

    /**
     * 新增时间
     */
    private Date createTime;

    /**
     * 更新时间
     */
    private Date updateTime;
}

PsmSerialNoRecordDao

import com.example.serialnodemo.po.PsmSerialNoRecord;
import org.apache.ibatis.annotations.Mapper;
import org.apache.ibatis.annotations.Param;

import java.util.List;

@Mapper
public interface PsmSerialNoRecordDao {

    int updateMaxId(@Param("key") String key, @Param("step") int step, @Param("limit") long limit);

    List<String> selectForLock(@Param("keyList") List<String> keyList);

    void insertRecord(PsmSerialNoRecord record);

    PsmSerialNoRecord selectRecordByKey(@Param("key") String key);
}

PsmSerialNoRecordMapper

<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE mapper PUBLIC "-//mybatis.org//DTD Mapper 3.0//EN" "http://mybatis.org/dtd/mybatis-3-mapper.dtd">
<mapper namespace="com.example.serialnodemo.dao.PsmSerialNoRecordDao">

    <resultMap id="BaseResultMap" type="com.example.serialnodemo.po.PsmSerialNoRecord">
        <id column="id" jdbcType="BIGINT" property="id"/>
        <result column="biz_source" jdbcType="VARCHAR" property="bizSource"/>
        <result column="max_id" jdbcType="BIGINT" property="maxId"/>
        <result column="is_delete" jdbcType="INTEGER" property="isDelete"/>
        <result column="create_time" jdbcType="TIMESTAMP" property="createTime"/>
        <result column="update_time" jdbcType="TIMESTAMP" property="updateTime"/>
    </resultMap>

    <sql id="Base_Column_List">
      id, biz_source, max_id, is_delete, create_time, update_time
    </sql>

    <update id="updateMaxId">
        UPDATE psm_serial_no_record
        SET max_id =
                CASE WHEN max_id = #{limit,jdbcType=BIGINT} THEN #{step,jdbcType=INTEGER}
                     WHEN max_id + #{step,jdbcType=INTEGER} &lt; #{limit,jdbcType=BIGINT} THEN max_id + #{step,jdbcType=INTEGER}
                     WHEN max_id + #{step,jdbcType=INTEGER} &gt;= #{limit,jdbcType=BIGINT} THEN #{limit,jdbcType=BIGINT}
                END
        WHERE biz_source = #{key,jdbcType=VARCHAR}
    </update>

    <select id="selectForLock" resultType="string">
        SELECT
        biz_source
        FROM psm_serial_no_record
        WHERE is_delete = 0
        <if test="keyList != null and keyList.size()>0">
            AND biz_source IN
            <foreach collection="keyList" item="key" open="(" close=")" separator=",">
                #{key}
            </foreach>
        </if>
        FOR UPDATE
    </select>

    <insert id="insertRecord" keyColumn="id" keyProperty="id"
            parameterType="com.example.serialnodemo.po.PsmSerialNoRecord" useGeneratedKeys="true">
        INSERT INTO psm_serial_no_record
        <trim prefix="(" suffix=")" suffixOverrides=",">
            <if test="bizSource != null">
                biz_source,
            </if>
            <if test="maxId != null">
                max_id,
            </if>
            <if test="isDelete != null">
                is_delete,
            </if>
            <if test="createTime != null">
                create_time,
            </if>
            <if test="updateTime != null">
                update_time,
            </if>
        </trim>
        <trim prefix="values (" suffix=")" suffixOverrides=",">
            <if test="bizSource != null">
                #{bizSource,jdbcType=VARCHAR},
            </if>
            <if test="maxId != null">
                #{maxId,jdbcType=BIGINT},
            </if>
            <if test="isDelete != null">
                #{isDelete,jdbcType=BOOLEAN},
            </if>
            <if test="createTime != null">
                #{createTime,jdbcType=TIMESTAMP},
            </if>
            <if test="updateTime != null">
                #{updateTime,jdbcType=TIMESTAMP},
            </if>
        </trim>
    </insert>

    <select id="selectRecordByKey" resultMap="BaseResultMap">
        SELECT
        <include refid="Base_Column_List"/>
        FROM psm_serial_no_record
        WHERE biz_source = #{key,jdbcType=VARCHAR}
    </select>

</mapper>

PsmSerialNoRecordService

import com.example.serialnodemo.po.PsmSerialNoRecord;

import java.util.List;

public interface PsmSerialNoRecordService {

    int modifyMaxId(String key, int step, long limit);


    List<String> queryForLock(List<String> keyList);

    void addRecord(PsmSerialNoRecord record);

    PsmSerialNoRecord findRecordByKey(String key);

    PsmSerialNoRecord modifyMaxIdAndGet(String key, int step, long limit) throws Exception;
}

PsmSerialNoRecordServiceImpl

import com.example.serialnodemo.po.PsmSerialNoRecord;
import com.example.serialnodemo.dao.PsmSerialNoRecordDao;
import com.example.serialnodemo.service.PsmSerialNoRecordService;
import com.google.common.collect.Lists;
import org.apache.commons.lang.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.util.CollectionUtils;

import java.util.List;

/**
 * @Description
 * @Program psm-report
 **/
@Service
public class PsmSerialNoRecordServiceImpl implements PsmSerialNoRecordService {

    @Autowired
    private PsmSerialNoRecordDao psmSerialNoRecordDao;

    @Override
    public int modifyMaxId(String key, int step, long limit) {
        if (StringUtils.isNotBlank(key) && step > 0 && limit > 0L) {
            return psmSerialNoRecordDao.updateMaxId(key, step, limit);
        }
        return -1;
    }

    @Override
    public List<String> queryForLock(List<String> keyList) {
        if (!CollectionUtils.isEmpty(keyList)) {
            return psmSerialNoRecordDao.selectForLock(keyList);
        }
        return Lists.newArrayList();
    }

    @Override
    public void addRecord(PsmSerialNoRecord record) {
        if (record != null) {
            psmSerialNoRecordDao.insertRecord(record);
        }
    }

    @Override
    public PsmSerialNoRecord findRecordByKey(String key) {
        if (StringUtils.isNotBlank(key)) {
            return psmSerialNoRecordDao.selectRecordByKey(key);
        }
        return null;
    }

    @Transactional(rollbackFor = Exception.class)
    @Override
    public PsmSerialNoRecord modifyMaxIdAndGet(String key, int step, long limit) throws Exception {
        int update = this.modifyMaxId(key, step, limit);
        // 若没有该业务类型的序列号记录，则报错
        if (update <= 0) {
            throw new Exception("更新序列号失败，key:" + key);
        }
        return this.findRecordByKey(key);
    }
}

Segment

import lombok.Data;

import java.util.concurrent.atomic.AtomicLong;

/**
 * @description: Segment
 */
@Data
public class Segment {

    /**
     * 步长
     */
    public volatile int step;

    /**
     * 自增值
     */
    private AtomicLong value = new AtomicLong(0);

    /**
     * 最大值
     */
    private volatile long max;

    /**
     * buffer
     */
    private SegmentBuffer buffer;

    public Segment(SegmentBuffer buffer) {
        this.buffer = buffer;
    }

    /**
     * @description: 获取空闲数
     *
     * @return long 空闲数
     */
    public long getIdle() {
        return this.getMax() - getValue().get();
    }

    @Override
    public String toString() {
        StringBuilder sb = new StringBuilder("Segment(");
        sb.append("value:");
        sb.append(value);
        sb.append(",max:");
        sb.append(max);
        sb.append(",step:");
        sb.append(step);
        sb.append(")");
        return sb.toString();
    }

}

SegmentBuffer

import lombok.Data;

import java.util.Arrays;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;

/**
 * @description: 双buffer
 */
@Data
public class SegmentBuffer {

    /**
     * key,唯一值
     */
    private String key;

    /**
     * 双buffer段数组，这里只有2段
     */
    private Segment[] segments;

    /**
     * 当前使用的segment的index
     */
    private volatile int currentPos;

    /**
     * 下一个segment是否处于可切换状态
     */
    private volatile boolean nextReady;

    /**
     * 线程是否在运行中
     */
    private final AtomicBoolean threadRunning;

    /**
     * 读写锁
     */
    private final ReadWriteLock lock;

    /**
     * 步长
     */
    private volatile int step;

    /**
     * 最大步长不超过100,0000
     */
    public static final int MAX_STEP = 100;

    /**
     * 最小步长不小于200
     */
    public static final int MIN_STEP = 10;

    /**
     * 更新Segment时时间戳
     */
    private volatile long updateTimestamp;

    public SegmentBuffer() {
        segments = new Segment[]{new Segment(this), new Segment(this)};
        currentPos = 0;
        nextReady = false;
        threadRunning = new AtomicBoolean(false);
        lock = new ReentrantReadWriteLock();
    }

    /**
     * 获取当前的segment
     */
    public Segment getCurrent() {
        return segments[currentPos];
    }

    /**
     * 获取下一个segment
     */
    public int nextPos() {
        return (currentPos + 1) % 2;
    }

    /**
     * 切换segment
     */
    public void switchPos() {
        currentPos = nextPos();
    }

    /**
     * segmentBuffer加读锁
     */
    public Lock rLock() {
        return lock.readLock();
    }

    /**
     * segmentBuffer加写锁
     */
    public Lock wLock() {
        return lock.writeLock();
    }

    @Override
    public String toString() {
        final StringBuilder sb = new StringBuilder("SegmentBuffer{");
        sb.append("key='").append(key).append('\'');
        sb.append(", segments=").append(Arrays.toString(segments));
        sb.append(", currentPos=").append(currentPos);
        sb.append(", nextReady=").append(nextReady);
        sb.append(", threadRunning=").append(threadRunning);
        sb.append(", step=").append(step);
        sb.append(", updateTimestamp=").append(updateTimestamp);
        sb.append('}');
        return sb.toString();
    }
}

BizSourceTypeEnum

import com.google.common.collect.Lists;
import lombok.AllArgsConstructor;
import lombok.Getter;

import java.util.List;

/**
 *@Description 业务类型枚举
 **/
@Getter
@AllArgsConstructor
public enum BizSourceTypeEnum {

    BTA("bta", "业务A"),

    BTB("btb", "业务B"),

    BTC("btc", "业务C");

    private String type;

    private String desc;

    /**
     * 根据类型获取对应的中文名字
     * @param type type
     * @return
     */
    public static String getDescByType(String type) {
        for (BizSourceTypeEnum e : BizSourceTypeEnum.values()) {
            if (e.getType().equals(type)) {
                return e.getDesc();
            }
        }
        return "";
    }

    public static List<String> getTypeList() {
        List<String> typeList = Lists.newArrayList();
        for (BizSourceTypeEnum e : BizSourceTypeEnum.values()) {
            typeList.add(e.getType());
        }
        return typeList;
    }
}

序列号生成器

import com.example.serialnodemo.common.enums.BizSourceTypeEnum;
import com.example.serialnodemo.po.PsmSerialNoRecord;
import com.example.serialnodemo.common.segment.model.Segment;
import com.example.serialnodemo.common.segment.model.SegmentBuffer;
import com.example.serialnodemo.service.PsmSerialNoRecordService;
import com.google.common.base.Stopwatch;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.transaction.annotation.Transactional;

import javax.annotation.PostConstruct;
import java.text.NumberFormat;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;


/**
 * @description: 序列号生成器
 */
@Slf4j
@Component
public class SerialNoGenerator {

    @Autowired
    private PsmSerialNoRecordService psmSerialNoRecordService;

    /**
     * 可以自定义线程池
     */
    @Autowired
    private ThreadPoolExecutor segmentUpdateThreadPool;

    private Map<String, SegmentBuffer> cache = new ConcurrentHashMap<>();


    /**
     * 一个Segment维持时间为15分钟
     */
    private static final long SEGMENT_DURATION = 15 * 60 * 1000L;

    /**
     * 最大序列号
     */
    public static final int SERIAL_NO_LIMIT = 9999999;

    /**
     * 序列号位数
     */
    public static final int SERIAL_LENGTH = 7;


    /**
     * @description: 初始化DB序列号
     */
    @PostConstruct
    @Transactional
    public void init() {
        log.info("Init ...");
        List<String> allTypeList = BizSourceTypeEnum.getTypeList();
        List<String> existTypeList = psmSerialNoRecordService.queryForLock(allTypeList);
        BizSourceTypeEnum[] bizSourceTypeEnums = BizSourceTypeEnum.values();
        int existTypeNum = bizSourceTypeEnums.length;
        if (allTypeList.size() == existTypeNum) return;
        for (BizSourceTypeEnum typeEnum : bizSourceTypeEnums) {
            String type = typeEnum.getType();
            if (!existTypeList.contains(type)) {
                PsmSerialNoRecord record = new PsmSerialNoRecord();
                record.setBizSource(type);
                record.setMaxId(0);
                psmSerialNoRecordService.addRecord(record);
            }
        }
    }


    /**
     * @description: 获取序列号
     *
     * @param key 唯一键
     * @return long 序列号
     */
    public String getSerialNo(final String key) throws Exception {
        SegmentBuffer buffer = cache.get(key);
        if (Objects.isNull(buffer)) {
            synchronized (key) {
                buffer = cache.get(key);
                if (Objects.isNull(buffer)) {
                    buffer = new SegmentBuffer();
                    buffer.setKey(key);
                    updateSegmentFromDb(key, buffer.getCurrent());
                    cache.put(key, buffer);
                }
            }
        }
        long seq = getSerialNoFromSegmentBuffer(buffer);
        NumberFormat serialFormat = NumberFormat.getNumberInstance();
        serialFormat.setMinimumIntegerDigits(SERIAL_LENGTH);
        serialFormat.setGroupingUsed(false);
        return serialFormat.format(seq);
    }

    /**
     * @description: 从DB获取数据更新segment
     *
     * @param key 唯一键
     * @param segment segment
     */
    private void updateSegmentFromDb(String key, Segment segment) throws Exception {
        Stopwatch sw = Stopwatch.createStarted();
        SegmentBuffer buffer = segment.getBuffer();
        PsmSerialNoRecord record;
        if (buffer.getUpdateTimestamp() == 0) {
            record = psmSerialNoRecordService.modifyMaxIdAndGet(key, SegmentBuffer.MIN_STEP, SERIAL_NO_LIMIT);
            buffer.setUpdateTimestamp(System.currentTimeMillis());
            buffer.setStep(SegmentBuffer.MIN_STEP);
        } else {
            long duration = System.currentTimeMillis() - buffer.getUpdateTimestamp();
            int nextStep = buffer.getStep();
            if (duration < SEGMENT_DURATION) {
                if (nextStep << 1 > SegmentBuffer.MAX_STEP) {
                    // do nothing
                } else {
                    nextStep = nextStep << 1;
                }
            } else if (duration < SEGMENT_DURATION << 1) {
                // do nothing with nextStep
            } else {
                nextStep = nextStep >> 1 >= SegmentBuffer.MIN_STEP ? nextStep >> 1 : nextStep;
            }
            log.info("SerialNoGenerator updateSegmentFromDb key:{}, step:{}, duration:{}min, nextStep{}", key, buffer.getStep(), String.format("%.2f", ((double) duration / (1000 * 60))), nextStep);
            record = psmSerialNoRecordService.modifyMaxIdAndGet(key, nextStep, SERIAL_NO_LIMIT);
            buffer.setUpdateTimestamp(System.currentTimeMillis());
            buffer.setStep(nextStep);
        }
        long value = SERIAL_NO_LIMIT != record.getMaxId() ? record.getMaxId() - buffer.getStep() + 1 : buffer.getCurrent().getMax() + 1;
        segment.getValue().set(value);
        segment.setMax(record.getMaxId());
        segment.setStep(buffer.getStep());
        log.info("SerialNoGenerator updateSegmentFromDb 数据库更新耗时:{}", sw.stop().elapsed(TimeUnit.MILLISECONDS));
    }

    /**
     * @description: 从buffer中获取序列号
     *
     * @param buffer buffer
     * @return long
     */
    private long getSerialNoFromSegmentBuffer(final SegmentBuffer buffer) throws Exception {
        while (true) {
            buffer.rLock().lock();
            try {
                final Segment segment = buffer.getCurrent();
                // 如果是不可切换状态 && 空闲数量 < 90% && 如果线程状态是没有运行，则将状态设置为运行状态
                if (!buffer.isNextReady() && (segment.getIdle() < 0.9 * segment.getStep())
                        && buffer.getThreadRunning().compareAndSet(false, true)) {
                    segmentUpdateThreadPool.execute(() -> {
                        Segment nextSegment = buffer.getSegments()[buffer.nextPos()];
                        // 数据是否加载完毕，默认false
                        boolean updateOk = false;
                        try {
                            updateSegmentFromDb(buffer.getKey(), nextSegment);
                            updateOk = true;
                            log.info("SerialNoGenerator getIdFromSegmentBuffer key:{}，更新segment:{}", buffer.getKey(), nextSegment);
                        } catch (Exception e) {
                            log.error("SerialNoGenerator getIdFromSegmentBuffer key:{}，更新segment异常:{}", buffer.getKey(), e.getMessage(), e);
                        } finally {
                            if (updateOk) {
                                buffer.wLock().lock();
                                buffer.setNextReady(true);
                                buffer.getThreadRunning().set(false);
                                buffer.wLock().unlock();
                            } else {
                                buffer.getThreadRunning().set(false);
                            }
                        }
                    });
                }
                long value = segment.getValue().getAndIncrement();
                if (value <= segment.getMax()) {
                    return value;
                }
            } finally {
                buffer.rLock().unlock();
            }
            // 如果上面当前段没有拿到序列号，说明当前段序列用完了，需要切换，下一段是在线程池中运行的，所以这里等待一小会儿
            waitAndSleep(buffer);
            buffer.wLock().lock();
            try {
                // 再次拿当前段，因为有可能前面一个线程已经切换好了
                final Segment segment = buffer.getCurrent();
                long value = segment.getValue().getAndIncrement();
                if (value <= segment.getMax()) {
                    return value;
                }
                // 如果是处于可以切换状态，就切换段并设置为不可切换状态，下次获取时就可以在线程池中加载下一段
                if (buffer.isNextReady()) {
                    buffer.switchPos();
                    buffer.setNextReady(false);
                } else {
                    log.error("SerialNoGenerator getIdFromSegmentBuffer 两个Segment都未从数据库中加载 buffer:{}!", buffer);
                    throw new Exception("SerialNoGenerator getIdFromSegmentBuffer 两个Segment都未从数据库中加载");
                }
            } finally {
                buffer.wLock().unlock();
            }
        }
    }

    /**
     * @description: 等待下一段加载完成
     *
     * @param buffer buffer
     */
    private void waitAndSleep(SegmentBuffer buffer) {
        int roll = 0;
        while (buffer.getThreadRunning().get()) {
            roll += 1;
            if (roll > 100) {
                try {
                    TimeUnit.MILLISECONDS.sleep(10);
                    break;
                } catch (InterruptedException e) {
                    log.error("SerialNoGenerator waitAndSleep 线程睡眠异常，Thread:{}, 异常:{}", Thread.currentThread().getName(), e);
                    break;
                }
            }
        }
    }

}

注意

需要注意的是，这种作为工具类的方式放到项目中的方式，若项目服务器宕机或重启会丢失缓存在segment中序列号，造成序列号黑洞，如每个segment的step为10，segment1缓存的是1-10，segment2缓存的是11-20，此时DB中maxid为20，当系统使用到序列号5时，系统宕机了，重启后，因DB中maxid为20，segment1缓存的是21-30，segment2缓存的是31-40，再次获取序列号获取到的是21，其中序列号6-20丢失了，并没有用到。所以序列号最大值可以适时比估计的每天的数据量大些。

参考：

Leaf——美团点评分布式ID生成系统

Leaf：美团分布式ID生成服务开源

Leaf项目Github地址