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

Adds distributed tracing instrumentation to GitLab-Shell using LabKit

1 files changed, 15 insertions, 17 deletions

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
 }