github.com/sentienttechnologies/studio-go-runner@v0.0.0-20201118202441-6d21f2ced8ee/internal/runner/taskqueue.go (about)

     1  // Copyright 2018-2020 (c) Cognizant Digital Business, Evolutionary AI. All rights reserved. Issued under the Apache 2.0 License.
     2  
     3  package runner
     4  
     5  // This file defines an interface for task queues used by the runner
     6  import (
     7  	"context"
     8  	"os"
     9  	"regexp"
    10  	"strings"
    11  
    12  	runnerReports "github.com/leaf-ai/studio-go-runner/internal/gen/dev.cognizant_dev.ai/genproto/studio-go-runner/reports/v1"
    13  
    14  	"github.com/go-stack/stack"
    15  	"github.com/jjeffery/kv" // MIT License
    16  )
    17  
    18  // QueueTask encapsulates the metadata needed to handle requests on a queue.
    19  //
    20  type QueueTask struct {
    21  	FQProject    string // A proprietary runner label for a project to uniquely identify it
    22  	Project      string
    23  	QueueType    string
    24  	Subscription string
    25  	ShortQName   string // The short queue name for the current task, will be used to retrieve signing keys
    26  	Credentials  string
    27  	Msg          []byte
    28  	Handler      MsgHandler
    29  	Wrapper      *Wrapper                   // A store of encryption related information for messages
    30  	ResponseQ    chan *runnerReports.Report // A response queue the runner can employ to send progress updates on
    31  }
    32  
    33  // MsgHandler defines the function signature for a generic message handler for a specified queue implementation
    34  //
    35  type MsgHandler func(ctx context.Context, qt *QueueTask) (resource *Resource, ack bool, err kv.Error)
    36  
    37  // TaskQueue is the interface definition for a queue message handling implementation.
    38  //
    39  type TaskQueue interface {
    40  	// Refresh is used to scan the catalog of queues work could arrive on and pass them back to the caller
    41  	Refresh(ctx context.Context, qNameMatch *regexp.Regexp, qNameMismatch *regexp.Regexp) (known map[string]interface{}, err kv.Error)
    42  
    43  	// Process a single unit of work if available on a queue, blocking operation on the queue and on the processing
    44  	// of the work itself
    45  	Work(ctx context.Context, qt *QueueTask) (msgProcessed bool, resource *Resource, err kv.Error)
    46  
    47  	// Check that the specified queue exists
    48  	Exists(ctx context.Context, subscription string) (exists bool, err kv.Error)
    49  
    50  	// HasWork is a probe to see if there is a potential for work to be available
    51  	HasWork(ctx context.Context, subscription string) (hasWork bool, err kv.Error)
    52  
    53  	// Responder is used to open a connection to an existing repsonse queue if
    54  	// one was made available and also to provision a channel into which the
    55  	// runner can place report messages
    56  	Responder(ctx context.Context, subscription string) (sender chan *runnerReports.Report, err kv.Error)
    57  }
    58  
    59  // NewTaskQueue is used to initiate processing for any of the types of queues
    60  // the runner supports.  It also performs some lazy initialization.
    61  //
    62  func NewTaskQueue(project string, creds string, wrapper *Wrapper) (tq TaskQueue, err kv.Error) {
    63  
    64  	switch {
    65  	case strings.HasPrefix(project, "amqp://"):
    66  		tq, err = NewRabbitMQ(project, creds, wrapper)
    67  	default:
    68  		// SQS uses a number of credential and config file names
    69  		files := strings.Split(creds, ",")
    70  		for _, file := range files {
    71  			_, errGo := os.Stat(file)
    72  			if errGo != nil {
    73  				return nil, kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime()).With("file", file).With("project", project)
    74  			}
    75  		}
    76  		tq, err = NewSQS(project, creds, wrapper)
    77  	}
    78  
    79  	return tq, err
    80  }