Merge pull request #20 from uc-cdis/feat/remove_mq

Feat/remove mq

Author: giangbui

handlers/handler.go

@@ -1,7 +1,6 @@
 package handlers
 
 import (
-	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -10,17 +9,14 @@ import (
 	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/service/sqs"
 	"github.com/golang/glog"
-	mq "github.com/remind101/mq-go"
 )
 
-const MAX_RETRIES = 3
-
 type SQSHandler struct {
 	QueueURL      string
 	Start         bool
 	JobConfigs    []JobConfig
-	Server        *mq.Server
 	MonitoredJobs []*JobInfo
 	Mu            sync.Mutex
 }
@@ -49,27 +45,84 @@ func NewSQSHandler(queueURL string) *SQSHandler {
 
 // StartServer starts a server
 func (handler *SQSHandler) StartServer() error {
-	// return nil if the server already start
-	if handler.Server != nil {
-		return nil
-	}
 
+	glog.Info("Starting a new server ...")
+
+	go handler.StartConsumingProcess()
+	go handler.StartMonitoringProcess()
+	go handler.RemoveCompletedJobsProcess()
+
+	glog.Info("The server is started")
+
+	return nil
+
+}
+
+// StartConsumingProcess starts consumming the queue
+func (handler *SQSHandler) StartConsumingProcess() error {
 	newClient, err := NewSQSClient()
 	if err != nil {
 		return err
 	}
 
-	glog.Info("Starting a new server...")
-	handler.Server = mq.NewServer(handler.QueueURL, mq.HandlerFunc(func(m *mq.Message) error {
-		return handler.HandleSQSMessage(aws.StringValue(m.SQSMessage.Body))
-	}), mq.WithClient(newClient))
-	handler.Server.Start()
-	glog.Info("The server is started")
+	receiveParams := &sqs.ReceiveMessageInput{
+		QueueUrl:            aws.String(handler.QueueURL),
+		MaxNumberOfMessages: aws.Int64(1),
+		VisibilityTimeout:   aws.Int64(30),
+		WaitTimeSeconds:     aws.Int64(20),
+	}
+	for {
+		time.Sleep(1 * time.Second)
+		receiveResp, err := newClient.ReceiveMessage(receiveParams)
+		if err != nil {
+			glog.Error(err)
+		}
 
-	go handler.StartMonitoringProcess()
-	go handler.RemoveCompletedJobsProcess()
+		for _, message := range receiveResp.Messages {
+			err := handler.HandleSQSMessage(message)
+			if err != nil {
+				glog.Errorf("Can not process the message. Error %s. Message %s", err, *message.Body)
+				continue
+			}
+
+			handler.RemoveSQSMessage(message)
+
+		}
 
+	}
+}
+
+// RemoveSQSMessage removes SQS message
+func (handler *SQSHandler) RemoveSQSMessage(message *sqs.Message) error {
+	newClient, err := NewSQSClient()
+	if err != nil {
+		return err
+	}
+	deleteParams := &sqs.DeleteMessageInput{
+		QueueUrl:      aws.String(handler.QueueURL), // Required
+		ReceiptHandle: message.ReceiptHandle,        // Required
+	}
+	_, err = newClient.DeleteMessage(deleteParams) // No response returned when successed.
+	if err != nil {
+		glog.Error(err)
+		return err
+	}
+	glog.Infof("Message ID: %s has beed deleted.\n\n", *message.MessageId)
 	return nil
+}
+
+// ResendSQSMessage resends the message
+func (handler *SQSHandler) ResendSQSMessage(queueURL string, message *sqs.Message) error {
+	newClient, err := NewSQSClient()
+	if err != nil {
+		return err
+	}
+	sentMessageInput := &sqs.SendMessageInput{
+		MessageBody: message.Body,
+		QueueUrl:    aws.String(queueURL),
+	}
+	_, err = newClient.SendMessage(sentMessageInput)
+	return err
 
 }
 
@@ -81,7 +134,8 @@ func (handler *SQSHandler) StartMonitoringProcess() {
 		for _, jobInfo := range handler.MonitoredJobs {
 			k8sJob, err := GetJobStatusByID(jobInfo.UID)
 			if err != nil {
-				glog.Errorf("Can not get k8s job %s. Detail %s", jobInfo.Name, err)
+				glog.Errorf("Can not get k8s job %s. Detail %s. Resend the message to the queue", jobInfo.Name, err)
+				handler.ResendSQSMessage(handler.QueueURL, jobInfo.SQSMessage)
 			} else {
 				glog.Infof("%s: %s", k8sJob.Name, k8sJob.Status)
 				if k8sJob.Status == "Unknown" || k8sJob.Status == "Running" {
@@ -107,17 +161,6 @@ func (handler *SQSHandler) RemoveCompletedJobsProcess() {
 	}
 }
 
-// ShutdownServer shutdowns a server
-func (handler *SQSHandler) ShutdownServer() error {
-	fmt.Println("Shutdown the server")
-	if handler.Server == nil {
-		return nil
-	}
-	err := handler.Server.Shutdown(context.Background())
-	handler.Server = nil
-	return err
-}
-
 /*
 getObjectFromSQSMessage returns s3 object from sqs message
 
@@ -202,8 +245,9 @@ to the queue and retry later (handled by `md` library). That makes sure
 the message is properly handle before it actually deleted
 
 */
-func (handler *SQSHandler) HandleSQSMessage(jsonBody string) error {
+func (handler *SQSHandler) HandleSQSMessage(message *sqs.Message) error {
 
+	jsonBody := *message.Body
 	objectPaths := getObjectsFromSQSMessage(jsonBody)
 
 	jobNameList := make([]string, 0)
@@ -220,21 +264,24 @@ func (handler *SQSHandler) HandleSQSMessage(jsonBody string) error {
 			}
 		}
 	}
+	glog.Info("message:", jsonBody)
 
 	glog.Infof("Start to run %d jobs", len(jobMap))
 
 	for objectPath, jobConfig := range jobMap {
 		for GetNumberRunningJobs() > GetMaxJobConfig() {
 			time.Sleep(5 * time.Second)
 		}
-		glog.Info("Processing: ", objectPath)
 		jobInfo, err := CreateK8sJob(objectPath, jobConfig)
 		if err != nil {
+			glog.Infof("Error :%s", err)
 			glog.Errorln(err)
 			return err
 		}
+		jobInfo.SQSMessage = message
 		out, err := json.Marshal(jobInfo)
 		if err != nil {
+			glog.Infof("Error :%s", err)
 			glog.Errorln(err)
 			return err
 		}
@@ -296,5 +343,8 @@ func (handler *SQSHandler) RetryCreateIndexingJob(jsonBytes []byte) error {
 	str := fmt.Sprintf(`{
 		"Type" : "Notification",
 		"Message" : "{\"Records\":[{\"eventSource\":\"aws:s3\",\"awsRegion\":\"us-east-1\",\"eventName\":\"ObjectCreated:Put\",\"s3\":{\"s3SchemaVersion\":\"1.0\",\"bucket\":{\"name\":\"%s\"},\"object\":{\"key\":\"%s\"}}}]}"}`, retryMessage.Bucket, retryMessage.Key)
-	return handler.HandleSQSMessage(str)
+	sqsMessage := sqs.Message{}
+	sqsMessage.SetBody(str)
+
+	return handler.HandleSQSMessage(&sqsMessage)
 }

handlers/handler_api.go

@@ -8,76 +8,10 @@ import (
 
 // RegisterSQSHandler registers endpoints
 func (handler *SQSHandler) RegisterSQSHandler() {
-	http.HandleFunc("/server", handler.ServiceHandler)
-	http.HandleFunc("/sqs", handler.SQSHandler)
 	http.HandleFunc("/jobConfig", handler.HandleJobConfig)
 	http.HandleFunc("/dispatchJob", handler.HandleDispatchJob)
 }
 
-// ServiceHandler handles stop/start/status the SQS querrying service
-// To start the server
-//		curl -X POST http://localhost/server?start=true
-// To stop the server
-//		curl -X POST http://localhost/server?start=false
-// To check the status
-// 		curl http://localhost/server
-func (handler *SQSHandler) ServiceHandler(w http.ResponseWriter, r *http.Request) {
-	if r.Method == "PUT" {
-		val := r.URL.Query().Get("start")
-		if val == "true" {
-			if err := handler.StartServer(); err != nil {
-				fmt.Fprintf(w, fmt.Sprintf("Can not start a server. Detail %s", err))
-			} else {
-				fmt.Fprintf(w, "Successfully start a server")
-			}
-		} else if val == "false" {
-			if err := handler.ShutdownServer(); err != nil {
-				fmt.Fprintf(w, fmt.Sprintf("Can not shutdown the server. Detail %s", err))
-			} else {
-				fmt.Fprintf(w, "Successfully shutdown the server")
-			}
-		}
-
-	} else if r.Method == "GET" {
-		if handler.Server != nil {
-			fmt.Fprintf(w, "SQS server is running")
-		} else {
-			fmt.Fprintf(w, "SQS server is not running")
-		}
-	} else {
-		http.Error(w, "Not supported request method.", 405)
-	}
-
-}
-
-// SQSHandler handles update/get SQS URL string
-// to update SQS url
-//		curl -X PUT http://localhost/sqs?url=http://sqs.aws.example
-// to get current sqs url
-//		curl http://localhost/sqs
-func (handler *SQSHandler) SQSHandler(w http.ResponseWriter, r *http.Request) {
-	if r.Method == "PUT" {
-		handler.QueueURL = r.URL.Query().Get("url")
-
-		if err := handler.ShutdownServer(); err != nil {
-			fmt.Fprintf(w, fmt.Sprintf("Can not shutdown the server. Detail %s", err))
-		} else {
-			fmt.Fprintf(w, "Successfully shutdown the server")
-		}
-
-		if err := handler.StartServer(); err != nil {
-			fmt.Fprintf(w, fmt.Sprintf("Can not start a server. Detail %s", err))
-		} else {
-			fmt.Fprintf(w, "Successfully start a server")
-		}
-
-	} else if r.Method == "GET" {
-		fmt.Fprintf(w, handler.QueueURL)
-	} else {
-		http.Error(w, "Not supported request method.", 405)
-	}
-}
-
 // HandleJobConfig handles job config endpoints
 // to add a jobConfig
 //		curl -X POST http://localhost/jobConfig -d `{"name": "usersync", "pattern": "s3://bucket/usersync.yaml", "image": "quay.io/cdis/fence:master", "imageConfig":{}}`

handlers/handler_test.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"os"
 
+	"github.com/aws/aws-sdk-go/service/sqs"
 	"github.com/golang/glog"
 
 	batchv1 "k8s.io/api/batch/v1"
@@ -26,9 +27,10 @@ type JobsArray struct {
 }
 
 type JobInfo struct {
-	UID    string `json:"uid"`
-	Name   string `json:"name"`
-	Status string `json:"status"`
+	UID        string `json:"uid"`
+	Name       string `json:"name"`
+	Status     string `json:"status"`
+	SQSMessage *sqs.Message
 }
 
 func getJobClient() batchtypev1.JobInterface {

handlers/jobs.go

@@ -14,7 +14,6 @@ curl -X POST http://localhost:8000/job -d '{"Name":"TEST3", "Pattern":"s3://test
 package main
 
 import (
-	"context"
 	"encoding/json"
 	"flag"
 	"fmt"
@@ -68,10 +67,10 @@ func main() {
 
 	// start an SQSHandler instance
 	SQSHandler := handlers.NewSQSHandler(sqsURL)
+
 	if err := SQSHandler.StartServer(); err != nil {
 		glog.Errorf("Can not start the server. Detail %s", err)
 	}
-	defer SQSHandler.Server.Shutdown(context.Background())
 
 	SQSHandler.JobConfigs = jobConfigs