使用redigo实现 pub/sub(代码示例)

package main

import (
	"context"
	"fmt"
	"log"
	"strconv"
	"time"

	"github.com/gomodule/redigo/redis"
)

// ConsumeFunc consumes message at the channel.
type ConsumeFunc func(channel string, message []byte) error

// RedisClient represents a redis client with connection pool.
type RedisClient struct {
	pool *redis.Pool
}

// NewRedisClient returns a RedisClient.
func NewRedisClient(addr string, db int, passwd string) *RedisClient {
	pool := &redis.Pool{
		MaxIdle:     10,
		IdleTimeout: 300 * time.Second,
		Dial: func() (redis.Conn, error) {
			c, err := redis.Dial("tcp", addr, redis.DialPassword(passwd), redis.DialDatabase(db))
			if err != nil {
				return nil, err
			}
			return c, nil
		},
		TestOnBorrow: func(c redis.Conn, t time.Time) error {
			if time.Since(t) < time.Minute {
				return nil
			}
			_, err := c.Do("PING")
			return err
		},
	}
	log.Printf("new redis pool at %s", addr)
	client := &RedisClient{
		pool: pool,
	}
	return client
}

// Close closes connection pool.
func (r *RedisClient) Close() error {
	err := r.pool.Close()
	return err
}

// Publish publishes message to channel.
func (r *RedisClient) Publish(channel, message string) (int, error) {
	c := r.pool.Get()
	defer c.Close()
	n, err := redis.Int(c.Do("PUBLISH", channel, message))
	if err != nil {
		return 0, fmt.Errorf("redis publish %s %s, err: %v", channel, message, err)
	}
	return n, nil
}

// Subscribe subscribes messages at the channels.
func (r *RedisClient) Subscribe(ctx context.Context, consume ConsumeFunc, channel ...string) error {
	psc := redis.PubSubConn{Conn: r.pool.Get()}

	log.Printf("redis pubsub subscribe channel: %v", channel)
	if err := psc.Subscribe(redis.Args{}.AddFlat(channel)...); err != nil {
		return err
	}
	done := make(chan error, 1)
	// start a new goroutine to receive message
	go func() {
		defer psc.Close()
		for {
			switch msg := psc.Receive().(type) {
			case error:
				done <- fmt.Errorf("redis pubsub receive err: %v", msg)
				return
			case redis.Message:
				if err := consume(msg.Channel, msg.Data); err != nil {
					done <- err
					return
				}
			case redis.Subscription:
				if msg.Count == 0 {
					// all channels are unsubscribed
					done <- nil
					return
				}
			}
		}
	}()

	ch <- 0

	// health check
	tick := time.NewTicker(time.Minute)
	defer tick.Stop()
	for {
		select {
		case <-ctx.Done():
			if err := psc.Unsubscribe(); err != nil {
				return fmt.Errorf("redis pubsub unsubscribe err: %v", err)
			}
			return nil
		case err := <-done:
			return err
		case <-tick.C:
			if err := psc.Ping(""); err != nil {
				return err
			}
		}
	}

}

func myConsumer(channel string, message []byte) error {
	log.Printf("receive message[%s] at the channel[%s]\n", string(message), channel)
	return nil
}

// ch 用于保证发布线程在订阅线程启动成功后才开始发布消息
var ch = make(chan int)

func main() {
	redisClient := NewRedisClient("127.0.0.1:6379", 0, "pj8888")

	go func() {
		var subscriber int
		<-ch
		for i := 0; i < 3; i++ {
			subscriber, _ = redisClient.Publish("testx", "hello world"+strconv.Itoa(i))
			log.Printf("there is %d subscriber.\n", subscriber)
		}

	}()

	ctx, cancel := context.WithCancel(context.Background())
	err := redisClient.Subscribe(ctx,
		func(channel string, message []byte) error {
			log.Printf("receive message[%s] at the channel[%s]\n", string(message), channel)
			if string(message) == "goodbye" {
				cancel()
			}
			return nil
		},
		"testx")
	if err != nil {
		fmt.Printf("get error: %v\n", err)
	}

	redisClient.Close()

	return
}

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值