使用雪花算法（Snowflake Algorithm）在Python中生成唯一ID

engchina

于 2024-06-30 06:44:01 发布

阅读量1.7k

点赞数 39

分类专栏： LINUX 文章标签： python 开发语言雪花算法唯一ID

本文链接：https://blog.csdn.net/engchina/article/details/140075526

版权

LINUX 专栏收录该内容

493 篇文章 60 订阅

订阅专栏

使用雪花算法Snowflake Algorithm在Python中生成唯一ID

使用雪花算法（Snowflake Algorithm）在Python中生成唯一ID

使用雪花算法（Snowflake Algorithm）在Python中生成唯一ID

在分布式系统中生成唯一ID是一个常见的需求。Twitter的雪花算法（Snowflake Algorithm）是一种优秀的解决方案，它可以生成64位的唯一ID，包含时间戳、工作机器ID和序列号。今天，我们将探讨如何在Python中实现雪花算法。

雪花算法简介

雪花算法生成的ID结构如下：

1位符号位（始终为0）
41位时间戳（毫秒级）
5位数据中心ID
5位工作机器ID
12位序列号

这种结构允许在同一毫秒内生成4096个唯一ID，并支持多个数据中心和工作机器。

Python实现

import time
import threading

class SnowflakeGenerator:
    def __init__(self, datacenter_id, worker_id):
        self.datacenter_id = datacenter_id
        self.worker_id = worker_id
        self.sequence = 0
        self.last_timestamp = -1

        # Bit lengths for different parts
        self.datacenter_id_bits = 5
        self.worker_id_bits = 5
        self.sequence_bits = 12

        # Maximum values
        self.max_datacenter_id = -1 ^ (-1 << self.datacenter_id_bits)
        self.max_worker_id = -1 ^ (-1 << self.worker_id_bits)
        self.max_sequence = -1 ^ (-1 << self.sequence_bits)

        # Shift amounts
        self.worker_id_shift = self.sequence_bits
        self.datacenter_id_shift = self.sequence_bits + self.worker_id_bits
        self.timestamp_shift = self.sequence_bits + self.worker_id_bits + self.datacenter_id_bits

        self.lock = threading.Lock()

    def _current_milliseconds(self):
        return int(time.time() * 1000)

    def _til_next_millis(self, last_timestamp):
        timestamp = self._current_milliseconds()
        while timestamp <= last_timestamp:
            timestamp = self._current_milliseconds()
        return timestamp

    def generate_id(self):
        with self.lock:
            timestamp = self._current_milliseconds()

            if timestamp < self.last_timestamp:
                raise ValueError("Clock moved backwards. Refusing to generate id.")

            if timestamp == self.last_timestamp:
                self.sequence = (self.sequence + 1) & self.max_sequence
                if self.sequence == 0:
                    timestamp = self._til_next_millis(self.last_timestamp)
            else:
                self.sequence = 0

            self.last_timestamp = timestamp

            return ((timestamp - 1288834974657) << self.timestamp_shift) | \
                   (self.datacenter_id << self.datacenter_id_shift) | \
                   (self.worker_id << self.worker_id_shift) | \
                   self.sequence

def generate_unique_id(prefix: str, datacenter_id: int, worker_id: int) -> str:
    generator = SnowflakeGenerator(datacenter_id, worker_id)
    snowflake_id = generator.generate_id()
    return f"{prefix}{snowflake_id}"

代码解析

SnowflakeGenerator类：
- 初始化方法设置数据中心ID和工作机器ID，并定义各部分的位长度和最大值。
- _current_milliseconds()方法获取当前时间戳（毫秒级）。
- _til_next_millis()方法等待直到下一毫秒。
- generate_id()方法是核心，它生成唯一的雪花ID。
generate_id()方法：
- 使用锁确保线程安全。
- 获取当前时间戳。
- 处理时钟回拨问题（抛出异常）。
- 处理同一毫秒内的序列号。
- 组合各部分生成最终的ID。
generate_unique_id()函数：
- 创建SnowflakeGenerator实例。
- 生成雪花ID并添加前缀。

使用示例

# 创建一个生成器，指定数据中心ID和工作机器ID
datacenter_id = 1
worker_id = 1

# 生成用户ID
user_id = generate_unique_id("USER_", datacenter_id, worker_id)
print(f"生成的用户ID: {user_id}")

# 生成订单ID
order_id = generate_unique_id("ORDER_", datacenter_id, worker_id)
print(f"生成的订单ID: {order_id}")

# 生成多个产品ID
for i in range(5):
    product_id = generate_unique_id("PROD_", datacenter_id, worker_id)
    print(f"产品 {i+1} ID: {product_id}")

输出可能如下：

生成的用户ID: USER_6791951648483729408
生成的订单ID: ORDER_6791951648483729409
产品 1 ID: PROD_6791951648483729410
产品 2 ID: PROD_6791951648483729411
产品 3 ID: PROD_6791951648483729412
产品 4 ID: PROD_6791951648483729413
产品 5 ID: PROD_6791951648483729414