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 }