refactor idgen so that it exits gracefully

buffer.go

@@ -156,7 +156,7 @@ func (s *Source) Setup(state *Peco) {
 			// Not a great thing to do, allowing nil to be passed
 			// as state, but for testing I couldn't come up with anything
 			// better for the moment
-			if state != nil && !state.ExecQuery() {
+			if state != nil {
 				state.Hub().SendDraw(false)
 			}
 		}
@@ -198,7 +198,7 @@ func (s *Source) Setup(state *Peco) {
 		for scanner.Scan() {
 			txt := scanner.Text()
 			readCount++
-			s.Append(NewRawLine(txt, s.enableSep))
+			s.Append(state.NewRawLine(txt, s.enableSep))
 			notify.Do(notifycb)
 		}
 
@@ -217,11 +217,15 @@ func (s *Source) Setup(state *Peco) {
 
 // Start starts
 func (s *Source) Start(ctx context.Context) {
+	// I should be the only one running this method until I bail out
 	if pdebug.Enabled {
 		g := pdebug.Marker("Source.Start")
 		defer g.End()
 		defer pdebug.Printf("Source sent %d lines", len(s.lines))
 	}
+	s.start<-struct{}{}
+	defer func() { <-s.start }()
+
 	defer s.OutputChannel.SendEndMark("end of input")
 	defer close(s.done)
 

filter.go

@@ -402,7 +402,7 @@ func NewSmartCaseFilter() *RegexpFilter {
 	return rf
 }
 
-func NewExternalCmdFilter(name string, cmd string, args []string, threshold int, enableSep bool) *ExternalCmdFilter {
+func NewExternalCmdFilter(state *Peco, name string, cmd string, args []string, threshold int, enableSep bool) *ExternalCmdFilter {
 	if len(args) == 0 {
 		args = []string{"$QUERY"}
 	}
@@ -417,6 +417,7 @@ func NewExternalCmdFilter(name string, cmd string, args []string, threshold int,
 		enableSep:       enableSep,
 		name:            name,
 		outCh:           pipeline.OutputChannel(make(chan interface{})),
+		state: state,
 		thresholdBufsiz: threshold,
 	}
 }
@@ -545,7 +546,7 @@ func (ecf *ExternalCmdFilter) launchExternalCmd(ctx context.Context, buf []Line)
 				// This is the ONLY location where we need to actually
 				// RECREATE a RawLine, and thus the only place where
 				// ctx.enableSep is required.
-				cmdCh <- NewMatchedLine(NewRawLine(string(b), ecf.enableSep), nil)
+				cmdCh <- NewMatchedLine(ecf.state.NewRawLine(string(b), ecf.enableSep), nil)
 			}
 			if err != nil {
 				break

interface.go

@@ -46,8 +46,8 @@ const (
 	RegexpMatch        = "Regexp"
 )
 
-type idGen struct {
-	genCh chan uint64
+type idgen struct {
+	ch chan uint64
 }
 
 // Peco is the global object containing everything required to run peco.
@@ -64,12 +64,13 @@ type Peco struct {
 	// Config contains the values read in from config file
 	config                  Config
 	currentLineBuffer       Buffer
-	filters                 FilterSet
-	keymap                  Keymap
 	enableSep               bool     // Enable parsing on separators
+	filters                 FilterSet
+	idgen                   *idgen
 	initialFilter           string   // populated if --initial-filter is specified
 	initialQuery            string   // populated if --query is specified
 	inputseq                Inputseq // current key sequence (just the names)
+	keymap                  Keymap
 	layoutType              string
 	location                Location
 	mutex                   sync.Mutex
@@ -428,6 +429,7 @@ type Source struct {
 	enableSep bool
 	done      chan struct{}
 	ready     chan struct{}
+	start     chan struct{}
 	setupOnce sync.Once
 }
 
@@ -518,6 +520,7 @@ type ExternalCmdFilter struct {
 	args            []string
 	name            string
 	query           string
+	state *Peco
 	thresholdBufsiz int
 	outCh           pipeline.OutputChannel
 }

line.go

@@ -7,34 +7,11 @@ import (
 	"github.com/peco/peco/internal/util"
 )
 
-func newIDGen() *idGen {
-	ch := make(chan uint64)
-	go func() {
-		var i uint64
-		for ; ; i++ {
-			ch <- i
-			if i >= uint64(1<<63)-1 {
-				i = 0
-			}
-		}
-	}()
-	return &idGen{
-		genCh: ch,
-	}
-}
-
-func (ig *idGen) create() uint64 {
-	return <-ig.genCh
-}
-
-var idGenerator = newIDGen()
-
 // NewRawLine creates a new RawLine. The `enableSep` flag tells
 // it if we should search for a null character to split the
 // string to display and the string to emit upon selection of
 // of said line
-func NewRawLine(v string, enableSep bool) *RawLine {
-	id := idGenerator.create()
+func NewRawLine(id uint64, v string, enableSep bool) *RawLine {
 	rl := &RawLine{
 		id:            id,
 		buf:           v,

peco.go

@@ -4,6 +4,7 @@ import (
 	"io"
 	"os"
 	"reflect"
+	"sync"
 	"time"
 	"unicode/utf8"
 
@@ -54,6 +55,34 @@ func (is *Inputseq) Reset() {
 	*is = []string(nil)
 }
 
+func newIDGen() *idgen {
+	return &idgen{
+		ch: make(chan uint64),
+	}
+}
+
+func (ig *idgen) Run(ctx context.Context) {
+	var i uint64
+	for ; ; i++ {
+		select {
+		case <-ctx.Done():
+			return
+		case ig.ch <- i:
+		}
+
+		if i >= uint64(1<<63)-1 {
+			// If this happens, it's a disaster, but what can we do...
+			i = 0
+		}
+	}
+}
+
+func (ig *idgen) next() uint64 {
+	return <-ig.ch
+}
+
+var idGenerator = newIDGen()
+
 func New() *Peco {
 	return &Peco{
 		Argv:              os.Args,
@@ -61,6 +90,7 @@ func New() *Peco {
 		Stdin:             os.Stdin,
 		Stdout:            os.Stdout,
 		currentLineBuffer: NewMemoryBuffer(), // XXX revisit this
+		idgen: newIDGen(),
 		queryExecDelay:    50 * time.Millisecond,
 		readyCh:           make(chan struct{}),
 		screen:            &Termbox{},
@@ -245,15 +275,21 @@ func (p *Peco) Run(ctx context.Context) (err error) {
 	p.screen.Init()
 	defer p.screen.Close()
 
+	var _cancelOnce sync.Once
 	var _cancel func()
 	ctx, _cancel = context.WithCancel(ctx)
 	cancel := func() {
+		_cancelOnce.Do(func() {
 		if pdebug.Enabled {
 			pdebug.Printf("Peco.Run cancel called")
 		}
 		_cancel()
+		})
 	}
 
+	// start the ID generator
+	go p.idgen.Run(ctx)
+
 	// remember this cancel func so p.Exit works (XXX requires locking?)
 	p.cancelFunc = cancel
 
@@ -467,7 +503,7 @@ func (p *Peco) populateFilters() error {
 	p.filters.Add(NewRegexpFilter())
 
 	for name, c := range p.config.CustomFilter {
-		f := NewExternalCmdFilter(name, c.Cmd, c.Args, c.BufferThreshold, p.enableSep)
+		f := NewExternalCmdFilter(p, name, c.Cmd, c.Args, c.BufferThreshold, p.enableSep)
 		p.filters.Add(f)
 	}
 
@@ -576,3 +612,7 @@ func (p *Peco) CollectResults() {
 	})
 	close(p.resultCh)
 }
+
+func (p *Peco) NewRawLine(s string, enableSep bool) *RawLine {
+	return NewRawLine(p.idgen.next(), s, enableSep)
+}