github.com/searKing/golang/go@v1.2.74/sync/subject.go (about) 1 // Copyright 2022 The searKing Author. All rights reserved. 2 // Use of this source code is governed by a BSD-style 3 // license that can be found in the LICENSE file. 4 5 package sync 6 7 import ( 8 "context" 9 "fmt" 10 "sync" 11 12 "github.com/searKing/golang/go/errors" 13 "github.com/searKing/golang/go/pragma" 14 "github.com/searKing/golang/go/sync/atomic" 15 ) 16 17 // Subject implements a condition variable like with channel, a rendezvous point 18 // for goroutines waiting for or announcing the occurrence 19 // of an event. 20 // 21 // The caller typically cannot assume that the condition is true when 22 // Subscribe chan returns. Instead, the caller should Wait in a loop: 23 // 24 // time.After(timeout, c.PublishBroadcast()) // for timeout or periodic event 25 // c.PublishBroadcast() // for async notify event directly 26 // eventC, cancel := c.Subscribe() 27 // for !condition() { 28 // select{ 29 // case event, closed := <- eventC: 30 // ... make use of event ... 31 // } 32 // } 33 // ... make use of condition ... 34 // 35 type Subject struct { 36 noCopy pragma.DoNotCopy 37 38 mu sync.Mutex 39 subscribers map[*subscriber]struct{} 40 41 inShutdown atomic.Bool // true when when server is in shutdown 42 } 43 44 type subscriber struct { 45 mu sync.Mutex // guard close of channel msgC 46 msgC chan interface{} 47 48 once sync.Once 49 doneC chan struct{} // closed when subscriber is in shutdown, like removed. 50 } 51 52 func (s *subscriber) Shutdown() { 53 if s == nil { 54 return 55 } 56 s.once.Do(func() { 57 close(s.doneC) 58 s.mu.Lock() 59 defer s.mu.Unlock() 60 close(s.msgC) 61 }) 62 } 63 64 // publish wakes a listener waiting on c to consume the event. 65 // event will be dropped if ctx is Done before event is received. 66 func (s *subscriber) publish(ctx context.Context, event interface{}) error { 67 // guard of msgC's close 68 s.mu.Lock() 69 defer s.mu.Unlock() 70 select { 71 case <-ctx.Done(): 72 // event dropped because of publisher 73 return ctx.Err() 74 case <-s.doneC: 75 // event dropped because of subscriber 76 return fmt.Errorf("event dropped because of subscriber unsubscribed") 77 default: 78 } 79 80 select { 81 case <-ctx.Done(): 82 // event dropped because of publisher 83 return ctx.Err() 84 case <-s.doneC: 85 // event dropped because of subscriber 86 return fmt.Errorf("event dropped because of subscriber unsubscribed") 87 case s.msgC <- event: 88 // event consumed 89 return nil 90 } 91 } 92 93 // Subscribe returns a channel that's closed when awoken by PublishSignal or PublishBroadcast. 94 // never be canceled. Successive calls to Subscribe return different values. 95 // The close of the Subscribe channel may happen asynchronously, 96 // after the cancel function returns. 97 func (s *Subject) Subscribe() (<-chan interface{}, context.CancelFunc) { 98 listener := &subscriber{ 99 msgC: make(chan interface{}), 100 doneC: make(chan struct{}), 101 } 102 s.trackChannel(listener, true) 103 return listener.msgC, func() { 104 s.trackChannel(listener, false) 105 } 106 } 107 108 // PublishSignal wakes one listener waiting on c, if there is any. 109 // PublishSignal blocks until event is received or dropped. 110 func (s *Subject) PublishSignal(ctx context.Context, event interface{}) error { 111 s.mu.Lock() 112 defer s.mu.Unlock() 113 var wg sync.WaitGroup 114 var errs []error 115 for listener := range s.subscribers { 116 wg.Add(1) 117 go func(listener *subscriber) { 118 defer wg.Done() 119 err := listener.publish(ctx, event) 120 if err != nil { 121 errs = append(errs, err) 122 } 123 }(listener) 124 break 125 } 126 wg.Wait() 127 return errors.Multi(errs...) 128 } 129 130 // PublishBroadcast wakes all listeners waiting on c. 131 // PublishBroadcast blocks until event is received or dropped. 132 // event will be dropped if ctx is Done before event is received. 133 func (s *Subject) PublishBroadcast(ctx context.Context, event interface{}) error { 134 var wg sync.WaitGroup 135 var errs []error 136 func() { 137 s.mu.Lock() 138 defer s.mu.Unlock() 139 for listener := range s.subscribers { 140 wg.Add(1) 141 go func(listener *subscriber) { 142 defer wg.Done() 143 err := listener.publish(ctx, event) 144 if err != nil { 145 errs = append(errs, err) 146 } 147 }(listener) 148 } 149 }() 150 wg.Wait() 151 return errors.Multi(errs...) 152 } 153 154 func (s *Subject) trackChannel(c *subscriber, add bool) { 155 func() { 156 s.mu.Lock() 157 defer s.mu.Unlock() 158 if s.subscribers == nil { 159 s.subscribers = make(map[*subscriber]struct{}) 160 } 161 _, has := s.subscribers[c] 162 if has { 163 if add { 164 return 165 } 166 delete(s.subscribers, c) 167 return 168 } 169 if add { 170 s.subscribers[c] = struct{}{} 171 } 172 }() 173 if !add { 174 c.Shutdown() 175 } 176 }