main.go

package main

import (
	"errors"
	"fmt"
	"runtime"
	"sync"
	"time"
)

/*
	Twitter Snowflake implementation
*/

const (
	workerIDBits     = 5
	dataCenterIDBits = 5
	sequenceBits     = 12

	maxWorkerID     = -1 ^ (-1 << workerIDBits)     // 32 mask:   0000 0000 0001 1111
	maxDataCenterID = -1 ^ (-1 << dataCenterIDBits) // 32 mask:   0000 0000 0001 1111
	maxSequence     = -1 ^ (-1 << sequenceBits)     // 4096 mas 1111 1111 1111

	timestampShift  = sequenceBits + workerIDBits + dataCenterIDBits // 22
	datacenterShift = sequenceBits + workerIDBits                    // 17
	workerShift     = sequenceBits                                   // 12

	customEpoch = 1735689600000 // time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC).UnixMilli()
)

type Generator struct {
	mt            sync.Mutex
	workerID      int64
	dataCenterID  int64
	sequence      int64
	lastTimestamp int64
}

func NewGenerator(workerID, dataCenterID int64) (*Generator, error) {
	if workerID < 0 || workerID > maxWorkerID {
		return nil, fmt.Errorf("worker ID must be between 0 and %d", maxWorkerID)
	}
	if dataCenterID < 0 || dataCenterID > maxDataCenterID {
		return nil, fmt.Errorf("worker ID must be between 0 and %d", maxDataCenterID)
	}

	return &Generator{
		workerID:     workerID,
		dataCenterID: dataCenterID,
	}, nil
}

func (g *Generator) Generate() (uint64, error) {
	g.mt.Lock()
	defer g.mt.Unlock()

	now := time.Now().UTC().UnixMilli()

	if now < g.lastTimestamp {
		return 0, errors.New("clock moved backwards, refusing to generate id")
	}

	if now == g.lastTimestamp {
		g.sequence = (g.sequence + 1) & maxSequence
		fmt.Printf("parallel generating ids %d\n", g.sequence)

		if g.sequence == 0 {
			fmt.Println("sequence overflow, wait for next millisecond")

			for now <= g.lastTimestamp {
				now = time.Now().UTC().UnixMilli()
			}
		}
	} else {
		g.sequence = 0
	}

	g.lastTimestamp = now

	timestamp := now - customEpoch

	return (uint64(timestamp) << timestampShift) |
			(uint64(g.dataCenterID) << datacenterShift) |
			(uint64(g.workerID) << workerShift) |
			uint64(g.sequence),
		nil
}

func main() {
	const (
		genPerThread = 10_000

		workerID     int64 = 5
		dataCenterID int64 = 9
	)

	generator, err := NewGenerator(workerID, dataCenterID)
	if err != nil {
		panic(err)
	}

	wg := sync.WaitGroup{}

	for range runtime.NumCPU() {
		wg.Add(1)

		go func() {
			defer wg.Done()

			for range genPerThread {
				id, err := generator.Generate()
				if err != nil {
					fmt.Println(fmt.Errorf("err id generating %w", err))
					continue
				}

				fmt.Printf("ID: %d binary: %s\n", id, sPrintfBits(id))
			}
		}()
	}

	wg.Wait()
}

const (
	ColorReset  = "\033[0m"
	ColorRed    = "\033[31m"
	ColorGreen  = "\033[32m"
	ColorYellow = "\033[33m"
	ColorBlue   = "\033[34m"
	ColorPurple = "\033[35m"
)

func sPrintfBits(id uint64) string {
	sequence := id & uint64(maxSequence)
	workerID := (id >> workerShift) & uint64(maxWorkerID)
	dataCenterID := (id >> datacenterShift) & uint64(maxDataCenterID)
	timestamp := id >> timestampShift
	signBit := id >> 63

	return fmt.Sprintf(
		"[%s%1b%s|%s%041b%s|%s%05b%s|%s%05b%s|%s%012b%s]",
		ColorPurple, signBit, ColorReset,
		ColorRed, timestamp, ColorReset,
		ColorGreen, dataCenterID, ColorReset,
		ColorYellow, workerID, ColorReset,
		ColorBlue, sequence, ColorReset,
	)
}