bramble-xray-plugin.go

package plugins

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"regexp"
	"strings"

	"github.com/aws/aws-xray-sdk-go/strategy/sampling"
	"github.com/aws/aws-xray-sdk-go/xray"
	"github.com/movio/bramble"
)

var (
	whitespaceRegexp = regexp.MustCompile(`\s+`)
)

func init() {
	bramble.RegisterPlugin(&AWSXrayPlugin{})
}

type AWSXrayPlugin struct {
	bramble.BasePlugin

	logGroup string
}

func (p *AWSXrayPlugin) ID() string {
	return "aws-xray"
}

func (p *AWSXrayPlugin) Configure(cfg *bramble.Config, data json.RawMessage) error {
	return nil
}

func (p *AWSXrayPlugin) WrapGraphQLClientTransport(t http.RoundTripper) http.RoundTripper {
	if t == nil {
		t = http.DefaultTransport
	}
	return xray.RoundTripper(t)
}

// InterceptRequest annotates the segment as per https://docs.aws.amazon.com/xray/latest/devguide/xray-api-segmentdocuments.html
func (p *AWSXrayPlugin) InterceptRequest(ctx context.Context, operationName, query string, variables map[string]interface{}) {
	seg := xray.GetSegment(ctx)
	if seg == nil {
		return
	}
	headers := bramble.GetOutgoingRequestHeadersFromContext(ctx)
	seg.AddAnnotation("operationName", operationName)
	if headers.Get("X-Movio-Tenant") != "" {
		seg.AddAnnotation("tenant", headers.Get("X-Movio-Tenant"))
	}
	if headers.Get("JWT-Claim-Email") != "" {
		seg.User = headers.Get("JWT-Claim-Email")
	}
	seg.AddMetadata("operationName", operationName)
	seg.AddMetadata("query", whitespaceRegexp.ReplaceAllString(query, " "))
	seg.AddMetadata("variables", variables)
	seg.AddMetadata("headers", headers)
	if p.logGroup != "" {
		aws := seg.GetAWS()
		aws["cloudwatch_logs"] = []map[string]interface{}{{"log_group": p.logGroup}}
	}
}

func (p *AWSXrayPlugin) Init(s *bramble.ExecutableSchema) {
	collector := os.Getenv("XRAY_DAEMON")
	p.logGroup = os.Getenv("XRAY_APP_CLOUDWATCH_LOG_GROUP")

	strat, _ := sampling.NewCentralizedStrategy()

	xray.Configure(xray.Config{
		DaemonAddr:       fmt.Sprintf("%s:2000", collector),
		ServiceVersion:   "0.0.1",
		SamplingStrategy: &samplingStrategy{strat},
	})
}

type samplingStrategy struct {
	proxy *sampling.CentralizedStrategy
}

func (s *samplingStrategy) ShouldTrace(r *sampling.Request) *sampling.Decision {
	if shouldFilter(r) {
		return &sampling.Decision{
			Sample: false,
		}
	}
	return s.proxy.ShouldTrace(r)
}

// We want to filter out healthchecks and internal requests
func shouldFilter(r *sampling.Request) bool {
	return r.Host == "" || strings.HasPrefix(r.Host, "localhost") || r.URL == "/health"
}

func (p *AWSXrayPlugin) ApplyMiddlewarePublicMux(h http.Handler) http.Handler {
	return xray.Handler(xray.NewFixedSegmentNamer("bramble"), h)
}

func (p *AWSXrayPlugin) ApplyMiddlewarePrivateMux(h http.Handler) http.Handler {
	return xray.Handler(xray.NewFixedSegmentNamer("bramble"), h)
}