github.com/andy2046/gopie@v0.7.0/pkg/pubsub/pubsub.go (about) 1 // Package pubsub provides a pubsub implementation. 2 package pubsub 3 4 import ( 5 "context" 6 "crypto/rand" 7 "errors" 8 "fmt" 9 "log" 10 "sync" 11 ) 12 13 type ( 14 // PubSub is a Pub/Sub instance for a single project. 15 PubSub struct { 16 projectID string 17 topics map[string]*Topic 18 mu sync.RWMutex 19 } 20 21 // Topic represents a PubSub topic. 22 Topic struct { 23 // The identifier for the topic, 24 // in the format "projects/<projid>/topics/<name>". 25 name string 26 27 inbox chan Message 28 29 stopped bool 30 31 pubSub *PubSub 32 33 subscriptions map[string]*Subscription 34 35 // Errors is the error output channel back to the user. You MUST read from this 36 // channel or the Publish will deadlock when the channel is full. 37 Errors chan PublishError 38 39 numGoroutines int 40 41 wgPublish sync.WaitGroup 42 43 once sync.Once 44 45 mu sync.RWMutex 46 47 wg sync.WaitGroup 48 } 49 50 // Subscription represents a PubSub subscription. 51 Subscription struct { 52 // The identifier for the subscription, 53 // in the format "projects/<projid>/topics/<name>/subscriptions/<name>". 54 name string 55 56 inbox chan Message 57 58 topic *Topic 59 60 stopped bool 61 62 mu sync.RWMutex 63 64 wg sync.WaitGroup 65 66 done chan struct{} 67 68 // numGoroutines is the number of goroutines it will spawn to pull msg concurrently. 69 numGoroutines int 70 } 71 72 // Message represents a Pub/Sub message. 73 Message struct { 74 ID string 75 Data []byte 76 } 77 78 // PublishError is the error generated when it fails to publish a message. 79 PublishError struct { 80 Msg *Message 81 Err error 82 } 83 ) 84 85 func (pe PublishError) Error() string { 86 return fmt.Sprintf("failed to publish message %s -> %s", pe.Msg.ID, pe.Err) 87 } 88 89 // UUID generates uuid. 90 func UUID() string { 91 b := make([]byte, 16) 92 rand.Read(b) 93 return fmt.Sprintf("%x-%x-%x-%x-%x", b[0:4], b[4:6], b[6:8], b[8:10], b[10:]) 94 } 95 96 // New creates a new PubSub. 97 func New(project string) *PubSub { 98 return &PubSub{ 99 projectID: project, 100 topics: make(map[string]*Topic), 101 } 102 } 103 104 // Name returns the full name for the PubSub. 105 func (p *PubSub) Name() string { 106 return fmt.Sprintf("projects/%s", p.projectID) 107 } 108 109 // NewTopic creates a new Topic with the given name, 110 // size is the channel buffer size for topic message chan, 111 // numGoroutines is the number of goroutines it will spawn to push msg concurrently. 112 func (p *PubSub) NewTopic(name string, size int, numGoroutines int) (*Topic, error) { 113 p.mu.RLock() 114 if _, ok := p.topics[name]; ok { 115 p.mu.RUnlock() 116 return nil, errors.New("duplicated topic name") 117 } 118 p.mu.RUnlock() 119 p.mu.Lock() 120 t := &Topic{ 121 name: name, 122 pubSub: p, 123 subscriptions: make(map[string]*Subscription), 124 inbox: make(chan Message, size), 125 Errors: make(chan PublishError, size), 126 numGoroutines: numGoroutines, 127 } 128 p.topics[name] = t 129 p.mu.Unlock() 130 return t, nil 131 } 132 133 // Topic returns the topic by name. 134 func (p *PubSub) Topic(name string) *Topic { 135 p.mu.RLock() 136 defer p.mu.RUnlock() 137 if t, ok := p.topics[name]; ok { 138 return t 139 } 140 return nil 141 } 142 143 // Topics list all the topics in the PubSub. 144 func (p *PubSub) Topics() []string { 145 p.mu.RLock() 146 defer p.mu.RUnlock() 147 ts := make([]string, 0, len(p.topics)) 148 for k := range p.topics { 149 ts = append(ts, k) 150 } 151 return ts 152 } 153 154 // Publish publishes msg to the topic asynchronously. 155 func (t *Topic) Publish(ctx context.Context, msg *Message) error { 156 t.mu.RLock() 157 defer t.mu.RUnlock() 158 if t.stopped { 159 return errors.New("topic stopped") 160 } 161 t.wgPublish.Add(1) 162 go func() { 163 select { 164 case <-ctx.Done(): 165 t.Errors <- PublishError{ 166 msg, 167 ctx.Err(), 168 } 169 case t.inbox <- *msg: 170 } 171 t.wgPublish.Done() 172 }() 173 return nil 174 } 175 176 // Delete removes itself from PubSuband stop it. 177 func (t *Topic) Delete() { 178 t.pubSub.mu.Lock() 179 for k := range t.pubSub.topics { 180 if k == t.name { 181 delete(t.pubSub.topics, k) 182 } 183 } 184 t.pubSub.mu.Unlock() 185 t.Stop() 186 } 187 188 // Name returns the full name for the topic. 189 func (t *Topic) Name() string { 190 return fmt.Sprintf("projects/%s/topics/%s", t.pubSub.projectID, t.name) 191 } 192 193 // Stop stops the topic. 194 func (t *Topic) Stop() { 195 t.mu.Lock() 196 t.stopped = true 197 t.wgPublish.Wait() 198 close(t.inbox) 199 t.mu.Unlock() 200 t.wg.Wait() 201 202 for _, v := range t.subscriptions { 203 go func(s *Subscription) { 204 s.mu.Lock() 205 if s.stopped { 206 s.mu.Unlock() 207 return 208 } 209 s.stopped = true 210 close(s.inbox) 211 s.wg.Wait() 212 close(s.done) 213 s.mu.Unlock() 214 }(v) 215 } 216 } 217 218 func (t *Topic) start() { 219 for { 220 m, open := <-t.inbox 221 if !open { 222 log.Printf("topic %s inbox closed, exit", t.Name()) 223 t.wg.Done() 224 return 225 } 226 t.mu.RLock() 227 subs := make(map[string]*Subscription, len(t.subscriptions)) 228 for k, v := range t.subscriptions { 229 subs[k] = v 230 } 231 t.mu.RUnlock() 232 233 for _, v := range subs { 234 go func(s *Subscription) { 235 s.mu.RLock() 236 if s.stopped { 237 s.mu.RUnlock() 238 return 239 } 240 s.inbox <- m 241 s.mu.RUnlock() 242 }(v) 243 } 244 } 245 } 246 247 // Subscriptions list all the subscriptions to this topic. 248 func (t *Topic) Subscriptions() []string { 249 t.mu.RLock() 250 defer t.mu.RUnlock() 251 sub := make([]string, 0, len(t.subscriptions)) 252 for k := range t.subscriptions { 253 sub = append(sub, k) 254 } 255 return sub 256 } 257 258 // NewSubscription creates a new Subscription to this topic, 259 // numGoroutines is the number of goroutines it will spawn to pull msg concurrently. 260 func (t *Topic) NewSubscription(numGoroutines int) (*Subscription, error) { 261 t.once.Do(func() { 262 for range make([]struct{}, t.numGoroutines) { 263 t.wg.Add(1) 264 go t.start() 265 } 266 }) 267 t.mu.Lock() 268 n := fmt.Sprintf("%s-sub-%s", t.name, UUID()) 269 s := &Subscription{ 270 name: n, 271 inbox: make(chan Message, 10*numGoroutines), 272 topic: t, 273 done: make(chan struct{}), 274 numGoroutines: numGoroutines, 275 } 276 t.subscriptions[n] = s 277 t.mu.Unlock() 278 return s, nil 279 } 280 281 // Subscription returns the subscription by name.. 282 func (t *Topic) Subscription(name string) *Subscription { 283 t.mu.RLock() 284 defer t.mu.RUnlock() 285 if s, ok := t.subscriptions[name]; ok { 286 return s 287 } 288 return nil 289 } 290 291 // Receive receives message for this subscription. 292 func (s *Subscription) Receive(f func(*Message)) { 293 for range make([]struct{}, s.numGoroutines) { 294 s.wg.Add(1) 295 d, i := s.done, s.inbox 296 go func() { 297 for { 298 select { 299 case <-d: 300 d = nil 301 return 302 case m, open := <-i: 303 if !open { 304 i = nil 305 s.wg.Done() 306 break 307 } 308 f(&m) 309 } 310 } 311 }() 312 } 313 } 314 315 // Delete unsubscribes itself from topic. 316 func (s *Subscription) Delete() { 317 s.topic.mu.Lock() 318 for k := range s.topic.subscriptions { 319 if k == s.name { 320 delete(s.topic.subscriptions, k) 321 } 322 } 323 s.topic.mu.Unlock() 324 s.mu.Lock() 325 if s.stopped { 326 s.mu.Unlock() 327 return 328 } 329 s.stopped = true 330 close(s.inbox) 331 s.wg.Wait() 332 close(s.done) 333 s.mu.Unlock() 334 }