Merge pull request #2034 from jaffee/worker-pool

Worker pool

executor.go

@@ -19,6 +19,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"sort"
+	"sync"
 	"time"
 
 	"github.com/pilosa/pilosa/pql"
@@ -55,6 +56,10 @@ type executor struct {
 
 	// Stores key/id translation data.
 	TranslateStore TranslateStore
+
+	workersWG      sync.WaitGroup
+	workerPoolSize int
+	work           chan job
 }
 
 // executorOption is a functional option type for pilosa.Executor
@@ -67,10 +72,18 @@ func optExecutorInternalQueryClient(c InternalQueryClient) executorOption {
 	}
 }
 
+func optExecutorWorkerPoolSize(size int) executorOption {
+	return func(e *executor) error {
+		e.workerPoolSize = size
+		return nil
+	}
+}
+
 // newExecutor returns a new instance of Executor.
 func newExecutor(opts ...executorOption) *executor {
 	e := &executor{
-		client: newNopInternalQueryClient(),
+		client:         newNopInternalQueryClient(),
+		workerPoolSize: 2,
 	}
 	for _, opt := range opts {
 		err := opt(e)
@@ -78,9 +91,27 @@ func newExecutor(opts ...executorOption) *executor {
 			panic(err)
 		}
 	}
+	// this channel cap doesn't necessarily have to be the same as
+	// workerPoolSize... any larger doesn't seem to have an effect in
+	// the few tests we've done at scale with concurrent query
+	// workloads. Possible that it could be smaller.
+	e.work = make(chan job, e.workerPoolSize)
+	for i := 0; i < e.workerPoolSize; i++ {
+		e.workersWG.Add(1)
+		go func() {
+			defer e.workersWG.Done()
+			worker(e.work)
+		}()
+	}
 	return e
 }
 
+func (e *executor) Close() error {
+	close(e.work)
+	e.workersWG.Wait()
+	return nil
+}
+
 // Execute executes a PQL query.
 func (e *executor) Execute(ctx context.Context, index string, q *pql.Query, shards []uint64, opt *execOptions) (QueryResponse, error) {
 	span, ctx := tracing.StartSpanFromContext(ctx, "Executor.Execute")
@@ -2516,6 +2547,24 @@ func (e *executor) mapper(ctx context.Context, ch chan mapResponse, nodes []*Nod
 	return nil
 }
 
+type job struct {
+	shard      uint64
+	mapFn      mapFunc
+	ctx        context.Context
+	resultChan chan mapResponse
+}
+
+func worker(work chan job) {
+	for j := range work {
+		result, err := j.mapFn(j.shard)
+
+		select {
+		case <-j.ctx.Done():
+		case j.resultChan <- mapResponse{result: result, err: err}:
+		}
+	}
+}
+
 // mapperLocal performs map & reduce entirely on the local node.
 func (e *executor) mapperLocal(ctx context.Context, shards []uint64, mapFn mapFunc, reduceFn reduceFunc) (interface{}, error) {
 	span, ctx := tracing.StartSpanFromContext(ctx, "Executor.mapperLocal")
@@ -2524,15 +2573,12 @@ func (e *executor) mapperLocal(ctx context.Context, shards []uint64, mapFn mapFu
 	ch := make(chan mapResponse, len(shards))
 
 	for _, shard := range shards {
-		go func(shard uint64) {
-			result, err := mapFn(shard)
-
-			// Return response to the channel.
-			select {
-			case <-ctx.Done():
-			case ch <- mapResponse{result: result, err: err}:
-			}
-		}(shard)
+		e.work <- job{
+			shard:      shard,
+			mapFn:      mapFn,
+			ctx:        ctx,
+			resultChan: ch,
+		}
 	}
 
 	// Reduce results

holder.go

@@ -234,8 +234,8 @@ func (h *Holder) Close() error {
 // This is used to determine if the rebalancing of data is necessary
 // when a node joins the cluster.
 func (h *Holder) HasData() (bool, error) {
-	h.mu.Lock()
-	defer h.mu.Unlock()
+	h.mu.RLock()
+	defer h.mu.RUnlock()
 	if len(h.indexes) > 0 {
 		return true, nil
 	}
@@ -385,15 +385,21 @@ func (h *Holder) CreateIndex(name string, opt IndexOptions) (*Index, error) {
 // CreateIndexIfNotExists returns an index by name.
 // The index is created if it does not already exist.
 func (h *Holder) CreateIndexIfNotExists(name string, opt IndexOptions) (*Index, error) {
-	h.mu.Lock()
-	defer h.mu.Unlock()
+	h.mu.RLock()
 
 	// Find index in cache first.
 	if index := h.indexes[name]; index != nil {
+		h.mu.RUnlock()
 		return index, nil
 	}
 
-	return h.createIndex(name, opt)
+	h.mu.RUnlock()
+
+	index, err := h.CreateIndex(name, opt)
+	if _, ok := err.(ConflictError); err != nil && !ok {
+		return nil, err
+	}
+	return index, nil
 }
 
 func (h *Holder) createIndex(name string, opt IndexOptions) (*Index, error) {

server.go

@@ -49,13 +49,14 @@ type Server struct { // nolint: maligned
 	closing chan struct{}
 
 	// Internal
-	holder          *Holder
-	cluster         *cluster
-	diagnostics     *diagnosticsCollector
-	executor        *executor
-	hosts           []string
-	clusterDisabled bool
-	serializer      Serializer
+	holder           *Holder
+	cluster          *cluster
+	diagnostics      *diagnosticsCollector
+	executor         *executor
+	executorPoolSize int
+	hosts            []string
+	clusterDisabled  bool
+	serializer       Serializer
 
 	// External
 	systemInfo SystemInfo
@@ -179,13 +180,19 @@ func OptServerGCNotifier(gcn GCNotifier) ServerOption {
 // used to set the implementation of InternalClient.
 func OptServerInternalClient(c InternalClient) ServerOption {
 	return func(s *Server) error {
-		s.executor = newExecutor(optExecutorInternalQueryClient(c))
 		s.defaultClient = c
 		s.cluster.InternalClient = c
 		return nil
 	}
 }
 
+func OptServerExecutorPoolSize(size int) ServerOption {
+	return func(s *Server) error {
+		s.executorPoolSize = size
+		return nil
+	}
+}
+
 // OptServerPrimaryTranslateStore has been deprecated.
 func OptServerPrimaryTranslateStore(store TranslateStore) ServerOption {
 	return func(s *Server) error {
@@ -306,7 +313,6 @@ func NewServer(opts ...ServerOption) (*Server, error) {
 
 		logger: logger.NopLogger,
 	}
-	s.executor = newExecutor(optExecutorInternalQueryClient(s.defaultClient))
 	s.cluster.InternalClient = s.defaultClient
 
 	s.diagnostics.server = s
@@ -317,6 +323,14 @@ func NewServer(opts ...ServerOption) (*Server, error) {
 			return nil, errors.Wrap(err, "applying option")
 		}
 	}
+
+	// set up executor after server opts have been processed
+	executorOpts := []executorOption{optExecutorInternalQueryClient(s.defaultClient)}
+	if s.executorPoolSize > 0 {
+		executorOpts = append(executorOpts, optExecutorWorkerPoolSize(s.executorPoolSize))
+	}
+	s.executor = newExecutor(executorOpts...)
+
 	s.holder.translateFile.logger = s.logger
 
 	path, err := expandDirName(s.dataDir)
@@ -426,6 +440,8 @@ func (s *Server) Open() error {
 
 // Close closes the server and waits for it to shutdown.
 func (s *Server) Close() error {
+	errE := s.executor.Close()
+
 	// Notify goroutines to stop.
 	close(s.closing)
 	s.wg.Wait()
@@ -445,7 +461,11 @@ func (s *Server) Close() error {
 	if errh != nil {
 		return errors.Wrap(errh, "closing holder")
 	}
-	return errors.Wrap(errc, "closing cluster")
+	if errc != nil {
+		return errors.Wrap(errc, "closing cluster")
+	}
+	return errors.Wrap(errE, "closing executor")
+
 }
 
 // loadNodeID gets NodeID from disk, or creates a new value.

server/config.go

@@ -19,6 +19,7 @@ import (
 	"fmt"
 	"log"
 	"net"
+	"runtime"
 	"strconv"
 	"strings"
 	"time"
@@ -85,6 +86,13 @@ type Config struct {
 	// TLS
 	TLS TLSConfig `toml:"tls"`
 
+	// WorkerPoolSize controls how many goroutines are created for
+	// processing queries. Defaults to runtime.NumCPU(). It is
+	// intentionally not defined as a flag... only exposed here so
+	// that we can limit the size while running tests in CI so we
+	// don't exhaust the goroutine limit.
+	WorkerPoolSize int
+
 	Cluster struct {
 		// Disabled controls whether clustering functionality is enabled.
 		Disabled    bool     `toml:"disabled"`
@@ -151,7 +159,10 @@ func NewConfig() *Config {
 		// a bit below your system limits.
 		MaxMapCount:  1000000,
 		MaxFileCount: 1000000,
-		TLS:          TLSConfig{},
+
+		TLS: TLSConfig{},
+
+		WorkerPoolSize: runtime.NumCPU(),
 	}
 
 	// Cluster config.

server/server.go

@@ -284,6 +284,7 @@ func (m *Command) SetupServer() error {
 		pilosa.OptServerMaxWritesPerRequest(m.Config.MaxWritesPerRequest),
 		pilosa.OptServerMetricInterval(time.Duration(m.Config.Metric.PollInterval)),
 		pilosa.OptServerDiagnosticsInterval(diagnosticsInterval),
+		pilosa.OptServerExecutorPoolSize(m.Config.WorkerPoolSize),
 
 		pilosa.OptServerLogger(m.logger),
 		pilosa.OptServerAttrStoreFunc(boltdb.NewAttrStore),

test/pilosa.go

@@ -72,6 +72,7 @@ func newCommand(opts ...server.CommandOption) *Command {
 	m.Config.Bind = "http://localhost:0"
 	m.Config.Cluster.Disabled = true
 	m.Config.Translation.MapSize = 140000
+	m.Config.WorkerPoolSize = 2
 
 	if testing.Verbose() {
 		m.Command.Stdout = os.Stdout