github.com/songzhibin97/gkit@v1.2.13/generator/snowflake.go (about)

     1  package generator
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  	"net"
     7  	"sync"
     8  	"time"
     9  )
    10  
    11  /*
    12  Total: 			64bit
    13  First: 			1bit
    14  TimeBit:		39bit
    15  NodeID: 		16bit(uint16)
    16  SequenceID: 	8bit(uint8)
    17  */
    18  
    19  const (
    20  	TimeBit     = 39
    21  	NodeBit     = 16
    22  	SequenceBit = 8
    23  )
    24  
    25  // Snowflake is a distributed unique ID generator.
    26  type Snowflake struct {
    27  	// 互斥锁
    28  	sync.Mutex
    29  
    30  	// startTime: 开始时间
    31  	startTime int64
    32  
    33  	// elapsedTime: 经过的时间
    34  	elapsedTime int64
    35  
    36  	// sequence: 流水号
    37  	sequence uint16
    38  
    39  	// node: 节点
    40  	node uint16
    41  }
    42  
    43  // NextID 获取下一个ID
    44  func (s *Snowflake) NextID() (uint64, error) {
    45  	const maskSequence = uint16(1<<SequenceBit - 1)
    46  
    47  	s.Lock()
    48  	defer s.Unlock()
    49  
    50  	current := currentElapsedTime(s.startTime)
    51  	if s.elapsedTime < current {
    52  		s.elapsedTime = current
    53  		s.sequence = 0
    54  	} else {
    55  		// s.elapsedTime >= current
    56  		s.sequence = (s.sequence + 1) & maskSequence
    57  		if s.sequence == 0 {
    58  			s.elapsedTime++
    59  			overtime := s.elapsedTime - current
    60  			time.Sleep(sleepTime(overtime))
    61  		}
    62  	}
    63  	return s.toID()
    64  }
    65  
    66  // timeUnit 时间单位,纳秒
    67  const timeUnit = 1e7
    68  
    69  func toSnowflakeTime(t time.Time) int64 {
    70  	return t.UTC().UnixNano() / timeUnit
    71  }
    72  
    73  func currentElapsedTime(startTime int64) int64 {
    74  	return toSnowflakeTime(time.Now()) - startTime
    75  }
    76  
    77  func sleepTime(overtime int64) time.Duration {
    78  	return time.Duration(overtime)*10*time.Millisecond -
    79  		time.Duration(time.Now().UTC().UnixNano()%timeUnit)*time.Nanosecond
    80  }
    81  
    82  func (s *Snowflake) toID() (uint64, error) {
    83  	if s.elapsedTime >= 1<<TimeBit {
    84  		return 0, errors.New("over the time limit")
    85  	}
    86  
    87  	return uint64(s.elapsedTime)<<(SequenceBit+NodeBit) |
    88  		uint64(s.sequence)<<NodeBit |
    89  		uint64(s.node), nil
    90  }
    91  
    92  // Decompose 根据id分解为mapping包含这组数据的所有信息
    93  func Decompose(id uint64) map[string]uint64 {
    94  	const maskSequence = uint64((1<<SequenceBit - 1) << NodeBit)
    95  	const maskMachineID = uint64(1<<NodeBit - 1)
    96  
    97  	msb := id >> 63
    98  	t := id >> (SequenceBit + NodeBit)
    99  	sequence := id & maskSequence >> NodeBit
   100  	node := id & maskMachineID
   101  	return map[string]uint64{
   102  		"id":       id,
   103  		"msb":      msb,
   104  		"time":     t,
   105  		"sequence": sequence,
   106  		"node":     node,
   107  	}
   108  }
   109  
   110  // localIPv4 获取本地IP
   111  func localIPv4() (net.IP, error) {
   112  	as, err := net.InterfaceAddrs()
   113  	if err != nil {
   114  		return nil, err
   115  	}
   116  
   117  	for _, a := range as {
   118  		i, ok := a.(*net.IPNet)
   119  		if !ok || i.IP.IsLoopback() {
   120  			continue
   121  		}
   122  
   123  		ip := i.IP.To4()
   124  		if isPrivateIPv4(ip) {
   125  			return ip, nil
   126  		}
   127  	}
   128  	return nil, errors.New("no private ip address")
   129  }
   130  
   131  // isPrivateIPv4 判断是否有效IP
   132  func isPrivateIPv4(ip net.IP) bool {
   133  	return ip != nil &&
   134  		(ip[0] == 10 || ip[0] == 172 && (ip[1] >= 16 && ip[1] < 32) || ip[0] == 192 && ip[1] == 168)
   135  }
   136  
   137  // IpToUint16 将IP地址转化为uint16
   138  func IpToUint16(ip net.IP) (uint16, error) {
   139  	return uint16(ip[2])<<8 + uint16(ip[3]), nil
   140  }
   141  
   142  // LocalIpToUint16 本地IP转化为uint16
   143  func LocalIpToUint16() (uint16, error) {
   144  	ip, err := localIPv4()
   145  	if err != nil {
   146  		return 0, err
   147  	}
   148  	fmt.Println(ip)
   149  	return uint16(ip[2])<<8 + uint16(ip[3]), nil
   150  }
   151  
   152  // NewSnowflake 初始化
   153  // StartTime 起始时间
   154  // NodeID 服务器ID
   155  func NewSnowflake(startTime time.Time, nodeID uint16) Generator {
   156  	sf := new(Snowflake)
   157  	sf.sequence = uint16(1<<SequenceBit - 1)
   158  	sf.node = nodeID
   159  	if startTime.After(time.Now()) {
   160  		return nil
   161  	}
   162  	if startTime.IsZero() {
   163  		sf.startTime = toSnowflakeTime(time.Date(2014, 9, 1, 0, 0, 0, 0, time.UTC))
   164  	} else {
   165  		sf.startTime = toSnowflakeTime(startTime)
   166  	}
   167  	return sf
   168  }