diff options
Diffstat (limited to 'portlist')
| -rw-r--r-- | portlist/poller.go | 127 | ||||
| -rw-r--r-- | portlist/portlist_test.go | 2 |
2 files changed, 79 insertions, 50 deletions
diff --git a/portlist/poller.go b/portlist/poller.go index 16dd8e74c..53a8e05b1 100644 --- a/portlist/poller.go +++ b/portlist/poller.go @@ -29,7 +29,7 @@ type Poller struct { // This field should only be changed before calling Run. IncludeLocalhost bool - c chan List // unbuffered + c chan Update // unbuffered // os, if non-nil, is an OS-specific implementation of the portlist getting // code. When non-nil, it's responsible for getting the complete list of @@ -52,6 +52,23 @@ type Poller struct { prev List // most recent data, not aliasing scratch } +// Update is sent by Poller to indicate +// an update has been made to the machine's +// open ports. Receiver of this struct must +// check the Err() method before calling List(). +type Update struct { + list List + err error +} + +func (u *Update) Err() error { + return u.err +} + +func (u *Update) List() List { + return u.list +} + // osImpl is the OS-specific implementation of getting the open listening ports. type osImpl interface { Close() error @@ -71,32 +88,6 @@ var newOSImpl func(includeLocalhost bool) osImpl var errUnimplemented = errors.New("portlist poller not implemented on " + runtime.GOOS) -// NewPoller returns a new portlist Poller. It returns an error -// if the portlist couldn't be obtained. -func NewPoller() (*Poller, error) { - if debugDisablePortlist() { - return nil, errors.New("portlist disabled by envknob") - } - p := &Poller{ - c: make(chan List), - runDone: make(chan struct{}), - } - p.closeCtx, p.closeCtxCancel = context.WithCancel(context.Background()) - p.osOnce.Do(p.initOSField) - if p.os == nil { - return nil, errUnimplemented - } - - // Do one initial poll synchronously so we can return an error - // early. - if pl, err := p.getList(); err != nil { - return nil, err - } else { - p.setPrev(pl) - } - return p, nil -} - func (p *Poller) setPrev(pl List) { // Make a copy, as the pass in pl slice aliases pl.scratch and we don't want // that to except to the caller. @@ -112,11 +103,15 @@ func (p *Poller) initOSField() { // Updates return the channel that receives port list updates. // // The channel is closed when the Poller is closed. -func (p *Poller) Updates() <-chan List { return p.c } +func (p *Poller) Updates() <-chan Update { return p.c } // Close closes the Poller. // Run will return with a nil error. func (p *Poller) Close() error { + // Skip if uninitialized. + if p.os == nil { + return nil + } p.closeCtxCancel() <-p.runDone if p.os != nil { @@ -126,14 +121,14 @@ func (p *Poller) Close() error { } // send sends pl to p.c and returns whether it was successfully sent. -func (p *Poller) send(ctx context.Context, pl List) (sent bool, err error) { +func (p *Poller) send(ctx context.Context, pl List, listErr error) (sent bool) { select { - case p.c <- pl: - return true, nil + case p.c <- Update{list: pl, err: listErr}: + return true case <-ctx.Done(): - return false, ctx.Err() + return false case <-p.closeCtx.Done(): - return false, nil + return false } } @@ -141,48 +136,82 @@ func (p *Poller) send(ctx context.Context, pl List) (sent bool, err error) { // is done, or the Close is called. // // Run may only be called once. -func (p *Poller) Run(ctx context.Context) error { +func (p *Poller) Run(ctx context.Context) (chan Update, error) { + if debugDisablePortlist() { + return nil, errors.New("portlist disabled by envknob") + } + if p.os != nil { + return nil, errors.New("method called more than once") + } + p.initOSField() + if p.os == nil { + return nil, errUnimplemented + } + + p.c = make(chan Update) + p.runDone = make(chan struct{}) + p.closeCtx, p.closeCtxCancel = context.WithCancel(context.Background()) + + // Do one initial poll synchronously so we can return an error + // early. + if pl, err := p.getList(); err != nil { + return nil, err + } else { + p.setPrev(pl) + } + tick := time.NewTicker(pollInterval) defer tick.Stop() - return p.runWithTickChan(ctx, tick.C) + go p.runWithTickChan(ctx, tick.C) + return p.c, nil } -func (p *Poller) runWithTickChan(ctx context.Context, tickChan <-chan time.Time) error { +func (p *Poller) runWithTickChan(ctx context.Context, tickChan <-chan time.Time) { defer close(p.runDone) defer close(p.c) // Send out the pre-generated initial value. - if sent, err := p.send(ctx, p.prev); !sent { - return err + if sent := p.send(ctx, p.prev, nil); !sent { + return } + // Order of events: + // 1. If the context is done, exit + // 2. If the user called p.Close(), exit. + // 3. If we received a tick, then get the list + // 3B. If that error'd, send an error to the user. + // 3C. If the context or p.Close where called in the meantime, exit. + // 3D. If getList succeeded, skip if there are no updates. + // 3E. If there are indeed updates, send them, or exit if 1/2 are true. + // We check 1 & 2 in 3 places: top of the for-loop, + // whenever we send (which is two places: sending an error, or sending a list). for { select { + case <-ctx.Done(): + return + case <-p.closeCtx.Done(): + return case <-tickChan: pl, err := p.getList() if err != nil { - return err + sent := p.send(ctx, nil, err) + if !sent { + return + } + continue } if pl.equal(p.prev) { continue } p.setPrev(pl) - if sent, err := p.send(ctx, p.prev); !sent { - return err + if sent := p.send(ctx, p.prev, nil); !sent { + return } - case <-ctx.Done(): - return ctx.Err() - case <-p.closeCtx.Done(): - return nil } } } func (p *Poller) getList() (List, error) { - if debugDisablePortlist() { - return nil, nil - } - p.osOnce.Do(p.initOSField) var err error p.scratch, err = p.os.AppendListeningPorts(p.scratch[:0]) return p.scratch, err diff --git a/portlist/portlist_test.go b/portlist/portlist_test.go index 6055a8426..ee480ea7f 100644 --- a/portlist/portlist_test.go +++ b/portlist/portlist_test.go @@ -192,7 +192,7 @@ func TestPoller(t *testing.T) { for pl := range p.Updates() { // Look at all the pl slice memory to maximize // chance of race detector seeing violations. - for _, v := range pl { + for _, v := range pl.List() { if v == (Port{}) { // Force use panic("empty port") |