id-dispatch.py

from math import ceil
from typing import List

import redis


def n1(n):
    # left shift then minus 1
    return (0b1 << n + 1) - 1


class IDDispatch:
    def __init__(
        self,
        db: redis.Redis,
        max_id: int,
        used: List[int] = [],
        prefix='',
        infix='',
        init_db=False,
        check_revoke_count=15,
        expire=10,
    ):
        self.db = db
        self.all_prefix = f'{prefix}:{infix}'
        self.dispatched_set_key = f'{self.all_prefix}:dispatched'
        self.used_key = f'{self.all_prefix}:used'
        self.max_id = max_id
        self.used = used
        self.check_revoke_count = check_revoke_count
        self.expire = expire

        self.all_ids_bitmap = None
        self.all_ids_bitmap_str = None
        if init_db:
            self.init_db()
        self._count = 0

    def init_db(self):
        self.db.delete(self.dispatched_set_key)
        self.db.delete(self.used_key)
        self.db.delete(self.all_prefix)
        self.all_ids_bitmap = n1(self.max_id)
        self.all_ids_bitmap_str = self.all_ids_bitmap.to_bytes(
            ceil(self.max_id / 8), 'little'
        )
        self.db.set(self.all_prefix, self.all_ids_bitmap_str)
        # print(hex(self.all_ids_bitmap), self.all_ids_bitmap_str)
        # 0 is not used
        self.db.setbit(self.all_prefix, 0, 0)
        for i in self.used:
            self.add_used(i)
        print(f'init_db ends')

    def check_revoke(self):
        self._count += 1
        if self._count > self.check_revoke_count:
            self._count = 0
            ks = self.db.smembers(self.dispatched_set_key)
            for k in ks:
                if self.db.get(k) is None:
                    self.revoke(int(k.decode().rsplit(':', 1)[-1]))
                    self.db.srem(self.dispatched_set_key, k)

    def get_next(self) -> int:
        self.check_revoke()
        i = self.db.bitpos(self.all_prefix, 1)
        if i >= 0:
            self.db.setbit(self.all_prefix, i, 0)
            self.add_dispatched(i, self.expire)
        return i

    def add_dispatched(self, i: int, expire: int):
        key = f'{self.all_prefix}:{i}'
        self.db.sadd(self.dispatched_set_key, key)
        self.db.set(key, 'dispatched', expire)

    def revoke(self, i: int):
        if not self.is_used(i):
            self.db.setbit(self.all_prefix, i, 1)

    def add_used(self, i: int):
        self.db.setbit(self.used_key, i, 1)

    def is_used(self, i: int):
        return bool(self.db.getbit(self.used_key, i))

    def get_remains_count(self) -> int:
        return self.db.bitcount(self.all_prefix)

    @staticmethod
    def ensure_int(v: list):
        return [int(x.decode()) for x in v]


class OrderIDDispatch:
    def __init__(self, n):
        self.redis = redis.Redis(**REDIS_SETTINGS)
        print('init all_set ...')
        self.max_id = 10 ** n - 1
        self.account_ids_map = {}

    def get_next_order_id(self, account_id):
        if account_id not in self.account_ids_map:
            self.account_ids_map[account_id] = IDDispatch(
                self.redis,
                self.max_id,
                prefix='order_id',
                infix=account_id,
                init_db=True,
            )
        dispatch = self.account_ids_map[account_id]
        id = dispatch.get_next()
        return f'{account_id}-R{id:08d}'


def test_dispatch_implement():
    import time

    r = redis.Redis(**REDIS_SETTINGS)
    d = IDDispatch(r, 10 ** 2 - 1, prefix='test', init_db=True)
    print('test_dispatch_implement inited')
    assert hex(d.all_ids_bitmap) == '0xfffffffffffffffffffffffff'

    assert d.get_remains_count() == d.max_id, (
        d.get_remains_count(),
        hex(d.all_ids_bitmap),
        d.all_ids_bitmap_str,
    )
    next = d.get_next()
    assert next == 1, next
    next = d.get_next()
    assert next == 2, next
    d.add_used(3)
    d.add_used(5)
    for _ in range(30):
        next = d.get_next()
        print(next)
        time.sleep(1)


order_id_dispatcher = OrderIDDispatch(n=8)


def test_get_order_id():
    for _ in range(100):
        x = order_id_dispatcher.get_next_order_id('XCTD-001234')
        print(x)


# test_dispatch_implement()
test_get_order_id()