Delayed queue deletion (#1952)

Author: vishalbollu

pkg/lib/cron/cron.go

@@ -32,7 +32,7 @@ func Run(f func() error, errHandler func(error), delay time.Duration) Cron {
 	cronCancel := make(chan struct{}, 1)
 
 	runCron := func() {
-		defer recoverer(errHandler)
+		defer Recoverer(errHandler)
 		err := f()
 		if err != nil && errHandler != nil {
 			errHandler(err)
@@ -69,7 +69,7 @@ func (c *Cron) Cancel() {
 	c.cronCancel <- struct{}{}
 }
 
-func recoverer(errHandler func(error)) {
+func Recoverer(errHandler func(error)) {
 	if errInterface := recover(); errInterface != nil {
 		err := errors.CastRecoverError(errInterface)
 		errors.PrintStacktrace(err)

pkg/operator/resources/job/batchapi/job.go

@@ -224,7 +224,7 @@ func handleJobSubmissionError(jobKey spec.JobKey, jobErr error) {
 func deleteJobRuntimeResources(jobKey spec.JobKey) error {
 	err := errors.FirstError(
 		deleteK8sJob(jobKey),
-		deleteQueueByJobKeyIfExists(jobKey),
+		deleteQueueWithDelay(jobKey),
 		saveMetricsToCloud(jobKey),
 	)
 

pkg/operator/resources/job/batchapi/queue.go

@@ -19,10 +19,12 @@ package batchapi
 import (
 	"fmt"
 	"strings"
+	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/sqs"
 	awslib "github.com/cortexlabs/cortex/pkg/lib/aws"
+	"github.com/cortexlabs/cortex/pkg/lib/cron"
 	"github.com/cortexlabs/cortex/pkg/lib/errors"
 	libjson "github.com/cortexlabs/cortex/pkg/lib/json"
 	s "github.com/cortexlabs/cortex/pkg/lib/strings"
@@ -31,6 +33,11 @@ import (
 	"github.com/cortexlabs/cortex/pkg/types/spec"
 )
 
+const (
+	_markForDeletion          = "cortex.dev/to-be-deleted"
+	_queueGraceKillTimePeriod = 5 * time.Minute
+)
+
 func apiQueueNamePrefix(apiName string) string {
 	return config.CoreConfig.SQSNamePrefix() + apiName + "-"
 }
@@ -121,21 +128,75 @@ func listQueueURLsForAllAPIs() ([]string, error) {
 	return queueURLs, nil
 }
 
-func deleteQueueByJobKey(jobKey spec.JobKey) error {
+func markForDeletion(queueURL string) error {
+	_, err := config.AWS.SQS().TagQueue(&sqs.TagQueueInput{
+		QueueUrl: aws.String(queueURL),
+		Tags: aws.StringMap(map[string]string{
+			_markForDeletion: time.Now().Format(time.RFC3339Nano),
+		}),
+	})
+	if err != nil {
+		return errors.WithStack(err)
+	}
+	return nil
+}
+
+func deleteQueueWithDelay(jobKey spec.JobKey) error {
 	queueURL, err := getJobQueueURL(jobKey)
 	if err != nil {
 		return err
 	}
 
-	return deleteQueueByURL(queueURL)
-}
-
-func deleteQueueByJobKeyIfExists(jobKey spec.JobKey) error {
-	err := deleteQueueByJobKey(jobKey)
-	if err != nil && awslib.IsNonExistentQueueErr(errors.CauseOrSelf(err)) {
+	output, err := config.AWS.SQS().ListQueueTags(&sqs.ListQueueTagsInput{
+		QueueUrl: aws.String(queueURL),
+	})
+	if err != nil {
+		if !awslib.IsNonExistentQueueErr(errors.CauseOrSelf(err)) {
+			operatorLogger.Error(err)
+		}
 		return nil
 	}
-	return err
+
+	if value, exists := output.Tags[_markForDeletion]; exists {
+		markedTime, err := time.Parse(time.RFC3339Nano, *value)
+		if err != nil {
+			err = deleteQueueByURL(queueURL)
+			if err != nil {
+				if !awslib.IsNonExistentQueueErr(errors.CauseOrSelf(err)) {
+					operatorLogger.Error(err)
+				}
+				return nil
+			}
+		}
+
+		if time.Since(markedTime) > _queueGraceKillTimePeriod {
+			err := deleteQueueByURL(queueURL)
+			if err != nil {
+				if !awslib.IsNonExistentQueueErr(errors.CauseOrSelf(err)) {
+					operatorLogger.Error(err)
+				}
+				return nil
+			}
+		}
+	} else {
+		operatorLogger.Info("scheduling deleting queue " + jobKey.UserString())
+		err = markForDeletion(queueURL)
+		if err != nil && awslib.IsNonExistentQueueErr(errors.CauseOrSelf(err)) {
+			return nil
+		}
+
+		time.AfterFunc(_queueGraceKillTimePeriod, func() {
+			defer cron.Recoverer(nil)
+			operatorLogger.Info("deleting queue " + jobKey.UserString())
+			err := deleteQueueByURL(queueURL)
+			// ignore non existent queue errors
+			if err != nil && !awslib.IsNonExistentQueueErr(errors.CauseOrSelf(err)) {
+				operatorLogger.Error(err)
+			}
+		})
+	}
+
+	return nil
 }
 
 func deleteQueueByURL(queueURL string) error {