github.com/songzhibin97/gkit@v1.2.13/window/array.go (about)

     1  package window
     2  
     3  import (
     4  	"sync/atomic"
     5  	"unsafe"
     6  
     7  	"github.com/songzhibin97/gkit/internal/clock"
     8  	"github.com/songzhibin97/gkit/internal/sys/safe"
     9  )
    10  
    11  // AtomicArray 封装原子操作, 底层维护 []*Bucket
    12  type AtomicArray struct {
    13  	// length: array长度
    14  	length uint64
    15  
    16  	// base: 数据基地址
    17  	base unsafe.Pointer
    18  	data []*Bucket
    19  }
    20  
    21  // offset 根据index获取底层的桶
    22  func (a *AtomicArray) offset(index uint64) (unsafe.Pointer, bool) {
    23  	if index < 0 {
    24  		return nil, false
    25  	}
    26  	index = index % a.length
    27  	base := a.base
    28  	return unsafe.Pointer(uintptr(base) + uintptr(index*PtrOffSize)), true
    29  }
    30  
    31  // getBucket 根据index获取底层bucket
    32  func (a *AtomicArray) getBucket(index uint64) *Bucket {
    33  	if ptr, ok := a.offset(index); ok {
    34  		return (*Bucket)(atomic.LoadPointer((*unsafe.Pointer)(ptr)))
    35  	}
    36  	return nil
    37  }
    38  
    39  // compareAndSwap 比较交换
    40  func (a *AtomicArray) compareAndSwap(index uint64, old, new *Bucket) bool {
    41  	if ptr, ok := a.offset(index); ok {
    42  		return atomic.CompareAndSwapPointer((*unsafe.Pointer)(ptr), unsafe.Pointer(old), unsafe.Pointer(new))
    43  	}
    44  	return false
    45  }
    46  
    47  // NewAtomicArrayWithTime 初始化 AtomicArray, 需要手动传入 startTime 作为时间戳
    48  func NewAtomicArrayWithTime(length uint64, bucketSize uint64, now uint64, Builder BucketBuilder) *AtomicArray {
    49  	array := &AtomicArray{
    50  		length: length,
    51  		data:   make([]*Bucket, length),
    52  	}
    53  	id := now / bucketSize
    54  	index := id % length
    55  	startTime := calculateStartTime(now, bucketSize)
    56  	for i := index; i < length; i++ {
    57  		b := &Bucket{
    58  			Start: startTime,
    59  			Value: atomic.Value{},
    60  		}
    61  		b.Value.Store(Builder.NewEmptyBucket())
    62  		array.data[i] = b
    63  		startTime += bucketSize
    64  	}
    65  	for i := (uint64)(0); i < index; i++ {
    66  		b := &Bucket{
    67  			Start: startTime,
    68  			Value: atomic.Value{},
    69  		}
    70  		b.Value.Store(Builder.NewEmptyBucket())
    71  		array.data[i] = b
    72  		startTime += bucketSize
    73  	}
    74  	header := (*safe.SliceModel)(unsafe.Pointer(&array.data))
    75  	array.base = header.Data
    76  	return array
    77  }
    78  
    79  // NewAtomicArray 初始化 AtomicArray, startTime 为当前时间
    80  func NewAtomicArray(length uint64, bucketSize uint64, Builder BucketBuilder) *AtomicArray {
    81  	return NewAtomicArrayWithTime(length, bucketSize, clock.GetTimeMillis(), Builder)
    82  }