Adds distributed tracing instrumentation to GitLab-Shell

Adds distributed tracing instrumentation to GitLab-Shell using LabKit

config.yml.example

@@ -56,3 +56,7 @@ audit_usernames: false
 migration:
   enabled: false
   features: []
+
+# Distributed Tracing. GitLab-Shell has distributed tracing instrumentation.
+# For more details, visit https://docs.gitlab.com/ee/development/distributed_tracing.html
+# gitlab_tracing: opentracing://driver

go/cmd/gitaly-receive-pack/main.go

 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
 }

go/cmd/gitaly-upload-archive/main.go

 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
 }

go/cmd/gitaly-upload-archive/main_test.go

-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)
-			}
-		})
-	}
-}

go/cmd/gitaly-upload-pack/main.go

 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
 }

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) {

go/internal/handler/exec.go

+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
+}

go/internal/handler/exec_test.go

+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
+}

go/internal/handler/handler.go

-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
-}

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)
 }

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)
 }

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)
 }