├── LICENSE ├── README.md ├── broadcast.go ├── flow.go └── flow_test.go /LICENSE: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2016 Chzyer 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # flow 2 | Help you to control the flow of goroutines 3 | 4 | 5 | 6 | ## Feature 7 | 8 | 9 | 10 | ##### Wait on all goroutines exited, just like `sync.WaitGroup` 11 | 12 | ```go 13 | func main() { 14 | f := flow.New() 15 | go func() { 16 | f.Add(1) 17 | defer f.Done() 18 | println("exit") 19 | }() 20 | f.Wait() // wait will capture signals 21 | } 22 | // output: 23 | // exit 24 | ``` 25 | 26 | 27 | 28 | ##### Notify goroutines to exit 29 | 30 | ```go 31 | func main() { 32 | f := flow.New() 33 | go func() { 34 | f.Add(1) 35 | defer f.Done() 36 | ticker := time.NewTicker(time.Second) 37 | defer ticker.Stop() 38 | 39 | loop: 40 | for { 41 | select { 42 | case <-ticker.C: 43 | println("tick") 44 | case <-f.IsClose(): 45 | println("break") 46 | break loop 47 | } 48 | } 49 | }() 50 | 51 | go func() { 52 | time.Sleep(2 * time.Second) 53 | f.Close() 54 | }() 55 | 56 | f.Wait() 57 | println("exited") 58 | } 59 | // output: 60 | // tick 61 | // tick 62 | // break 63 | // exited 64 | ``` 65 | 66 | If we kill this process by `Ctrl+C` in 1s, we will got output as follows: 67 | 68 | ```go 69 | // output: 70 | // tick 71 | // break 72 | // exited 73 | ``` 74 | 75 | 76 | 77 | ##### Multiple goroutines can all run or all die 78 | 79 | ```go 80 | func main() { 81 | f := flow.New() 82 | 83 | ch := make(chan string) 84 | // read 85 | go func() { 86 | f.Add(1) 87 | defer f.DoneAndClose() // Done and also close this flow 88 | ticker := time.NewTicker(time.Second) 89 | defer ticker.Stop() 90 | exitTime := time.Now().Add(3 * time.Second) 91 | 92 | loop: 93 | for { 94 | select { 95 | case now := <-ticker.C: 96 | if now.After(exitTime) { 97 | break loop 98 | } 99 | ch <- now.String() 100 | case <-f.IsClose(): 101 | break loop 102 | } 103 | } 104 | 105 | println("readloop exit") 106 | }() 107 | 108 | go func() { 109 | f.Add(1) 110 | defer f.DoneAndClose() // Done and also close this flow 111 | 112 | loop: 113 | for { 114 | select { 115 | case text := <-ch: 116 | fmt.Fprintln(os.Stdout, text) 117 | case <-f.IsClose(): 118 | break loop 119 | } 120 | } 121 | 122 | println("writeloop exit") 123 | }() 124 | 125 | f.Wait() 126 | println("all exit") 127 | } 128 | 129 | // output: 130 | // 2016-05-31 18:10:18.525209975 +0800 CST 131 | // 2016-05-31 18:10:19.525009926 +0800 CST 132 | // readloop exit 133 | // writeloop exit 134 | ``` 135 | 136 | 137 | 138 | ##### Indicate leaking goroutine 139 | 140 | ```Go 141 | func goroutine1(f *flow.Flow) { 142 | f.Add(1) 143 | defer f.DoneAndClose() 144 | for _ = range time.Tick(time.Second) { 145 | 146 | } 147 | println("goroutine 1 exit") 148 | } 149 | 150 | func goroutine2(f *flow.Flow) { 151 | f.Add(1) 152 | defer f.DoneAndClose() 153 | loop: 154 | for { 155 | select { 156 | case <-f.IsClose(): 157 | break loop 158 | } 159 | } 160 | println("goroutine 2 exit") 161 | } 162 | 163 | func main() { 164 | f := flow.New() 165 | go goroutine1(f) 166 | go goroutine2(f) 167 | f.Wait() 168 | } 169 | // output: 170 | // (press Ctrl+C) 171 | // goroutine 2 exit 172 | 173 | // 31 18:18:59 flow-wait.go:124 main.main - init 174 | // 31 18:18:59 flow-wait.go:103 main.goroutine1 - add: 1, ref: 1 175 | // 31 18:18:59 flow-wait.go:111 main.goroutine2 - add: 1, ref: 2 176 | // 31 18:19:00 flow-wait.go:127 main.main - got signal 177 | // 31 18:19:00 flow-wait.go:127 main.main - stop 178 | // 31 18:19:00 flow-wait.go:127 main.main - wait 179 | // 31 18:19:00 flow-wait.go:121 main.goroutine2 - done and close, ref: 1 180 | 181 | // (press Ctrl+C again) 182 | // force close 183 | ``` 184 | 185 | If the `flow.Wait()` waiting too long, it will print some debug info to let you indicate which goroutine is leaked. For example above, we know that `main.goroutine1` is leaked. 186 | 187 | 188 | 189 | ##### Raise error out to main goroutine from sub-goroutine 190 | 191 | ```Go 192 | func main() { 193 | f := flow.New() 194 | 195 | go func() { 196 | err := http.ListenAndServe(":80", nil) 197 | if err != nil { 198 | f.Error(err) 199 | return 200 | } 201 | }() 202 | 203 | if err := f.Wait(); err != nil { 204 | println("error:", err.Error()) 205 | os.Exit(1) 206 | } 207 | } 208 | 209 | // output: 210 | // error: listen tcp :80: bind: permission denied 211 | // exit status 1 212 | ``` 213 | 214 | -------------------------------------------------------------------------------- /broadcast.go: -------------------------------------------------------------------------------- 1 | package flow 2 | 3 | import ( 4 | "errors" 5 | "sync/atomic" 6 | ) 7 | 8 | var ErrCanceled = errors.New("operation is canceled") 9 | 10 | type Broadcast struct { 11 | notify atomic.Value 12 | } 13 | 14 | func NewBroadcast() *Broadcast { 15 | b := &Broadcast{} 16 | b.init() 17 | return b 18 | } 19 | 20 | func (b *Broadcast) init() { 21 | b.notify.Store(make(chan struct{})) 22 | } 23 | 24 | func (b *Broadcast) Wait() chan struct{} { 25 | return b.notify.Load().(chan struct{}) 26 | } 27 | 28 | func (b *Broadcast) Notify() { 29 | ch := b.Wait() 30 | b.init() 31 | close(ch) 32 | } 33 | 34 | func (b *Broadcast) Close() { 35 | close(b.Wait()) 36 | } 37 | -------------------------------------------------------------------------------- /flow.go: -------------------------------------------------------------------------------- 1 | package flow 2 | 3 | import ( 4 | "bytes" 5 | "fmt" 6 | "os" 7 | "os/signal" 8 | "path" 9 | "reflect" 10 | "runtime" 11 | "strings" 12 | "sync" 13 | "sync/atomic" 14 | "syscall" 15 | "time" 16 | ) 17 | 18 | var ( 19 | pkgPath = reflect.TypeOf(Flow{}).PkgPath() 20 | DefaultDebug = true 21 | ) 22 | 23 | type debugInfo struct { 24 | Time string 25 | FileInfo string 26 | Stack string 27 | Info string 28 | } 29 | 30 | func (d *debugInfo) String() string { 31 | return d.Stack + " - " + d.Info 32 | } 33 | 34 | type Flow struct { 35 | errChan chan error 36 | stopChan chan struct{} 37 | ref *int32 38 | wg sync.WaitGroup 39 | Parent *Flow 40 | Children []*Flow 41 | stoped int32 42 | exited int32 43 | onClose []func() 44 | id uintptr 45 | 46 | mutex sync.Mutex 47 | debug []debugInfo 48 | printed int32 49 | } 50 | 51 | func NewEx(n int) *Flow { 52 | f := &Flow{ 53 | errChan: make(chan error, 1), 54 | stopChan: make(chan struct{}), 55 | ref: new(int32), 56 | } 57 | f.appendDebug("init") 58 | if n > 0 { 59 | f.Add(n) 60 | } 61 | return f 62 | } 63 | 64 | func New() *Flow { 65 | return NewEx(0) 66 | } 67 | 68 | func (f *Flow) WaitNotify(ch chan struct{}) bool { 69 | select { 70 | case <-ch: 71 | return true 72 | case <-f.IsClose(): 73 | return false 74 | } 75 | } 76 | 77 | func (f *Flow) MarkExit() bool { 78 | return atomic.CompareAndSwapInt32(&f.exited, 0, 1) 79 | } 80 | 81 | func (f *Flow) IsExit() bool { 82 | return atomic.LoadInt32(&f.exited) == 1 83 | } 84 | 85 | func (f *Flow) GetDebug() []byte { 86 | buf := bytes.NewBuffer(nil) 87 | var stackML, fileML, timeML int 88 | for _, d := range f.debug { 89 | if stackML < len(d.Stack) { 90 | stackML = len(d.Stack) 91 | } 92 | if fileML < len(d.FileInfo) { 93 | fileML = len(d.FileInfo) 94 | } 95 | if timeML < len(d.Time) { 96 | timeML = len(d.Time) 97 | } 98 | } 99 | fill := func(a string, n int) string { 100 | return a + strings.Repeat(" ", n-len(a)) 101 | } 102 | buf.WriteString("\n") 103 | for _, d := range f.debug { 104 | buf.WriteString(fmt.Sprintf("%v %v %v - %v\n", 105 | fill(d.Time, timeML), 106 | fill(d.FileInfo, fileML), 107 | fill(d.Stack, stackML), d.Info, 108 | )) 109 | } 110 | return buf.Bytes() 111 | } 112 | 113 | func (f *Flow) printDebug() { 114 | println(string(f.GetDebug())) 115 | } 116 | 117 | func (f *Flow) appendDebug(info string) { 118 | pc, fp, line, _ := runtime.Caller(f.getCaller()) 119 | name := runtime.FuncForPC(pc).Name() 120 | f.debug = append(f.debug, debugInfo{ 121 | Time: time.Now().Format("02 15:04:05"), 122 | FileInfo: fmt.Sprintf("%v:%v", path.Base(fp), line), 123 | Stack: path.Base(name), 124 | Info: info, 125 | }) 126 | } 127 | 128 | func (f *Flow) SetOnClose(exit func()) *Flow { 129 | if f.IsClosed() { 130 | f.appendDebug("set close after closed") 131 | exit() 132 | return f 133 | } 134 | 135 | f.onClose = []func(){exit} 136 | return f 137 | } 138 | 139 | func (f *Flow) AddOnClose(exit func()) *Flow { 140 | if f.IsClosed() { 141 | f.appendDebug("add close after closed") 142 | exit() 143 | return f 144 | } 145 | 146 | f.onClose = append(f.onClose, exit) 147 | return f 148 | } 149 | 150 | const ( 151 | F_CLOSED = true 152 | F_TIMEOUT = false 153 | ) 154 | 155 | func (f *Flow) Tick(t *time.Ticker) bool { 156 | select { 157 | case <-t.C: 158 | return F_TIMEOUT 159 | case <-f.IsClose(): 160 | return F_CLOSED 161 | } 162 | } 163 | 164 | func (f *Flow) CloseOrWait(duration time.Duration) bool { 165 | select { 166 | case <-time.After(duration): 167 | return F_TIMEOUT 168 | case <-f.IsClose(): 169 | return F_CLOSED 170 | } 171 | } 172 | 173 | func (f *Flow) Errorf(layout string, obj ...interface{}) { 174 | f.Error(fmt.Errorf(layout, obj...)) 175 | } 176 | 177 | func (f *Flow) Error(err error) { 178 | f.errChan <- err 179 | } 180 | 181 | func (f *Flow) ForkTo(ref **Flow, exit func()) { 182 | child := f.Fork(0) 183 | *ref = child 184 | child.AddOnClose(exit) 185 | } 186 | 187 | func (f *Flow) Fork(n int) *Flow { 188 | f2 := NewEx(n) 189 | f2.Parent = f 190 | // TODO(chzyer): test it ! 191 | f2.errChan = f.errChan 192 | f.Children = append(f.Children, f2) 193 | f.Add(1) // for f2 194 | 195 | if f.IsClosed() { 196 | f.appendDebug("fork when closed") 197 | // stop-wait 198 | // ->fork 199 | // done 200 | f2.Stop() 201 | } 202 | 203 | return f2 204 | } 205 | 206 | func (f *Flow) StopAll() { 207 | flow := f 208 | for flow.Parent != nil { 209 | flow = flow.Parent 210 | } 211 | flow.Stop() 212 | } 213 | 214 | func (f *Flow) Close() { 215 | f.appendDebug("close") 216 | f.close() 217 | } 218 | 219 | func (f *Flow) close() { 220 | f.Stop() 221 | f.wait() 222 | } 223 | 224 | func (f *Flow) Stop() { 225 | if !atomic.CompareAndSwapInt32(&f.stoped, 0, 1) { 226 | return 227 | } 228 | f.appendDebug("stop") 229 | 230 | close(f.stopChan) 231 | for _, cf := range f.Children { 232 | cf.Stop() 233 | } 234 | if len(f.onClose) > 0 { 235 | go func() { 236 | for _, f := range f.onClose { 237 | f() 238 | } 239 | }() 240 | } 241 | } 242 | 243 | func (f *Flow) IsClosed() bool { 244 | return atomic.LoadInt32(&f.stoped) == 1 245 | } 246 | 247 | func (f *Flow) IsClose() chan struct{} { 248 | return f.stopChan 249 | } 250 | 251 | func (f *Flow) Add(n int) { 252 | ref := atomic.AddInt32(f.ref, int32(n)) 253 | f.appendDebug(fmt.Sprintf("add: %v, ref: %v", n, ref)) 254 | f.wg.Add(n) 255 | } 256 | 257 | func (f *Flow) getCaller() int { 258 | for i := 0; ; i++ { 259 | pc, _, _, ok := runtime.Caller(i) 260 | if !ok { 261 | break 262 | } 263 | f := runtime.FuncForPC(pc).Name() 264 | if !strings.HasPrefix(f, pkgPath) { 265 | return i - 1 266 | } 267 | } 268 | return 0 269 | } 270 | 271 | func (f *Flow) getRef() int32 { 272 | return atomic.LoadInt32(f.ref) 273 | } 274 | 275 | func (f *Flow) Done() { 276 | f.wg.Done() 277 | ref := atomic.AddInt32(f.ref, -1) 278 | f.appendDebug(fmt.Sprintf("done, ref: %v", ref)) 279 | if ref == 0 { 280 | f.Stop() 281 | } 282 | } 283 | 284 | func (f *Flow) DoneAndClose() { 285 | f.wg.Done() 286 | ref := atomic.AddInt32(f.ref, -1) 287 | f.appendDebug(fmt.Sprintf("done and close, ref: %v", ref)) 288 | f.Stop() 289 | } 290 | 291 | func (f *Flow) wait() { 292 | f.appendDebug("wait") 293 | 294 | done := make(chan struct{}) 295 | printed := int32(0) 296 | if DefaultDebug && atomic.CompareAndSwapInt32(&f.printed, 0, 1) { 297 | go func() { 298 | select { 299 | case <-done: 300 | return 301 | case <-time.After(1000 * time.Millisecond): 302 | f.printDebug() 303 | atomic.StoreInt32(&printed, 1) 304 | } 305 | }() 306 | } 307 | <-f.stopChan 308 | f.wg.Wait() 309 | close(done) 310 | if atomic.LoadInt32(&printed) == 1 { 311 | println(fmt.Sprint(&f) + " - exit") 312 | } 313 | 314 | if f.Parent != nil { 315 | f.Parent.Done() 316 | f.Parent = nil 317 | } 318 | } 319 | 320 | func (f *Flow) Wait() error { 321 | signalChan := make(chan os.Signal) 322 | signal.Notify(signalChan, 323 | os.Interrupt, os.Kill, syscall.SIGTERM, syscall.SIGHUP) 324 | var err error 325 | select { 326 | case <-f.IsClose(): 327 | f.appendDebug("got closed") 328 | case <-signalChan: 329 | f.appendDebug("got signal") 330 | f.Stop() 331 | case err = <-f.errChan: 332 | f.appendDebug(fmt.Sprintf("got error: %v", err)) 333 | 334 | if err != nil { 335 | f.Stop() 336 | } 337 | } 338 | 339 | go func() { 340 | <-signalChan 341 | // force close 342 | println("force close") 343 | os.Exit(1) 344 | }() 345 | 346 | f.wait() 347 | return err 348 | } 349 | -------------------------------------------------------------------------------- /flow_test.go: -------------------------------------------------------------------------------- 1 | package flow 2 | 3 | import ( 4 | "sync/atomic" 5 | "testing" 6 | "time" 7 | ) 8 | 9 | func TestStopWaitFork(t *testing.T) { 10 | f := New() 11 | 12 | done := make(chan struct{}) 13 | 14 | { // worker 15 | go func() { 16 | f.Add(1) 17 | defer f.DoneAndClose() 18 | <-done 19 | }() 20 | } 21 | 22 | go func() { 23 | for _ = range time.Tick(time.Second) { 24 | println(string(f.GetDebug())) 25 | break 26 | } 27 | }() 28 | 29 | go func() { 30 | f.Stop() 31 | { 32 | f2 := f.Fork(0) 33 | done2 := make(chan struct{}) 34 | go func() { 35 | f2.Add(1) 36 | <-done2 37 | f2.Done() 38 | }() 39 | close(done2) 40 | f2.Wait() 41 | } 42 | time.Sleep(100 * time.Millisecond) 43 | close(done) 44 | }() 45 | f.Wait() 46 | } 47 | 48 | func TestFlow(t *testing.T) { 49 | { 50 | f := NewEx(1) 51 | done := 0 52 | go func() { 53 | defer f.Done() 54 | for i := 0; i < 3; i++ { 55 | done++ 56 | } 57 | }() 58 | 59 | f.Wait() 60 | if done != 3 { 61 | t.Fatal(done) 62 | } 63 | } 64 | 65 | { 66 | f := NewEx(2) 67 | var done int64 68 | go func() { 69 | defer f.Done() 70 | atomic.AddInt64(&done, 1) 71 | }() 72 | go func() { 73 | defer f.Done() 74 | atomic.AddInt64(&done, 1) 75 | }() 76 | f.Wait() 77 | if done != 2 { 78 | t.Fatal(done) 79 | } 80 | } 81 | } 82 | --------------------------------------------------------------------------------