Function TestForwarder_shutdown

modules/distributor/forwarder_test.go:75–118

(t *testing.T)

}

func TestForwarder_shutdown(t *testing.T) {
	oCfg := overrides.Config{}
	oCfg.RegisterFlagsAndApplyDefaults(&flag.FlagSet{})
	oCfg.Defaults.MetricsGenerator.Forwarder.QueueSize = 200

	id, err := util.HexStringToTraceID("1234567890abcdef")
	require.NoError(t, err)

	b := test.MakeBatch(10, id)
	keys, rebatchedTraces, _, _, err := requestsByTraceID([]*v1.ResourceSpans{b}, tenantID, 10, 1000)
	require.NoError(t, err)

	o, err := overrides.NewOverrides(oCfg, nil, prometheus.DefaultRegisterer)
	require.NoError(t, err)

	signalCh := make(chan struct{})
	f := newGeneratorForwarder(
		log.NewNopLogger(),
		func(_ context.Context, userID string, k []uint32, traces []*rebatchedTrace, noGenerateMetrics bool) error {
			<-signalCh

			assert.Equal(t, tenantID, userID)
			assert.Equal(t, keys, k)
			assert.Equal(t, rebatchedTraces, traces)
			assert.False(t, noGenerateMetrics)
			return nil
		},
		o,
	)

	require.NoError(t, f.start(context.Background()))
	defer func() {
		go func() {
			// Wait to unblock processing of requests so shutdown and draining the queue is done in parallel
			time.Sleep(time.Second)
			close(signalCh)
		}()
		require.NoError(t, f.stop(nil))
	}()

	for i := 0; i < 100; i++ {
		f.SendTraces(context.Background(), tenantID, keys, rebatchedTraces)
	}
}