github.com/matrixorigin/matrixone@v0.7.0/pkg/fileservice/lru.go (about)

     1  // Copyright 2022 Matrix Origin
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package fileservice
    16  
    17  import (
    18  	"container/list"
    19  	"sync"
    20  )
    21  
    22  type LRU struct {
    23  	sync.Mutex
    24  	capacity int64
    25  	size     int64
    26  	evicts   *list.List
    27  	kv       map[any]*list.Element
    28  }
    29  
    30  type lruItem struct {
    31  	Key   any
    32  	Value any
    33  	Size  int64
    34  }
    35  
    36  func NewLRU(capacity int64) *LRU {
    37  	return &LRU{
    38  		capacity: capacity,
    39  		evicts:   list.New(),
    40  		kv:       make(map[any]*list.Element),
    41  	}
    42  }
    43  
    44  func (l *LRU) Set(key any, value any, size int64) {
    45  	l.Lock()
    46  	defer l.Unlock()
    47  
    48  	if elem, ok := l.kv[key]; ok {
    49  		// replace
    50  		item := elem.Value.(*lruItem)
    51  		l.size -= item.Size
    52  		l.size += size
    53  		l.evicts.MoveToFront(elem)
    54  		item.Size = size
    55  		item.Key = key
    56  		item.Value = value
    57  
    58  	} else {
    59  		// insert
    60  		item := &lruItem{
    61  			Key:   key,
    62  			Value: value,
    63  			Size:  size,
    64  		}
    65  		elem := l.evicts.PushFront(item)
    66  		l.kv[key] = elem
    67  		l.size += size
    68  	}
    69  
    70  	l.evict()
    71  }
    72  
    73  func (l *LRU) evict() {
    74  	for {
    75  		if l.size <= l.capacity {
    76  			return
    77  		}
    78  		if len(l.kv) == 0 {
    79  			return
    80  		}
    81  
    82  		elem := l.evicts.Back()
    83  		for {
    84  			if elem == nil {
    85  				return
    86  			}
    87  			item := elem.Value.(*lruItem)
    88  
    89  			// RC value
    90  			if rc, ok := item.Value.(interface {
    91  				RefCount() int64
    92  			}); ok {
    93  				if rc.RefCount() > 0 {
    94  					// skip
    95  					elem = elem.Prev()
    96  					continue
    97  				}
    98  			}
    99  
   100  			l.size -= item.Size
   101  			l.evicts.Remove(elem)
   102  			delete(l.kv, item.Key)
   103  			break
   104  		}
   105  
   106  	}
   107  }
   108  
   109  func (l *LRU) Get(key any) (value any, size int64, ok bool) {
   110  	l.Lock()
   111  	defer l.Unlock()
   112  	if elem, ok := l.kv[key]; ok {
   113  		l.evicts.MoveToFront(elem)
   114  		item := elem.Value.(*lruItem)
   115  		return item.Value, item.Size, true
   116  	}
   117  	return nil, 0, false
   118  }
   119  
   120  func (l *LRU) Flush() {
   121  	l.Lock()
   122  	defer l.Unlock()
   123  	l.size = 0
   124  	l.evicts = list.New()
   125  	l.kv = make(map[any]*list.Element)
   126  }