github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/lokifrontend/frontend/v1/queue_test.go (about)

     1  package v1
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"strconv"
     7  	"strings"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/go-kit/log"
    12  	"github.com/grafana/dskit/flagext"
    13  	"github.com/grafana/dskit/services"
    14  	"github.com/pkg/errors"
    15  	"github.com/stretchr/testify/require"
    16  	"github.com/weaveworks/common/httpgrpc"
    17  	"github.com/weaveworks/common/user"
    18  	"google.golang.org/grpc/metadata"
    19  
    20  	"github.com/grafana/loki/pkg/lokifrontend/frontend/v1/frontendv1pb"
    21  )
    22  
    23  func setupFrontend(t *testing.T, config Config) *Frontend {
    24  	logger := log.NewNopLogger()
    25  
    26  	frontend, err := New(config, limits{queriers: 3}, logger, nil)
    27  	require.NoError(t, err)
    28  
    29  	t.Cleanup(func() {
    30  		require.NoError(t, services.StopAndAwaitTerminated(context.Background(), frontend))
    31  	})
    32  	return frontend
    33  }
    34  
    35  func testReq(ctx context.Context, reqID, user string) *request {
    36  	return &request{
    37  		originalCtx: ctx,
    38  		err:         make(chan error, 1),
    39  		request: &httpgrpc.HTTPRequest{
    40  			// Good enough for testing.
    41  			Method: user,
    42  			Url:    reqID,
    43  		},
    44  		response: make(chan *httpgrpc.HTTPResponse, 1),
    45  	}
    46  }
    47  
    48  func TestDequeuesExpiredRequests(t *testing.T) {
    49  	var config Config
    50  	flagext.DefaultValues(&config)
    51  	config.MaxOutstandingPerTenant = 10
    52  	userID := "1"
    53  
    54  	f := setupFrontend(t, config)
    55  
    56  	ctx := user.InjectOrgID(context.Background(), userID)
    57  	expired, cancel := context.WithCancel(ctx)
    58  	cancel()
    59  
    60  	good := 0
    61  	for i := 0; i < config.MaxOutstandingPerTenant; i++ {
    62  		var err error
    63  		if i%5 == 0 {
    64  			good++
    65  			err = f.queueRequest(ctx, testReq(ctx, fmt.Sprintf("good-%d", i), userID))
    66  		} else {
    67  			err = f.queueRequest(ctx, testReq(expired, fmt.Sprintf("expired-%d", i), userID))
    68  		}
    69  
    70  		require.Nil(t, err)
    71  	}
    72  
    73  	// Calling Process will only return when client disconnects or context is finished.
    74  	// We use context timeout to stop Process call.
    75  	ctx2, cancel2 := context.WithTimeout(context.Background(), 1*time.Second)
    76  	defer cancel2()
    77  
    78  	m := &processServerMock{ctx: ctx2, querierID: "querier"}
    79  	err := f.Process(m)
    80  	require.EqualError(t, err, context.DeadlineExceeded.Error())
    81  
    82  	// Verify that only non-expired requests were forwarded to querier.
    83  	for _, r := range m.requests {
    84  		require.True(t, strings.HasPrefix(r.Url, "good-"), r.Url)
    85  	}
    86  	require.Len(t, m.requests, good)
    87  }
    88  
    89  func TestRoundRobinQueues(t *testing.T) {
    90  	var config Config
    91  	flagext.DefaultValues(&config)
    92  
    93  	const (
    94  		requests = 100
    95  		tenants  = 10
    96  	)
    97  
    98  	config.MaxOutstandingPerTenant = requests
    99  
   100  	f := setupFrontend(t, config)
   101  
   102  	for i := 0; i < requests; i++ {
   103  		userID := fmt.Sprint(i / tenants)
   104  		ctx := user.InjectOrgID(context.Background(), userID)
   105  
   106  		err := f.queueRequest(ctx, testReq(ctx, fmt.Sprintf("%d", i), userID))
   107  		require.NoError(t, err)
   108  	}
   109  
   110  	// Calling Process will only return when client disconnects or context is finished.
   111  	// We use context timeout to stop Process call.
   112  	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
   113  	defer cancel()
   114  
   115  	m := &processServerMock{ctx: ctx, querierID: "querier"}
   116  	err := f.Process(m)
   117  	require.EqualError(t, err, context.DeadlineExceeded.Error())
   118  
   119  	require.Len(t, m.requests, requests)
   120  	for i, r := range m.requests {
   121  		intUserID, err := strconv.Atoi(r.Method)
   122  		require.NoError(t, err)
   123  
   124  		require.Equal(t, i%tenants, intUserID)
   125  	}
   126  }
   127  
   128  // This mock behaves as connected querier worker to frontend. It will remember each request
   129  // that frontend sends, and reply with 200 HTTP status code.
   130  type processServerMock struct {
   131  	ctx       context.Context
   132  	querierID string
   133  
   134  	response *frontendv1pb.ClientToFrontend
   135  
   136  	requests []*httpgrpc.HTTPRequest
   137  }
   138  
   139  func (p *processServerMock) Send(client *frontendv1pb.FrontendToClient) error {
   140  	switch {
   141  	case client.GetType() == frontendv1pb.GET_ID:
   142  		p.response = &frontendv1pb.ClientToFrontend{ClientID: p.querierID}
   143  		return nil
   144  
   145  	case client.GetType() == frontendv1pb.HTTP_REQUEST:
   146  		p.requests = append(p.requests, client.HttpRequest)
   147  		p.response = &frontendv1pb.ClientToFrontend{HttpResponse: &httpgrpc.HTTPResponse{Code: 200}}
   148  		return nil
   149  
   150  	default:
   151  		return errors.New("unknown message")
   152  	}
   153  }
   154  
   155  func (p *processServerMock) Recv() (*frontendv1pb.ClientToFrontend, error) {
   156  	if p.response != nil {
   157  		m := p.response
   158  		p.response = nil
   159  		return m, nil
   160  	}
   161  	return nil, errors.New("no message")
   162  }
   163  
   164  func (p *processServerMock) SetHeader(_ metadata.MD) error  { return nil }
   165  func (p *processServerMock) SendHeader(_ metadata.MD) error { return nil }
   166  func (p *processServerMock) SetTrailer(md metadata.MD)      {}
   167  func (p *processServerMock) Context() context.Context       { return p.ctx }
   168  func (p *processServerMock) SendMsg(m interface{}) error    { return nil }
   169  func (p *processServerMock) RecvMsg(m interface{}) error    { return nil }