workpool.go

type WorkPool struct {
	work     chan Runnable
	shutdown chan struct{}
	o        sync.Once
}

type Runnable func()

func NewWorkPool(tasks int) *WorkPool {
	work := make(chan Runnable, tasks)
	shutdown := make(chan struct{})
	for i := 0; i < tasks; i++ {
		go func(i int) {
			s := shutdown
			for s != nil || len(work) > 0 {
				select {
				case r := <-work:
					r()
				case <-s:
					fmt.Println("got shutdown in", i)
					s = nil
				}
			}
			fmt.Println("closing", i)
		}(i)
	}
	return &WorkPool{
		work:     work,
		shutdown: shutdown,
	}
}

func (wp *WorkPool) SubmitWithTimeout(t time.Duration, r Runnable) error {
	select {
	case wp.work <- r:
	case <-time.After(t):
		return errors.New("Task submit timed out")
	case <-wp.shutdown:
		return errors.New("Work Pool shut down")
	}
	return nil
}

func (wp *WorkPool) Shutdown() {
	wp.o.Do(func() {
		fmt.Println("shutting down")
		close(wp.shutdown)
	})
}

func (wp *WorkPool) Empty() bool {
	return len(wp.work) == 0
}