github.com/duomi520/utils@v0.0.0-20240430123446-e03a4cddd6ec/snowflake_test.go (about)

     1  package utils
     2  
     3  import (
     4  	"errors"
     5  	"sync"
     6  	"testing"
     7  	"time"
     8  )
     9  
    10  //
    11  // 工作组中心分配工作机器id
    12  //
    13  
    14  // ErrFailureAllocID 定义错误
    15  var ErrFailureAllocID = errors.New("utils.IdWorker.GetId|工作机器id耗尽。")
    16  
    17  // IDWorker 工作组用于分配工作机器id
    18  type IDWorker struct {
    19  	//2018-6-1 00:00:00 UTC  ，时间戳启动计算时间零点
    20  	SystemCenterStartupTime int64
    21  	queue                   []int
    22  	queueMap                []int
    23  	//消费
    24  	consumption int
    25  	//生产
    26  	production int
    27  	mutex      sync.Mutex
    28  }
    29  
    30  // NewIDWorker 初始化
    31  func NewIDWorker(st int64) *IDWorker {
    32  	w := &IDWorker{
    33  		SystemCenterStartupTime: st,
    34  		queue:                   make([]int, MaxWorkNumber),
    35  		queueMap:                make([]int, MaxWorkNumber),
    36  		production:              0,
    37  		consumption:             0,
    38  	}
    39  	for i := 0; i < MaxWorkNumber; i++ {
    40  		w.queue[i] = i
    41  		w.queueMap[i] = i
    42  	}
    43  	return w
    44  }
    45  
    46  // GetID 取id
    47  func (w *IDWorker) GetID() (int, error) {
    48  	n := 0
    49  	var err error
    50  	w.mutex.Lock()
    51  	if w.queue[w.consumption] == -1 {
    52  		n = -1
    53  		err = ErrFailureAllocID
    54  	} else {
    55  		n = w.queue[w.consumption]
    56  		w.queue[w.consumption] = -1
    57  		w.consumption++
    58  		if w.consumption >= MaxWorkNumber {
    59  			w.consumption = 0
    60  		}
    61  		w.queueMap[n] = -1
    62  		err = nil
    63  	}
    64  	w.mutex.Unlock()
    65  	return n, err
    66  }
    67  
    68  //
    69  //当工作机器与工作组中心的时间不同步时，释放后再利用的workID与之前释放的workID的snowflake id会重复，逻辑上产生bug。
    70  //
    71  
    72  // PutID 还id
    73  func (w *IDWorker) PutID(n int) {
    74  	if n < 0 || n >= MaxWorkNumber {
    75  		return
    76  	}
    77  	w.mutex.Lock()
    78  	if w.queueMap[n] != -1 || w.queue[w.production] != -1 {
    79  		w.mutex.Unlock()
    80  		return
    81  	}
    82  	w.queue[w.production] = n
    83  	w.queueMap[n] = w.production
    84  	w.production++
    85  	if w.production >= MaxWorkNumber {
    86  		w.production = 0
    87  	}
    88  	w.mutex.Unlock()
    89  }
    90  
    91  func TestIDWorker(t *testing.T) {
    92  	// Test table
    93  	idWorkerTests := []int{2, 3, 8, 10, -1, MaxWorkNumber + 1, -10}
    94  	// Verify table
    95  	idWorkerVerify := []int{2, 3, 8, 10}
    96  	w := NewIDWorker(time.Date(2017, time.January, 1, 0, 0, 0, 0, time.UTC).UnixNano())
    97  	for i := 0; i < MaxWorkNumber; i++ {
    98  		f, err := w.GetID()
    99  		if f != i || err != nil {
   100  			t.Error(f, err, w)
   101  		}
   102  	}
   103  	g0, err := w.GetID()
   104  	if err == nil {
   105  		t.Error(g0, err, w)
   106  	}
   107  	for _, l := range idWorkerTests {
   108  		w.PutID(l)
   109  	}
   110  	for _, l := range idWorkerVerify {
   111  		g, err := w.GetID()
   112  		if l != g || err != nil {
   113  			t.Error(g, err, w)
   114  		}
   115  
   116  	}
   117  }
   118  func TestNextID(t *testing.T) {
   119  	var v [30000]int64
   120  	s1 := NewSnowFlakeIDPlus(1, time.Date(2017, time.January, 1, 0, 0, 0, 0, time.UTC).UnixNano())
   121  	for i := 0; i < 10000; i++ {
   122  		v[i] = s1.NextID()
   123  	}
   124  	time.Sleep(time.Millisecond)
   125  	s2 := NewSnowFlakeIDPlus(1, time.Date(2017, time.January, 1, 0, 0, 0, 0, time.UTC).UnixNano())
   126  	for i := 10000; i < 20000; i++ {
   127  		v[i] = s2.NextID()
   128  	}
   129  	s3 := NewSnowFlakeIDPlus(3, time.Date(2017, time.January, 1, 0, 0, 0, 0, time.UTC).UnixNano())
   130  	for i := 20000; i < 30000; i++ {
   131  		v[i] = s3.NextID()
   132  	}
   133  	//验证
   134  	for i := 0; i < (30000 - 1); i++ {
   135  		if v[i] >= v[i+1] {
   136  			t.Error("失败：i:", i, "v[i]:", v[i], "v[i+1]:", v[i+1])
   137  			t.FailNow()
   138  		}
   139  	}
   140  }
   141  
   142  func TestGetWorkID(t *testing.T) {
   143  	s1 := NewSnowFlakeIDPlus(1, time.Date(2017, time.January, 1, 0, 0, 0, 0, time.UTC).UnixNano())
   144  	s2 := NewSnowFlakeIDPlus(555, time.Date(2017, time.January, 1, 0, 0, 0, 0, time.UTC).UnixNano())
   145  	s3 := NewSnowFlakeIDPlus(1022, time.Date(2017, time.January, 1, 0, 0, 0, 0, time.UTC).UnixNano())
   146  	n1 := s1.NextID()
   147  	n2 := s2.NextID()
   148  	n3 := s3.NextID()
   149  	if GetWorkID(n1) != 1 || s1.GetWorkID() != 1 {
   150  		t.Error("失败:", n1, GetWorkID(n1), s1.GetWorkID())
   151  	}
   152  	if GetWorkID(n2) != 555 || s2.GetWorkID() != 555 {
   153  		t.Error("失败:", n2, GetWorkID(n2), s2.GetWorkID())
   154  	}
   155  	if GetWorkID(n3) != 1022 || s3.GetWorkID() != 1022 {
   156  		t.Error("失败:", n3, GetWorkID(n3), s3.GetWorkID())
   157  	}
   158  }