gitee.com/h79/goutils@v1.22.10/mq/pull_consumer.go (about)

     1  package mq
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"fmt"
     7  	"gitee.com/h79/goutils/common/result"
     8  	"gitee.com/h79/goutils/common/system"
     9  	"github.com/apache/rocketmq-client-go/v2"
    10  	"github.com/apache/rocketmq-client-go/v2/consumer"
    11  	"github.com/apache/rocketmq-client-go/v2/primitive"
    12  	"time"
    13  )
    14  
    15  var _ Consumer = (*pullConsumer)(nil)
    16  
    17  type pullConsumer struct {
    18  	cfg      Config
    19  	stop     bool
    20  	consumer rocketmq.PullConsumer
    21  }
    22  
    23  func NewPullConsumer(cfg Config) Consumer {
    24  	return &pullConsumer{cfg: cfg, stop: false}
    25  }
    26  
    27  func (con *pullConsumer) Start() error {
    28  
    29  	p, e := rocketmq.NewPullConsumer(
    30  		consumer.WithGroupName(con.cfg.Client.GroupName),
    31  		consumer.WithNsResolver(primitive.NewPassthroughResolver([]string{con.cfg.Client.Server})),
    32  		consumer.WithCredentials(primitive.Credentials{
    33  			AccessKey:     con.cfg.Credentials.AccessKey,
    34  			SecretKey:     con.cfg.Credentials.SecretKey,
    35  			SecurityToken: con.cfg.Credentials.SecretToken},
    36  		),
    37  		consumer.WithInstance(con.cfg.Client.InstanceName),
    38  		consumer.WithNameServerDomain(con.cfg.Client.Domain),
    39  	)
    40  	if e != nil {
    41  		return result.Errorf(result.ErrMqInitInternal, "[MQ] Init pull consumer failed, error: %v", e).Log()
    42  	}
    43  
    44  	if e := p.Start(); e != nil {
    45  		return result.Errorf(result.ErrMqStartInternal, "[MQ] Start pull consumer failed, error: %v", e).Log()
    46  	}
    47  	con.consumer = p
    48  	return nil
    49  }
    50  
    51  func (con *pullConsumer) Stop() {
    52  	con.stop = true
    53  	_ = con.consumer.Shutdown()
    54  }
    55  
    56  func (con *pullConsumer) Subscribe(topic Topic, exp Expression, call func(msg *Message) Status) error {
    57  
    58  	ctx := context.Background()
    59  	queue := primitive.MessageQueue{
    60  		Topic:      topic.Topic,
    61  		BrokerName: topic.BrokerName, // replace with your broker name. otherwise, pull will failed.
    62  		QueueId:    topic.QueueId,
    63  	}
    64  	offset := int64(0)
    65  	for {
    66  		if con.stop {
    67  			return nil
    68  		}
    69  		resp, err := con.consumer.PullFrom(ctx, queue, offset, 10)
    70  		if err != nil {
    71  			if errors.Is(err, rocketmq.ErrRequestTimeout) {
    72  				fmt.Printf("timeout \n")
    73  				time.Sleep(1 * time.Second)
    74  				continue
    75  			}
    76  			fmt.Printf("unexpectable err: %v \n", err)
    77  			return err
    78  		}
    79  		if resp.Status == primitive.PullFound {
    80  			callMsg(resp.GetMessageExts(), call)
    81  		}
    82  		offset = resp.NextBeginOffset
    83  
    84  		select {
    85  		case <-system.Closed():
    86  			return nil
    87  		default:
    88  		}
    89  	}
    90  }
    91  
    92  func callMsg(ext []*primitive.MessageExt, call func(msg *Message) Status) {
    93  	for _, msg := range ext {
    94  		m := Message{
    95  			MsgId: msg.MsgId,
    96  		}
    97  		call(&m)
    98  	}
    99  }