github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/querier/worker/frontend_processor_test.go (about)

     1  package worker
     2  
     3  import (
     4  	"context"
     5  	"net"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/go-kit/log"
    10  	"github.com/stretchr/testify/assert"
    11  	"github.com/stretchr/testify/require"
    12  	"go.uber.org/atomic"
    13  	"google.golang.org/grpc"
    14  	"google.golang.org/grpc/credentials/insecure"
    15  	"google.golang.org/grpc/test/bufconn"
    16  
    17  	"github.com/grafana/loki/pkg/util/test"
    18  )
    19  
    20  const bufConnSize = 1024 * 1024
    21  
    22  func TestRecvFailDoesntCancelProcess(t *testing.T) {
    23  	ctx, cancel := context.WithCancel(context.Background())
    24  	defer cancel()
    25  
    26  	listener := bufconn.Listen(bufConnSize)
    27  	defer listener.Close()
    28  	cc, err := grpc.DialContext(ctx, "bufnet", grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) {
    29  		return listener.Dial()
    30  	}), grpc.WithTransportCredentials(insecure.NewCredentials()))
    31  
    32  	require.NoError(t, err)
    33  
    34  	cfg := Config{}
    35  	mgr := newFrontendProcessor(cfg, nil, log.NewNopLogger())
    36  	running := atomic.NewBool(false)
    37  	go func() {
    38  		running.Store(true)
    39  		defer running.Store(false)
    40  
    41  		mgr.processQueriesOnSingleStream(ctx, cc, "test:12345")
    42  	}()
    43  
    44  	test.Poll(t, time.Second, true, func() interface{} {
    45  		return running.Load()
    46  	})
    47  
    48  	// Wait a bit, and verify that processQueriesOnSingleStream is still running, and hasn't stopped
    49  	// just because it cannot contact frontend.
    50  	time.Sleep(100 * time.Millisecond)
    51  	assert.Equal(t, true, running.Load())
    52  
    53  	cancel()
    54  	test.Poll(t, time.Second, false, func() interface{} {
    55  		return running.Load()
    56  	})
    57  }
    58  
    59  func TestContextCancelStopsProcess(t *testing.T) {
    60  	ctx, cancel := context.WithCancel(context.Background())
    61  	defer cancel()
    62  
    63  	listener := bufconn.Listen(bufConnSize)
    64  	defer listener.Close()
    65  	cc, err := grpc.DialContext(ctx, "bufnet", grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) {
    66  		return listener.Dial()
    67  	}), grpc.WithTransportCredentials(insecure.NewCredentials()))
    68  	require.NoError(t, err)
    69  
    70  	pm := newProcessorManager(ctx, &mockProcessor{}, cc, "test")
    71  	pm.concurrency(1)
    72  
    73  	test.Poll(t, time.Second, 1, func() interface{} {
    74  		return int(pm.currentProcessors.Load())
    75  	})
    76  
    77  	cancel()
    78  
    79  	test.Poll(t, time.Second, 0, func() interface{} {
    80  		return int(pm.currentProcessors.Load())
    81  	})
    82  
    83  	pm.stop()
    84  	test.Poll(t, time.Second, 0, func() interface{} {
    85  		return int(pm.currentProcessors.Load())
    86  	})
    87  }