uid/uid.go

// Copyright 2022 @qoobing.com. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//
// Unique id creator tools, It can be used to create 'userid' or
// 'transaction id' or something else need global unique id in a
// service cluster.
package uid

import (
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"github.com/gammazero/deque"
)

// IdCreator
type IdCreator interface {
	GetId(parts ...interface{}) string
}

// IdCreatorHelper
type IdCreatorHelper interface {
	// CreatePrefix create id prefix
	CreatePrefix(parts ...interface{})
	// CreateRange create id range from database table t_id.
	CreateRange(prefix string) (start, length int64, err error)
	// CreateMixedId create finnal transaction id
	CreateMixedId(prefix string, id int64, parts ...interface{}) string
}

// idCreater
type idCreater struct {
	lock   sync.Mutex
	Type   string
	Ranges map[string]*deque.Deque
	Helper IdCreateHelper
}

// NewIdCreater to new a id creator by helper.
func NewIdCreater(typ string, helper IdCreatorHelper) IdCreater {
	// Step 1. basic value
	var idg = &idCreater{
		lock:   sync.Mutex,
		Type:   typ,
		Ranges: map[string]*deque.Deque{},
		Helper: helper,
	}

	// Step 2. generate the first range.
	if prefix, err = idg.createDefaultPrefix(); err == nil {
		if err := idg.createNewRange(prefix); err != nil {
			log.Errorf("generate the first range failed: %s", err)
		}
	}

	// Step 3. backgroud task for generate range
	go func() {
		for {
			for prefix, queue := range idg.Ranges {
				var qlen = queue.Len()
				var timeoutMs = 1

				if queueLen <= 0 {
					timeoutMs = 1
				} else if queueLen <= 2 {
					timeoutMs = 50
				} else {
					timeoutMs = -1000
				}

				if timeoutMs > 0 {
					idg.createNewRange(prefix)
					time.Sleep(timeoutMs * time.Millisecond)
				} else if timeoutMs < 0 {
					time.Sleep((-timeoutMs) * time.Millisecond)
				}
			}
			time.Sleep(1000 * time.Millisecond)
		}
	}()

	return idg
}

func (idg *idCreater) GetId(idparts ...interface{}) string {
	var prefix = idg.Creator.CreatePrefix(idparts...)
	for {
		var id int64 = 0
		var cur = idg.getCurrentRange(prefix)
		// Step 1. Optimistic[no lock] to add id in current range.
		for {
			id = atomic.LoadInt64(&cur.Next)
			if id > cur.End {
				//next range
				break
			}
			if atomic.CompareAndSwapInt64(&cur.Next, id, id+1) {
				//success
				break
			}
		}

		// Step 2. No lock add failed, try add id in next range.
		if id > cur.End {
			//next range
			idg.turnToNextRange(cur)
			continue
		}

		// Step 3. Success, return mixed id.
		return idg.MixId(cur, id)
	}
}

// Range Range define the range [start, end] of ids.
type Range struct {
	Id     int64  // id
	Status int    // 0:invlid; 1:ok; 2:used
	Prefix string // prefix
	Start  int64  // begin id of this range(include)
	End    int64  // end id of this range(include)
	Length int64  // lenght of this range
}

func (idg *idCreater) createDefaultPrefix() (prefix string, err error) {
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("default prefix not surport by user")
		}
	}()
	if prefix = idg.Helper.CreatePrefix(); len(prefix) == 0 {
		err = fmt.Errorf("default prefix not surport by user")
	}
	return prefix, err
}

func (idg *idCreater) createNewRange(prefix string) error {
	// Step 1. create range from db or something else.
	start, length, err := idg.Creator.CreateRange(idg.Type, prefix)
	if err != nil {
		return fmt.Errorf(
			"idCreater '%s' generate range err:'%s'",
			idg.Type, err.Error())
	}

	// Step 2. lock generator.
	idg.lock.Lock()
	defer idg.lock.Unlock()
	lock, ok := idg.RangesLocks[prefix]
	if !ok {
		lock = &sync.Mutex{}
		queue = deque.New[*Range]()
		idg.Ranges[prefix] = queue
		idg.RangesLocks[prefix] = lock
	}
	lock.Lock()
	defer lock.Unlock()

	// Step 3. push to ranges queue.
	queue.PushBack(
		&Range{
			Id:     0,
			Status: 1,
			Prefix: prefix,
			Start:  start,
			End:    start + length - 1,
			Length: length,
		})
	return nil
}

func (idg *idCreater) getCurrentRange(prefix string) (r *Range) {
	// Step 0. get ranges queue by prefix.
	var queue, ok = idg.Ranges[prefix]
	if !ok {
		idg.createNewRange(prefix)
	}
	var lock, ok = idg.RangesLocks[prefix]
	if !ok {
		panic("unreachable code: lock is not exist")
	}
	lock.Lock()
	defer lock.Unlock()

	// Step 1. try read from queue directly.
	if queue.Len() >= 1 {
		return queue.Front()
	}

	// Step 2. try create a new range & try again.
	idg.createNewRange(prefix)
	if queue.Len() >= 1 {
		return queue.Front()
	}
	panic("get current range failed, too many pressure???")
}

func (idg *idCreater) turnToNextRange(r *Range) {
	// We donot lock the range, because there is no too much
	// side effect even if 'createNewRange' execute twice.
	switch r.Status {
	case 1:
		r.Status = 2
		idg.createNewRange(r.Prefix)
	case 2:
		//have turn by others, we do nothing.
	default:
		panic("invalid status, maybe somethong bug")
	}
}

// demo demo range creator, return current timestamp % 86400 * lenght.
// **DONOT** use in product enviroment.
// **DONOT** use in product enviroment.
// **DONOT** use in product enviroment.
type demo struct{}

func (d *demo) GetRange(typ, prefix string) (start, length int64, err error) {
	var curdate string = time.Now().Format("20060102")
	if prefix != curdate {
		return 0, 0, fmt.Errorf("demo just support prefix is current date")
	}

	start = uint64((time.Now().Unix() % 86400) * 10000)
	length = 10000
	return start, length, nil
}