github.com/zooyer/miskit@v1.0.71/smq/rabbit.go (about)

     1  package smq
     2  
     3  import (
     4  	"fmt"
     5  	"github.com/streadway/amqp"
     6  	"sync"
     7  )
     8  
     9  type Exchange struct {
    10  	Type  string // 交换机类型
    11  	Name  string // 交换机名称
    12  	Key   string // key值
    13  	Queue string // 队列名称
    14  }
    15  
    16  type rabbit struct {
    17  	Exchange
    18  	mutex      sync.RWMutex
    19  	broker     string
    20  	channel    *amqp.Channel
    21  	connection *amqp.Connection
    22  }
    23  
    24  func NewRabbitMQ(broker string, exchange Exchange) *rabbit {
    25  	return &rabbit{
    26  		broker:   broker,
    27  		Exchange: exchange,
    28  	}
    29  }
    30  
    31  func (r *rabbit) Connect() (err error) {
    32  	r.mutex.Lock()
    33  	defer r.mutex.Unlock()
    34  
    35  	if r.connection == nil {
    36  		if r.connection, err = amqp.Dial(r.broker); err != nil {
    37  			return
    38  		}
    39  	}
    40  
    41  	if r.channel == nil {
    42  		if r.channel, err = r.connection.Channel(); err != nil {
    43  			return
    44  		}
    45  	}
    46  
    47  	return
    48  }
    49  
    50  func (r *rabbit) Pub(topic string, message []byte, option ...Option) (err error) {
    51  	options := options(option...)
    52  
    53  	// 用于检查队列是否存在,已经存在不需要重复声明
    54  	if _, err = r.channel.QueueDeclarePassive(r.Queue, true, false, false, true, nil); err != nil {
    55  		// 队列不存在,声明队列
    56  		// name:队列名称;durable:是否持久化,队列存盘,true服务重启后信息不会丢失,影响性能;autoDelete:是否自动删除;noWait:是否非阻塞,
    57  		// true为是,不等待RMQ返回信息;args:参数,传nil即可;exclusive:是否设置排他
    58  		if _, err = r.channel.QueueDeclare(r.Queue, true, false, false, true, nil); err != nil {
    59  			return
    60  		}
    61  	}
    62  
    63  	// 队列绑定
    64  	if err = r.channel.QueueBind(r.Queue, r.Key, r.Name, true, nil); err != nil {
    65  		return
    66  	}
    67  
    68  	// 用于检查交换机是否存在,已经存在不需要重复声明
    69  	if err = r.channel.ExchangeDeclarePassive(r.Name, r.Type, true, false, false, true, nil); err != nil {
    70  		// 注册交换机
    71  		// name:交换机名称,kind:交换机类型,durable:是否持久化,队列存盘,true服务重启后信息不会丢失,影响性能;autoDelete:是否自动删除;
    72  		// noWait:是否非阻塞, true为是,不等待RMQ返回信息;args:参数,传nil即可; internal:是否为内部
    73  		if err = r.channel.ExchangeDeclare(r.Name, r.Type, true, false, false, true, nil); err != nil {
    74  			return
    75  		}
    76  	}
    77  
    78  	// 发送任务消息
    79  	var publish = amqp.Publishing{
    80  		ContentType: "text/plain",
    81  		Body:        message,
    82  		Expiration:  fmt.Sprint(options.TTL.Milliseconds()),
    83  	}
    84  	if err = r.channel.Publish(r.Name, r.Key, false, false, publish); err != nil {
    85  		return
    86  	}
    87  
    88  	return
    89  }
    90  
    91  func (r *rabbit) Sub(topic string, subscriber Subscriber) (err error) {
    92  	// 用于检查队列是否存在,已经存在不需要重复声明
    93  	if _, err = r.channel.QueueDeclarePassive(r.Queue, true, false, false, true, nil); err != nil {
    94  		// 队列不存在,声明队列
    95  		// name:队列名称;durable:是否持久化,队列存盘,true服务重启后信息不会丢失,影响性能;autoDelete:是否自动删除;noWait:是否非阻塞,
    96  		// true为是,不等待RMQ返回信息;args:参数,传nil即可;exclusive:是否设置排他
    97  		if _, err = r.channel.QueueDeclare(r.Queue, true, false, false, true, nil); err != nil {
    98  			return
    99  		}
   100  	}
   101  
   102  	// 队列绑定
   103  	if err = r.channel.QueueBind(r.Queue, r.Key, r.Name, true, nil); err != nil {
   104  		return
   105  	}
   106  
   107  	// 获取消费通道,确保rabbitMQ一个一个发送消息
   108  	if err = r.channel.Qos(1, 0, true); err != nil {
   109  		return
   110  	}
   111  	consume, err := r.channel.Consume(r.Queue, "", false, false, false, false, nil)
   112  	if err != nil {
   113  		return
   114  	}
   115  
   116  	go func() {
   117  		for msg := range consume {
   118  			var multiple bool
   119  			if err = subscriber(r, topic, msg.Body, Options{}); err != nil {
   120  				multiple = true
   121  				// TODO callback error
   122  			}
   123  			if err = msg.Ack(multiple); err != nil {
   124  				// TODO ack error
   125  			}
   126  		}
   127  	}()
   128  
   129  	return
   130  }
   131  
   132  func (r *rabbit) Close() (err error) {
   133  	if err = r.channel.Close(); err != nil {
   134  		return
   135  	}
   136  
   137  	if err = r.connection.Close(); err != nil {
   138  		return
   139  	}
   140  
   141  	return
   142  }