Add test, more logging, better cancellation.

Signed-off-by: Tom Wilkie <[email protected]>

Author: tomwilkie

pkg/querier/frontend/frontend.go

@@ -1,6 +1,7 @@
 package frontend
 
 import (
+	"context"
 	"flag"
 	"math/rand"
 	"net/http"
@@ -139,28 +140,38 @@ func (f *Frontend) serveHTTP(w http.ResponseWriter, r *http.Request) error {
 
 		case resp = <-request.response:
 		case lastErr = <-request.err:
-			level.Error(f.log).Log("msg", "error processing request", "try", tries, "err", lastErr)
 			resp, _ = httpgrpc.HTTPResponseFromError(lastErr)
 		}
 
-		// Only fail is we get a valid HTTP non-500; otherwise retry.
-		if resp != nil && resp.Code/100 != 5 {
-			retries.Observe(float64(tries))
-			server.WriteResponse(w, resp)
-			return nil
+		// Only retry is we get a HTTP 500 or non-HTTP error.
+		if resp == nil || resp.Code/100 == 5 {
+			level.Error(f.log).Log("msg", "error processing request", "try", tries, "err", lastErr, "resp", resp)
+			continue
 		}
+
+		retries.Observe(float64(tries))
+		server.WriteResponse(w, resp)
+		return nil
 	}
 
 	return lastErr
 }
 
 // Process allows backends to pull requests from the frontend.
 func (f *Frontend) Process(server Frontend_ProcessServer) error {
+
+	// If this request is canceled, ping the condition to unblock. This is done
+	// once, here (instead of in getNextRequest) as we expect calls to Process to
+	// process many requests.
+	go func() {
+		<-server.Context().Done()
+		f.cond.Broadcast()
+	}()
+
 	for {
-		request := f.getNextRequest()
-		if request == nil {
-			// Occurs when server is shutting down.
-			return nil
+		request, err := f.getNextRequest(server.Context())
+		if err != nil {
+			return err
 		}
 
 		if err := server.Send(&ProcessRequest{
@@ -208,16 +219,20 @@ func (f *Frontend) queueRequest(userID string, req *request) error {
 
 // getQueue picks a random queue and takes the next request off of it, so we
 // faily process users queries.  Will block if there are no requests.
-func (f *Frontend) getNextRequest() *request {
+func (f *Frontend) getNextRequest(ctx context.Context) (*request, error) {
 	f.mtx.Lock()
 	defer f.mtx.Unlock()
 
-	for len(f.queues) == 0 && !f.closed {
+	for len(f.queues) == 0 && !f.closed && ctx.Err() == nil {
 		f.cond.Wait()
 	}
 
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+
 	if f.closed {
-		return nil
+		return nil, errServerClosing
 	}
 
 	i, n := 0, rand.Intn(len(f.queues))
@@ -234,7 +249,7 @@ func (f *Frontend) getNextRequest() *request {
 
 		queueDutation.Observe(time.Now().Sub(request.enqueueTime).Seconds())
 		queueLength.Add(-1)
-		return request
+		return request, nil
 	}
 
 	panic("should never happen")

pkg/querier/frontend/frontend_test.go

@@ -0,0 +1,110 @@
+package frontend
+
+import (
+	"context"
+	"fmt"
+	"io/ioutil"
+	"net"
+	"net/http"
+	"sync/atomic"
+	"testing"
+
+	"github.com/go-kit/kit/log"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	httpgrpc_server "github.com/weaveworks/common/httpgrpc/server"
+	"github.com/weaveworks/common/middleware"
+	"google.golang.org/grpc"
+
+	"github.com/weaveworks/common/user"
+	"github.com/weaveworks/cortex/pkg/util"
+)
+
+func TestFrontend(t *testing.T) {
+	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Write([]byte("Hello World"))
+	})
+	test := func(addr string) {
+		req, err := http.NewRequest("GET", fmt.Sprintf("http://%s/", addr), nil)
+		require.NoError(t, err)
+		err = user.InjectOrgIDIntoHTTPRequest(user.InjectOrgID(context.Background(), "1"), req)
+		require.NoError(t, err)
+
+		resp, err := http.DefaultClient.Do(req)
+		require.NoError(t, err)
+		require.Equal(t, 200, resp.StatusCode)
+
+		body, err := ioutil.ReadAll(resp.Body)
+		require.NoError(t, err)
+
+		assert.Equal(t, "Hello World", string(body))
+	}
+	testFrontend(t, handler, test)
+}
+
+func TestFrontendRetries(t *testing.T) {
+	try := int32(0)
+	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if atomic.AddInt32(&try, 1) == 5 {
+			w.Write([]byte("Hello World"))
+			return
+		}
+
+		w.WriteHeader(http.StatusInternalServerError)
+	})
+	test := func(addr string) {
+		req, err := http.NewRequest("GET", fmt.Sprintf("http://%s/", addr), nil)
+		require.NoError(t, err)
+		err = user.InjectOrgIDIntoHTTPRequest(user.InjectOrgID(context.Background(), "1"), req)
+		require.NoError(t, err)
+
+		resp, err := http.DefaultClient.Do(req)
+		require.NoError(t, err)
+		require.Equal(t, 200, resp.StatusCode)
+
+		body, err := ioutil.ReadAll(resp.Body)
+		require.NoError(t, err)
+
+		assert.Equal(t, "Hello World", string(body))
+	}
+	testFrontend(t, handler, test)
+}
+
+func testFrontend(t *testing.T, handler http.Handler, test func(addr string)) {
+	logger := log.NewNopLogger() //log.NewLogfmtLogger(log.NewSyncWriter(os.Stderr))
+
+	var (
+		config       Config
+		workerConfig WorkerConfig
+	)
+	util.DefaultValues(&config, &workerConfig)
+
+	grpcListen, err := net.Listen("tcp", "")
+	require.NoError(t, err)
+	workerConfig.Address = grpcListen.Addr().String()
+
+	httpListen, err := net.Listen("tcp", "")
+	require.NoError(t, err)
+
+	frontend, err := New(config, logger)
+	require.NoError(t, err)
+
+	grpcServer := grpc.NewServer()
+	defer grpcServer.GracefulStop()
+
+	RegisterFrontendServer(grpcServer, frontend)
+
+	httpServer := http.Server{
+		Handler: middleware.AuthenticateUser.Wrap(frontend),
+	}
+	defer httpServer.Shutdown(context.Background())
+
+	go httpServer.Serve(httpListen)
+	go grpcServer.Serve(grpcListen)
+
+	worker, err := NewWorker(workerConfig, httpgrpc_server.NewServer(handler), logger)
+	require.NoError(t, err)
+	defer worker.Stop()
+
+	test(httpListen.Addr().String())
+}

pkg/querier/frontend/worker.go

@@ -12,12 +12,12 @@ import (
 	"github.com/grpc-ecosystem/grpc-opentracing/go/otgrpc"
 	"github.com/mwitkow/go-grpc-middleware"
 	"github.com/opentracing/opentracing-go"
-	"github.com/weaveworks/common/httpgrpc"
-	"github.com/weaveworks/common/httpgrpc/server"
-	"github.com/weaveworks/common/middleware"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/naming"
 
+	"github.com/weaveworks/common/httpgrpc"
+	"github.com/weaveworks/common/httpgrpc/server"
+	"github.com/weaveworks/common/middleware"
 	"github.com/weaveworks/cortex/pkg/util"
 )
 
@@ -38,7 +38,7 @@ type WorkerConfig struct {
 // RegisterFlags adds the flags required to config this to the given FlagSet.
 func (cfg *WorkerConfig) RegisterFlags(f *flag.FlagSet) {
 	f.StringVar(&cfg.Address, "querier.frontend-address", "", "Address of query frontend service.")
-	f.IntVar(&cfg.Parallelism, "querier.worker-parallelism", 1, "Number of simultaneous queries to process.")
+	f.IntVar(&cfg.Parallelism, "querier.worker-parallelism", 10, "Number of simultaneous queries to process.")
 	f.DurationVar(&cfg.DNSLookupDuration, "querier.dns-lookup-period", 10*time.Second, "How often to query DNS.")
 }
 
@@ -66,6 +66,7 @@ func (noopWorker) Stop() {}
 // NewWorker creates a new Worker.
 func NewWorker(cfg WorkerConfig, server *server.Server, log log.Logger) (Worker, error) {
 	if cfg.Address == "" {
+		level.Info(log).Log("msg", "no address specified, not starting worker")
 		return noopWorker{}, nil
 	}
 
@@ -113,15 +114,23 @@ func (w *worker) watchDNSLoop() {
 		}
 	}()
 
-	for updates, err := w.watcher.Next(); err != nil; {
+	for {
+		updates, err := w.watcher.Next()
+		if err != nil {
+			level.Error(w.log).Log("msg", "error from DNS watcher", "err", err)
+			return
+		}
+
 		for _, update := range updates {
 			switch update.Op {
 			case naming.Add:
+				level.Debug(w.log).Log("msg", "adding connection", "addr", update.Addr)
 				ctx, cancel := context.WithCancel(w.ctx)
 				cancels[update.Addr] = cancel
 				w.runMany(ctx, update.Addr)
 
 			case naming.Delete:
+				level.Debug(w.log).Log("msg", "removing connection", "addr", update.Addr)
 				if cancel, ok := cancels[update.Addr]; ok {
 					cancel()
 				}
@@ -134,17 +143,17 @@ func (w *worker) watchDNSLoop() {
 }
 
 // runMany starts N runOne loops for a given address.
-func (w *worker) runMany(ctx context.Context, address string) error {
+func (w *worker) runMany(ctx context.Context, address string) {
 	client, err := connect(address)
 	if err != nil {
-		return err
+		level.Error(w.log).Log("msg", "error connecting", "addr", address, "err", err)
+		return
 	}
 
 	w.wg.Add(w.cfg.Parallelism)
 	for i := 0; i < w.cfg.Parallelism; i++ {
 		go w.runOne(ctx, client)
 	}
-	return nil
 }
 
 // runOne loops, trying to establish a stream to the frontend to begin