github.com/TeaOSLab/EdgeNode@v1.3.8/internal/utils/expires/list.go (about)

     1  package expires
     2  
     3  import (
     4  	"github.com/TeaOSLab/EdgeNode/internal/zero"
     5  	"sync"
     6  )
     7  
     8  type ItemMap = map[uint64]zero.Zero
     9  
    10  type List struct {
    11  	expireMap map[int64]ItemMap // expires timestamp => map[id]ItemMap
    12  	itemsMap  map[uint64]int64  // itemId => timestamp
    13  
    14  	mu sync.RWMutex
    15  
    16  	gcCallback      func(itemId uint64)
    17  	gcBatchCallback func(itemIds ItemMap)
    18  
    19  	lastTimestamp int64
    20  }
    21  
    22  func NewList() *List {
    23  	var list = &List{
    24  		expireMap: map[int64]ItemMap{},
    25  		itemsMap:  map[uint64]int64{},
    26  	}
    27  
    28  	SharedManager.Add(list)
    29  
    30  	return list
    31  }
    32  
    33  func NewSingletonList() *List {
    34  	var list = &List{
    35  		expireMap: map[int64]ItemMap{},
    36  		itemsMap:  map[uint64]int64{},
    37  	}
    38  
    39  	return list
    40  }
    41  
    42  // Add 添加条目
    43  // 如果条目已经存在，则覆盖
    44  func (this *List) Add(itemId uint64, expiresAt int64) {
    45  	this.mu.Lock()
    46  	defer this.mu.Unlock()
    47  
    48  	if this.lastTimestamp == 0 || this.lastTimestamp > expiresAt {
    49  		this.lastTimestamp = expiresAt
    50  	}
    51  
    52  	// 是否已经存在
    53  	oldExpiresAt, ok := this.itemsMap[itemId]
    54  	if ok {
    55  		if oldExpiresAt == expiresAt {
    56  			return
    57  		}
    58  		delete(this.expireMap[oldExpiresAt], itemId)
    59  		if len(this.expireMap[oldExpiresAt]) == 0 {
    60  			delete(this.expireMap, oldExpiresAt)
    61  		}
    62  	}
    63  
    64  	expireItemMap, ok := this.expireMap[expiresAt]
    65  	if ok {
    66  		expireItemMap[itemId] = zero.New()
    67  	} else {
    68  		this.expireMap[expiresAt] = ItemMap{
    69  			itemId: zero.New(),
    70  		}
    71  	}
    72  
    73  	this.itemsMap[itemId] = expiresAt
    74  }
    75  
    76  func (this *List) Remove(itemId uint64) {
    77  	this.mu.Lock()
    78  	defer this.mu.Unlock()
    79  	this.removeItem(itemId)
    80  }
    81  
    82  func (this *List) ExpiresAt(itemId uint64) int64 {
    83  	this.mu.RLock()
    84  	defer this.mu.RUnlock()
    85  	return this.itemsMap[itemId]
    86  }
    87  
    88  func (this *List) GC(timestamp int64) ItemMap {
    89  	if this.lastTimestamp > timestamp+1 {
    90  		return nil
    91  	}
    92  	var itemMap = this.gcItems(timestamp)
    93  	if len(itemMap) == 0 {
    94  		return itemMap
    95  	}
    96  
    97  	if this.gcCallback != nil {
    98  		for itemId := range itemMap {
    99  			this.gcCallback(itemId)
   100  		}
   101  	}
   102  	if this.gcBatchCallback != nil {
   103  		this.gcBatchCallback(itemMap)
   104  	}
   105  
   106  	return itemMap
   107  }
   108  
   109  func (this *List) Clean() {
   110  	this.mu.Lock()
   111  	this.itemsMap = map[uint64]int64{}
   112  	this.expireMap = map[int64]ItemMap{}
   113  	this.mu.Unlock()
   114  }
   115  
   116  func (this *List) Count() int {
   117  	this.mu.RLock()
   118  	var count = len(this.itemsMap)
   119  	this.mu.RUnlock()
   120  	return count
   121  }
   122  
   123  func (this *List) OnGC(callback func(itemId uint64)) *List {
   124  	this.gcCallback = callback
   125  	return this
   126  }
   127  
   128  func (this *List) OnGCBatch(callback func(itemMap ItemMap)) *List {
   129  	this.gcBatchCallback = callback
   130  	return this
   131  }
   132  
   133  func (this *List) ExpireMap() map[int64]ItemMap {
   134  	return this.expireMap
   135  }
   136  
   137  func (this *List) ItemsMap() map[uint64]int64 {
   138  	return this.itemsMap
   139  }
   140  
   141  func (this *List) LastTimestamp() int64 {
   142  	return this.lastTimestamp
   143  }
   144  
   145  func (this *List) removeItem(itemId uint64) {
   146  	expiresAt, ok := this.itemsMap[itemId]
   147  	if !ok {
   148  		return
   149  	}
   150  	delete(this.itemsMap, itemId)
   151  
   152  	expireItemMap, ok := this.expireMap[expiresAt]
   153  	if ok {
   154  		delete(expireItemMap, itemId)
   155  		if len(expireItemMap) == 0 {
   156  			delete(this.expireMap, expiresAt)
   157  		}
   158  	}
   159  }
   160  
   161  func (this *List) gcItems(timestamp int64) ItemMap {
   162  	this.mu.RLock()
   163  	expireItemsMap, ok := this.expireMap[timestamp]
   164  	this.mu.RUnlock()
   165  
   166  	if ok {
   167  		this.mu.Lock()
   168  		for itemId := range expireItemsMap {
   169  			delete(this.itemsMap, itemId)
   170  		}
   171  		delete(this.expireMap, timestamp)
   172  		this.mu.Unlock()
   173  	}
   174  
   175  	return expireItemsMap
   176  }