github.com/tumi8/quic-go@v0.37.4-tum/token_store.go (about)

     1  package quic
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/tumi8/quic-go/noninternal/utils"
     7  
     8  	list "github.com/tumi8/quic-go/noninternal/utils/linkedlist"
     9  )
    10  
    11  type singleOriginTokenStore struct {
    12  	tokens []*ClientToken
    13  	len    int
    14  	p      int
    15  }
    16  
    17  func newSingleOriginTokenStore(size int) *singleOriginTokenStore {
    18  	return &singleOriginTokenStore{tokens: make([]*ClientToken, size)}
    19  }
    20  
    21  func (s *singleOriginTokenStore) Add(token *ClientToken) {
    22  	s.tokens[s.p] = token
    23  	s.p = s.index(s.p + 1)
    24  	s.len = utils.Min(s.len+1, len(s.tokens))
    25  }
    26  
    27  func (s *singleOriginTokenStore) Pop() *ClientToken {
    28  	s.p = s.index(s.p - 1)
    29  	token := s.tokens[s.p]
    30  	s.tokens[s.p] = nil
    31  	s.len = utils.Max(s.len-1, 0)
    32  	return token
    33  }
    34  
    35  func (s *singleOriginTokenStore) Len() int {
    36  	return s.len
    37  }
    38  
    39  func (s *singleOriginTokenStore) index(i int) int {
    40  	mod := len(s.tokens)
    41  	return (i + mod) % mod
    42  }
    43  
    44  type lruTokenStoreEntry struct {
    45  	key   string
    46  	cache *singleOriginTokenStore
    47  }
    48  
    49  type lruTokenStore struct {
    50  	mutex sync.Mutex
    51  
    52  	m                map[string]*list.Element[*lruTokenStoreEntry]
    53  	q                *list.List[*lruTokenStoreEntry]
    54  	capacity         int
    55  	singleOriginSize int
    56  }
    57  
    58  var _ TokenStore = &lruTokenStore{}
    59  
    60  // NewLRUTokenStore creates a new LRU cache for tokens received by the client.
    61  // maxOrigins specifies how many origins this cache is saving tokens for.
    62  // tokensPerOrigin specifies the maximum number of tokens per origin.
    63  func NewLRUTokenStore(maxOrigins, tokensPerOrigin int) TokenStore {
    64  	return &lruTokenStore{
    65  		m:                make(map[string]*list.Element[*lruTokenStoreEntry]),
    66  		q:                list.New[*lruTokenStoreEntry](),
    67  		capacity:         maxOrigins,
    68  		singleOriginSize: tokensPerOrigin,
    69  	}
    70  }
    71  
    72  func (s *lruTokenStore) Put(key string, token *ClientToken) {
    73  	s.mutex.Lock()
    74  	defer s.mutex.Unlock()
    75  
    76  	if el, ok := s.m[key]; ok {
    77  		entry := el.Value
    78  		entry.cache.Add(token)
    79  		s.q.MoveToFront(el)
    80  		return
    81  	}
    82  
    83  	if s.q.Len() < s.capacity {
    84  		entry := &lruTokenStoreEntry{
    85  			key:   key,
    86  			cache: newSingleOriginTokenStore(s.singleOriginSize),
    87  		}
    88  		entry.cache.Add(token)
    89  		s.m[key] = s.q.PushFront(entry)
    90  		return
    91  	}
    92  
    93  	elem := s.q.Back()
    94  	entry := elem.Value
    95  	delete(s.m, entry.key)
    96  	entry.key = key
    97  	entry.cache = newSingleOriginTokenStore(s.singleOriginSize)
    98  	entry.cache.Add(token)
    99  	s.q.MoveToFront(elem)
   100  	s.m[key] = elem
   101  }
   102  
   103  func (s *lruTokenStore) Pop(key string) *ClientToken {
   104  	s.mutex.Lock()
   105  	defer s.mutex.Unlock()
   106  
   107  	var token *ClientToken
   108  	if el, ok := s.m[key]; ok {
   109  		s.q.MoveToFront(el)
   110  		cache := el.Value.cache
   111  		token = cache.Pop()
   112  		if cache.Len() == 0 {
   113  			s.q.Remove(el)
   114  			delete(s.m, key)
   115  		}
   116  	}
   117  	return token
   118  }