Report errors to sentry for the async gateway / enqueuer (#2138)

Co-authored-by: Miguel Varela Ramos <miguel@cortexlabs.com>

Author: RobertLucian

async-gateway/go.mod

@@ -20,110 +20,115 @@ import (
 	"flag"
 	"net/http"
 	"os"
-	"strings"
 
-	"github.com/aws/aws-sdk-go/aws"
-	"github.com/aws/aws-sdk-go/aws/session"
+	gateway "github.com/cortexlabs/cortex/pkg/async-gateway"
+	"github.com/cortexlabs/cortex/pkg/lib/aws"
+	"github.com/cortexlabs/cortex/pkg/lib/errors"
+	"github.com/cortexlabs/cortex/pkg/lib/logging"
+	"github.com/cortexlabs/cortex/pkg/lib/telemetry"
+	"github.com/cortexlabs/cortex/pkg/types/clusterconfig"
+	"github.com/cortexlabs/cortex/pkg/types/userconfig"
 	"github.com/gorilla/handlers"
 	"github.com/gorilla/mux"
-	"go.uber.org/zap"
-	"go.uber.org/zap/zapcore"
 )
 
 const (
 	_defaultPort = "8080"
 )
 
-func createLogger() (*zap.Logger, error) {
-	logLevelEnv := os.Getenv("CORTEX_LOG_LEVEL")
-	disableJSONLogging := os.Getenv("CORTEX_DISABLE_JSON_LOGGING")
-
-	var logLevelZap zapcore.Level
-	switch logLevelEnv {
-	case "DEBUG":
-		logLevelZap = zapcore.DebugLevel
-	case "WARNING":
-		logLevelZap = zapcore.WarnLevel
-	case "ERROR":
-		logLevelZap = zapcore.ErrorLevel
-	default:
-		logLevelZap = zapcore.InfoLevel
+var (
+	gatewayLogger = logging.GetLogger()
+)
+
+func Exit(err error, wrapStrs ...string) {
+	for _, str := range wrapStrs {
+		err = errors.Wrap(err, str)
 	}
 
-	encoderConfig := zap.NewProductionEncoderConfig()
-	encoderConfig.MessageKey = "message"
+	if err != nil && !errors.IsNoTelemetry(err) {
+		telemetry.Error(err)
+	}
 
-	encoding := "json"
-	if strings.ToLower(disableJSONLogging) == "true" {
-		encoding = "console"
+	if err != nil && !errors.IsNoPrint(err) {
+		gatewayLogger.Error(err)
 	}
 
-	return zap.Config{
-		Level:            zap.NewAtomicLevelAt(logLevelZap),
-		Encoding:         encoding,
-		EncoderConfig:    encoderConfig,
-		OutputPaths:      []string{"stdout"},
-		ErrorOutputPaths: []string{"stderr"},
-	}.Build()
+	telemetry.Close()
+
+	os.Exit(1)
 }
 
 // usage: ./gateway -bucket <bucket> -region <region> -port <port> -queue queue <apiName>
 func main() {
-	log, err := createLogger()
-	if err != nil {
-		panic(err)
-	}
+	log := logging.GetLogger()
 	defer func() {
 		_ = log.Sync()
 	}()
 
 	var (
-		port       = flag.String("port", _defaultPort, "port on which the gateway server runs on")
-		queueURL   = flag.String("queue", "", "SQS queue URL")
-		region     = flag.String("region", "", "AWS region")
-		bucket     = flag.String("bucket", "", "AWS bucket")
-		clusterUID = flag.String("cluster-uid", "", "cluster UID")
+		port              = flag.String("port", _defaultPort, "port on which the gateway server runs on")
+		queueURL          = flag.String("queue", "", "SQS queue URL")
+		clusterConfigPath = flag.String("cluster-config", "", "cluster config path")
 	)
 	flag.Parse()
 
 	switch {
 	case *queueURL == "":
 		log.Fatal("missing required option: -queue")
-	case *region == "":
-		log.Fatal("missing required option: -region")
-	case *bucket == "":
-		log.Fatal("missing required option: -bucket")
-	case *clusterUID == "":
-		log.Fatal("missing required option: -cluster-uid")
+	case *clusterConfigPath == "":
+		log.Fatal("missing required option: -cluster-config")
 	}
 
 	apiName := flag.Arg(0)
 	if apiName == "" {
 		log.Fatal("apiName argument was not provided")
 	}
 
-	sess, err := session.NewSessionWithOptions(session.Options{
-		Config: aws.Config{
-			Region: region,
+	clusterConfig, err := clusterconfig.NewForFile(*clusterConfigPath)
+	if err != nil {
+		Exit(err)
+	}
+
+	awsClient, err := aws.NewForRegion(clusterConfig.Region)
+	if err != nil {
+		Exit(err)
+	}
+
+	_, userID, err := awsClient.CheckCredentials()
+	if err != nil {
+		Exit(err)
+	}
+
+	err = telemetry.Init(telemetry.Config{
+		Enabled: clusterConfig.Telemetry,
+		UserID:  userID,
+		Properties: map[string]string{
+			"kind":       userconfig.AsyncAPIKind.String(),
+			"image_type": "async-gateway",
 		},
-		SharedConfigState: session.SharedConfigEnable,
+		Environment: "api",
+		LogErrors:   true,
+		BackoffMode: telemetry.BackoffDuplicateMessages,
 	})
 	if err != nil {
-		log.Fatal("failed to create AWS session: %s", zap.Error(err))
+		Exit(err)
 	}
 
-	s3Storage := NewS3(sess, *bucket)
-	sqsQueue := NewSQS(*queueURL, sess)
+	sess := awsClient.Session()
+	s3Storage := gateway.NewS3(sess, clusterConfig.Bucket)
+	sqsQueue := gateway.NewSQS(*queueURL, sess)
 
-	svc := NewService(*clusterUID, apiName, sqsQueue, s3Storage, log)
-	ep := NewEndpoint(svc, log)
+	svc := gateway.NewService(clusterConfig.ClusterUID, apiName, sqsQueue, s3Storage, log)
+	ep := gateway.NewEndpoint(svc, log)
 
 	router := mux.NewRouter()
 	router.HandleFunc("/", ep.CreateWorkload).Methods("POST")
 	router.HandleFunc(
 		"/healthz",
 		func(w http.ResponseWriter, r *http.Request) {
-			respondPlainText(w, http.StatusOK, "ok")
+			w.WriteHeader(http.StatusOK)
+			w.Header().Set("Content-Type", "text/plain")
+			_, _ = w.Write([]byte("ok"))
 		},
 	)
 	router.HandleFunc("/{id}", ep.GetWorkload).Methods("GET")
@@ -140,6 +145,6 @@ func main() {
 
 	log.Info("Running on port " + *port)
 	if err = http.ListenAndServe(":"+*port, handlers.CORS(corsOptions...)(router)); err != nil {
-		log.Fatal("failed to start server", zap.Error(err))
+		Exit(err)
 	}
 }

async-gateway/go.sum

@@ -22,6 +22,7 @@ import (
 	"strings"
 
 	"github.com/cortexlabs/cortex/pkg/consts"
+	"github.com/cortexlabs/cortex/pkg/enqueuer"
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
 )
@@ -105,7 +106,7 @@ func main() {
 		log.Fatal("-jobID is a required option")
 	}
 
-	envConfig := EnvConfig{
+	envConfig := enqueuer.EnvConfig{
 		ClusterUID: clusterUID,
 		Region:     region,
 		Version:    version,
@@ -114,17 +115,17 @@ func main() {
 		JobID:      jobID,
 	}
 
-	enqueuer, err := NewEnqueuer(envConfig, queueURL, log)
+	eqr, err := enqueuer.NewEnqueuer(envConfig, queueURL, log)
 	if err != nil {
 		log.Fatal("failed to create enqueuer", zap.Error(err))
 	}
 
-	totalBatches, err := enqueuer.Enqueue()
+	totalBatches, err := eqr.Enqueue()
 	if err != nil {
 		log.Fatal("failed to enqueue batches", zap.Error(err))
 	}
 
-	if err = enqueuer.UploadBatchCount(totalBatches); err != nil {
+	if err = eqr.UploadBatchCount(totalBatches); err != nil {
 		log.Fatal("failed to upload batch count", zap.Error(err))
 	}
 

async-gateway/main.go renamed to cmd/async-gateway/main.go

@@ -23,10 +23,10 @@ import (
 	"github.com/cortexlabs/cortex/pkg/config"
 	"github.com/cortexlabs/cortex/pkg/lib/cron"
 	"github.com/cortexlabs/cortex/pkg/lib/errors"
+	"github.com/cortexlabs/cortex/pkg/lib/logging"
 	"github.com/cortexlabs/cortex/pkg/lib/telemetry"
 	"github.com/cortexlabs/cortex/pkg/operator/endpoints"
 	"github.com/cortexlabs/cortex/pkg/operator/lib/exit"
-	"github.com/cortexlabs/cortex/pkg/operator/lib/logging"
 	"github.com/cortexlabs/cortex/pkg/operator/operator"
 	"github.com/cortexlabs/cortex/pkg/operator/resources/asyncapi"
 	"github.com/cortexlabs/cortex/pkg/operator/resources/job/taskapi"
@@ -37,7 +37,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 )
 
-var operatorLogger = logging.GetOperatorLogger()
+var operatorLogger = logging.GetLogger()
 
 const _operatorPortStr = "8888"
 

enqueuer/main.go renamed to cmd/enqueuer/main.go

@@ -76,7 +76,7 @@ fi
 export CORTEX_OPERATOR_IN_CLUSTER=false
 export CORTEX_CLUSTER_CONFIG_PATH=~/.cortex/cluster-dev.yaml
 export CORTEX_DISABLE_JSON_LOGGING=true
-export CORTEX_OPERATOR_LOG_LEVEL=debug
+export CORTEX_LOG_LEVEL=debug
 export CORTEX_PROMETHEUS_URL="http://localhost:9090"
 
 portForwardCMD="kubectl port-forward -n default prometheus-prometheus-0 9090"
@@ -89,17 +89,17 @@ mkdir -p $ROOT/bin
 
 if [ "$operator_only" = "true" ]; then
   kill $(pgrep -f rerun) >/dev/null 2>&1 || true
-  rerun -watch $ROOT/pkg $ROOT/dev/config -run sh -c \
-  "clear && echo 'building operator...' && go build -o $ROOT/bin/operator $ROOT/pkg/operator && echo 'starting local operator...' && $ROOT/bin/operator"
+  rerun -watch $ROOT/pkg $ROOT/dev/config $ROOT/cmd/operator -run sh -c \
+  "clear && echo 'building operator...' && go build -o $ROOT/bin/operator $ROOT/cmd/operator && echo 'starting local operator...' && $ROOT/bin/operator"
 elif [ "$debug" = "true" ]; then
-  DEBUG_CMD="dlv --listen=:2345 --headless=true --api-version=2 debug $ROOT/pkg/operator --output ${ROOT}/bin/__debug_bin"
+  DEBUG_CMD="dlv --listen=:2345 --headless=true --api-version=2 debug $ROOT/cmd/operator --output ${ROOT}/bin/__debug_bin"
   kill $(pgrep -f "${DEBUG_CMD}") >/dev/null 2>&1 || true
   kill $(pgrep -f __debug_bin) >/dev/null 2>&1 || true
   echo 'starting local operator in debug mode...' && eval "${DEBUG_CMD}"
 else
   kill $(pgrep -f rerun) >/dev/null 2>&1 || true
-  rerun -watch $ROOT/pkg $ROOT/cli $ROOT/dev/config -run sh -c \
-  "clear && echo 'building cli...' && go build -o $ROOT/bin/cortex $ROOT/cli && echo 'building operator...' && go build -o $ROOT/bin/operator $ROOT/pkg/operator && echo 'starting local operator...' && $ROOT/bin/operator"
+  rerun -watch $ROOT/pkg $ROOT/cli $ROOT/dev/config $ROOT/cmd/operator -run sh -c \
+  "clear && echo 'building cli...' && go build -o $ROOT/bin/cortex $ROOT/cli && echo 'building operator...' && go build -o $ROOT/bin/operator $ROOT/cmd/operator && echo 'starting local operator...' && $ROOT/bin/operator"
 fi
 
-# go run -race $ROOT/pkg/operator/main.go  # Check for race conditions. Doesn't seem to catch them all?
+# go run -race $ROOT/cmd/operator # Check for race conditions. Doesn't seem to catch them all?