github.com/whoyao/protocol@v0.0.0-20230519045905-2d8ace718ca5/utils/messaging.go (about)

     1  package utils
     2  
     3  import (
     4  	"context"
     5  	"crypto/sha256"
     6  	"encoding/base64"
     7  	"math/rand"
     8  	"time"
     9  
    10  	"github.com/eapache/channels"
    11  	"github.com/prometheus/client_golang/prometheus"
    12  	"github.com/redis/go-redis/v9"
    13  	"google.golang.org/protobuf/proto"
    14  )
    15  
    16  const lockExpiration = time.Second * 5
    17  
    18  var (
    19  	PromMessageBusCounter = prometheus.NewCounterVec(
    20  		prometheus.CounterOpts{
    21  			Namespace: "livekit",
    22  			Subsystem: "messagebus",
    23  			Name:      "messages",
    24  		},
    25  		[]string{"type", "status"},
    26  	)
    27  )
    28  
    29  func init() {
    30  	prometheus.MustRegister(PromMessageBusCounter)
    31  }
    32  
    33  type MessageBus interface {
    34  	Subscribe(ctx context.Context, channel string) (PubSub, error)
    35  	// SubscribeQueue is like subscribe, but ensuring only a single instance gets to process the message
    36  	SubscribeQueue(ctx context.Context, channel string) (PubSub, error)
    37  	Publish(ctx context.Context, channel string, msg proto.Message) error
    38  }
    39  
    40  type PubSub interface {
    41  	Channel() <-chan interface{}
    42  	Payload(msg interface{}) []byte
    43  	Close() error
    44  }
    45  
    46  type RedisMessageBus struct {
    47  	rc redis.UniversalClient
    48  }
    49  
    50  func NewRedisMessageBus(rc redis.UniversalClient) MessageBus {
    51  	return &RedisMessageBus{rc: rc}
    52  }
    53  
    54  func (r *RedisMessageBus) Lock(ctx context.Context, key string, expiration time.Duration) (bool, error) {
    55  	return r.rc.SetNX(ctx, key, rand.Int(), expiration).Result()
    56  }
    57  
    58  func (r *RedisMessageBus) Subscribe(ctx context.Context, channel string) (PubSub, error) {
    59  	ps := r.rc.Subscribe(ctx, channel)
    60  	return &RedisPubSub{
    61  		ps:   ps,
    62  		c:    channels.Wrap(ps.Channel()).Out(),
    63  		done: make(chan struct{}, 1),
    64  	}, nil
    65  }
    66  
    67  func (r *RedisMessageBus) SubscribeQueue(ctx context.Context, channel string) (PubSub, error) {
    68  	sub := r.rc.Subscribe(ctx, channel)
    69  	c := make(chan *redis.Message, 100) // same chan size as redis pubsub
    70  	ps := &RedisPubSub{
    71  		ps:   sub,
    72  		c:    channels.Wrap(c).Out(),
    73  		done: make(chan struct{}, 1),
    74  	}
    75  
    76  	go func() {
    77  		for {
    78  			select {
    79  			case <-ps.done:
    80  				return
    81  			case msg := <-sub.Channel():
    82  				sha := sha256.Sum256([]byte(msg.Payload))
    83  				hash := base64.StdEncoding.EncodeToString(sha[:])
    84  				acquired, _ := r.Lock(ctx, hash, lockExpiration)
    85  				if acquired {
    86  					PromMessageBusCounter.WithLabelValues("in", "success").Add(1)
    87  					c <- msg
    88  				}
    89  			}
    90  		}
    91  	}()
    92  
    93  	return ps, nil
    94  }
    95  
    96  func (r *RedisMessageBus) Publish(ctx context.Context, channel string, msg proto.Message) error {
    97  	b, err := proto.Marshal(msg)
    98  	if err != nil {
    99  		PromMessageBusCounter.WithLabelValues("out", "failure").Add(1)
   100  		return err
   101  	}
   102  
   103  	err = r.rc.Publish(ctx, channel, b).Err()
   104  	if err == nil {
   105  		PromMessageBusCounter.WithLabelValues("out", "success").Add(1)
   106  	} else {
   107  		PromMessageBusCounter.WithLabelValues("out", "failure").Add(1)
   108  	}
   109  
   110  	return err
   111  }
   112  
   113  type RedisPubSub struct {
   114  	ps   *redis.PubSub
   115  	c    <-chan interface{}
   116  	done chan struct{}
   117  }
   118  
   119  func (r *RedisPubSub) Channel() <-chan interface{} {
   120  	return r.c
   121  }
   122  
   123  func (r *RedisPubSub) Payload(msg interface{}) []byte {
   124  	return []byte(msg.(*redis.Message).Payload)
   125  }
   126  
   127  func (r *RedisPubSub) Close() error {
   128  	r.done <- struct{}{}
   129  	return r.ps.Close()
   130  }