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  }