Fix async api responses' statuses (status would be stuck in in_queue state) (#2072)

Co-authored-by: vishal <vishalbollu@users.noreply.github.com>

Author: RobertLucian

async-gateway/service.go

@@ -66,9 +66,9 @@ func (s *service) CreateWorkload(id string, payload io.Reader, contentType strin
 		return "", err
 	}
 
-	statusPath := fmt.Sprintf("%s/%s/status", prefix, id)
+	statusPath := fmt.Sprintf("%s/%s/%s", prefix, id, StatusInQueue)
 	log.Debug(fmt.Sprintf("setting status to %s", StatusInQueue))
-	if err := s.storage.Upload(statusPath, strings.NewReader(string(StatusInQueue)), "text/plain"); err != nil {
+	if err := s.storage.Upload(statusPath, strings.NewReader(""), "text/plain"); err != nil {
 		return "", err
 	}
 
@@ -77,32 +77,24 @@ func (s *service) CreateWorkload(id string, payload io.Reader, contentType strin
 
 // GetWorkload retrieves the status and result, if available, of a given workload
 func (s *service) GetWorkload(id string) (GetWorkloadResponse, error) {
-	prefix := s.workloadStoragePrefix()
 	log := s.logger.With(zap.String("id", id))
 
-	// download workload status
-	statusPath := fmt.Sprintf("%s/%s/status", prefix, id)
-	log.Debug("downloading status file", zap.String("path", statusPath))
-	statusBuf, err := s.storage.Download(statusPath)
+	status, err := s.getStatus(id)
 	if err != nil {
 		return GetWorkloadResponse{}, err
 	}
 
-	status := Status(statusBuf[:])
-	switch status {
-	case StatusFailed, StatusInProgress, StatusInQueue:
+	if status != StatusCompleted {
 		return GetWorkloadResponse{
 			ID:     id,
 			Status: status,
 		}, nil
-	case StatusCompleted: // continues execution after switch/case, below
-	default:
-		return GetWorkloadResponse{}, fmt.Errorf("invalid workload status: %s", status)
 	}
 
 	// attempt to download user result
+	prefix := s.workloadStoragePrefix()
 	resultPath := fmt.Sprintf("%s/%s/result.json", prefix, id)
-	log.Debug("donwloading user result", zap.String("path", resultPath))
+	log.Debug("downloading user result", zap.String("path", resultPath))
 	resultBuf, err := s.storage.Download(resultPath)
 	if err != nil {
 		return GetWorkloadResponse{}, err
@@ -127,6 +119,38 @@ func (s *service) GetWorkload(id string) (GetWorkloadResponse, error) {
 	}, nil
 }
 
+func (s *service) getStatus(id string) (Status, error) {
+	prefix := s.workloadStoragePrefix()
+	log := s.logger.With(zap.String("id", id))
+
+	// download workload status
+	log.Debug("checking status", zap.String("path", fmt.Sprintf("%s/%s/*", prefix, id)))
+	files, err := s.storage.List(fmt.Sprintf("%s/%s", prefix, id))
+	if err != nil {
+		return "", err
+	}
+
+	// determine request status
+	status := StatusInQueue
+	for _, file := range files {
+		fileStatus := Status(file)
+
+		if !fileStatus.Valid() {
+			status = fileStatus
+			return "", fmt.Errorf("invalid workload status: %s", status)
+		}
+		if fileStatus == StatusInProgress {
+			status = fileStatus
+		}
+		if fileStatus == StatusCompleted || fileStatus == StatusFailed {
+			status = fileStatus
+			break
+		}
+	}
+
+	return status, nil
+}
+
 func (s *service) workloadStoragePrefix() string {
 	return fmt.Sprintf("%s/apis/%s/workloads", s.clusterName, s.apiName)
 }

async-gateway/storage.go

@@ -18,6 +18,8 @@ package main
 
 import (
 	"io"
+	"path"
+	"strings"
 	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
@@ -30,6 +32,7 @@ import (
 type Storage interface {
 	Upload(key string, payload io.Reader, contentType string) error
 	Download(key string) ([]byte, error)
+	List(key string) ([]string, error)
 	GetLastModified(key string) (time.Time, error)
 }
 
@@ -80,6 +83,30 @@ func (s *s3) Download(key string) ([]byte, error) {
 	return buff.Bytes(), nil
 }
 
+// List lists a set of files from a given S3 path.
+// Works only for one level deep sub-paths.
+func (s *s3) List(key string) ([]string, error) {
+	if key != "" && !strings.HasSuffix(key, "/") {
+		key += "/"
+	}
+
+	result, err := s.client.ListObjectsV2(&awss3.ListObjectsV2Input{
+		Prefix: aws.String(key),
+		Bucket: aws.String(s.bucket),
+	})
+
+	if err != nil {
+		return nil, err
+	}
+
+	files := []string{}
+	for _, obj := range result.Contents {
+		_, file := path.Split(*obj.Key)
+		files = append(files, file)
+	}
+	return files, nil
+}
+
 // GetLastModified retrieves the last modified timestamp of an S3 object
 func (s *s3) GetLastModified(key string) (time.Time, error) {
 	input := awss3.GetObjectInput{

async-gateway/types.go

@@ -32,6 +32,19 @@ const (
 	StatusCompleted  Status = "completed"
 )
 
+func (status Status) String() string {
+	return string(status)
+}
+
+func (status Status) Valid() bool {
+	switch status {
+	case StatusFailed, StatusInProgress, StatusInQueue, StatusCompleted:
+		return true
+	default:
+		return false
+	}
+}
+
 //CreateWorkloadResponse represents the response returned to the user on workload creation
 type CreateWorkloadResponse struct {
 	ID string `json:"id"`

pkg/cortex/serve/cortex_internal/lib/api/async_api.py

@@ -94,7 +94,7 @@ def statsd(self):
         return self.__statsd
 
     def update_status(self, request_id: str, status: str):
-        self.storage.put_str(status, f"{self.storage_path}/{request_id}/status")
+        self.storage.put_str("", f"{self.storage_path}/{request_id}/{status}")
 
     def upload_result(self, request_id: str, result: Dict[str, Any]):
         if not isinstance(result, dict):