mirror of https://github.com/tailscale/tailscale/
net/flowtrack: add new package to specialize groupcache/lru key type
Reduces allocs. Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>bradfitz/umaskless_permissions
Brad Fitzpatrick
3 years ago
parent
f69e46175d
commit
4d15e954bd
@ -0,0 +1,99 @@
|
||||
// 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) }
|
||||
@ -0,0 +1,82 @@
|
||||
// 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.
|
||||
|
||||
package flowtrack
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"inet.af/netaddr"
|
||||
)
|
||||
|
||||
func TestCache(t *testing.T) {
|
||||
c := &Cache{MaxEntries: 2}
|
||||
|
||||
k1 := Tuple{Src: netaddr.MustParseIPPort("1.1.1.1:1"), Dst: netaddr.MustParseIPPort("1.1.1.1:1")}
|
||||
k2 := Tuple{Src: netaddr.MustParseIPPort("1.1.1.1:1"), Dst: netaddr.MustParseIPPort("2.2.2.2:2")}
|
||||
k3 := Tuple{Src: netaddr.MustParseIPPort("1.1.1.1:1"), Dst: netaddr.MustParseIPPort("3.3.3.3:3")}
|
||||
k4 := Tuple{Src: netaddr.MustParseIPPort("1.1.1.1:1"), Dst: netaddr.MustParseIPPort("4.4.4.4:4")}
|
||||
|
||||
wantLen := func(want int) {
|
||||
t.Helper()
|
||||
if got := c.Len(); got != want {
|
||||
t.Fatalf("Len = %d; want %d", got, want)
|
||||
}
|
||||
}
|
||||
wantVal := func(key Tuple, want interface{}) {
|
||||
t.Helper()
|
||||
got, ok := c.Get(key)
|
||||
if !ok {
|
||||
t.Fatalf("Get(%q) failed; want value %v", key, want)
|
||||
}
|
||||
if got != want {
|
||||
t.Fatalf("Get(%q) = %v; want %v", key, got, want)
|
||||
}
|
||||
}
|
||||
wantMissing := func(key Tuple) {
|
||||
t.Helper()
|
||||
if got, ok := c.Get(key); ok {
|
||||
t.Fatalf("Get(%q) = %v; want absent from cache", key, got)
|
||||
}
|
||||
}
|
||||
|
||||
wantLen(0)
|
||||
c.RemoveOldest() // shouldn't panic
|
||||
c.Remove(k4) // shouldn't panic
|
||||
|
||||
c.Add(k1, 1)
|
||||
wantLen(1)
|
||||
c.Add(k2, 2)
|
||||
wantLen(2)
|
||||
c.Add(k3, 3)
|
||||
wantLen(2) // hit the max
|
||||
|
||||
wantMissing(k1)
|
||||
c.Remove(k1)
|
||||
wantLen(2) // no change; k1 should've been the deleted one per LRU
|
||||
|
||||
wantVal(k3, 3)
|
||||
|
||||
wantVal(k2, 2)
|
||||
c.Remove(k2)
|
||||
wantLen(1)
|
||||
wantMissing(k2)
|
||||
|
||||
c.Add(k3, 30)
|
||||
wantVal(k3, 30)
|
||||
wantLen(1)
|
||||
|
||||
allocs := int(testing.AllocsPerRun(1000, func() {
|
||||
got, ok := c.Get(k3)
|
||||
if !ok {
|
||||
t.Fatal("missing k3")
|
||||
}
|
||||
if got != 30 {
|
||||
t.Fatalf("got = %d; want 30", got)
|
||||
}
|
||||
}))
|
||||
if allocs != 0 {
|
||||
t.Errorf("allocs = %v; want 0", allocs)
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue