github.com/livekit/protocol@v1.16.1-0.20240517185851-47e4c6bba773/utils/messaging.go (about) 1 // Copyright 2023 LiveKit, Inc. 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 15 package utils 16 17 import ( 18 "context" 19 "crypto/sha256" 20 "encoding/base64" 21 "math/rand" 22 "time" 23 24 "github.com/eapache/channels" 25 "github.com/prometheus/client_golang/prometheus" 26 "github.com/redis/go-redis/v9" 27 "google.golang.org/protobuf/proto" 28 ) 29 30 const lockExpiration = time.Second * 5 31 32 var ( 33 PromMessageBusCounter = prometheus.NewCounterVec( 34 prometheus.CounterOpts{ 35 Namespace: "livekit", 36 Subsystem: "messagebus", 37 Name: "messages", 38 }, 39 []string{"type", "status"}, 40 ) 41 ) 42 43 func init() { 44 prometheus.MustRegister(PromMessageBusCounter) 45 } 46 47 type MessageBus interface { 48 Subscribe(ctx context.Context, channel string) (PubSub, error) 49 // SubscribeQueue is like subscribe, but ensuring only a single instance gets to process the message 50 SubscribeQueue(ctx context.Context, channel string) (PubSub, error) 51 Publish(ctx context.Context, channel string, msg proto.Message) error 52 } 53 54 type PubSub interface { 55 Channel() <-chan interface{} 56 Payload(msg interface{}) []byte 57 Close() error 58 } 59 60 type RedisMessageBus struct { 61 rc redis.UniversalClient 62 } 63 64 func NewRedisMessageBus(rc redis.UniversalClient) MessageBus { 65 return &RedisMessageBus{rc: rc} 66 } 67 68 func (r *RedisMessageBus) Lock(ctx context.Context, key string, expiration time.Duration) (bool, error) { 69 return r.rc.SetNX(ctx, key, rand.Int(), expiration).Result() 70 } 71 72 func (r *RedisMessageBus) Subscribe(ctx context.Context, channel string) (PubSub, error) { 73 ps := r.rc.Subscribe(ctx, channel) 74 return &RedisPubSub{ 75 ps: ps, 76 c: channels.Wrap(ps.Channel()).Out(), 77 done: make(chan struct{}, 1), 78 }, nil 79 } 80 81 func (r *RedisMessageBus) SubscribeQueue(ctx context.Context, channel string) (PubSub, error) { 82 sub := r.rc.Subscribe(ctx, channel) 83 c := make(chan *redis.Message, 100) // same chan size as redis pubsub 84 ps := &RedisPubSub{ 85 ps: sub, 86 c: channels.Wrap(c).Out(), 87 done: make(chan struct{}, 1), 88 } 89 90 go func() { 91 for { 92 select { 93 case <-ps.done: 94 return 95 case msg := <-sub.Channel(): 96 sha := sha256.Sum256([]byte(msg.Payload)) 97 hash := base64.StdEncoding.EncodeToString(sha[:]) 98 acquired, _ := r.Lock(ctx, hash, lockExpiration) 99 if acquired { 100 PromMessageBusCounter.WithLabelValues("in", "success").Add(1) 101 c <- msg 102 } 103 } 104 } 105 }() 106 107 return ps, nil 108 } 109 110 func (r *RedisMessageBus) Publish(ctx context.Context, channel string, msg proto.Message) error { 111 b, err := proto.Marshal(msg) 112 if err != nil { 113 PromMessageBusCounter.WithLabelValues("out", "failure").Add(1) 114 return err 115 } 116 117 err = r.rc.Publish(ctx, channel, b).Err() 118 if err == nil { 119 PromMessageBusCounter.WithLabelValues("out", "success").Add(1) 120 } else { 121 PromMessageBusCounter.WithLabelValues("out", "failure").Add(1) 122 } 123 124 return err 125 } 126 127 type RedisPubSub struct { 128 ps *redis.PubSub 129 c <-chan interface{} 130 done chan struct{} 131 } 132 133 func (r *RedisPubSub) Channel() <-chan interface{} { 134 return r.c 135 } 136 137 func (r *RedisPubSub) Payload(msg interface{}) []byte { 138 return []byte(msg.(*redis.Message).Payload) 139 } 140 141 func (r *RedisPubSub) Close() error { 142 r.done <- struct{}{} 143 return r.ps.Close() 144 }