takahe/core/snowflake.py

import secrets
import time


class Snowflake:
    """
    Snowflake ID generator and parser.
    """

    # Epoch is 2022/1/1 at midnight, as these are used for _created_ times in our
    # own database, not original publish times (which would need an earlier one)
    EPOCH = 1641020400

    TYPE_POST = 0b000
    TYPE_POST_INTERACTION = 0b001
    TYPE_IDENTITY = 0b010
    TYPE_REPORT = 0b011
    TYPE_FOLLOW = 0b100

    @classmethod
    def generate(cls, type_id: int) -> int:
        """
        Generates a snowflake-style ID for the given "type". They are designed
        to fit inside 63 bits (a signed bigint)

        ID layout is:
        * 41 bits of millisecond-level timestamp (enough for EPOCH + 69 years)
        * 19 bits of random data (1% chance of clash at 10000 per millisecond)
        * 3 bits of type information

        We use random data rather than a sequence ID to try and avoid pushing
        this job onto the DB - we may do that in future. If a clash does
        occur, the insert will fail and Stator will retry the work for anything
        that's coming in remotely, leaving us to just handle that scenario for
        our own posts, likes, etc.
        """
        # Get the current time in milliseconds
        now: int = int((time.time() - cls.EPOCH) * 1000)
        # Generate random data
        rand_seq: int = secrets.randbits(19)
        # Compose them together
        return (now << 22) | (rand_seq << 3) | type_id

    @classmethod
    def get_type(cls, snowflake: int) -> int:
        """
        Returns the type of a given snowflake ID
        """
        if snowflake < (1 << 22):
            raise ValueError("Not a valid Snowflake ID")
        return snowflake & 0b111

    @classmethod
    def get_time(cls, snowflake: int) -> float:
        """
        Returns the generation time (in UNIX timestamp seconds) of the ID
        """
        if snowflake < (1 << 22):
            raise ValueError("Not a valid Snowflake ID")
        return ((snowflake >> 22) / 1000) + cls.EPOCH

    # Handy pre-baked methods for django model defaults
    @classmethod
    def generate_post(cls) -> int:
        return cls.generate(cls.TYPE_POST)

    @classmethod
    def generate_post_interaction(cls) -> int:
        return cls.generate(cls.TYPE_POST_INTERACTION)

    @classmethod
    def generate_identity(cls) -> int:
        return cls.generate(cls.TYPE_IDENTITY)

    @classmethod
    def generate_report(cls) -> int:
        return cls.generate(cls.TYPE_REPORT)

    @classmethod
    def generate_follow(cls) -> int:
        return cls.generate(cls.TYPE_FOLLOW)