Add tracing to bundle/discovery download

Signed-off-by: Magnus Jungsbluth <magnus@jungsbluth.de>
This commit is contained in:
Magnus Jungsbluth
2023-06-05 21:00:44 +02:00
committed by Ashutosh Narkar
parent 67529bf428
commit 22619d26c0
5 changed files with 98 additions and 21 deletions
+8 -6
View File
@@ -18,15 +18,17 @@ import (
"github.com/open-policy-agent/opa/keys"
"github.com/open-policy-agent/opa/logging"
"github.com/open-policy-agent/opa/plugins/rest"
"github.com/open-policy-agent/opa/tracing"
"github.com/open-policy-agent/opa/util"
)
// ServiceOptions stores the options passed to ParseServicesConfig
type ServiceOptions struct {
Raw json.RawMessage
AuthPlugin rest.AuthPluginLookupFunc
Keys map[string]*keys.Config
Logger logging.Logger
Raw json.RawMessage
AuthPlugin rest.AuthPluginLookupFunc
Keys map[string]*keys.Config
Logger logging.Logger
DistributedTacingOpts tracing.Options
}
// ParseServicesConfig returns a set of named service clients. The service
@@ -42,7 +44,7 @@ func ParseServicesConfig(opts ServiceOptions) (map[string]rest.Client, error) {
if err := util.Unmarshal(opts.Raw, &arr); err == nil {
for _, s := range arr {
client, err := rest.New(s, opts.Keys, rest.AuthPluginLookup(opts.AuthPlugin), rest.Logger(opts.Logger))
client, err := rest.New(s, opts.Keys, rest.AuthPluginLookup(opts.AuthPlugin), rest.Logger(opts.Logger), rest.DistributedTracingOpts(opts.DistributedTacingOpts))
if err != nil {
return nil, err
}
@@ -50,7 +52,7 @@ func ParseServicesConfig(opts ServiceOptions) (map[string]rest.Client, error) {
}
} else if util.Unmarshal(opts.Raw, &obj) == nil {
for k := range obj {
client, err := rest.New(obj[k], opts.Keys, rest.Name(k), rest.AuthPluginLookup(opts.AuthPlugin), rest.Logger(opts.Logger))
client, err := rest.New(obj[k], opts.Keys, rest.Name(k), rest.AuthPluginLookup(opts.AuthPlugin), rest.Logger(opts.Logger), rest.DistributedTracingOpts(opts.DistributedTacingOpts))
if err != nil {
return nil, err
}
+18 -7
View File
@@ -29,6 +29,7 @@ import (
"github.com/open-policy-agent/opa/storage"
"github.com/open-policy-agent/opa/topdown/cache"
"github.com/open-policy-agent/opa/topdown/print"
"github.com/open-policy-agent/opa/tracing"
)
// Factory defines the interface OPA uses to instantiate your plugin.
@@ -192,6 +193,7 @@ type Manager struct {
router *mux.Router
prometheusRegister prometheus.Registerer
tracerProvider *trace.TracerProvider
distributedTacingOpts tracing.Options
registeredNDCacheTriggers []func(bool)
bootstrapConfigLabels map[string]string
}
@@ -365,6 +367,13 @@ func WithTracerProvider(tracerProvider *trace.TracerProvider) func(*Manager) {
}
}
// WithDistributedTracingOpts sets the options to be used by distributed tracing.
func WithDistributedTracingOpts(tr tracing.Options) func(*Manager) {
return func(m *Manager) {
m.distributedTacingOpts = tr
}
}
// New creates a new Manager using config.
func New(raw []byte, id string, store storage.Store, opts ...func(*Manager)) (*Manager, error) {
@@ -409,10 +418,11 @@ func New(raw []byte, id string, store storage.Store, opts ...func(*Manager)) (*M
}
serviceOpts := cfg.ServiceOptions{
Raw: parsedConfig.Services,
AuthPlugin: m.AuthPlugin,
Keys: keys,
Logger: m.logger,
Raw: parsedConfig.Services,
AuthPlugin: m.AuthPlugin,
Keys: keys,
Logger: m.logger,
DistributedTacingOpts: m.distributedTacingOpts,
}
services, err := cfg.ParseServicesConfig(serviceOpts)
@@ -647,9 +657,10 @@ func (m *Manager) Stop(ctx context.Context) {
// Reconfigure updates the configuration on the manager.
func (m *Manager) Reconfigure(config *config.Config) error {
opts := cfg.ServiceOptions{
Raw: config.Services,
AuthPlugin: m.AuthPlugin,
Logger: m.logger,
Raw: config.Services,
AuthPlugin: m.AuthPlugin,
Logger: m.logger,
DistributedTacingOpts: m.distributedTacingOpts,
}
keys, err := keys.ParseKeysConfig(config.Keys)
+20 -7
View File
@@ -20,6 +20,7 @@ import (
"github.com/open-policy-agent/opa/internal/version"
"github.com/open-policy-agent/opa/keys"
"github.com/open-policy-agent/opa/logging"
"github.com/open-policy-agent/opa/tracing"
"github.com/open-policy-agent/opa/util"
)
@@ -125,13 +126,14 @@ func (c *Config) authPrepare(req *http.Request, lookup AuthPluginLookupFunc) err
// Client implements an HTTP/REST client for communicating with remote
// services.
type Client struct {
bytes *[]byte
json *interface{}
config Config
headers map[string]string
authPluginLookup AuthPluginLookupFunc
logger logging.Logger
loggerFields map[string]interface{}
bytes *[]byte
json *interface{}
config Config
headers map[string]string
authPluginLookup AuthPluginLookupFunc
logger logging.Logger
loggerFields map[string]interface{}
distributedTacingOpts tracing.Options
}
// Name returns an option that overrides the service name on the client.
@@ -158,6 +160,13 @@ func Logger(l logging.Logger) func(*Client) {
}
}
// DistributedTracingOpts sets the options to be used by distributed tracing.
func DistributedTracingOpts(tr tracing.Options) func(*Client) {
return func(c *Client) {
c.distributedTacingOpts = tr
}
}
// New returns a new Client for config.
func New(config []byte, keys map[string]*keys.Config, opts ...func(*Client)) (Client, error) {
var parsedConfig Config
@@ -260,6 +269,10 @@ func (c Client) Do(ctx context.Context, method, path string) (*http.Response, er
return nil, err
}
if len(c.distributedTacingOpts) > 0 {
httpClient.Transport = tracing.NewTransport(httpClient.Transport, c.distributedTacingOpts)
}
path = strings.Trim(path, "/")
var body io.Reader
+51
View File
@@ -36,6 +36,7 @@ import (
"github.com/open-policy-agent/opa/internal/jwx/jws"
"github.com/open-policy-agent/opa/keys"
"github.com/open-policy-agent/opa/logging"
"github.com/open-policy-agent/opa/tracing"
"github.com/open-policy-agent/opa/internal/version"
"github.com/open-policy-agent/opa/util/test"
@@ -802,6 +803,56 @@ func TestDoWithResponseHeaderTimeout(t *testing.T) {
}
}
type tracemock struct {
called int
}
func (m *tracemock) NewTransport(rt http.RoundTripper, _ tracing.Options) http.RoundTripper {
m.called++
return rt
}
func (*tracemock) NewHandler(http.Handler, string, tracing.Options) http.Handler {
panic("unreachable")
}
func TestDoWithDistributedTracingOpts(t *testing.T) {
ctx := context.Background()
mock := tracemock{}
tracing.RegisterHTTPTracing(&mock)
body := "Some Bad Request was received"
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusBadRequest)
fmt.Fprintln(w, body)
}))
defer ts.Close()
buf := bytes.Buffer{}
logger := logging.New()
logger.SetOutput(&buf)
logger.SetLevel(logging.Debug)
config := fmt.Sprintf(`{
"name": "foo",
"url": %q,
}`, ts.URL)
ks := map[string]*keys.Config{}
client, err := New([]byte(config), ks, DistributedTracingOpts(tracing.Options{"testoption"}))
if err != nil {
t.Fatalf("Unexpected error: %v", err)
}
_, err = client.Do(ctx, "GET", ts.URL)
if err != nil {
t.Fatalf("Unexpected error: %v", err)
}
if exp, act := 1, mock.called; exp != act {
t.Errorf("calls to NewTransport: expected %d, got %d", exp, act)
}
}
func TestDoWithResponseInClientLog(t *testing.T) {
ctx := context.Background()
+1 -1
View File
@@ -3217,7 +3217,7 @@ func TestDistributedTracingEnabled(t *testing.T) {
}
if exp, act := 1, mock.called; exp != act {
t.Errorf("calls to NewTransported: expected %d, got %d", exp, act)
t.Errorf("calls to NewTransport: expected %d, got %d", exp, act)
}
}