package goprocess import ( "sync" ) // process implements Process type process struct { children map[*processLink]struct{} // process to close with us waitfors map[*processLink]struct{} // process to only wait for waiters []*processLink // processes that wait for us. for gc. teardown TeardownFunc // called to run the teardown logic. closing chan struct{} // closed once close starts. closed chan struct{} // closed once close is done. closeErr error // error to return to clients of Close() sync.Mutex } // newProcess constructs and returns a Process. // It will call tf TeardownFunc exactly once: // **after** all children have fully Closed, // **after** entering <-Closing(), and // **before** <-Closed(). func newProcess(tf TeardownFunc) *process { return &process{ teardown: tf, closed: make(chan struct{}), closing: make(chan struct{}), waitfors: make(map[*processLink]struct{}), children: make(map[*processLink]struct{}), } } func (p *process) WaitFor(q Process) { if q == nil { panic("waiting for nil process") } p.Lock() defer p.Unlock() select { case <-p.Closed(): panic("Process cannot wait after being closed") default: } pl := newProcessLink(p, q) if p.waitfors == nil { // This may be nil when we're closing. In close, we'll keep // reading this map till it stays nil. p.waitfors = make(map[*processLink]struct{}, 1) } p.waitfors[pl] = struct{}{} go pl.AddToChild() } func (p *process) AddChildNoWait(child Process) { if child == nil { panic("adding nil child process") } p.Lock() defer p.Unlock() select { case <-p.Closing(): // Either closed or closing, close child immediately. This is // correct because we aren't asked to _wait_ on this child. go child.Close() // Wait for the child to start closing so the child is in the // "correct" state after this function finishes (see #17). <-child.Closing() return default: } pl := newProcessLink(p, child) p.children[pl] = struct{}{} go pl.AddToChild() } func (p *process) AddChild(child Process) { if child == nil { panic("adding nil child process") } p.Lock() defer p.Unlock() pl := newProcessLink(p, child) select { case <-p.Closed(): // AddChild must not be called on a dead process. Maybe that's // too strict? panic("Process cannot add children after being closed") default: } select { case <-p.Closing(): // Already closing, close child in background. go child.Close() // Wait for the child to start closing so the child is in the // "correct" state after this function finishes (see #17). <-child.Closing() default: // Only add the child when not closing. When closing, just add // it to the "waitfors" list. p.children[pl] = struct{}{} } if p.waitfors == nil { // This may be be nil when we're closing. In close, we'll keep // reading this map till it stays nil. p.waitfors = make(map[*processLink]struct{}, 1) } p.waitfors[pl] = struct{}{} go pl.AddToChild() } func (p *process) Go(f ProcessFunc) Process { child := newProcess(nil) waitFor := newProcess(nil) child.WaitFor(waitFor) // prevent child from closing // add child last, to prevent a closing parent from // closing all of them prematurely, before running the func. p.AddChild(child) go func() { f(child) waitFor.Close() // allow child to close. child.CloseAfterChildren() // close to tear down. }() return child } // SetTeardown to assign a teardown function func (p *process) SetTeardown(tf TeardownFunc) { if tf == nil { panic("cannot set nil TeardownFunc") } p.Lock() if p.teardown != nil { panic("cannot SetTeardown twice") } p.teardown = tf select { case <-p.Closed(): // Call the teardown function, but don't set the error. We can't // change that after we shut down. tf() default: } p.Unlock() } // Close is the external close function. // it's a wrapper around internalClose that waits on Closed() func (p *process) Close() error { p.Lock() // if already closing, or closed, get out. (but wait!) select { case <-p.Closing(): p.Unlock() <-p.Closed() return p.closeErr default: } p.doClose() p.Unlock() return p.closeErr } func (p *process) Closing() <-chan struct{} { return p.closing } func (p *process) Closed() <-chan struct{} { return p.closed } func (p *process) Err() error { <-p.Closed() return p.closeErr } // the _actual_ close process. func (p *process) doClose() { // this function is only be called once (protected by p.Lock()). // and it will panic (on closing channels) otherwise. close(p.closing) // signal that we're shutting down (Closing) // We won't add any children after we start closing so we can do this // once. for plc, _ := range p.children { child := plc.Child() if child != nil { // check because child may already have been removed. go child.Close() // force all children to shut down } // safe to call multiple times per link plc.ParentClear() } p.children = nil // clear them. release memory. // We may repeatedly continue to add waiters while we wait to close so // we have to do this in a loop. for len(p.waitfors) > 0 { // we must be careful not to iterate over waitfors directly, as it may // change under our feet. wf := p.waitfors p.waitfors = nil // clear them. release memory. for w, _ := range wf { // Here, we wait UNLOCKED, so that waitfors who are in the middle of // adding a child to us can finish. we will immediately close the child. p.Unlock() <-w.ChildClosed() // wait till all waitfors are fully closed (before teardown) p.Lock() // safe to call multiple times per link w.ParentClear() } } if p.teardown != nil { p.closeErr = p.teardown() // actually run the close logic (ok safe to teardown) } close(p.closed) // signal that we're shut down (Closed) // go remove all the parents from the process links. optimization. go func(waiters []*processLink) { for _, pl := range waiters { pl.ClearChild() pr, ok := pl.Parent().(*process) if !ok { // parent has already been called to close continue } pr.Lock() delete(pr.waitfors, pl) delete(pr.children, pl) pr.Unlock() } }(p.waiters) // pass in so p.waiters = nil // clear them. release memory. } // We will only wait on the children we have now. // We will not wait on children added subsequently. // this may change in the future. func (p *process) CloseAfterChildren() error { p.Lock() select { case <-p.Closed(): p.Unlock() return p.Close() // get error. safe, after p.Closed() default: } p.Unlock() // here only from one goroutine. nextToWaitFor := func() Process { p.Lock() defer p.Unlock() for e, _ := range p.waitfors { c := e.Child() if c == nil { continue } select { case <-c.Closed(): default: return c } } return nil } // wait for all processes we're waiting for are closed. // the semantics here are simple: we will _only_ close // if there are no processes currently waiting for. for next := nextToWaitFor(); next != nil; next = nextToWaitFor() { <-next.Closed() } // YAY! we're done. close return p.Close() }