Grace-full restart and cancelation

install/app-engine/Dockerfile

@@ -1 +1 @@
-from dekartxyz/dekart:0.9
+FROM dekartxyz/dekart:0.9

src/server/http/http.go → src/server/app/app.go

@@ -1,4 +1,4 @@
-package http
+package app
 
 import (
 	"dekart/src/proto"

src/server/http/staticfilehandler.go → src/server/app/staticfilehandler.go

@@ -1,4 +1,4 @@
-package http
+package app
 
 import (
 	"bytes"

src/server/athenajob/athenajob.go

@@ -1,29 +1,25 @@
 package athenajob
 
 import (
-	"context"
 	"dekart/src/proto"
-	"dekart/src/server/dekart"
+	"dekart/src/server/job"
 	"dekart/src/server/storage"
-	"dekart/src/server/uuid"
 	"fmt"
 	"os"
-	"sync"
 	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/athena"
-	"github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
 )
 
-// Store implements dekart.Store interface for athena
+// Store implements job.Store interface for athena
 type Store struct {
-	mutex          sync.Mutex
+	job.BasicStore
 	session        *session.Session
 	outputLocation string
-	jobs           []*Job
+	Jobs           []*Job
 }
 
 func NewStore(storage storage.Storage) *Store {
@@ -40,142 +36,43 @@ func NewStore(storage storage.Storage) *Store {
 		session:        session,
 		outputLocation: fmt.Sprintf("s3://%s", outputLocation),
 	}
-	store.jobs = make([]*Job, 0)
+	store.Jobs = make([]*Job, 0)
 	return store
 
 }
 
-func (s *Store) removeJobWhenDone(job *Job) {
-	<-job.ctx.Done()
-	s.mutex.Lock()
-	for i, j := range s.jobs {
-		if job.id == j.id {
-			// removing job from slice
-			last := len(s.jobs) - 1
-			s.jobs[i] = s.jobs[last]
-			s.jobs = s.jobs[:last]
-			break
-		}
-	}
-	s.mutex.Unlock()
-}
-
 // Create a new Athena job within the store
-func (s *Store) Create(reportID string, queryID string, queryText string) (dekart.Job, chan int32, error) {
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
-	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
-	status := make(chan int32)
+func (s *Store) Create(reportID string, queryID string, queryText string) (job.Job, chan int32, error) {
+	s.Lock()
+	defer s.Unlock()
 	client := athena.New(s.session)
 	job := &Job{
-		id:             uuid.GetUUID(),
-		reportID:       reportID,
-		queryID:        queryID,
-		ctx:            ctx,
-		cancel:         cancel,
-		status:         status,
-		logger:         log.With().Str("reportID", reportID).Str("queryID", queryID).Str("outputLocation", s.outputLocation).Logger(),
-		queryText:      queryText,
+		BasicJob: job.BasicJob{
+			ReportID:  reportID,
+			QueryID:   queryID,
+			QueryText: queryText,
+			Logger:    log.With().Str("reportID", reportID).Str("queryID", queryID).Logger(),
+		},
 		session:        s.session,
 		client:         client,
 		outputLocation: s.outputLocation,
 	}
-	s.jobs = append(s.jobs, job)
-	go s.removeJobWhenDone(job)
-	return job, status, nil
-}
-
-func (s *Store) Cancel(queryID string) {
-	s.mutex.Lock()
-	for _, job := range s.jobs {
-		if job.queryID == queryID {
-			job.status <- int32(proto.Query_JOB_STATUS_UNSPECIFIED)
-			job.logger.Info().Msg("Canceling Job Context")
-			job.cancel()
-		}
-	}
-	s.mutex.Unlock()
+	job.Init()
+	s.Jobs = append(s.Jobs, job)
+	go s.RemoveJobWhenDone(job)
+	return job, job.Status(), nil
 }
 
 // Job implements dekart.Job interface for Athena
 type Job struct {
-	id               string
-	queryID          string
-	reportID         string
-	ctx              context.Context
-	cancel           context.CancelFunc
-	status           chan int32
-	err              string
-	queryText        string
-	totalRows        int64
-	processedBytes   int64
-	resultSize       int64
-	resultID         *string
-	mutex            sync.Mutex
-	logger           zerolog.Logger
+	job.BasicJob
 	session          *session.Session
 	queryExecutionId string
 	client           *athena.Athena
 	outputLocation   string
 	storageObject    storage.StorageObject
 }
 
-func (j *Job) GetID() string {
-	return j.id
-}
-
-func (j *Job) GetReportID() string {
-	return j.reportID
-}
-
-func (j *Job) GetQueryID() string {
-	return j.queryID
-}
-
-func (j *Job) GetResultID() *string {
-	j.mutex.Lock()
-	defer j.mutex.Unlock()
-	return j.resultID
-}
-
-func (j *Job) GetTotalRows() int64 {
-	j.mutex.Lock()
-	defer j.mutex.Unlock()
-	return j.totalRows
-}
-
-func (j *Job) GetProcessedBytes() int64 {
-	j.mutex.Lock()
-	defer j.mutex.Unlock()
-	return j.processedBytes
-}
-
-func (j *Job) GetResultSize() int64 {
-	j.mutex.Lock()
-	defer j.mutex.Unlock()
-	return j.resultSize
-}
-
-func (j *Job) GetCtx() context.Context {
-	return j.ctx
-}
-
-func (j *Job) Err() string {
-	j.mutex.Lock()
-	defer j.mutex.Unlock()
-	return j.err
-}
-
-func (job *Job) cancelWithError(err error) {
-	if err != context.Canceled {
-		job.mutex.Lock()
-		job.err = err.Error()
-		job.mutex.Unlock()
-	}
-	job.status <- int32(proto.Query_JOB_STATUS_UNSPECIFIED)
-	job.cancel()
-}
-
 func (j *Job) pullQueryExecutionStatus() (*athena.QueryExecution, error) {
 	var err error
 	ticker := time.NewTicker(1 * time.Second)
@@ -186,17 +83,17 @@ func (j *Job) pullQueryExecutionStatus() (*athena.QueryExecution, error) {
 
 	for {
 		select {
-		case <-j.ctx.Done():
-			err = j.ctx.Err()
+		case <-j.GetCtx().Done():
+			err = j.GetCtx().Err()
 			return nil, err
 		case <-ticker.C:
-			out, err := j.client.GetQueryExecutionWithContext(j.ctx, input)
+			out, err := j.client.GetQueryExecutionWithContext(j.GetCtx(), input)
 			if err != nil {
-				j.cancelWithError(err)
+				j.CancelWithError(err)
 				return nil, err
 			}
 			status := *out.QueryExecution.Status.State
-			j.logger.Debug().Str("status", status).Send()
+			j.Logger.Debug().Str("status", status).Send()
 			switch status {
 			case "RUNNING":
 				continue
@@ -214,64 +111,67 @@ func (j *Job) pullQueryExecutionStatus() (*athena.QueryExecution, error) {
 	}
 
 }
+
 func (j *Job) wait() {
 	queryExecution, err := j.pullQueryExecutionStatus()
 	if err != nil {
-		j.cancelWithError(err)
+		j.CancelWithError(err)
 		return
 	}
-	j.logger.Debug().Msg("job done")
+	j.Logger.Debug().Msg("job done")
 	{
-		j.mutex.Lock()
-		j.processedBytes = *queryExecution.Statistics.DataScannedInBytes
-		j.mutex.Unlock()
+		j.Lock()
+		j.ProcessedBytes = *queryExecution.Statistics.DataScannedInBytes
+		j.Unlock()
 	}
-	j.status <- int32(proto.Query_JOB_STATUS_READING_RESULTS)
-	err = j.storageObject.CopyFromS3(j.ctx, *queryExecution.ResultConfiguration.OutputLocation)
+	j.Status() <- int32(proto.Query_JOB_STATUS_READING_RESULTS)
+	err = j.storageObject.CopyFromS3(j.GetCtx(), *queryExecution.ResultConfiguration.OutputLocation)
 	if err != nil {
-		j.cancelWithError(err)
+		j.CancelWithError(err)
 		return
 	}
-	size, err := j.storageObject.GetSize(j.ctx)
+	size, err := j.storageObject.GetSize(j.GetCtx())
 	if err != nil {
-		j.cancelWithError(err)
+		j.CancelWithError(err)
 		return
 	}
 	{
-		j.mutex.Lock()
-		j.resultSize = *size
-		j.resultID = &j.id
-		j.mutex.Unlock()
+		j.Lock()
+		j.ResultSize = *size
+		resultID := j.GetID()
+		j.ResultID = &resultID
+		j.Unlock()
 	}
-	j.status <- int32(proto.Query_JOB_STATUS_DONE)
-	j.cancel()
+	j.Status() <- int32(proto.Query_JOB_STATUS_DONE)
+	j.Cancel()
 }
 
 func (j *Job) Run(storageObject storage.StorageObject) error {
-	j.mutex.Lock()
+	j.Lock()
 	j.storageObject = storageObject
-	j.mutex.Unlock()
+	j.Unlock()
 
-	j.status <- int32(proto.Query_JOB_STATUS_PENDING)
-	out, err := j.client.StartQueryExecutionWithContext(j.ctx, &athena.StartQueryExecutionInput{
-		QueryString: &j.queryText,
+	j.Status() <- int32(proto.Query_JOB_STATUS_PENDING)
+	queryString := j.GetQueryText()
+	out, err := j.client.StartQueryExecutionWithContext(j.GetCtx(), &athena.StartQueryExecutionInput{
+		QueryString: &queryString,
 		ResultConfiguration: &athena.ResultConfiguration{
 			OutputLocation: &j.outputLocation,
 		},
 	})
 	if err != nil {
-		j.logger.Error().Err(err).Msg("Error starting query execution")
-		j.cancelWithError(err)
+		j.Logger.Error().Err(err).Msg("Error starting query execution")
+		j.CancelWithError(err)
 		return nil
 	}
 
-	j.mutex.Lock()
+	j.Lock()
 	j.queryExecutionId = *out.QueryExecutionId
-	j.mutex.Unlock()
+	j.Unlock()
 
-	j.status <- int32(proto.Query_JOB_STATUS_RUNNING)
+	j.Status() <- int32(proto.Query_JOB_STATUS_RUNNING)
 
-	j.logger.Debug().Str("queryExecutionId", j.queryExecutionId).Msg("waiting")
+	j.Logger.Debug().Str("queryExecutionId", j.queryExecutionId).Msg("waiting")
 	go j.wait()
 	return nil
 }