Implement response cache

This commit is contained in:
Martin Polden
2019-12-25 21:04:26 +01:00
parent 66719b9932
commit 78116f69ad
6 changed files with 111 additions and 8 deletions

54
http/cache.go Normal file
View File

@ -0,0 +1,54 @@
package http
import (
"hash/fnv"
"net"
"sync"
)
type Cache struct {
capacity int
mu sync.RWMutex
entries map[uint64]*Response
keys []uint64
}
func NewCache(capacity int) *Cache {
if capacity < 0 {
capacity = 0
}
return &Cache{
capacity: capacity,
entries: make(map[uint64]*Response),
keys: make([]uint64, 0, capacity),
}
}
func key(ip net.IP) uint64 {
h := fnv.New64a()
h.Write(ip)
return h.Sum64()
}
func (c *Cache) Set(ip net.IP, resp *Response) {
if c.capacity == 0 {
return
}
k := key(ip)
c.mu.Lock()
defer c.mu.Unlock()
if len(c.entries) == c.capacity && c.capacity > 0 {
delete(c.entries, c.keys[0])
c.keys = c.keys[1:]
}
c.entries[k] = resp
c.keys = append(c.keys, k)
}
func (c *Cache) Get(ip net.IP) (*Response, bool) {
k := key(ip)
c.mu.RLock()
defer c.mu.RUnlock()
r, ok := c.entries[k]
return r, ok
}

41
http/cache_test.go Normal file
View File

@ -0,0 +1,41 @@
package http
import (
"fmt"
"net"
"testing"
)
func TestCacheCapacity(t *testing.T) {
var tests = []struct {
addCount, capacity, size int
}{
{1, 0, 0},
{1, 2, 1},
{2, 2, 2},
{3, 2, 2},
}
for i, tt := range tests {
c := NewCache(tt.capacity)
var responses []*Response
for i := 0; i < tt.addCount; i++ {
ip := net.ParseIP(fmt.Sprintf("192.0.2.%d", i))
r := &Response{IP: ip}
responses = append(responses, r)
c.Set(ip, r)
}
if got := len(c.entries); got != tt.size {
t.Errorf("#%d: len(entries) = %d, want %d", i, got, tt.size)
}
if tt.capacity > 0 && tt.addCount > tt.capacity && tt.capacity == tt.size {
lastAdded := responses[tt.addCount-1]
if _, ok := c.Get(lastAdded.IP); !ok {
t.Errorf("#%d: Get(%s) = (_, %t), want (_, %t)", i, lastAdded.IP.String(), ok, !ok)
}
firstAdded := responses[0]
if _, ok := c.Get(firstAdded.IP); ok {
t.Errorf("#%d: Get(%s) = (_, %t), want (_, %t)", i, firstAdded.IP.String(), ok, !ok)
}
}
}
}

View File

@ -27,6 +27,7 @@ type Server struct {
IPHeaders []string
LookupAddr func(net.IP) (string, error)
LookupPort func(net.IP, uint64) error
cache *Cache
gr geo.Reader
}
@ -51,8 +52,8 @@ type PortResponse struct {
Reachable bool `json:"reachable"`
}
func New(db geo.Reader) *Server {
return &Server{gr: db}
func New(db geo.Reader, cache *Cache) *Server {
return &Server{cache: cache, gr: db}
}
func ipFromForwardedForHeader(v string) string {
@ -93,6 +94,10 @@ func (s *Server) newResponse(r *http.Request) (Response, error) {
if err != nil {
return Response{}, err
}
response, ok := s.cache.Get(ip)
if ok {
return *response, nil
}
ipDecimal := iputil.ToDecimal(ip)
country, _ := s.gr.Country(ip)
city, _ := s.gr.City(ip)
@ -111,7 +116,7 @@ func (s *Server) newResponse(r *http.Request) (Response, error) {
parsed := useragent.Parse(userAgentRaw)
userAgent = &parsed
}
return Response{
response = &Response{
IP: ip,
IPDecimal: ipDecimal,
Country: country.Name,
@ -124,7 +129,9 @@ func (s *Server) newResponse(r *http.Request) (Response, error) {
ASN: autonomousSystemNumber,
ASNOrg: asn.AutonomousSystemOrganization,
UserAgent: userAgent,
}, nil
}
s.cache.Set(ip, response)
return *response, nil
}
func (s *Server) newPortResponse(r *http.Request) (PortResponse, error) {

View File

@ -31,7 +31,7 @@ func (t *testDb) ASN(net.IP) (geo.ASN, error) {
func (t *testDb) IsEmpty() bool { return false }
func testServer() *Server {
return &Server{gr: &testDb{}, LookupAddr: lookupAddr, LookupPort: lookupPort}
return &Server{cache: NewCache(100), gr: &testDb{}, LookupAddr: lookupAddr, LookupPort: lookupPort}
}
func httpGet(url string, acceptMediaType string, userAgent string) (string, int, error) {