open-consul/agent/rpc/middleware/interceptors_test.go

313 lines
8.0 KiB
Go

package middleware
import (
"errors"
"net"
"net/netip"
"strings"
"sync"
"testing"
"time"
"github.com/armon/go-metrics"
"github.com/hashicorp/consul/agent/consul/rate"
"github.com/hashicorp/go-hclog"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
)
// obs holds all the things we want to assert on that we recorded correctly in our tests.
type obs struct {
key []string
elapsed float32
labels []metrics.Label
}
// recorderStore acts as an in-mem mock storage for all the RequestRecorder.Record() RecorderFunc calls.
type recorderStore struct {
lock sync.Mutex
store map[string]obs
}
func (rs *recorderStore) put(key []string, o obs) {
rs.lock.Lock()
defer rs.lock.Unlock()
actualKey := strings.Join(append(key, o.labels[0].Value), "")
rs.store[actualKey] = o
}
func (rs *recorderStore) get(key []string) obs {
rs.lock.Lock()
defer rs.lock.Unlock()
actualKey := strings.Join(key, "")
return rs.store[actualKey]
}
var store = recorderStore{store: make(map[string]obs)}
var simpleRecorderFunc = func(key []string, val float32, labels []metrics.Label) {
o := obs{key: key, elapsed: val, labels: labels}
store.put(key, o)
}
type readRequest struct{}
type writeRequest struct{}
type readReqWithTD struct{}
type writeReqWithTD struct{}
func (rr readRequest) IsRead() bool {
return true
}
func (wr writeRequest) IsRead() bool {
return false
}
func (r readReqWithTD) IsRead() bool {
return true
}
func (r readReqWithTD) RequestDatacenter() string {
return "dc3"
}
func (r readReqWithTD) GetMinQueryIndex() uint64 {
return 1
}
func (r readReqWithTD) AllowStaleRead() bool {
return false
}
func (w writeReqWithTD) IsRead() bool {
return false
}
func (w writeReqWithTD) RequestDatacenter() string {
return "dc2"
}
type testCase struct {
name string
// description is meant for human friendliness
description string
// requestName is encouraged to be unique across tests to
// avoid lock contention
requestName string
requestI interface{}
rpcType string
errored bool
isLeader func() bool
dc string
// the first element in expectedLabels should be the method name
expectedLabels []metrics.Label
}
var testCases = []testCase{
{
name: "simple ok",
description: "This is a simple happy path test case. We check for pass through and normal request processing",
requestName: "A.B",
requestI: struct{}{},
rpcType: RPCTypeInternal,
errored: false,
dc: "dc1",
expectedLabels: []metrics.Label{
{Name: "method", Value: "A.B"},
{Name: "errored", Value: "false"},
{Name: "request_type", Value: "unreported"},
{Name: "rpc_type", Value: RPCTypeInternal},
{Name: "leader", Value: "unreported"},
},
},
{
name: "simple ok errored",
description: "Checks that the errored value is populated right.",
requestName: "A.C",
requestI: struct{}{},
rpcType: "test",
errored: true,
dc: "dc1",
expectedLabels: []metrics.Label{
{Name: "method", Value: "A.C"},
{Name: "errored", Value: "true"},
{Name: "request_type", Value: "unreported"},
{Name: "rpc_type", Value: "test"},
{Name: "leader", Value: "unreported"},
},
},
{
name: "read request, rpc type internal",
description: "Checks for read request interface parsing",
requestName: "B.C",
requestI: readRequest{},
rpcType: RPCTypeInternal,
errored: false,
dc: "dc1",
expectedLabels: []metrics.Label{
{Name: "method", Value: "B.C"},
{Name: "errored", Value: "false"},
{Name: "request_type", Value: "read"},
{Name: "rpc_type", Value: RPCTypeInternal},
{Name: "leader", Value: "unreported"},
},
},
{
name: "write request, rpc type net/rpc",
description: "Checks for write request interface, different RPC type",
requestName: "D.E",
requestI: writeRequest{},
rpcType: RPCTypeNetRPC,
errored: false,
dc: "dc1",
expectedLabels: []metrics.Label{
{Name: "method", Value: "D.E"},
{Name: "errored", Value: "false"},
{Name: "request_type", Value: "write"},
{Name: "rpc_type", Value: RPCTypeNetRPC},
{Name: "leader", Value: "unreported"},
},
},
{
name: "read request with blocking stale and target dc",
description: "Checks for locality, blocking status and target dc",
requestName: "E.F",
requestI: readReqWithTD{},
rpcType: RPCTypeNetRPC,
errored: false,
dc: "dc1",
expectedLabels: []metrics.Label{
{Name: "method", Value: "E.F"},
{Name: "errored", Value: "false"},
{Name: "request_type", Value: "read"},
{Name: "rpc_type", Value: RPCTypeNetRPC},
{Name: "leader", Value: "unreported"},
{Name: "allow_stale", Value: "false"},
{Name: "blocking", Value: "true"},
{Name: "target_datacenter", Value: "dc3"},
{Name: "locality", Value: "forwarded"},
},
},
{
name: "write request with TD, locality local",
description: "Checks for write request with local forwarding and target dc",
requestName: "F.G",
requestI: writeReqWithTD{},
rpcType: RPCTypeNetRPC,
errored: false,
dc: "dc2",
expectedLabels: []metrics.Label{
{Name: "method", Value: "F.G"},
{Name: "errored", Value: "false"},
{Name: "request_type", Value: "write"},
{Name: "rpc_type", Value: RPCTypeNetRPC},
{Name: "leader", Value: "unreported"},
{Name: "target_datacenter", Value: "dc2"},
{Name: "locality", Value: "local"},
},
},
{
name: "is leader",
description: "checks for is leader",
requestName: "G.H",
requestI: struct{}{},
rpcType: "test",
errored: false,
isLeader: func() bool {
return true
},
expectedLabels: []metrics.Label{
{Name: "method", Value: "G.H"},
{Name: "errored", Value: "false"},
{Name: "request_type", Value: "unreported"},
{Name: "rpc_type", Value: "test"},
{Name: "leader", Value: "true"},
},
},
{
name: "is not leader",
description: "checks for is not leader",
requestName: "H.I",
requestI: struct{}{},
rpcType: "test",
errored: false,
isLeader: func() bool {
return false
},
expectedLabels: []metrics.Label{
{Name: "method", Value: "H.I"},
{Name: "errored", Value: "false"},
{Name: "request_type", Value: "unreported"},
{Name: "rpc_type", Value: "test"},
{Name: "leader", Value: "false"},
},
},
}
// TestRequestRecorder goes over all the parsing and reporting that RequestRecorder
// is expected to perform.
func TestRequestRecorder(t *testing.T) {
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
r := RequestRecorder{
Logger: hclog.NewInterceptLogger(&hclog.LoggerOptions{}),
RecorderFunc: simpleRecorderFunc,
serverIsLeader: tc.isLeader,
localDC: tc.dc,
}
start := time.Now()
r.Record(tc.requestName, tc.rpcType, start, tc.requestI, tc.errored)
key := append(metricRPCRequest, tc.expectedLabels[0].Value)
o := store.get(key)
require.Equal(t, o.key, metricRPCRequest)
require.LessOrEqual(t, o.elapsed, float32(time.Now().Sub(start).Microseconds())/1000)
require.Equal(t, o.labels, tc.expectedLabels)
})
}
}
func TestGetNetRPCRateLimitingInterceptor(t *testing.T) {
limiter := rate.NewMockRequestLimitsHandler(t)
logger := hclog.NewNullLogger()
rateLimitInterceptor := GetNetRPCRateLimitingInterceptor(limiter, NewPanicHandler(logger))
addr := net.TCPAddrFromAddrPort(netip.MustParseAddrPort("1.2.3.4:5678"))
t.Run("allow operation", func(t *testing.T) {
limiter.On("Allow", mock.Anything).
Return(nil).
Once()
err := rateLimitInterceptor("Status.Leader", addr)
require.NoError(t, err)
})
t.Run("allow returns error", func(t *testing.T) {
limiter.On("Allow", mock.Anything).
Return(errors.New("uh oh")).
Once()
err := rateLimitInterceptor("Status.Leader", addr)
require.Error(t, err)
require.Equal(t, "uh oh", err.Error())
})
t.Run("allow panics", func(t *testing.T) {
limiter.On("Allow", mock.Anything).
Panic("uh oh").
Once()
err := rateLimitInterceptor("Status.Leader", addr)
require.Error(t, err)
require.Equal(t, "rpc: panic serving request", err.Error())
})
}