github.com/GuanceCloud/cliutils@v1.1.21/pipeline/ptinput/plcache/cache.go (about)

     1  // Unless explicitly stated otherwise all files in this repository are licensed
     2  // under the MIT License.
     3  // This product includes software developed at Guance Cloud (https://www.guance.com/).
     4  // Copyright 2021-present Guance, Inc.
     5  
     6  // Package plcache implements cache in pipeline
     7  package plcache
     8  
     9  import (
    10  	"container/list"
    11  	"errors"
    12  	"sync"
    13  	"time"
    14  )
    15  
    16  var (
    17  	ErrClosed = errors.New("cache is closed")
    18  	ErrArg    = errors.New("illegal arguments")
    19  )
    20  
    21  type cacheItem struct {
    22  	key        string
    23  	value      any
    24  	expiration time.Duration
    25  	circle     int
    26  }
    27  
    28  type Cache struct {
    29  	mu        sync.RWMutex
    30  	interval  time.Duration
    31  	ticker    *time.Ticker
    32  	slots     []*list.List
    33  	numSlots  int
    34  	tickedPos int
    35  
    36  	// All cache items for fast query
    37  	items map[string]*list.Element
    38  
    39  	setChannel  chan cacheItem
    40  	stopChannel chan struct{}
    41  }
    42  
    43  // NewCache returns a Cache.
    44  func NewCache(interval time.Duration, numSlots int) (*Cache, error) {
    45  	if interval <= 0 || numSlots <= 0 {
    46  		return nil, ErrArg
    47  	}
    48  
    49  	return NewCacheWithTicker(interval, numSlots, time.NewTicker(interval))
    50  }
    51  
    52  // NewCacheWithTicker returns a Cache with the given ticker.
    53  func NewCacheWithTicker(interval time.Duration, numSlots int, ticker *time.Ticker) (*Cache, error) {
    54  	c := &Cache{
    55  		interval:  interval,
    56  		ticker:    ticker,
    57  		slots:     make([]*list.List, numSlots),
    58  		numSlots:  numSlots,
    59  		tickedPos: numSlots - 1,
    60  		items:     make(map[string]*list.Element),
    61  
    62  		setChannel:  make(chan cacheItem),
    63  		stopChannel: make(chan struct{}),
    64  	}
    65  	c.initSlot()
    66  	go c.run()
    67  
    68  	return c, nil
    69  }
    70  
    71  func (c *Cache) initSlot() {
    72  	for i := 0; i < c.numSlots; i++ {
    73  		c.slots[i] = list.New()
    74  	}
    75  }
    76  
    77  func (c *Cache) run() {
    78  	for {
    79  		select {
    80  		case <-c.ticker.C:
    81  			c.onTick()
    82  		case ci := <-c.setChannel:
    83  			c.setCacheItem(ci)
    84  		case <-c.stopChannel:
    85  			c.ticker.Stop()
    86  			return
    87  		}
    88  	}
    89  }
    90  
    91  func (c *Cache) onTick() {
    92  	c.tickedPos = (c.tickedPos + 1) % c.numSlots
    93  	lst := c.slots[c.tickedPos]
    94  	c.scanAndRemoveExpiredCache(lst)
    95  }
    96  
    97  func (c *Cache) scanAndRemoveExpiredCache(lst *list.List) {
    98  	c.mu.Lock()
    99  	defer c.mu.Unlock()
   100  
   101  	for e := lst.Front(); e != nil; {
   102  		ci := e.Value.(cacheItem)
   103  		if ci.circle > 0 {
   104  			ci.circle--
   105  			e = e.Next()
   106  			continue
   107  		}
   108  		lst.Remove(e)
   109  		delete(c.items, e.Value.(cacheItem).key)
   110  		e = e.Next()
   111  	}
   112  }
   113  
   114  func (c *Cache) setCacheItem(ci cacheItem) {
   115  	c.mu.Lock()
   116  	defer c.mu.Unlock()
   117  
   118  	if ci.expiration < c.interval {
   119  		ci.expiration = c.interval
   120  	}
   121  
   122  	pos, circle := c.getPosAndCircle(ci.expiration)
   123  	ci.circle = circle
   124  	c.slots[pos].PushBack(ci)
   125  	c.items[ci.key] = c.slots[pos].Back()
   126  }
   127  
   128  func (c *Cache) getPosAndCircle(d time.Duration) (pos, circle int) {
   129  	step := int(d / c.interval)
   130  	pos = (c.tickedPos + step) % c.numSlots
   131  	circle = (step - 1) / c.numSlots
   132  
   133  	return
   134  }
   135  
   136  func (c *Cache) Get(key string) (any, bool, error) {
   137  	select {
   138  	case <-c.stopChannel:
   139  		return nil, false, ErrClosed
   140  	default:
   141  		if elem, exists := c.items[key]; exists {
   142  			item := elem.Value.(cacheItem)
   143  
   144  			return item.value, true, nil
   145  		}
   146  		return nil, false, nil
   147  	}
   148  }
   149  
   150  func (c *Cache) Set(key string, value any, expiration time.Duration) error {
   151  	if expiration <= 0 {
   152  		return ErrArg
   153  	}
   154  
   155  	select {
   156  	case c.setChannel <- cacheItem{
   157  		key:        key,
   158  		value:      value,
   159  		expiration: expiration,
   160  	}:
   161  		return nil
   162  	case <-c.stopChannel:
   163  		return ErrClosed
   164  	}
   165  }
   166  
   167  func (c *Cache) Stop() {
   168  	close(c.stopChannel)
   169  }