github.com/andy2046/gopie@v0.7.0/pkg/sequence/memflake.go (about)

     1  package sequence
     2  
     3  import (
     4  	"errors"
     5  	"sync"
     6  	"time"
     7  )
     8  
     9  // MemFlake is an implementation of in-memory Iceflake.
    10  type MemFlake struct {
    11  	sync.Mutex
    12  	startTime      time.Time
    13  	bitLenSequence uint8
    14  	elapsedTime    uint64
    15  	startEpoch     int64
    16  	machineID      uint64
    17  }
    18  
    19  var (
    20  	errLessThanOne          = errors.New("`n` should not be less than 1")
    21  	errTimeDrift            = errors.New("time drifts too much")
    22  	_              Iceflake = &MemFlake{}
    23  )
    24  
    25  // NewMemFlake creates a MemFlake.
    26  func NewMemFlake(startTime time.Time, bitLenSequence uint8, machineID uint64) *MemFlake {
    27  	if startTime.After(time.Now()) {
    28  		// no future time
    29  		return nil
    30  	}
    31  	if startTime.IsZero() {
    32  		startTime = time.Date(2019, 10, 9, 0, 0, 0, 0, time.UTC)
    33  	}
    34  	startEpoch := toIceflakeTime(startTime)
    35  	if startEpoch < 0 {
    36  		// should not be too far away from now
    37  		// should be after 1970/1/1
    38  		return nil
    39  	}
    40  	if bitLenSequence < 10 || bitLenSequence > 21 {
    41  		// keep control
    42  		bitLenSequence = 18
    43  	}
    44  
    45  	return &MemFlake{
    46  		startTime:      startTime,
    47  		bitLenSequence: bitLenSequence,
    48  		startEpoch:     startEpoch,
    49  		machineID:      machineID,
    50  	}
    51  }
    52  
    53  // Next ...
    54  func (m *MemFlake) Next() (uint64, error) {
    55  	return m.NextN(1)
    56  }
    57  
    58  // NextN ...
    59  func (m *MemFlake) NextN(n int) (uint64, error) {
    60  	if n < 1 {
    61  		return 0, errLessThanOne
    62  	}
    63  
    64  	current := currentElapsedTime(m.startEpoch)
    65  	if current < 0 {
    66  		return 0, errTimeDrift
    67  	}
    68  
    69  	nextTime, un := uint64(m.toID(current)), uint64(n)
    70  	m.Lock()
    71  	defer m.Unlock()
    72  	// [elapsedTime - n + 1, elapsedTime] will be returned
    73  	if m.elapsedTime < nextTime {
    74  		m.elapsedTime = nextTime + un - 1
    75  	} else {
    76  		m.elapsedTime += un
    77  	}
    78  	// return the first available timestamp
    79  	return m.elapsedTime - un + 1, nil
    80  }
    81  
    82  // MachineID ...
    83  func (m *MemFlake) MachineID() uint64 {
    84  	return m.machineID
    85  }
    86  
    87  // StartTime ...
    88  func (m *MemFlake) StartTime() time.Time {
    89  	return m.startTime
    90  }
    91  
    92  // BitLenSequence ...
    93  func (m *MemFlake) BitLenSequence() uint8 {
    94  	return m.bitLenSequence
    95  }
    96  
    97  func (m *MemFlake) toID(t int64) int64 {
    98  	return t << m.bitLenSequence
    99  }
   100  
   101  func toIceflakeTime(t time.Time) int64 {
   102  	return t.UTC().Unix() * 1000 // in milliseconds
   103  }
   104  
   105  func currentElapsedTime(startEpoch int64) int64 {
   106  	return toIceflakeTime(time.Now()) - startEpoch
   107  }