github.com/sentienttechnologies/studio-go-runner@v0.0.0-20201118202441-6d21f2ced8ee/internal/runner/rmq.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 contains the implementation of a RabbitMQ (rmq) client that will 6 // be used to retrieve work from RMQ and to query RMQ for extant queues 7 // within an StudioML Exchange 8 9 import ( 10 "context" 11 "fmt" 12 "net/http" 13 "net/url" 14 "os" 15 "regexp" 16 "strconv" 17 "strings" 18 "sync" 19 "time" 20 21 runnerReports "github.com/leaf-ai/studio-go-runner/internal/gen/dev.cognizant_dev.ai/genproto/studio-go-runner/reports/v1" 22 "google.golang.org/protobuf/encoding/prototext" 23 24 rh "github.com/michaelklishin/rabbit-hole" 25 26 "github.com/rs/xid" 27 "github.com/streadway/amqp" 28 29 "github.com/go-stack/stack" 30 "github.com/jjeffery/kv" // MIT License 31 ) 32 33 // RabbitMQ encapsulated the configuration and extant extant client for a 34 // queue server 35 // 36 type RabbitMQ struct { 37 url *url.URL // amqp URL to be used for the rmq Server 38 Identity string // A URL stripped of the user name and password, making it safe for logging etc 39 exchange string 40 mgmt *url.URL // URL for the management interface on the rmq 41 host string // The hostname that was specified for the RMQ server 42 user string // user name for the management interface on rmq 43 pass string // password for the management interface on rmq 44 transport *http.Transport // Custom transport to allow for connections to be actively closed 45 wrapper *Wrapper // Decryption infoprmation for messages with encrypted payloads 46 } 47 48 // DefaultStudioRMQExchange is the topic name used within RabbitMQ for StudioML based message queuing 49 const DefaultStudioRMQExchange = "StudioML.topic" 50 51 // NewRabbitMQ takes the uri identifing a server and will configure the client 52 // data structure needed to call methods against the server 53 // 54 // The order of these two parameters needs to reflect key, value pair that 55 // the GetKnown function returns 56 // 57 func NewRabbitMQ(uri string, creds string, wrapper *Wrapper) (rmq *RabbitMQ, err kv.Error) { 58 59 ampq, errGo := url.Parse(os.ExpandEnv(uri)) 60 if errGo != nil { 61 return nil, kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime()).With("uri", os.ExpandEnv(uri)) 62 } 63 64 rmq = &RabbitMQ{ 65 // "amqp://guest:guest@localhost:5672/%2F?connection_attempts=50", 66 // "http://127.0.0.1:15672", 67 exchange: DefaultStudioRMQExchange, 68 user: "guest", 69 pass: "guest", 70 host: ampq.Hostname(), 71 wrapper: wrapper, 72 } 73 74 // The Path will have a vhost that has been escaped. The identity does not require a valid URL just a unique 75 // label 76 ampq.Path, _ = url.PathUnescape(ampq.Path) 77 ampq.User = nil 78 ampq.RawQuery = "" 79 ampq.Fragment = "" 80 rmq.Identity = ampq.String() 81 82 userPass := strings.Split(creds, ":") 83 if len(userPass) != 2 { 84 return nil, kv.NewError("Username password missing or malformed").With("stack", stack.Trace().TrimRuntime()).With("creds", creds, "uri", ampq.String()) 85 } 86 ampq.User = url.UserPassword(userPass[0], userPass[1]) 87 88 // Update the fully qualified URL with the credentials 89 rmq.url = ampq 90 91 rmq.user = userPass[0] 92 rmq.pass = userPass[1] 93 rmq.mgmt = &url.URL{ 94 Scheme: "http", 95 User: url.UserPassword(userPass[0], userPass[1]), 96 Host: fmt.Sprintf("%s:%d", rmq.host, 15672), 97 } 98 99 return rmq, nil 100 } 101 func (rmq *RabbitMQ) IsEncrypted() (encrypted bool) { 102 return nil != rmq.wrapper 103 } 104 105 func (rmq *RabbitMQ) URL() (url string) { 106 return rmq.url.String() 107 } 108 109 func (rmq *RabbitMQ) attachQ(name string) (conn *amqp.Connection, ch *amqp.Channel, err kv.Error) { 110 111 conn, errGo := amqp.Dial(rmq.url.String()) 112 if errGo != nil { 113 return nil, nil, kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime()).With("uri", rmq.Identity) 114 } 115 116 if ch, errGo = conn.Channel(); errGo != nil { 117 return nil, nil, kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime()).With("uri", rmq.Identity) 118 } 119 120 if errGo := ch.ExchangeDeclare(name, "topic", true, true, false, false, nil); errGo != nil { 121 return nil, nil, kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime()).With("uri", rmq.Identity).With("exchange", rmq.exchange) 122 } 123 return conn, ch, nil 124 } 125 126 func (rmq *RabbitMQ) attachMgmt(timeout time.Duration) (mgmt *rh.Client, err kv.Error) { 127 user := rmq.mgmt.User.Username() 128 pass, _ := rmq.mgmt.User.Password() 129 130 mgmt, errGo := rh.NewClient(rmq.mgmt.String(), user, pass) 131 if errGo != nil { 132 return nil, kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime()).With("user", user).With("uri", rmq.mgmt).With("exchange", rmq.exchange) 133 } 134 135 if rmq.transport == nil { 136 rmq.transport = &http.Transport{ 137 MaxIdleConns: 1, 138 IdleConnTimeout: timeout, 139 } 140 } 141 mgmt.SetTransport(rmq.transport) 142 143 return mgmt, nil 144 } 145 146 // Refresh will examine the RMQ exchange a extract a list of the queues that relate to 147 // StudioML work from the rmq exchange. 148 // 149 func (rmq *RabbitMQ) Refresh(ctx context.Context, matcher *regexp.Regexp, mismatcher *regexp.Regexp) (known map[string]interface{}, err kv.Error) { 150 151 timeout := time.Duration(time.Minute) 152 if deadline, isPresent := ctx.Deadline(); isPresent { 153 timeout = time.Until(deadline) 154 } 155 156 known = map[string]interface{}{} 157 158 mgmt, err := rmq.attachMgmt(timeout) 159 if err != nil { 160 return known, err 161 } 162 defer func() { 163 rmq.transport.CloseIdleConnections() 164 }() 165 166 binds, errGo := mgmt.ListBindings() 167 if errGo != nil { 168 return known, kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime()).With("uri", rmq.mgmt) 169 } 170 171 for _, b := range binds { 172 if b.Source == DefaultStudioRMQExchange && strings.HasPrefix(b.RoutingKey, "StudioML.") { 173 // Make sure any retrieved Q names match the caller supplied regular expression 174 if matcher != nil { 175 if !matcher.MatchString(b.Destination) { 176 continue 177 } 178 } 179 if mismatcher != nil { 180 // We cannot allow an excluded queue 181 if mismatcher.MatchString(b.Destination) { 182 continue 183 } 184 } 185 queue := fmt.Sprintf("%s?%s", url.PathEscape(b.Vhost), url.PathEscape(b.Destination)) 186 known[queue] = b.Vhost 187 } 188 } 189 190 return known, nil 191 } 192 193 // GetKnown will connect to the rabbitMQ server identified in the receiver, rmq, and will 194 // query it for any queues that match the matcher regular expression 195 // 196 // found contains a map of keys that have an uncredentialed URL, and the value which is the user name and password for the URL 197 // 198 // The URL path is going to be the vhost and the queue name 199 // 200 func (rmq *RabbitMQ) GetKnown(ctx context.Context, matcher *regexp.Regexp, mismatcher *regexp.Regexp) (found map[string]string, err kv.Error) { 201 known, err := rmq.Refresh(ctx, matcher, mismatcher) 202 if err != nil { 203 return nil, err 204 } 205 206 creds := rmq.user + ":" + rmq.pass 207 208 // Construct the found queue reference prefix 209 qURL := rmq.url 210 rmq.url.User = nil 211 qURL.RawQuery = "" 212 213 found = make(map[string]string, len(known)) 214 215 for hostQueue := range known { 216 // Copy the credentials into the value portion of the returned collection 217 // and the uncredentialed URL and queue name into the key portion 218 found[qURL.String()+"?"+strings.TrimPrefix(hostQueue, "%2F?")] = creds 219 } 220 return found, nil 221 } 222 223 // Exists will connect to the rabbitMQ server identified in the receiver, rmq, and will 224 // query it to see if the queue identified by the studio go runner subscription exists 225 // 226 func (rmq *RabbitMQ) Exists(ctx context.Context, subscription string) (exists bool, err kv.Error) { 227 destHost := strings.Split(subscription, "?") 228 if len(destHost) != 2 { 229 return false, kv.NewError("subscription supplied was not question-mark separated").With("stack", stack.Trace().TrimRuntime()).With("subscription", subscription) 230 } 231 232 vhost, errGo := url.PathUnescape(destHost[0]) 233 if errGo != nil { 234 return false, kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime()).With("subscription", subscription).With("vhost", destHost[0]) 235 } 236 queue, errGo := url.PathUnescape(destHost[1]) 237 if errGo != nil { 238 return false, kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime()).With("subscription", subscription).With("queue", destHost[1]) 239 } 240 241 mgmt, err := rmq.attachMgmt(15 * time.Second) 242 if err != nil { 243 return false, err 244 } 245 defer func() { 246 rmq.transport.CloseIdleConnections() 247 }() 248 249 if _, errGo = mgmt.GetQueue(vhost, queue); errGo != nil { 250 if response, ok := errGo.(rh.ErrorResponse); ok && response.StatusCode == 404 { 251 return false, nil 252 } 253 return false, kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime()).With("uri", rmq.mgmt) 254 } 255 256 return true, nil 257 } 258 259 // Work will connect to the rabbitMQ server identified in the receiver, rmq, and will see if any work 260 // can be found on the queue identified by the go runner subscription and present work 261 // to the handler for processing 262 // 263 func (rmq *RabbitMQ) Work(ctx context.Context, qt *QueueTask) (msgProcessed bool, resource *Resource, err kv.Error) { 264 265 splits := strings.SplitN(qt.Subscription, "?", 2) 266 if len(splits) != 2 { 267 return false, nil, kv.NewError("malformed rmq subscription").With("stack", stack.Trace().TrimRuntime()).With("subscription", qt.Subscription) 268 } 269 270 conn, ch, err := rmq.attachQ(rmq.exchange) 271 if err != nil { 272 return false, nil, err 273 } 274 defer func() { 275 ch.Close() 276 conn.Close() 277 }() 278 279 queue, errGo := url.PathUnescape(splits[1]) 280 if errGo != nil { 281 return false, nil, kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime()).With("subscription", qt.Subscription) 282 } 283 queue = strings.Trim(queue, "/") 284 285 msg, ok, errGo := ch.Get(queue, false) 286 if errGo != nil { 287 return false, nil, kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime()).With("queue", queue) 288 } 289 if !ok { 290 return false, nil, nil 291 } 292 293 qt.Msg = msg.Body 294 qt.ShortQName = queue 295 296 rsc, ack, err := qt.Handler(ctx, qt) 297 if ack { 298 if errGo := msg.Ack(false); errGo != nil { 299 return false, rsc, kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime()).With("subscription", qt.Subscription) 300 } 301 } else { 302 msg.Nack(false, true) 303 } 304 305 return true, rsc, err 306 } 307 308 // This file contains the implementation of a test subsystem 309 // for deploying rabbitMQ in test scenarios where it 310 // has been installed for the purposes of running end-to-end 311 // tests related to queue handling and state management 312 313 var ( 314 testQErr = kv.NewError("uninitialized").With("stack", stack.Trace().TrimRuntime()) 315 qCheck sync.Once 316 ) 317 318 // PingRMQServer is used to validate the a RabbitMQ server is alive and active on the administration port. 319 // 320 // amqpURL is the standard client amqp uri supplied by a caller. amqpURL will be parsed and converted into 321 // the administration endpoint and then tested. 322 // 323 func PingRMQServer(amqpURL string) (err kv.Error) { 324 325 qCheck.Do(func() { 326 327 if len(amqpURL) == 0 { 328 testQErr = kv.NewError("amqpURL was not specified on the command line, or as an env var, cannot start rabbitMQ").With("stack", stack.Trace().TrimRuntime()) 329 return 330 } 331 332 q := os.ExpandEnv(amqpURL) 333 334 uri, errGo := amqp.ParseURI(q) 335 if errGo != nil { 336 testQErr = kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime()) 337 return 338 } 339 uri.Port += 10000 340 341 // Start by making sure that when things were started we saw a rabbitMQ configured 342 // on the localhost. If so then check that the rabbitMQ started automatically as a result of 343 // the Dockerfile_standalone, or Dockerfile_workstation setup 344 // 345 rmqc, errGo := rh.NewClient("http://"+uri.Host+":"+strconv.Itoa(uri.Port), uri.Username, uri.Password) 346 if errGo != nil { 347 testQErr = kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime()) 348 return 349 } 350 351 rmqc.SetTransport(&http.Transport{ 352 ResponseHeaderTimeout: time.Duration(15 * time.Second), 353 }) 354 rmqc.SetTimeout(time.Duration(15 * time.Second)) 355 356 // declares an exchange for the queues 357 exhangeSettings := rh.ExchangeSettings{ 358 Type: "topic", 359 Durable: true, 360 AutoDelete: true, 361 } 362 resp, errGo := rmqc.DeclareExchange("/", DefaultStudioRMQExchange, exhangeSettings) 363 if errGo != nil { 364 testQErr = kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime()) 365 return 366 } 367 resp.Body.Close() 368 369 // declares a queue 370 qn := "rmq_runner_test_" + xid.New().String() 371 if resp, errGo = rmqc.DeclareQueue("/", qn, rh.QueueSettings{Durable: false}); errGo != nil { 372 testQErr = kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime()) 373 return 374 } 375 resp.Body.Close() 376 377 bi := rh.BindingInfo{ 378 Source: DefaultStudioRMQExchange, 379 Destination: qn, 380 DestinationType: "queue", 381 RoutingKey: "StudioML." + qn, 382 Arguments: map[string]interface{}{}, 383 } 384 385 if resp, errGo = rmqc.DeclareBinding("/", bi); errGo != nil { 386 testQErr = kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime()) 387 return 388 } 389 resp.Body.Close() 390 391 testQErr = nil 392 }) 393 394 return testQErr 395 } 396 397 // QueueDeclare is a shim method for creating a queue within the rabbitMQ 398 // server defined by the receiver 399 // 400 func (rmq *RabbitMQ) QueueDeclare(qName string) (err kv.Error) { 401 conn, ch, err := rmq.attachQ(rmq.exchange) 402 if err != nil { 403 return err 404 } 405 defer func() { 406 ch.Close() 407 conn.Close() 408 }() 409 410 _, errGo := ch.QueueDeclare( 411 qName, // name 412 false, // durable 413 false, // delete when unused 414 false, // exclusive 415 false, // no-wait 416 nil, // arguments 417 ) 418 if errGo != nil { 419 return kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime()).With("qName", qName).With("uri", rmq.mgmt).With("exchange", rmq.exchange) 420 } 421 422 if errGo = ch.QueueBind(qName, "StudioML."+qName, "StudioML.topic", false, nil); errGo != nil { 423 return kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime()).With("qName", qName).With("uri", rmq.mgmt).With("exchange", rmq.exchange) 424 } 425 426 return nil 427 } 428 429 // QueueDestroy is a shim method for creating a queue within the rabbitMQ 430 // server defined by the receiver 431 // 432 func (rmq *RabbitMQ) QueueDestroy(qName string) (err kv.Error) { 433 conn, ch, err := rmq.attachQ(rmq.exchange) 434 if err != nil { 435 return err 436 } 437 defer func() { 438 ch.Close() 439 conn.Close() 440 }() 441 442 _, errGo := ch.QueueDelete( 443 qName, // name 444 false, // ifUnused 445 false, // ifEmpty 446 false, // noWait 447 ) 448 if errGo != nil { 449 return kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime()).With("qName", qName).With("uri", rmq.mgmt).With("exchange", rmq.exchange) 450 } 451 452 return nil 453 } 454 455 // One would typically keep a channel of publishings, a sequence number, and a 456 // set of unacknowledged sequence numbers and loop until the publishing channel 457 // is closed. 458 func confirmOne(confirms <-chan amqp.Confirmation) { 459 460 if confirmed := <-confirms; !confirmed.Ack { 461 fmt.Println("failed delivery of delivery tag: ", confirmed, "stack", stack.Trace().TrimRuntime()) 462 } 463 } 464 465 // Publish is a shim method for tests to use for sending requeues to a queue 466 // 467 func (rmq *RabbitMQ) Publish(routingKey string, contentType string, msg []byte) (err kv.Error) { 468 conn, ch, err := rmq.attachQ(rmq.exchange) 469 if err != nil { 470 return err 471 } 472 defer func() { 473 ch.Close() 474 conn.Close() 475 }() 476 477 errGo := ch.Confirm(false) 478 if errGo != nil { 479 return kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime()).With("routingKey", routingKey).With("uri", rmq.mgmt).With("exchange", rmq.exchange) 480 } 481 482 confirms := ch.NotifyPublish(make(chan amqp.Confirmation, 1)) 483 484 defer confirmOne(confirms) 485 486 errGo = ch.Publish( 487 rmq.exchange, // exchange 488 routingKey, // routing key 489 false, // mandatory 490 false, // immediate 491 amqp.Publishing{ 492 ContentType: contentType, 493 Body: msg, 494 }) 495 if errGo != nil { 496 return kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime()).With("routingKey", routingKey).With("uri", rmq.mgmt).With("exchange", rmq.exchange) 497 } 498 return nil 499 } 500 501 // HasWork will look at the SQS queue to see if there is any pending work. The function 502 // is called in an attempt to see if there is any point in processing new work without a 503 // lot of overhead. In the case of RabbitMQ at the moment we always assume there is work. 504 // 505 func (rmq *RabbitMQ) HasWork(ctx context.Context, subscription string) (hasWork bool, err kv.Error) { 506 return true, nil 507 } 508 509 // Responder is used to open a connection to an existing response queue if 510 // one was made available and also to provision a channel into which the 511 // runner can place report messages 512 func (rmq *RabbitMQ) Responder(ctx context.Context, subscription string) (sender chan *runnerReports.Report, err kv.Error) { 513 exists, err := rmq.Exists(ctx, subscription) 514 if !exists { 515 return nil, err 516 } 517 518 // Open the queue and if this cannot be done exit with the error 519 conn, ch, err := rmq.attachQ(subscription) 520 if err != nil { 521 return nil, err 522 } 523 524 sender = make(chan *runnerReports.Report, 1) 525 go func() { 526 defer conn.Close() 527 for { 528 select { 529 case data := <-sender: 530 if data == nil { 531 // If the responder channel is closed then there is nothing left 532 // to report so we stop 533 return 534 } 535 buf, errGo := prototext.Marshal(data) 536 if errGo != nil { 537 fmt.Println(kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime()).Error()) 538 continue 539 } 540 msg := amqp.Publishing{ 541 DeliveryMode: amqp.Persistent, 542 Timestamp: time.Now(), 543 ContentType: "text/plain", 544 Body: buf, 545 } 546 if err := ch.Publish(subscription, subscription, true, true, msg); err != nil { 547 fmt.Println(err.Error()) 548 } 549 continue 550 case <-ctx.Done(): 551 return 552 } 553 } 554 }() 555 return sender, err 556 }