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

Adds distributed tracing instrumentation to GitLab-Shell using LabKit

4 files changed, 44 insertions, 128 deletions

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
 }