From 210a5c141c9d76bc9718860d67d77d73997b1534 Mon Sep 17 00:00:00 2001 From: Andrew Newdigate Date: Mon, 25 Feb 2019 16:19:08 +0200 Subject: Adds distributed tracing instrumentation to GitLab-Shell Adds distributed tracing instrumentation to GitLab-Shell using LabKit --- go/cmd/gitaly-receive-pack/main.go | 32 +++---- go/cmd/gitaly-upload-archive/main.go | 38 +++----- go/cmd/gitaly-upload-archive/main_test.go | 70 -------------- go/cmd/gitaly-upload-pack/main.go | 32 +++---- go/internal/config/config.go | 11 ++- go/internal/handler/exec.go | 104 ++++++++++++++++++++ go/internal/handler/exec_test.go | 153 ++++++++++++++++++++++++++++++ go/internal/handler/handler.go | 51 ---------- go/internal/handler/receive_pack.go | 18 +--- go/internal/handler/upload_archive.go | 18 +--- go/internal/handler/upload_pack.go | 18 +--- 11 files changed, 322 insertions(+), 223 deletions(-) delete mode 100644 go/cmd/gitaly-upload-archive/main_test.go create mode 100644 go/internal/handler/exec.go create mode 100644 go/internal/handler/exec_test.go delete mode 100644 go/internal/handler/handler.go (limited to 'go') diff --git a/go/cmd/gitaly-receive-pack/main.go b/go/cmd/gitaly-receive-pack/main.go index a79eb63..85decd9 100644 --- a/go/cmd/gitaly-receive-pack/main.go +++ b/go/cmd/gitaly-receive-pack/main.go @@ -1,12 +1,12 @@ package main import ( + "context" "encoding/json" - "fmt" - "os" "gitlab.com/gitlab-org/gitlab-shell/go/internal/handler" "gitlab.com/gitlab-org/gitlab-shell/go/internal/logger" + "google.golang.org/grpc" pb "gitlab.com/gitlab-org/gitaly-proto/go/gitalypb" ) @@ -16,22 +16,20 @@ func init() { } func main() { - if err := handler.Prepare(); err != nil { - logger.Fatal("preparation failed", err) - } - - if n := len(os.Args); n != 3 { - logger.Fatal("wrong number of arguments", fmt.Errorf("expected 2 arguments, got %v", os.Args)) - } + handler.RunGitalyCommand(func(ctx context.Context, conn *grpc.ClientConn, requestJSON string) (int32, error) { + request, err := deserialize(requestJSON) + if err != nil { + return 1, err + } + + return handler.ReceivePack(ctx, conn, request) + }) +} +func deserialize(requestJSON string) (*pb.SSHReceivePackRequest, error) { var request pb.SSHReceivePackRequest - if err := json.Unmarshal([]byte(os.Args[2]), &request); err != nil { - logger.Fatal("unmarshaling request json failed", err) - } - - code, err := handler.ReceivePack(os.Args[1], &request) - if err != nil { - logger.Fatal("receive-pack failed", err) + if err := json.Unmarshal([]byte(requestJSON), request); err != nil { + return nil, err } - os.Exit(int(code)) + return &request, nil } diff --git a/go/cmd/gitaly-upload-archive/main.go b/go/cmd/gitaly-upload-archive/main.go index 10b10d3..fb07613 100644 --- a/go/cmd/gitaly-upload-archive/main.go +++ b/go/cmd/gitaly-upload-archive/main.go @@ -1,12 +1,12 @@ package main import ( + "context" "encoding/json" - "fmt" - "os" "gitlab.com/gitlab-org/gitlab-shell/go/internal/handler" "gitlab.com/gitlab-org/gitlab-shell/go/internal/logger" + "google.golang.org/grpc" pb "gitlab.com/gitlab-org/gitaly-proto/go/gitalypb" ) @@ -15,31 +15,21 @@ func init() { logger.ProgName = "gitaly-upload-archive" } -type uploadArchiveHandler func(gitalyAddress string, request *pb.SSHUploadArchiveRequest) (int32, error) - func main() { - if err := handler.Prepare(); err != nil { - logger.Fatal("preparation failed", err) - } - - code, err := uploadArchive(handler.UploadArchive, os.Args) - - if err != nil { - logger.Fatal("upload-archive failed", err) - } - - os.Exit(int(code)) + handler.RunGitalyCommand(func(ctx context.Context, conn *grpc.ClientConn, requestJSON string) (int32, error) { + request, err := deserialize(requestJSON) + if err != nil { + return 1, err + } + + return handler.UploadArchive(ctx, conn, request) + }) } -func uploadArchive(handler uploadArchiveHandler, args []string) (int32, error) { - if n := len(args); n != 3 { - return 0, fmt.Errorf("wrong number of arguments: expected 2 arguments, got %v", args) - } - +func deserialize(argumentJSON string) (*pb.SSHUploadArchiveRequest, error) { var request pb.SSHUploadArchiveRequest - if err := json.Unmarshal([]byte(args[2]), &request); err != nil { - return 0, fmt.Errorf("unmarshaling request json failed: %v", err) + if err := json.Unmarshal([]byte(argumentJSON), request); err != nil { + return nil, err } - - return handler(args[1], &request) + return &request, nil } diff --git a/go/cmd/gitaly-upload-archive/main_test.go b/go/cmd/gitaly-upload-archive/main_test.go deleted file mode 100644 index 1f30e56..0000000 --- a/go/cmd/gitaly-upload-archive/main_test.go +++ /dev/null @@ -1,70 +0,0 @@ -package main - -import ( - "fmt" - "strings" - "testing" - - pb "gitlab.com/gitlab-org/gitaly-proto/go/gitalypb" -) - -var testGitalyAddress = "unix:gitaly.socket" - -func TestUploadArchiveSuccess(t *testing.T) { - testRelativePath := "myrepo.git" - requestJSON := fmt.Sprintf(`{"repository":{"relative_path":"%s"}}`, testRelativePath) - mockHandler := func(gitalyAddress string, request *pb.SSHUploadArchiveRequest) (int32, error) { - if gitalyAddress != testGitalyAddress { - t.Fatalf("Expected gitaly address %s got %v", testGitalyAddress, gitalyAddress) - } - if relativePath := request.Repository.RelativePath; relativePath != testRelativePath { - t.Fatalf("Expected repository with relative path %s got %v", testRelativePath, request) - } - return 0, nil - } - - code, err := uploadArchive(mockHandler, []string{"git-upload-archive", testGitalyAddress, requestJSON}) - - if err != nil { - t.Fatal(err) - } - - if code != 0 { - t.Fatalf("Expected exit code 0, got %v", code) - } -} - -func TestUploadArchiveFailure(t *testing.T) { - mockHandler := func(_ string, _ *pb.SSHUploadArchiveRequest) (int32, error) { - t.Fatal("Expected handler not to be called") - - return 0, nil - } - - tests := []struct { - desc string - args []string - err string - }{ - { - desc: "With an invalid request json", - args: []string{"git-upload-archive", testGitalyAddress, "hello"}, - err: "unmarshaling request json failed", - }, - { - desc: "With an invalid argument count", - args: []string{"git-upload-archive", testGitalyAddress, "{}", "extra arg"}, - err: "wrong number of arguments: expected 2 arguments", - }, - } - - for _, test := range tests { - t.Run(test.desc, func(t *testing.T) { - _, err := uploadArchive(mockHandler, test.args) - - if !strings.Contains(err.Error(), test.err) { - t.Fatalf("Expected error %v, got %v", test.err, err) - } - }) - } -} diff --git a/go/cmd/gitaly-upload-pack/main.go b/go/cmd/gitaly-upload-pack/main.go index ba4618f..3fb721d 100644 --- a/go/cmd/gitaly-upload-pack/main.go +++ b/go/cmd/gitaly-upload-pack/main.go @@ -1,12 +1,12 @@ package main import ( + "context" "encoding/json" - "fmt" - "os" "gitlab.com/gitlab-org/gitlab-shell/go/internal/handler" "gitlab.com/gitlab-org/gitlab-shell/go/internal/logger" + "google.golang.org/grpc" pb "gitlab.com/gitlab-org/gitaly-proto/go/gitalypb" ) @@ -16,22 +16,20 @@ func init() { } func main() { - if err := handler.Prepare(); err != nil { - logger.Fatal("preparation failed", err) - } - - if n := len(os.Args); n != 3 { - logger.Fatal("wrong number of arguments", fmt.Errorf("expected 2 arguments, got %v", os.Args)) - } + handler.RunGitalyCommand(func(ctx context.Context, conn *grpc.ClientConn, requestJSON string) (int32, error) { + request, err := deserialize(requestJSON) + if err != nil { + return 1, err + } + + return handler.UploadPack(ctx, conn, request) + }) +} +func deserialize(requestJSON string) (*pb.SSHUploadPackRequest, error) { var request pb.SSHUploadPackRequest - if err := json.Unmarshal([]byte(os.Args[2]), &request); err != nil { - logger.Fatal("unmarshaling request json failed", err) - } - - code, err := handler.UploadPack(os.Args[1], &request) - if err != nil { - logger.Fatal("upload-pack failed", err) + if err := json.Unmarshal([]byte(requestJSON), request); err != nil { + return nil, err } - os.Exit(int(code)) + return &request, nil } diff --git a/go/internal/config/config.go b/go/internal/config/config.go index e745a25..f951fe6 100644 --- a/go/internal/config/config.go +++ b/go/internal/config/config.go @@ -21,11 +21,12 @@ type MigrationConfig struct { } type Config struct { - RootDir string - LogFile string `yaml:"log_file"` - LogFormat string `yaml:"log_format"` - Migration MigrationConfig `yaml:"migration"` - GitlabUrl string `yaml:"gitlab_url"` + RootDir string + LogFile string `yaml:"log_file"` + LogFormat string `yaml:"log_format"` + Migration MigrationConfig `yaml:"migration"` + GitlabUrl string `yaml:"gitlab_url"` + GitlabTracing string `yaml:"gitlab_tracing"` } func New() (*Config, error) { diff --git a/go/internal/handler/exec.go b/go/internal/handler/exec.go new file mode 100644 index 0000000..ee7b4a8 --- /dev/null +++ b/go/internal/handler/exec.go @@ -0,0 +1,104 @@ +package handler + +import ( + "context" + "fmt" + "os" + + "gitlab.com/gitlab-org/gitaly/auth" + "gitlab.com/gitlab-org/gitaly/client" + + "gitlab.com/gitlab-org/gitlab-shell/go/internal/config" + "gitlab.com/gitlab-org/gitlab-shell/go/internal/logger" + "gitlab.com/gitlab-org/labkit/tracing" + "google.golang.org/grpc" +) + +// GitalyHandlerFunc implementations are responsible for deserializing +// the request JSON into a GRPC request message, making an appropriate Gitaly +// call with the request, using the provided client, and returning the exit code +// or error from the Gitaly call. +type GitalyHandlerFunc func(ctx context.Context, client *grpc.ClientConn, requestJSON string) (int32, error) + +// RunGitalyCommand provides a bootstrap for Gitaly commands executed +// through GitLab-Shell. It ensures that logging, tracing and other +// common concerns are configured before executing the `handler`. +// RunGitalyCommand will handle errors internally and call +// `os.Exit()` on completion. This method will never return to +// the caller. +func RunGitalyCommand(handler GitalyHandlerFunc) { + exitCode, err := internalRunGitalyCommand(os.Args, handler) + + if err != nil { + logger.Fatal("error: %v", err) + } + + os.Exit(exitCode) +} + +// internalRunGitalyCommand is like RunGitalyCommand, except that since it doesn't +// call os.Exit, we can rely on its deferred handlers executing correctly +func internalRunGitalyCommand(args []string, handler GitalyHandlerFunc) (int, error) { + + if len(args) != 3 { + return 1, fmt.Errorf("expected 2 arguments, got %v", args) + } + + cfg, err := config.New() + if err != nil { + return 1, err + } + + if err := logger.Configure(cfg); err != nil { + return 1, err + } + + // Use a working directory that won't get removed or unmounted. + if err := os.Chdir("/"); err != nil { + return 1, err + } + + // Configure distributed tracing + serviceName := fmt.Sprintf("gitlab-shell-%v", args[0]) + closer := tracing.Initialize( + tracing.WithServiceName(serviceName), + + // For GitLab-Shell, we explicitly initialize tracing from a config file + // instead of the default environment variable (using GITLAB_TRACING) + // This decision was made owing to the difficulty in passing environment + // variables into GitLab-Shell processes. + // + // Processes are spawned as children of the SSH daemon, which tightly + // controls environment variables; doing this means we don't have to + // enable PermitUserEnvironment + tracing.WithConnectionString(cfg.GitlabTracing), + ) + defer closer.Close() + + ctx, finished := tracing.ExtractFromEnv(context.Background()) + defer finished() + + gitalyAddress := args[1] + if gitalyAddress == "" { + return 1, fmt.Errorf("no gitaly_address given") + } + + conn, err := client.Dial(gitalyAddress, dialOpts()) + if err != nil { + return 1, err + } + defer conn.Close() + + requestJSON := string(args[2]) + exitCode, err := handler(ctx, conn, requestJSON) + return int(exitCode), err +} + +func dialOpts() []grpc.DialOption { + connOpts := client.DefaultDialOpts + if token := os.Getenv("GITALY_TOKEN"); token != "" { + connOpts = append(client.DefaultDialOpts, grpc.WithPerRPCCredentials(gitalyauth.RPCCredentialsV2(token))) + } + + return connOpts +} diff --git a/go/internal/handler/exec_test.go b/go/internal/handler/exec_test.go new file mode 100644 index 0000000..e19ebc4 --- /dev/null +++ b/go/internal/handler/exec_test.go @@ -0,0 +1,153 @@ +package handler + +import ( + "context" + "fmt" + "io/ioutil" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" + "google.golang.org/grpc" +) + +func TestInteralRunHandler(t *testing.T) { + type testCase struct { + name string + args []string + handler func(context.Context, *grpc.ClientConn, string) (int32, error) + want int + wantErr bool + } + + var currentTest *testCase + makeHandler := func(r1 int32, r2 error) func(context.Context, *grpc.ClientConn, string) (int32, error) { + return func(ctx context.Context, client *grpc.ClientConn, requestJSON string) (int32, error) { + require.NotNil(t, ctx) + require.NotNil(t, client) + require.Equal(t, currentTest.args[2], requestJSON) + return r1, r2 + } + } + tests := []testCase{ + { + name: "expected", + args: []string{"test", "tcp://localhost:9999", "{}"}, + handler: makeHandler(0, nil), + want: 0, + wantErr: false, + }, + { + name: "handler_error", + args: []string{"test", "tcp://localhost:9999", "{}"}, + handler: makeHandler(0, fmt.Errorf("error")), + want: 0, + wantErr: true, + }, + { + name: "handler_exitcode", + args: []string{"test", "tcp://localhost:9999", "{}"}, + handler: makeHandler(1, nil), + want: 1, + wantErr: false, + }, + { + name: "handler_error_exitcode", + args: []string{"test", "tcp://localhost:9999", "{}"}, + handler: makeHandler(1, fmt.Errorf("error")), + want: 1, + wantErr: true, + }, + { + name: "too_few_arguments", + args: []string{"test"}, + handler: makeHandler(10, nil), + want: 1, + wantErr: true, + }, + { + name: "too_many_arguments", + args: []string{"test", "1", "2", "3"}, + handler: makeHandler(10, nil), + want: 1, + wantErr: true, + }, + { + name: "empty_gitaly_address", + args: []string{"test", "", "{}"}, + handler: makeHandler(10, nil), + want: 1, + wantErr: true, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + currentTest = &tt + defer func() { + currentTest = nil + }() + + done, err := createEnv() + defer done() + require.NoError(t, err) + + got, err := internalRunGitalyCommand(tt.args, tt.handler) + if tt.wantErr { + require.Error(t, err) + } else { + require.NoError(t, err) + } + + require.Equal(t, tt.want, got) + }) + } +} + +// createEnv sets up an environment for `config.New()`. +func createEnv() (func(), error) { + var dir string + var oldWd string + closer := func() { + if oldWd != "" { + err := os.Chdir(oldWd) + if err != nil { + panic(err) + } + } + + if dir != "" { + err := os.RemoveAll(dir) + if err != nil { + panic(err) + } + } + } + + dir, err := ioutil.TempDir("", "test") + if err != nil { + return closer, err + } + + err = ioutil.WriteFile(filepath.Join(dir, "config.yml"), []byte{}, 0644) + if err != nil { + return closer, err + } + + err = ioutil.WriteFile(filepath.Join(dir, "gitlab-shell.log"), []byte{}, 0644) + if err != nil { + return closer, err + } + + oldWd, err = os.Getwd() + if err != nil { + return closer, err + } + + err = os.Chdir(dir) + if err != nil { + return closer, err + } + + return closer, err +} diff --git a/go/internal/handler/handler.go b/go/internal/handler/handler.go deleted file mode 100644 index f8e8bee..0000000 --- a/go/internal/handler/handler.go +++ /dev/null @@ -1,51 +0,0 @@ -package handler - -import ( - "os" - "os/exec" - "syscall" - - "google.golang.org/grpc" - - "gitlab.com/gitlab-org/gitaly/auth" - "gitlab.com/gitlab-org/gitaly/client" - "gitlab.com/gitlab-org/gitlab-shell/go/internal/config" - "gitlab.com/gitlab-org/gitlab-shell/go/internal/logger" -) - -func Prepare() error { - cfg, err := config.New() - if err != nil { - return err - } - - if err := logger.Configure(cfg); err != nil { - return err - } - - // Use a working directory that won't get removed or unmounted. - if err := os.Chdir("/"); err != nil { - return err - } - - return nil -} - -func execCommand(command string, args ...string) error { - binPath, err := exec.LookPath(command) - if err != nil { - return err - } - - args = append([]string{binPath}, args...) - return syscall.Exec(binPath, args, os.Environ()) -} - -func dialOpts() []grpc.DialOption { - connOpts := client.DefaultDialOpts - if token := os.Getenv("GITALY_TOKEN"); token != "" { - connOpts = append(client.DefaultDialOpts, grpc.WithPerRPCCredentials(gitalyauth.RPCCredentialsV2(token))) - } - - return connOpts -} diff --git a/go/internal/handler/receive_pack.go b/go/internal/handler/receive_pack.go index 3496af0..bb6ca89 100644 --- a/go/internal/handler/receive_pack.go +++ b/go/internal/handler/receive_pack.go @@ -2,25 +2,17 @@ package handler import ( "context" - "fmt" "os" pb "gitlab.com/gitlab-org/gitaly-proto/go/gitalypb" "gitlab.com/gitlab-org/gitaly/client" + "google.golang.org/grpc" ) -func ReceivePack(gitalyAddress string, request *pb.SSHReceivePackRequest) (int32, error) { - if gitalyAddress == "" { - return 0, fmt.Errorf("no gitaly_address given") - } - - conn, err := client.Dial(gitalyAddress, dialOpts()) - if err != nil { - return 0, err - } - defer conn.Close() - - ctx, cancel := context.WithCancel(context.Background()) +// ReceivePack issues a Gitaly receive-pack rpc to the provided address +func ReceivePack(ctx context.Context, conn *grpc.ClientConn, request *pb.SSHReceivePackRequest) (int32, error) { + ctx, cancel := context.WithCancel(ctx) defer cancel() + return client.ReceivePack(ctx, conn, os.Stdin, os.Stdout, os.Stderr, request) } diff --git a/go/internal/handler/upload_archive.go b/go/internal/handler/upload_archive.go index 2175300..eb2fa90 100644 --- a/go/internal/handler/upload_archive.go +++ b/go/internal/handler/upload_archive.go @@ -2,25 +2,17 @@ package handler import ( "context" - "fmt" "os" pb "gitlab.com/gitlab-org/gitaly-proto/go/gitalypb" "gitlab.com/gitlab-org/gitaly/client" + "google.golang.org/grpc" ) -func UploadArchive(gitalyAddress string, request *pb.SSHUploadArchiveRequest) (int32, error) { - if gitalyAddress == "" { - return 0, fmt.Errorf("no gitaly_address given") - } - - conn, err := client.Dial(gitalyAddress, dialOpts()) - if err != nil { - return 0, err - } - defer conn.Close() - - ctx, cancel := context.WithCancel(context.Background()) +// UploadArchive issues a Gitaly upload-archive rpc to the provided address +func UploadArchive(ctx context.Context, conn *grpc.ClientConn, request *pb.SSHUploadArchiveRequest) (int32, error) { + ctx, cancel := context.WithCancel(ctx) defer cancel() + return client.UploadArchive(ctx, conn, os.Stdin, os.Stdout, os.Stderr, request) } diff --git a/go/internal/handler/upload_pack.go b/go/internal/handler/upload_pack.go index dd146fb..1ce5ea3 100644 --- a/go/internal/handler/upload_pack.go +++ b/go/internal/handler/upload_pack.go @@ -2,25 +2,17 @@ package handler import ( "context" - "fmt" "os" pb "gitlab.com/gitlab-org/gitaly-proto/go/gitalypb" "gitlab.com/gitlab-org/gitaly/client" + "google.golang.org/grpc" ) -func UploadPack(gitalyAddress string, request *pb.SSHUploadPackRequest) (int32, error) { - if gitalyAddress == "" { - return 0, fmt.Errorf("no gitaly_address given") - } - - conn, err := client.Dial(gitalyAddress, dialOpts()) - if err != nil { - return 0, err - } - defer conn.Close() - - ctx, cancel := context.WithCancel(context.Background()) +// UploadPack issues a Gitaly upload-pack rpc to the provided address +func UploadPack(ctx context.Context, conn *grpc.ClientConn, request *pb.SSHUploadPackRequest) (int32, error) { + ctx, cancel := context.WithCancel(ctx) defer cancel() + return client.UploadPack(ctx, conn, os.Stdin, os.Stdout, os.Stderr, request) } -- cgit v1.2.1