From 4105a1bce05ee16d2cea334c72fb717cb9190b39 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Peter=20=C5=A0tibran=C3=BD?= Date: Thu, 26 Oct 2023 10:22:49 +0200 Subject: [PATCH] Improve distributor's TestStartFinishRequest test. (#6481) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Check for inflight requests and inflight bytes while request is executing. Signed-off-by: Peter Štibraný * Use custom type for context keys to make lint happy. Signed-off-by: Peter Štibraný * Comment. Signed-off-by: Peter Štibraný --------- Signed-off-by: Peter Štibraný --- pkg/distributor/distributor_test.go | 29 ++++++++++++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/pkg/distributor/distributor_test.go b/pkg/distributor/distributor_test.go index 406a0327e59..eb2e672dab1 100644 --- a/pkg/distributor/distributor_test.go +++ b/pkg/distributor/distributor_test.go @@ -4993,9 +4993,30 @@ func TestStartFinishRequest(t *testing.T) { } } - // Pretend push went OK, make sure to call CleanUp. + type ctxType string + const ( + distributorKey ctxType = "dist" + expectedInflightRequestsKey ctxType = "req" + expectedInflightBytesKey ctxType = "bytes" + ) + + // Pretend push went OK, make sure to call CleanUp. Also check for expected values of inflight requests and inflight request size. finishPush := func(ctx context.Context, pushReq *Request) error { defer pushReq.CleanUp() + + distrib := ctx.Value(distributorKey).(*Distributor) + expReq := ctx.Value(expectedInflightRequestsKey).(int64) + expBytes := ctx.Value(expectedInflightBytesKey).(int64) + + reqs := distrib.inflightPushRequests.Load() + if expReq != reqs { + return errors.Errorf("unexpected number of inflight requests: %d, expected: %d", reqs, expReq) + } + + bs := distrib.inflightPushRequestsBytes.Load() + if expBytes != bs { + return errors.Errorf("unexpected number of inflight request bytes: %d, expected: %d", bs, expBytes) + } return nil } @@ -5123,6 +5144,12 @@ func TestStartFinishRequest(t *testing.T) { ds[0].ingestionRate.Tick() ctx := user.InjectOrgID(context.Background(), "user") + + // Set values that are checked by test handler. + ctx = context.WithValue(ctx, distributorKey, ds[0]) + ctx = context.WithValue(ctx, expectedInflightRequestsKey, int64(tc.inflightRequestsBeforePush)+1) + ctx = context.WithValue(ctx, expectedInflightBytesKey, tc.inflightRequestsSizeBeforePush+tc.httpgrpcRequestSize+int64(pushReq.Size())) + if tc.externalCheck { var err error ctx, err = ds[0].StartPushRequest(ctx, tc.httpgrpcRequestSize)