add close() method to api

api.go

@@ -25,6 +25,7 @@ import (
 	"io/ioutil"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/pilosa/pilosa/pql"
@@ -42,6 +43,7 @@ type API struct {
 	cluster *cluster
 	server  *Server
 
+	importWorkersWG      sync.WaitGroup
 	importWorkerPoolSize int
 	importWork           chan importJob
 
@@ -83,8 +85,10 @@ func NewAPI(opts ...apiOption) (*API, error) {
 
 	api.importWork = make(chan importJob, api.importWorkerPoolSize)
 	for i := 0; i < api.importWorkerPoolSize; i++ {
+		api.importWorkersWG.Add(1)
 		go func() {
 			importWorker(api.importWork)
+			defer api.importWorkersWG.Done()
 		}()
 	}
 
@@ -119,6 +123,13 @@ func (api *API) validate(f apiMethod) error {
 	return newAPIMethodNotAllowedError(errors.Errorf("api method %s not allowed in state %s", f, state))
 }
 
+// Close closes the api and waits for it to shutdown.
+func (api *API) Close() error {
+	close(api.importWork)
+	api.importWorkersWG.Wait()
+	return nil
+}
+
 // Query parses a PQL query out of the request and executes it.
 func (api *API) Query(ctx context.Context, req *QueryRequest) (QueryResponse, error) {
 	span, ctx := tracing.StartSpanFromContext(ctx, "API.Query")

server/server.go

@@ -378,6 +378,7 @@ func (m *Command) Close() error {
 	eg := errgroup.Group{}
 	eg.Go(m.Handler.Close)
 	eg.Go(m.Server.Close)
+	eg.Go(m.API.Close)
 	if m.gossipMemberSet != nil {
 		eg.Go(m.gossipMemberSet.Close)
 	}