Adds distributed tracing instrumentation to GitLab-Shellan-distributed-tracing

Adds distributed tracing instrumentation to GitLab-Shell using LabKit

7 files changed, 278 insertions, 95 deletions

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