github.com/turingchain2020/turingchain@v1.1.21/p2p/utils/blockcache.go (about)

     1  // Copyright Turing Corp. 2018 All Rights Reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package utils
     6  
     7  import (
     8  	"sync"
     9  
    10  	lru "github.com/hashicorp/golang-lru"
    11  )
    12  
    13  // SpaceLimitCache lru缓存封装, 控制占用空间大小
    14  type SpaceLimitCache struct {
    15  	capacity int
    16  	maxSize  int
    17  	currSize int
    18  	sizeMap  map[interface{}]int
    19  	data     *lru.Cache
    20  	lock     *sync.RWMutex
    21  }
    22  
    23  // NewSpaceLimitCache new space limit cache
    24  func NewSpaceLimitCache(num, maxByteSize int) *SpaceLimitCache {
    25  
    26  	cache := &SpaceLimitCache{capacity: num, maxSize: maxByteSize}
    27  	cache.sizeMap = make(map[interface{}]int)
    28  	cache.lock = &sync.RWMutex{}
    29  	var err error
    30  	cache.data, err = lru.New(num)
    31  	if err != nil {
    32  		panic(err)
    33  	}
    34  	return cache
    35  }
    36  
    37  // Add add key val
    38  func (c *SpaceLimitCache) Add(key interface{}, val interface{}, size int) bool {
    39  
    40  	c.lock.Lock()
    41  	defer c.lock.Unlock()
    42  
    43  	//如果存在先删除
    44  	if c.data.Contains(key) {
    45  		c.data.Remove(key)
    46  		c.currSize -= c.sizeMap[key]
    47  		delete(c.sizeMap, key)
    48  	}
    49  
    50  	//单个值超过最大大小
    51  	if size > c.maxSize {
    52  		return false
    53  	}
    54  	c.currSize += size
    55  
    56  	//超过最大大小, 移除最早的值
    57  	for c.currSize > c.maxSize || c.data.Len() >= c.capacity {
    58  		k, _, ok := c.data.RemoveOldest()
    59  		if !ok {
    60  			break
    61  		}
    62  		c.currSize -= c.sizeMap[k]
    63  		delete(c.sizeMap, k)
    64  	}
    65  
    66  	c.data.Add(key, val)
    67  	c.sizeMap[key] = size
    68  	return true
    69  }
    70  
    71  // Get get key
    72  func (c *SpaceLimitCache) Get(key interface{}) interface{} {
    73  	v, _ := c.data.Get(key)
    74  	return v
    75  }
    76  
    77  // Remove remove key
    78  func (c *SpaceLimitCache) Remove(key interface{}) (interface{}, bool) {
    79  
    80  	c.lock.Lock()
    81  	defer c.lock.Unlock()
    82  	val, exist := c.data.Get(key)
    83  	if exist {
    84  		c.data.Remove(key)
    85  		c.currSize -= c.sizeMap[key]
    86  		delete(c.sizeMap, key)
    87  	}
    88  	return val, exist
    89  }
    90  
    91  // Contains check if exist
    92  func (c *SpaceLimitCache) Contains(key interface{}) bool {
    93  
    94  	return c.data.Contains(key)
    95  }