Reduces allocs. Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>main
Brad Fitzpatrick
5 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