func (rc RequestChan) publish() error {
var (
err error
out *amqp.Channel
)
for r := range rc {
if r == nil {
break
}
for out == nil {
out, err = dial()
if err != nil {
out = nil
log.Println(err)
time.Sleep(*backoff)
}
}
err = out.Publish(*exchange, *key, false, false,
amqp.Publishing{
Headers: amqp.Table(r.Headers),
ContentType: r.ContentType,
Body: r.Body,
},
)
if err != nil {
return err
}
}
return nil
}
func (rc requestCh) publish() error {
var (
err error
out *amqp.Channel
)
for r := range rc {
if r == nil {
break
}
for out == nil {
out, err = dial()
if err != nil {
out = nil
log.Println(err)
time.Sleep(*backoff)
}
}
err = out.Publish(*exchange, *key, false, false,
amqp.Publishing{
Headers: amqp.Table(r.headers),
DeliveryMode: amqp.Persistent,
Body: []byte(r.body),
},
)
if err != nil {
return err
}
}
return nil
}
func NewRequestFromProto(req proto.Message) *Request {
protoBytes := make([]byte, 0)
if req != nil {
protoBytes, _ = proto.Marshal(req)
}
result := NewRequestFromDelivery(amqp.Delivery{
Body: protoBytes,
ContentType: "application/octetstream",
Headers: amqp.Table(make(map[string]interface{})),
})
result.unmarshaledData = req
return result
}
// Publish places a new message on the default queue
func (amqpBroker *AMQPBroker) Publish(signature *signatures.TaskSignature) error {
conn, channel, _, confirmsChan, err := amqpBroker.open()
if err != nil {
return err
}
defer amqpBroker.close(channel, conn)
message, err := json.Marshal(signature)
if err != nil {
return fmt.Errorf("JSON Encode Message: %v", err)
}
signature.AdjustRoutingKey(
amqpBroker.config.ExchangeType,
amqpBroker.config.BindingKey,
amqpBroker.config.DefaultQueue,
)
if err := channel.Publish(
amqpBroker.config.Exchange, // exchange
signature.RoutingKey, // routing key
false, // mandatory
false, // immediate
amqp.Publishing{
Headers: amqp.Table(signature.Headers),
ContentType: "application/json",
Body: message,
DeliveryMode: amqp.Persistent,
},
); err != nil {
return err
}
confirmed := <-confirmsChan
if confirmed.Ack {
return nil
}
return fmt.Errorf("Failed delivery of delivery tag: %v", confirmed.DeliveryTag)
}
// Connects to the message queue, opens a channel, declares a queue
func (amqpBroker *AMQPBroker) open() (*amqp.Connection, *amqp.Channel, amqp.Queue, <-chan amqp.Confirmation, error) {
var (
conn *amqp.Connection
channel *amqp.Channel
queue amqp.Queue
err error
)
// Connect
// From amqp docs: DialTLS will use the provided tls.Config when it encounters an amqps:// scheme
// and will dial a plain connection when it encounters an amqp:// scheme.
conn, err = amqp.DialTLS(amqpBroker.config.Broker, amqpBroker.config.TLSConfig)
if err != nil {
return conn, channel, queue, nil, fmt.Errorf("Dial: %s", err)
}
// Open a channel
channel, err = conn.Channel()
if err != nil {
return conn, channel, queue, nil, fmt.Errorf("Channel: %s", err)
}
// Declare an exchange
if err := channel.ExchangeDeclare(
amqpBroker.config.Exchange, // name of the exchange
amqpBroker.config.ExchangeType, // type
true, // durable
false, // delete when complete
false, // internal
false, // noWait
nil, // arguments
); err != nil {
return conn, channel, queue, nil, fmt.Errorf("Exchange Declare: %s", err)
}
// Declare a queue
queue, err = channel.QueueDeclare(
amqpBroker.config.DefaultQueue, // name
true, // durable
false, // delete when unused
false, // exclusive
false, // no-wait
nil, // arguments
)
if err != nil {
return conn, channel, queue, nil, fmt.Errorf("Queue Declare: %s", err)
}
// Bind the queue
if err := channel.QueueBind(
queue.Name, // name of the queue
amqpBroker.config.BindingKey, // binding key
amqpBroker.config.Exchange, // source exchange
false, // noWait
amqp.Table(amqpBroker.config.QueueBindingArguments), // arguments
); err != nil {
return conn, channel, queue, nil, fmt.Errorf("Queue Bind: %s", err)
}
// Enable publish confirmations
if err := channel.Confirm(false); err != nil {
return conn, channel, queue, nil, fmt.Errorf("Channel could not be put into confirm mode: %s", err)
}
return conn, channel, queue, channel.NotifyPublish(make(chan amqp.Confirmation, 1)), nil
}
func main() {
server := os.Getenv("AMQP_URL")
if server == "" {
server = os.ExpandEnv("amqp://$AMQP_USER:[email protected]:5672")
}
timeout := 5 * time.Second
clientID, _ := os.Hostname()
queue := "scanner"
mainCmd := &cobra.Command{
Use: "amqpc",
}
p := mainCmd.PersistentFlags()
p.StringVarP(&server, "server", "S", server, "server address")
p.DurationVarP(&timeout, "timeout", "", timeout, "timeout for commands")
p.StringVarP(&clientID, "id", "", clientID, "client ID")
p.StringVarP(&queue, "queue", "q", queue, "queue name to publish")
appID := queue
var noCompress bool
pubCmd := &cobra.Command{
Use: "pub",
Aliases: []string{"publish", "send", "write"},
Run: func(_ *cobra.Command, args []string) {
c, err := newClient(server, queue)
if err != nil {
log.Fatal(err)
}
defer c.Close()
if err = c.Confirm(false); err != nil {
log.Fatal(err)
}
confirms := c.NotifyPublish(make(chan amqp.Confirmation, 1))
returns := c.NotifyReturn(make(chan amqp.Return, 1))
var sendCount int
tbl := make(map[string]interface{}, 1)
for _, arg := range args {
for k := range tbl {
delete(tbl, k)
}
var r io.ReadCloser
mimeType, contentEncoding := "text/plain", ""
if strings.HasPrefix(arg, "@") {
arg = arg[1:]
if arg == "-" {
r = os.Stdin
} else if fh, err := os.Open(arg); err != nil {
log.Fatal(err)
} else {
if noCompress {
r = fh
} else {
pr, pw := io.Pipe()
go func() {
defer fh.Close()
gw := gzip.NewWriter(pw)
if _, err := io.Copy(gw, fh); err != nil {
pw.CloseWithError(err)
return
}
pw.CloseWithError(gw.Close())
}()
r = pr
contentEncoding = "application/gzip"
}
tbl["FileName"] = arg
if err := amqp.Table(tbl).Validate(); err != nil {
log.Fatal(err)
}
mimeType = mime.TypeByExtension(filepath.Ext(arg))
}
} else {
r = ioutil.NopCloser(strings.NewReader(arg))
}
b, err := ioutil.ReadAll(&io.LimitedReader{R: r, N: 256 << 20})
r.Close()
if err != nil {
log.Fatal(err)
}
if mimeType == "" {
if mimeType = magic.MIMEType(b); mimeType == "" {
mimeType = "application/octet-stream"
}
}
if err := c.Publish("", c.Queue.Name, false, false,
amqp.Publishing{
Headers: tbl,
DeliveryMode: amqp.Persistent,
ContentType: mimeType,
ContentEncoding: contentEncoding,
AppId: appID,
Body: b,
},
); err != nil {
log.Fatalf("Publish: %v", err)
}
log.Printf("Sent %q", arg)
sendCount++
}
Loop:
for i := 0; i < sendCount; {
select {
case c, ok := <-confirms:
if !ok {
break Loop
}
if !c.Ack {
log.Printf("couldn't deliver %d", c.DeliveryTag)
} else {
log.Printf("Delivered %d.", c.DeliveryTag)
i++
}
case r, ok := <-returns:
if !ok {
break Loop
}
log.Printf("RETURN: %#v", r)
i++
}
}
},
}
f := pubCmd.Flags()
f.StringVarP(&appID, "app-id", "", appID, "appID")
f.BoolVarP(&noCompress, "no-compress", "", noCompress, "disable file data compression (for slow devices)")
var keepFiles bool
subCmd := &cobra.Command{
Use: "sub",
Aliases: []string{"subscribe", "recv", "receive", "read"},
Run: func(_ *cobra.Command, args []string) {
c, err := newClient(server, queue)
if err != nil {
log.Fatal(err)
}
defer c.Close()
d, err := c.Consume(c.Queue.Name, clientID, false, false, false, false, nil)
if err != nil {
log.Fatal(err)
}
tempDir, err := ioutil.TempDir("", "amqpc-")
if err != nil {
log.Fatal(err)
}
defer os.RemoveAll(tempDir)
var i uint64
for msg := range d {
i++
log.Printf("Received %s with %q from %s@%s.",
msg.MessageId, msg.Headers, msg.UserId, msg.AppId)
var fn string
if fnI := msg.Headers["FileName"]; fnI != nil {
if fn, _ = fnI.(string); fn != "" {
fn = filepath.Base(fn)
}
}
if fn == "" {
var ext string
if exts, err := mime.ExtensionsByType(msg.ContentType); err != nil {
log.Printf("Extension for %q: %v", msg.ContentType)
} else if len(ext) > 0 {
ext = exts[0]
}
fn = fmt.Sprintf("%09d%s", i, ext)
}
fn = filepath.Join(tempDir, fn)
err = receive(fn, msg, args)
if !keepFiles {
os.Remove(fn)
}
if err != nil {
msg.Nack(false, true)
log.Fatal(err)
}
if err := msg.Ack(false); err != nil {
log.Printf("cannot ACK %q: %v", msg, err)
}
}
},
}
subCmd.Flags().BoolVarP(&keepFiles, "keep-files", "x", keepFiles, "keep temporary files")
mainCmd.AddCommand(pubCmd, subCmd)
mainCmd.Execute()
}
