github.com/livekit/protocol@v1.39.3/utils/time_size_cache.go (about)

     1  package utils
     2  
     3  import (
     4  	"slices"
     5  	"sync"
     6  	"time"
     7  
     8  	"github.com/livekit/protocol/utils/mono"
     9  )
    10  
    11  const (
    12  	defaultTTL = 3 * time.Second
    13  )
    14  
    15  type timeSizeCacheItem[T any] struct {
    16  	addedAt int64
    17  	item    *T
    18  	size    int
    19  }
    20  
    21  type TimeSizeCacheParams struct {
    22  	TTL     time.Duration
    23  	MaxSize int
    24  }
    25  
    26  type TimeSizeCache[T any] struct {
    27  	params TimeSizeCacheParams
    28  
    29  	lock      sync.Mutex
    30  	items     []timeSizeCacheItem[T]
    31  	itemsSize int
    32  }
    33  
    34  func NewTimeSizeCache[T any](params TimeSizeCacheParams) *TimeSizeCache[T] {
    35  	return &TimeSizeCache[T]{
    36  		params: params,
    37  	}
    38  }
    39  
    40  func (t *TimeSizeCache[T]) Add(item *T, size int) {
    41  	t.AddAt(item, size, mono.UnixNano())
    42  }
    43  
    44  func (t *TimeSizeCache[T]) AddAt(item *T, size int, atNano int64) {
    45  	t.lock.Lock()
    46  	defer t.lock.Unlock()
    47  
    48  	t.items = append(t.items, timeSizeCacheItem[T]{
    49  		addedAt: atNano,
    50  		item:    item,
    51  		size:    size,
    52  	})
    53  	t.itemsSize += size
    54  
    55  	// prune if size exceeds threshold
    56  	if t.params.MaxSize != 0 && t.itemsSize > t.params.MaxSize {
    57  		idx := 0
    58  		for t.itemsSize > t.params.MaxSize {
    59  			t.itemsSize -= t.items[idx].size
    60  			idx++
    61  		}
    62  		t.items = slices.Delete(t.items, 0, idx)
    63  	}
    64  }
    65  
    66  func (t *TimeSizeCache[T]) Get() []*T {
    67  	t.lock.Lock()
    68  	defer t.lock.Unlock()
    69  
    70  	items := make([]*T, len(t.items))
    71  	for idx, tsci := range t.items {
    72  		items[idx] = tsci.item
    73  	}
    74  
    75  	return items
    76  }
    77  
    78  func (t *TimeSizeCache[T]) Prune() {
    79  	t.PruneAt(mono.UnixNano())
    80  }
    81  
    82  func (t *TimeSizeCache[T]) PruneAt(atNano int64) {
    83  	t.lock.Lock()
    84  	defer t.lock.Unlock()
    85  
    86  	// age based pruning
    87  	threshold := atNano - t.getTTL().Nanoseconds()
    88  	for idx, item := range t.items {
    89  		if item.addedAt >= threshold {
    90  			t.items = slices.Delete(t.items, 0, idx)
    91  			return
    92  		}
    93  		t.itemsSize -= item.size
    94  	}
    95  
    96  	// all items have aged
    97  	t.items = nil
    98  	t.itemsSize = 0
    99  }
   100  
   101  func (t *TimeSizeCache[T]) getTTL() time.Duration {
   102  	ttl := t.params.TTL
   103  	if ttl == 0 {
   104  		ttl = defaultTTL
   105  	}
   106  	return ttl
   107  }