Skip to content

Commit

Permalink
Add xray.AWSSession to install handlers on session
Browse files Browse the repository at this point in the history
An application has to call xray.AWS for each AWS client it constructs.
This creates opportunities for blind spots if someone forgets to
configure a new client.

The xray.AWSSession installs the same handlers at the Session level.
Clients inherit handlers from the session they're created with. As long
as the application systematically reuses the same session to create
clients, it only needs to install X-Ray handlers once.
  • Loading branch information
logan committed Apr 11, 2019
1 parent 5ffa743 commit 0b43af7
Showing 1 changed file with 25 additions and 13 deletions.
38 changes: 25 additions & 13 deletions xray/aws.go
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ import (

"github.com/aws/aws-sdk-go/aws/client"
"github.com/aws/aws-sdk-go/aws/request"
"github.com/aws/aws-sdk-go/aws/session"
"github.com/aws/aws-xray-sdk-go/internal/logger"
"github.com/aws/aws-xray-sdk-go/resources"
)
Expand Down Expand Up @@ -150,33 +151,44 @@ var xRayAfterRetryHandler = request.NamedHandler{
},
}

func pushHandlers(c *client.Client) {
c.Handlers.Validate.PushFrontNamed(xRayBeforeValidateHandler)
c.Handlers.Build.PushBackNamed(xRayAfterBuildHandler)
c.Handlers.Sign.PushFrontNamed(xRayBeforeSignHandler)
c.Handlers.Send.PushBackNamed(xRayAfterSendHandler)
c.Handlers.Unmarshal.PushFrontNamed(xRayBeforeUnmarshalHandler)
c.Handlers.Unmarshal.PushBackNamed(xRayAfterUnmarshalHandler)
c.Handlers.Retry.PushFrontNamed(xRayBeforeRetryHandler)
c.Handlers.AfterRetry.PushBackNamed(xRayAfterRetryHandler)
func pushHandlers(handlers *request.Handlers, completionWhitelistFilename string) {
handlers.Validate.PushFrontNamed(xRayBeforeValidateHandler)
handlers.Build.PushBackNamed(xRayAfterBuildHandler)
handlers.Sign.PushFrontNamed(xRayBeforeSignHandler)
handlers.Send.PushBackNamed(xRayAfterSendHandler)
handlers.Unmarshal.PushFrontNamed(xRayBeforeUnmarshalHandler)
handlers.Unmarshal.PushBackNamed(xRayAfterUnmarshalHandler)
handlers.Retry.PushFrontNamed(xRayBeforeRetryHandler)
handlers.AfterRetry.PushBackNamed(xRayAfterRetryHandler)
handlers.Complete.PushFrontNamed(xrayCompleteHandler(completionWhitelistFilename))
}

// AWS adds X-Ray tracing to an AWS client.
func AWS(c *client.Client) {
if c == nil {
panic("Please initialize the provided AWS client before passing to the AWS() method.")
}
pushHandlers(c)
c.Handlers.Complete.PushFrontNamed(xrayCompleteHandler(""))
pushHandlers(&c.Handlers, "")
}

// AWSWithWhitelist allows a custom parameter whitelist JSON file to be defined.
func AWSWithWhitelist(c *client.Client, filename string) {
if c == nil {
panic("Please initialize the provided AWS client before passing to the AWSWithWhitelist() method.")
}
pushHandlers(c)
c.Handlers.Complete.PushFrontNamed(xrayCompleteHandler(filename))
pushHandlers(&c.Handlers, filename)
}

// AWSSession adds X-Ray tracing to an AWS session. Clients created under this
// session will inherit X-Ray tracing.
func AWSSession(s *session.Session) {
pushHandlers(&s.Handlers, "")
}

// AWSSessionWithWhitelist allows a custom parameter whitelist JSON file to be
// defined.
func AWSSessionWithWhitelist(s *session.Session, filename string) {
pushHandlers(&s.Handlers, filename)
}

func xrayCompleteHandler(filename string) request.NamedHandler {
Expand Down

0 comments on commit 0b43af7

Please sign in to comment.