yuid

A small, unique, URL-safe ID generator written in Go.
git clone git://git.lair.cx/yuid
Log | Files | Refs | README | LICENSE

generator.go (1842B)

      1 package yuid
      2 
      3 import (
      4 	"io"
      5 	"sync"
      6 	"time"
      7 )
      8 
      9 const (
     10 	Characters = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz-_" // Exactly 64 characters
     11 )
     12 
     13 type Generator struct {
     14 	sync.Mutex
     15 	Source    io.Reader
     16 	Now       func() int64
     17 	sequence  uint8
     18 	timestamp int64
     19 }
     20 
     21 type GeneratorOption func(*Generator)
     22 
     23 func NewGenerator(opts ...GeneratorOption) *Generator {
     24 	g := Generator{
     25 		Source: InsecureSource{},
     26 		Now:    timeNow,
     27 	}
     28 
     29 	for _, opt := range opts {
     30 		opt(&g)
     31 	}
     32 
     33 	return &g
     34 }
     35 
     36 func WithSource(source io.Reader) GeneratorOption {
     37 	return func(g *Generator) {
     38 		g.Source = source
     39 	}
     40 }
     41 
     42 func WithNowFunc(now func() int64) GeneratorOption {
     43 	return func(g *Generator) {
     44 		g.Now = now
     45 	}
     46 }
     47 
     48 func (g *Generator) Next(id *ID) error {
     49 	var ts int64
     50 	var seq uint8
     51 
     52 	now := g.Now()
     53 
     54 	g.Lock()
     55 
     56 	if now > g.timestamp {
     57 		g.timestamp = now
     58 		g.sequence = 0
     59 	} else if g.sequence == 255 {
     60 		g.timestamp++
     61 		g.sequence = 0
     62 	} else {
     63 		g.sequence++
     64 	}
     65 
     66 	ts = g.timestamp
     67 	seq = g.sequence
     68 
     69 	_, err := io.ReadFull(g.Source, id[12:])
     70 
     71 	g.Unlock()
     72 
     73 	if err != nil {
     74 		return err
     75 	}
     76 
     77 	putHeader(id, ts, seq)
     78 	setBody(id)
     79 
     80 	return nil
     81 }
     82 
     83 func putHeader(id *ID, ts int64, seq uint8) {
     84 	id[0] = Characters[(ts>>58)&63]
     85 	id[1] = Characters[(ts>>52)&63]
     86 	id[2] = Characters[(ts>>46)&63]
     87 	id[3] = Characters[(ts>>40)&63]
     88 	id[4] = Characters[(ts>>34)&63]
     89 	id[5] = Characters[(ts>>28)&63]
     90 	id[6] = Characters[(ts>>22)&63]
     91 	id[7] = Characters[(ts>>16)&63]
     92 	id[8] = Characters[(ts>>10)&63]
     93 	id[9] = Characters[(ts>>4)&63]
     94 	id[10] = Characters[uint8((ts&15)<<2)|(seq>>6)]
     95 	id[11] = Characters[seq&63]
     96 }
     97 
     98 func setBody(id *ID) {
     99 	for i := 12; i < IDSize; i++ {
    100 		id[i] = Characters[id[i]&63]
    101 	}
    102 }
    103 
    104 func (g *Generator) MustNext(id *ID) {
    105 	err := g.Next(id)
    106 	if err != nil {
    107 		panic(err)
    108 	}
    109 }
    110 
    111 func timeNow() int64 {
    112 	return time.Now().UnixMilli()
    113 }