tailscale/net/flowtrack/flowtrack.go
Brad Fitzpatrick 4d15e954bd net/flowtrack: add new package to specialize groupcache/lru key type
Reduces allocs.

Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
2021-01-11 13:08:03 -08:00

100 lines
2.5 KiB
Go

// Copyright (c) 2021 Tailscale Inc & AUTHORS All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//
// Original implementation (from same author) from which this was derived was:
// https://github.com/golang/groupcache/blob/5b532d6fd5efaf7fa130d4e859a2fde0fc3a9e1b/lru/lru.go
// ... which was Apache licensed:
// https://github.com/golang/groupcache/blob/master/LICENSE
// Package flowtrack contains types for tracking TCP/UDP flows by 4-tuples.
package flowtrack
import (
"container/list"
"inet.af/netaddr"
)
// Tuple is a 4-tuple of source and destination IP and port.
type Tuple struct {
Src netaddr.IPPort
Dst netaddr.IPPort
}
// Cache is an LRU cache keyed by Tuple.
//
// The zero value is valid to use.
//
// It is not safe for concurrent access.
type Cache struct {
// MaxEntries is the maximum number of cache entries before
// an item is evicted. Zero means no limit.
MaxEntries int
ll *list.List
m map[Tuple]*list.Element // of *entry
}
// entry is the container/list element type.
type entry struct {
key Tuple
value interface{}
}
// Add adds a value to the cache, set or updating its assoicated
// value.
//
// If MaxEntries is non-zero and the length of the cache is greater
// after any addition, the least recently used value is evicted.
func (c *Cache) Add(key Tuple, value interface{}) {
if c.m == nil {
c.m = make(map[Tuple]*list.Element)
c.ll = list.New()
}
if ee, ok := c.m[key]; ok {
c.ll.MoveToFront(ee)
ee.Value.(*entry).value = value
return
}
ele := c.ll.PushFront(&entry{key, value})
c.m[key] = ele
if c.MaxEntries != 0 && c.Len() > c.MaxEntries {
c.RemoveOldest()
}
}
// Get looks up a key's value from the cache, also reporting
// whether it was present.
func (c *Cache) Get(key Tuple) (value interface{}, ok bool) {
if ele, hit := c.m[key]; hit {
c.ll.MoveToFront(ele)
return ele.Value.(*entry).value, true
}
return nil, false
}
// Remove removes the provided key from the cache if it was present.
func (c *Cache) Remove(key Tuple) {
if ele, hit := c.m[key]; hit {
c.removeElement(ele)
}
}
// RemoveOldest removes the oldest item from the cache, if any.
func (c *Cache) RemoveOldest() {
if c.ll != nil {
if ele := c.ll.Back(); ele != nil {
c.removeElement(ele)
}
}
}
func (c *Cache) removeElement(e *list.Element) {
c.ll.Remove(e)
delete(c.m, e.Value.(*entry).key)
}
// Len returns the number of items in the cache.
func (c *Cache) Len() int { return len(c.m) }