diff options
| author | Ash McKenzie <amckenzie@gitlab.com> | 2020-09-21 04:47:01 +0000 |
|---|---|---|
| committer | Ash McKenzie <amckenzie@gitlab.com> | 2020-09-21 04:47:01 +0000 |
| commit | 1a2bfecd2f0ebb8e31f9833e0522c4643797041b (patch) | |
| tree | d17cf7bff45492a587027851bb6e0bcb493cff58 /internal/command/command.go | |
| parent | f100e7e83943b3bb5db232f5bf79a616fdba88f1 (diff) | |
| parent | a487572a904cc149840488eefdfe121173d8bcb5 (diff) | |
| download | gitlab-shell-1a2bfecd2f0ebb8e31f9833e0522c4643797041b.tar.gz | |
Merge branch 'sh-extract-context-from-env' into 'master'
Make it possible to propagate correlation ID across processes
Closes #474
See merge request gitlab-org/gitlab-shell!413
Diffstat (limited to 'internal/command/command.go')
| -rw-r--r-- | internal/command/command.go | 29 |
1 files changed, 28 insertions, 1 deletions
diff --git a/internal/command/command.go b/internal/command/command.go index 283b4a1..c69219b 100644 --- a/internal/command/command.go +++ b/internal/command/command.go @@ -1,6 +1,8 @@ package command import ( + "context" + "gitlab.com/gitlab-org/gitlab-shell/internal/command/authorizedkeys" "gitlab.com/gitlab-org/gitlab-shell/internal/command/authorizedprincipals" "gitlab.com/gitlab-org/gitlab-shell/internal/command/commandargs" @@ -16,10 +18,13 @@ import ( "gitlab.com/gitlab-org/gitlab-shell/internal/command/uploadpack" "gitlab.com/gitlab-org/gitlab-shell/internal/config" "gitlab.com/gitlab-org/gitlab-shell/internal/executable" + "gitlab.com/gitlab-org/labkit/correlation" + "gitlab.com/gitlab-org/labkit/log" + "gitlab.com/gitlab-org/labkit/tracing" ) type Command interface { - Execute() error + Execute(ctx context.Context) error } func New(e *executable.Executable, arguments []string, config *config.Config, readWriter *readwriter.ReadWriter) (Command, error) { @@ -35,6 +40,28 @@ func New(e *executable.Executable, arguments []string, config *config.Config, re return nil, disallowedcommand.Error } +// ContextWithCorrelationID() will always return a background Context +// with a correlation ID. It will first attempt to extract the ID from +// an environment variable. If is not available, a random one will be +// generated. +func ContextWithCorrelationID() (context.Context, func()) { + ctx, finished := tracing.ExtractFromEnv(context.Background()) + defer finished() + + correlationID := correlation.ExtractFromContext(ctx) + if correlationID == "" { + correlationID, err := correlation.RandomID() + if err != nil { + log.WithError(err).Warn("unable to generate correlation ID") + } else { + log.Info("generated random correlation ID") + ctx = correlation.ContextWithCorrelation(ctx, correlationID) + } + } + + return ctx, finished +} + func buildCommand(e *executable.Executable, args commandargs.CommandArgs, config *config.Config, readWriter *readwriter.ReadWriter) Command { switch e.Name { case executable.GitlabShell: |